From 1e5a5ccd284a1c33ce52b94195166be5a50e32c2 Mon Sep 17 00:00:00 2001
From: Hemang <hemang@invariantlabs.ai>
Date: Fri, 7 Mar 2025 13:45:00 +0100
Subject: [PATCH] Add endpoint to push to explorer. Update image handling with
 gemini.

---
 gateway/converters/gemini_to_invariant.py | 21 +++++++++++----------
 gateway/routes/gemini.py                  |  1 +
 2 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/gateway/converters/gemini_to_invariant.py b/gateway/converters/gemini_to_invariant.py
index 8bb9d0b..241f664 100644
--- a/gateway/converters/gemini_to_invariant.py
+++ b/gateway/converters/gemini_to_invariant.py
@@ -1,5 +1,8 @@
 """Converts the request and response formats from Gemini to Invariant API format."""
 
+import base64
+
+
 def convert_request(request: dict) -> list[dict]:
     """Converts the request from Gemini API to Invariant API format."""
     openai_messages = []
@@ -21,11 +24,16 @@ def convert_request(request: dict) -> list[dict]:
                     if "text" in part:
                         message_content.append({"type": "text", "text": part["text"]})
                     elif "inlineData" in part:
+                        # TODO: Handle other types of inline data.
+                        # 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("_", "/")
                         message_content.append(
                             {
-                                "type": "image",
+                                "type": "image_url",
                                 "image_url": {
-                                    "url": f"data:{part['inlineData']['mime_type']};base64,{part['inlineData']['data']}"
+                                    "url": f"data:{part['inlineData']['mime_type']};base64,{inline_data}"
                                 },
                             }
                         )
@@ -40,14 +48,7 @@ def convert_request(request: dict) -> list[dict]:
                             }
                         )
                 if message_content:
-                    openai_messages.append(
-                        {
-                            "role": "user",
-                            "content": message_content
-                            if len(message_content) > 1
-                            else message_content[0],
-                        }
-                    )
+                    openai_messages.append({"role": "user", "content": message_content})
 
             elif role == "model":
                 for part in content.get("parts", []):
diff --git a/gateway/routes/gemini.py b/gateway/routes/gemini.py
index fdb7a82..cc5dc64 100644
--- a/gateway/routes/gemini.py
+++ b/gateway/routes/gemini.py
@@ -21,6 +21,7 @@ GEMINI_AUTHORIZATION_HEADER = "x-goog-api-key"
 
 
 @gateway.post("/gemini/{api_version}/models/{model}:{endpoint}")
+@gateway.post("/{dataset_name}/gemini/{api_version}/models/{model}:{endpoint}")
 async def gemini_generate_content_gateway(
     request: Request,
     api_version: str,