From 5a6dd5c813bc22d21f229e0650182385471f4b85 Mon Sep 17 00:00:00 2001 From: Hemang Date: Mon, 10 Mar 2025 10:35:18 +0100 Subject: [PATCH] Add push to explorer support for Gemini while streaming. --- gateway/converters/gemini_to_invariant.py | 14 +++-- gateway/routes/gemini.py | 71 +++++++++++++++++++++-- 2 files changed, 76 insertions(+), 9 deletions(-) diff --git a/gateway/converters/gemini_to_invariant.py b/gateway/converters/gemini_to_invariant.py index 241f664..b7d8f12 100644 --- a/gateway/converters/gemini_to_invariant.py +++ b/gateway/converters/gemini_to_invariant.py @@ -28,7 +28,11 @@ def convert_request(request: dict) -> list[dict]: # Currently, only images are supported. # Gemini’s API returns URL-safe base64 (uses _ and -). # We need to convert it to standard base64 (uses + and /). - inline_data = part["inlineData"]["data"].replace("-", "+").replace("_", "/") + inline_data = ( + part["inlineData"]["data"] + .replace("-", "+") + .replace("_", "/") + ) message_content.append( { "type": "image_url", @@ -38,13 +42,15 @@ def convert_request(request: dict) -> list[dict]: } ) elif "functionResponse" in part: + result = part["functionResponse"]["response"].get("result", {}) + # TODO: Fix this once Explorer rendering is fixed. + if not isinstance(result, dict): + result = str(result) openai_messages.append( { "role": "tool", "tool_name": part["functionResponse"]["name"], - "content": part["functionResponse"]["response"].get( - "result", {} - ), + "content": result, } ) if message_content: diff --git a/gateway/routes/gemini.py b/gateway/routes/gemini.py index cc5dc64..a7bc1dd 100644 --- a/gateway/routes/gemini.py +++ b/gateway/routes/gemini.py @@ -53,7 +53,6 @@ async def gemini_generate_content_gateway( request_body_bytes = await request.body() request_body_json = json.loads(request_body_bytes) - print("Here is the request: ", request_body_json) client = httpx.AsyncClient(timeout=httpx.Timeout(CLIENT_TIMEOUT)) gemini_api_url = f"https://generativelanguage.googleapis.com/{api_version}/models/{model}:{endpoint}" @@ -71,6 +70,8 @@ async def gemini_generate_content_gateway( client, gemini_request, dataset_name, + request_body_json, + invariant_authorization, ) response = await client.send(gemini_request) return await handle_non_streaming_response( @@ -82,6 +83,8 @@ async def stream_response( client: httpx.AsyncClient, gemini_request: httpx.Request, dataset_name: Optional[str], + request_body_json: dict[str, Any], + invariant_authorization: Optional[str], ) -> Response: """Handles streaming the Gemini response to the client""" @@ -96,23 +99,82 @@ async def stream_response( raise HTTPException(status_code=response.status_code, detail=error_detail) async def event_generator() -> Any: + # Store the progressively merged response + merged_response = { + "candidates": [{"content": {"parts": []}, "finishReason": None}] + } + async for chunk in response.aiter_bytes(): chunk_text = chunk.decode().strip() if not chunk_text: continue # Yield chunk immediately to the client - print("Here is the response chunk: ", chunk) yield chunk - # Send full merged response to the explorer + # Parse and update merged_response incrementally + process_chunk_text(merged_response, chunk_text) + if dataset_name: # Push to Explorer - pass + await push_to_explorer( + dataset_name, + merged_response, + request_body_json, + invariant_authorization, + ) return StreamingResponse(event_generator(), media_type="text/event-stream") +def process_chunk_text( + merged_response: dict[str, Any], + chunk_text: str, +) -> None: + """Processes the chunk text and updates the merged_response to be sent to the explorer""" + # Split the chunk text into individual JSON strings + # A single chunk can contain multiple "data: " sections + for json_string in chunk_text.split("data: "): + json_string = json_string.replace("data: ", "").strip() + + if not json_string: + continue + + try: + json_chunk = json.loads(json_string) + except json.JSONDecodeError: + print("Warning: Could not parse chunk:", json_string) + + update_merged_response(merged_response, json_chunk) + + +def update_merged_response(merged_response: dict[str, Any], chunk_json: dict) -> None: + """Updates the merged response incrementally with a new chunk.""" + candidates = chunk_json.get("candidates", []) + + for candidate in candidates: + content = candidate.get("content", {}) + parts = content.get("parts", []) + + for part in parts: + if "text" in part: + existing_parts = merged_response["candidates"][0]["content"]["parts"] + if existing_parts and "text" in existing_parts[-1]: + existing_parts[-1]["text"] += part["text"] + else: + existing_parts.append({"text": part["text"]}) + + if "functionCall" in part: + merged_response["candidates"][0]["content"]["parts"].append( + {"functionCall": part["functionCall"]} + ) + + if "role" in content: + merged_response["candidates"][0]["content"]["role"] = content["role"] + + if "finishReason" in candidate: + merged_response["candidates"][0]["finishReason"] = candidate["finishReason"] + async def push_to_explorer( dataset_name: str, merged_response: dict[str, Any], @@ -138,7 +200,6 @@ async def handle_non_streaming_response( """Handles non-streaming Gemini responses""" try: json_response = response.json() - print("Here is the response: ", json_response) except json.JSONDecodeError as e: raise HTTPException( status_code=response.status_code,