feat: add vision support for Google (#141)

block · Oct 13, 2024 · 9679f07 · 9679f07
1 parent 6ea6d14
commit 9679f07
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/packages/exchange/src/exchange/providers/google.py b/packages/exchange/src/exchange/providers/google.py
@@ -7,7 +7,7 @@
 from exchange.content import Text, ToolResult, ToolUse
 from exchange.providers.base import Provider, Usage
 from tenacity import retry, wait_fixed, stop_after_attempt
-from exchange.providers.utils import raise_for_status, retry_if_status
+from exchange.providers.utils import raise_for_status, retry_if_status, encode_image
 
 GOOGLE_HOST = "https://generativelanguage.googleapis.com/v1beta"
 
@@ -111,9 +111,20 @@ def messages_to_google_spec(messages: List[Message]) -> List[Dict[str, Any]]:
                 elif isinstance(content, ToolUse):
                     converted["parts"].append({"functionCall": {"name": content.name, "args": content.parameters}})
                 elif isinstance(content, ToolResult):
-                    converted["parts"].append(
-                        {"functionResponse": {"name": content.tool_use_id, "response": {"content": content.output}}}
-                    )
+                    if content.output.startswith('"image:'):
+                        image_path = content.output.replace('"image:', "").replace('"', "")
+                        converted["parts"].append(
+                            {
+                                "inline_data": {
+                                    "mime_type": "image/png",
+                                    "data": f"{encode_image(image_path)}",
+                                }
+                            }
+                        )
+                    else:
+                        converted["parts"].append(
+                            {"functionResponse": {"name": content.tool_use_id, "response": {"content": content.output}}}
+                        )
             messages_spec.append(converted)
 
         if not messages_spec:

diff --git a/packages/exchange/tests/test_integration_vision.py b/packages/exchange/tests/test_integration_vision.py
@@ -9,6 +9,7 @@
 
 cases = [
     (get_provider("openai"), os.getenv("OPENAI_MODEL", "gpt-4o-mini")),
+    (get_provider("google"), os.getenv("GOOGLE_MODEL", "gemini-1.5-flash")),
 ]