Support returning ContentToolResultImage from a tool with ChatAnthropic (#118)

* wip support ContentToolResultImage/ContentToolResultResource where we can

* Update chatlas/_google.py

Co-authored-by: Copilot <[email protected]>

* Simplify

---------

Co-authored-by: Copilot <[email protected]>

Author: cpsievert

chatlas/__init__.py

@@ -2,7 +2,7 @@
 from ._anthropic import ChatAnthropic, ChatBedrockAnthropic
 from ._auto import ChatAuto
 from ._chat import Chat
-from ._content import ContentToolRequest, ContentToolResult
+from ._content import ContentToolRequest, ContentToolResult, ContentToolResultImage
 from ._content_image import content_image_file, content_image_plot, content_image_url
 from ._content_pdf import content_pdf_file, content_pdf_url
 from ._databricks import ChatDatabricks
@@ -46,6 +46,7 @@
     "content_pdf_url",
     "ContentToolRequest",
     "ContentToolResult",
+    "ContentToolResultImage",
     "interpolate",
     "interpolate_file",
     "Provider",

chatlas/_anthropic.py

@@ -17,6 +17,8 @@
     ContentText,
     ContentToolRequest,
     ContentToolResult,
+    ContentToolResultImage,
+    ContentToolResultResource,
 )
 from ._logging import log_model_default
 from ._provider import Provider, StandardModelParamNames, StandardModelParams
@@ -515,8 +517,26 @@ def _as_content_block(content: Content) -> "ContentBlockParam":
                 "tool_use_id": content.id,
                 "is_error": content.error is not None,
             }
-            # Anthropic supports non-text contents like ImageBlockParam
-            res["content"] = content.get_model_value()  # type: ignore
+
+            if isinstance(content, ContentToolResultImage):
+                res["content"] = [
+                    {
+                        "type": "image",
+                        "source": {
+                            "type": "base64",
+                            "media_type": content.mime_type,
+                            "data": content.value,
+                        },
+                    }
+                ]
+            elif isinstance(content, ContentToolResultResource):
+                raise NotImplementedError(
+                    "ContentToolResultResource is not currently supported by Anthropic."
+                )
+            else:
+                # Anthropic supports non-text contents like ImageBlockParam
+                res["content"] = content.get_model_value()  # type: ignore
+
             return res
 
         raise ValueError(f"Unknown content type: {type(content)}")

chatlas/_google.py

@@ -16,6 +16,8 @@
     ContentText,
     ContentToolRequest,
     ContentToolResult,
+    ContentToolResultImage,
+    ContentToolResultResource,
 )
 from ._logging import log_model_default
 from ._merge import merge_dicts
@@ -424,6 +426,13 @@ def _as_part_type(self, content: Content) -> "Part":
                 )
             )
         elif isinstance(content, ContentToolResult):
+            if isinstance(
+                content, (ContentToolResultImage, ContentToolResultResource)
+            ):
+                raise NotImplementedError(
+                    "Tool results with images or resources aren't supported by Google (Gemini). "
+                )
+
             if content.error:
                 resp = {"error": content.error}
             else:

chatlas/_openai.py

@@ -18,6 +18,8 @@
     ContentText,
     ContentToolRequest,
     ContentToolResult,
+    ContentToolResultImage,
+    ContentToolResultResource,
 )
 from ._logging import log_model_default
 from ._merge import merge_dicts
@@ -482,6 +484,12 @@ def _as_message_param(turns: list[Turn]) -> list["ChatCompletionMessageParam"]:
                             }
                         )
                     elif isinstance(x, ContentToolResult):
+                        if isinstance(
+                            x, (ContentToolResultImage, ContentToolResultResource)
+                        ):
+                            raise NotImplementedError(
+                                "OpenAI does not support tool results with images or resources."
+                            )
                         tool_results.append(
                             ChatCompletionToolMessageParam(
                                 # Currently, OpenAI only allows for text content in tool results

tests/test_provider_anthropic.py

@@ -2,8 +2,7 @@
 
 import httpx
 import pytest
-
-from chatlas import ChatAnthropic, ContentToolResult
+from chatlas import ChatAnthropic, ContentToolResultImage
 
 from .conftest import (
     assert_data_extraction,
@@ -108,17 +107,10 @@ def get_picture():
         # Local copy of https://upload.wikimedia.org/wikipedia/commons/4/47/PNG_transparency_demonstration_1.png
         with open(test_images_dir / "dice.png", "rb") as image:
             bytez = image.read()
-        res = [
-            {
-                "type": "image",
-                "source": {
-                    "type": "base64",
-                    "media_type": "image/png",
-                    "data": base64.b64encode(bytez).decode("utf-8"),
-                },
-            }
-        ]
-        return ContentToolResult(value=res, model_format="as_is")
+        return ContentToolResultImage(
+            value=base64.b64encode(bytez).decode("utf-8"),
+            mime_type="image/png",
+        )
 
     chat = ChatAnthropic()
     chat.register_tool(get_picture)