From e6ed42de1dd65fc26596a2169f868213038f6977 Mon Sep 17 00:00:00 2001 From: Hemang Date: Mon, 3 Mar 2025 13:14:10 +0100 Subject: [PATCH] Correct gemini API route and add logic to stream response. This is a pass through with no push to explorer. --- proxy/routes/anthropic.py | 9 ++- proxy/routes/gemini.py | 140 +++++++++++++++++++++++++++----------- proxy/routes/open_ai.py | 9 ++- 3 files changed, 109 insertions(+), 49 deletions(-) diff --git a/proxy/routes/anthropic.py b/proxy/routes/anthropic.py index c30461d..f62c044 100644 --- a/proxy/routes/anthropic.py +++ b/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( diff --git a/proxy/routes/gemini.py b/proxy/routes/gemini.py index 4475a06..da6d060 100644 --- a/proxy/routes/gemini.py +++ b/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 = /models/:generateContent` - reconstructed_endpoint = "/".join(endpoint_parts) - elif endpoint_parts[2] == "models": - # Case 2: With dataset_name - # `endpoint = //models/: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//models/:generateContent or \ - /api/v1/proxy/gemini//models/: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//models/:generateContent or \ + /api/v1/proxy//gemini/models/: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), + ) diff --git a/proxy/routes/open_ai.py b/proxy/routes/open_ai.py index 5357045..8053b71 100644 --- a/proxy/routes/open_ai.py +++ b/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(