Correct gemini API route and add logic to stream response. This is a pass through with no push to explorer.

proxy/routes/anthropic.py

@@ -82,11 +82,10 @@ async def anthropic_v1_messages_proxy(
         return await handle_streaming_response(
             client, anthropic_request, dataset_name, invariant_authorization
         )
-    else:
-        response = await client.send(anthropic_request)
-        return await handle_non_streaming_response(
-            response, dataset_name, request_body_json, invariant_authorization
-        )
+    response = await client.send(anthropic_request)
+    return await handle_non_streaming_response(
+        response, dataset_name, request_body_json, invariant_authorization
+    )
 
 
 async def push_to_explorer(

proxy/routes/gemini.py

@@ -1,61 +1,123 @@
 """Proxy service to forward requests to the Gemini APIs"""
 
 import json
+from typing import Any, Optional
 
+import httpx
 from common.config_manager import ProxyConfig, ProxyConfigManager
-from fastapi import APIRouter, Depends, Request, Response
-from utils.constants import IGNORED_HEADERS
+from fastapi import APIRouter, Depends, HTTPException, Query, Request, Response
+from fastapi.responses import StreamingResponse
+from utils.constants import CLIENT_TIMEOUT, IGNORED_HEADERS
 
 proxy = APIRouter()
 
 
-def _extract_dataset_name_and_endpoint(endpoint: str):
-    """Extracts the dataset name and endpoint from the given endpoint."""
-    endpoint_parts = endpoint.split("/")
-    dataset_name = None
-    if endpoint_parts[1] == "models":
-        # Case 1: Without dataset_name
-        # `endpoint = <version>/models/<model-name>:generateContent`
-        reconstructed_endpoint = "/".join(endpoint_parts)
-    elif endpoint_parts[2] == "models":
-        # Case 2: With dataset_name
-        # `endpoint = <dataset-name>/<version>/models/<model-name>:generateContent`
-        dataset_name = endpoint_parts[0]
-        reconstructed_endpoint = "/".join(endpoint_parts[1:])
-    else:
-        # Case 3: Invalid endpoint
-        return Response(
-            content=f"Invalid endpoint: {endpoint} - the endpoint should be in the format: \
-            /api/v1/proxy/gemini/<version>/models/<model-name>:generateContent or \
-            /api/v1/proxy/gemini/<dataset-name>/<version>models/<model-name>:generateContent",
-            status_code=400,
-        )
-    return dataset_name, reconstructed_endpoint
-
-
-@proxy.post(
-    "/gemini/{endpoint:path}",
-)
+@proxy.post("/gemini/{api_version}/models/{model}:{endpoint}")
+@proxy.post("/{dataset_name}/gemini/{api_version}/models/{model}:{endpoint}")
 async def gemini_generate_content_proxy(
     request: Request,
+    api_version: str,
+    model: str,
     endpoint: str,
+    dataset_name: str = None,
+    alt: str = Query(
+        None, title="Response Format", description="Set to 'sse' for streaming"
+    ),
     config: ProxyConfig = Depends(ProxyConfigManager.get_config),  # pylint: disable=unused-argument
 ) -> Response:
-    """Proxy calls to the OpenAI APIs"""
+    """Proxy calls to the Gemini GenerateContent API"""
+    if "generateContent" != endpoint and "streamGenerateContent" != endpoint:
+        return Response(
+            content="Invalid endpoint - the only endpoints supported are: \
+            /api/v1/proxy/gemini/<version>/models/<model-name>:generateContent or \
+            /api/v1/proxy/<dataset-name>/gemini/<version>models/<model-name>:generateContent",
+            status_code=400,
+        )
     headers = {
         k: v for k, v in request.headers.items() if k.lower() not in IGNORED_HEADERS
     }
     headers["accept-encoding"] = "identity"
 
     request_body_bytes = await request.body()
-    request_body_json = json.loads(request_body_bytes)
-    api_key = headers.get("x-goog-api-key")
-    print(f"API Key: {api_key}")
-    print("request body json: ", request_body_json)
-    dataset_name, reconstructed_endpoint = _extract_dataset_name_and_endpoint(endpoint)
+    client = httpx.AsyncClient(timeout=httpx.Timeout(CLIENT_TIMEOUT))
+    gemini_api_url = f"https://generativelanguage.googleapis.com/{api_version}/models/{model}:{endpoint}"
+    if alt == "sse":
+        gemini_api_url += "?alt=sse"
+    gemini_request = client.build_request(
+        "POST",
+        gemini_api_url,
+        content=request_body_bytes,
+        headers=headers,
+    )
 
-    print(f"API Key: {api_key}")
-    print("Processed Endpoint: ", reconstructed_endpoint)
-    print("Dataset Name: ", dataset_name)
+    if alt == "sse" or endpoint == "streamGenerateContent":
+        return await stream_response(
+            client,
+            gemini_request,
+            dataset_name,
+        )
+    response = await client.send(gemini_request)
+    return await handle_non_streaming_response(response, dataset_name)
 
-    return {}
+
+async def stream_response(
+    client: httpx.AsyncClient,
+    gemini_request: httpx.Request,
+    dataset_name: Optional[str],
+) -> Response:
+    """Handles streaming the Gemini response to the client"""
+
+    response = await client.send(gemini_request, stream=True)
+    if response.status_code != 200:
+        error_content = await response.aread()
+        try:
+            error_json = json.loads(error_content.decode("utf-8"))
+            error_detail = error_json.get("error", "Unknown error from Gemini API")
+        except json.JSONDecodeError:
+            error_detail = {"error": "Failed to parse Gemini error response"}
+        raise HTTPException(status_code=response.status_code, detail=error_detail)
+
+    async def event_generator() -> Any:
+        async for chunk in response.aiter_bytes():
+            chunk_text = chunk.decode().strip()
+            if not chunk_text:
+                continue
+
+            # Yield chunk immediately to the client (proxy behavior)
+            yield chunk
+
+        # Send full merged response to the explorer
+        if dataset_name:
+            # Push to Explorer
+            pass
+
+    return StreamingResponse(event_generator(), media_type="text/event-stream")
+
+
+async def handle_non_streaming_response(
+    response: httpx.Response,
+    dataset_name: Optional[str],
+) -> Response:
+    """Handles non-streaming Gemini responses"""
+    try:
+        json_response = response.json()
+    except json.JSONDecodeError as e:
+        raise HTTPException(
+            status_code=response.status_code,
+            detail="Invalid JSON response received from Gemini API",
+        ) from e
+    if response.status_code != 200:
+        raise HTTPException(
+            status_code=response.status_code,
+            detail=json_response.get("error", "Unknown error from Gemini API"),
+        )
+    if dataset_name:
+        # Push to Explorer
+        pass
+
+    return Response(
+        content=json.dumps(json_response),
+        status_code=response.status_code,
+        media_type="application/json",
+        headers=dict(response.headers),
+    )

proxy/routes/open_ai.py

@@ -87,11 +87,10 @@ async def openai_chat_completions_proxy(
             request_body_json,
             invariant_authorization,
         )
-    async with client:
-        response = await client.send(open_ai_request)
-        return await handle_non_streaming_response(
-            response, dataset_name, request_body_json, invariant_authorization
-        )
+    response = await client.send(open_ai_request)
+    return await handle_non_streaming_response(
+        response, dataset_name, request_body_json, invariant_authorization
+    )
 
 
 async def stream_response(