feat!(ContentToolResult): add a new model_format parameter (#87)

* feat!(ContentToolResult): add a new model_format parameter

* Update changelog

* Use orjson over json

* More sophisticated json dumping approach

* Improve docstring of ContentToolResult

* Address feedback

* Always dump to a string

* Small tweak to changelog

Author: cpsievert

CHANGELOG.md

@@ -12,10 +12,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### New features
 
 * Added `ChatDatabricks()`, for chatting with Databrick's [foundation models](https://docs.databricks.com/aws/en/machine-learning/model-serving/score-foundation-models). (#82)
-* `.stream()` and `.stream_async()` gain a `content` argument. Set this to `"all"` to include `ContentToolRequest` and `ContentToolResponse` instances in the stream. (#75)
-* `ContentToolRequest` and `ContentToolResponse` are now exported to `chatlas` namespace. (#75)
-* `ContentToolRequest` and `ContentToolResponse` now have `.tagify()` methods, making it so they can render automatically in a Shiny chatbot. (#75)
-* `ContentToolResult` instances can be returned from tools. This allows for custom rendering of the tool result. (#75)
+* `.stream()` and `.stream_async()` gain a `content` argument. Set this to `"all"` to include `ContentToolResult`/`ContentToolRequest` objects in the stream. (#75)
+* `ContentToolResult`/`ContentToolRequest` are now exported to `chatlas` namespace. (#75)
+* `ContentToolResult`/`ContentToolRequest` gain a `.tagify()` method so they render sensibly in a Shiny app. (#75)
+* A tool can now return a `ContentToolResult`. This is useful for: 
+    * Specifying the format used for sending the tool result to the chat model (`model_format`). (#87)
+    * Custom rendering of the tool result (by overriding relevant methods in a subclass). (#75)
 * `Chat` gains a new `.current_display` property. When a `.chat()` or `.stream()` is currently active, this property returns an object with a `.echo()` method (to echo new content to the display). This is primarily useful for displaying custom content during a tool call. (#79)
 
 ### Improvements
@@ -25,11 +27,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   * `.extract_data()` is now supported.
   *  `async` methods are now supported. (#81)
   * Fixed an issue with more than one session being active at once. (#83)
-* `ChatAnthropic()` no longer choke after receiving an output that consists only of whitespace. (#86)
+* `ChatAnthropic()` no longer chokes after receiving an output that consists only of whitespace. (#86)
+* `orjson` is now used for JSON loading and dumping. (#87)
 
 ### Changes
 
 * The `echo` argument of the `.chat()` method defaults to a new value of `"output"`. As a result, tool requests and results are now echoed by default. To revert to the previous behavior, set `echo="text"`. (#78)
+* Tool results are now dumped to JSON by default before being sent to the model. To revert to the previous behavior, have the tool return a `ContentToolResult` with `model_format="str"`. (#87)
 
 ### Breaking changes
 

chatlas/_anthropic.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
 
 import base64
-import json
 import warnings
 from typing import TYPE_CHECKING, Any, Literal, Optional, Union, cast, overload
 
+import orjson
 from pydantic import BaseModel
 
 from ._chat import Chat
@@ -366,8 +366,8 @@ def stream_merge_chunks(self, completion, chunk):
             this_content = completion.content[chunk.index]
             if this_content.type == "tool_use" and isinstance(this_content.input, str):
                 try:
-                    this_content.input = json.loads(this_content.input or "{}")
-                except json.JSONDecodeError as e:
+                    this_content.input = orjson.loads(this_content.input or "{}")
+                except orjson.JSONDecodeError as e:
                     raise ValueError(f"Invalid JSON input: {e}")
         elif chunk.type == "message_delta":
             completion.stop_reason = chunk.delta.stop_reason
@@ -488,12 +488,15 @@ def _as_content_block(content: Content) -> "ContentBlockParam":
                 "input": content.arguments,
             }
         elif isinstance(content, ContentToolResult):
-            return {
+            res: ToolResultBlockParam = {
                 "type": "tool_result",
                 "tool_use_id": content.id,
-                "content": content.get_final_value(),
                 "is_error": content.error is not None,
             }
+            # 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)}")
 
     @staticmethod

chatlas/_auto.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
-import json
 import os
 from typing import Callable, Literal, Optional
 
+import orjson
+
 from ._anthropic import ChatAnthropic, ChatBedrockAnthropic
 from ._chat import Chat
 from ._databricks import ChatDatabricks
@@ -178,7 +179,7 @@ def ChatAuto(
 
     env_kwargs = {}
     if env_kwargs_str := os.environ.get("CHATLAS_CHAT_ARGS"):
-        env_kwargs = json.loads(env_kwargs_str)
+        env_kwargs = orjson.loads(env_kwargs_str)
 
     kwargs = {**kwargs, **env_kwargs, **base_args}
     kwargs = {k: v for k, v in kwargs.items() if v is not None}

chatlas/_content.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
 
-import json
 import textwrap
 from pprint import pformat
 from typing import TYPE_CHECKING, Any, Literal, Optional, Union
 
+import orjson
 from pydantic import BaseModel, ConfigDict
 
 if TYPE_CHECKING:
@@ -218,27 +218,50 @@ class ContentToolResult(Content):
     """
     The result of calling a tool/function
 
-    This content type isn't meant to be used directly. Instead, it's
-    automatically generated by [](`~chatlas.Chat`) when a tool/function is
-    called (in response to a [](`~chatlas.ContentToolRequest`)).
+    A content type representing the result of a tool function call. When a model
+    requests a tool function, [](`~chatlas.Chat`) will create, (optionally)
+    echo, (optionally) yield, and store this content type in the chat history.
+
+    A tool function may also construct an instance of this class and return it.
+    This is useful for a tool that wishes to customize how the result is handled
+    (e.g., the format of the value sent to the model).
 
     Parameters
     ----------
     value
-        The value returned by the tool/function (to be sent to the model).
+        The return value of the tool/function.
+    model_format
+        The format used for sending the value to the model. The default,
+        `"auto"`, first attempts to format the value as a JSON string. If that
+        fails, it gets converted to a string via `str()`. To force
+        `orjson.dumps()` or `str()`, set to `"json"` or `"str"`. Finally,
+        `"as_is"` is useful for doing your own formatting and/or passing a
+        non-string value (e.g., a list or dict) straight to the model.
+        Non-string values are useful for tools that return images or other
+        'known' non-text content types.
     error
-        An exception that occurred during the tool request. If this is set, the
+        An exception that occurred while invoking the tool. If this is set, the
         error message sent to the model and the value is ignored.
     extra
        Additional data associated with the tool result that isn't sent to the
        model.
     request
         Not intended to be used directly. It will be set when the
         :class:`~chatlas.Chat` invokes the tool.
+
+    Note
+    ----
+    When `model_format` is `"json"` (or `"auto"`), and the value has a
+    `.to_json()`/`.to_dict()` method, those methods are called to obtain the
+    JSON representation of the value. This is convenient for classes, like
+    `pandas.DataFrame`, that have a `.to_json()` method, but don't necessarily
+    dump to JSON directly. If this happens to not be the desired behavior, set
+    `model_format="as_is"` return the desired value as-is.
     """
 
     # public
     value: Any
+    model_format: Literal["auto", "json", "str", "as_is"] = "auto"
     error: Optional[Exception] = None
     extra: Any = None
 
@@ -266,22 +289,11 @@ def arguments(self):
             )
         return self.request.arguments
 
-    def _get_value(self, pretty: bool = False) -> str:
-        if self.error:
-            return f"Tool call failed with error: '{self.error}'"
-        if not pretty:
-            return str(self.value)
-        try:
-            json_val = json.loads(self.value)  # type: ignore
-            return pformat(json_val, indent=2, sort_dicts=False)
-        except:  # noqa
-            return str(self.value)
-
     # Primarily used for `echo="all"`...
     def __str__(self):
         prefix = "✅ tool result" if not self.error else "❌ tool error"
         comment = f"# {prefix} ({self.id})"
-        value = self._get_value(pretty=True)
+        value = self._get_display_value()
         return f"""```python\n{comment}\n{value}\n```"""
 
     # ... and for displaying in the notebook
@@ -295,9 +307,62 @@ def __repr__(self, indent: int = 0):
             res += f" error='{self.error}'"
         return res + ">"
 
-    # The actual value to send to the model
-    def get_final_value(self) -> str:
-        return self._get_value()
+    # Format the value for display purposes
+    def _get_display_value(self) -> object:
+        if self.error:
+            return f"Tool call failed with error: '{self.error}'"
+
+        val = self.value
+
+        # If value is already a dict or list, format it directly
+        if isinstance(val, (dict, list)):
+            return pformat(val, indent=2, sort_dicts=False)
+
+        # For string values, try to parse as JSON
+        if isinstance(val, str):
+            try:
+                json_val = orjson.loads(val)
+                return pformat(json_val, indent=2, sort_dicts=False)
+            except orjson.JSONDecodeError:
+                # Not valid JSON, return as string
+                return val
+
+        return val
+
+    def get_model_value(self) -> object:
+        "Get the actual value sent to the model."
+
+        if self.error:
+            return f"Tool call failed with error: '{self.error}'"
+
+        val, mode = (self.value, self.model_format)
+
+        if isinstance(val, str):
+            return val
+
+        if mode == "auto":
+            try:
+                return self._to_json(val)
+            except Exception:
+                return str(val)
+        elif mode == "json":
+            return self._to_json(val)
+        elif mode == "str":
+            return str(val)
+        elif mode == "as_is":
+            return val
+        else:
+            raise ValueError(f"Unknown format mode: {mode}")
+
+    @staticmethod
+    def _to_json(value: Any) -> object:
+        if hasattr(value, "to_json") and callable(value.to_json):
+            return value.to_json()
+
+        if hasattr(value, "to_dict") and callable(value.to_dict):
+            value = value.to_dict()
+
+        return orjson.dumps(value).decode("utf-8")
 
     def tagify(self) -> "TagChild":
         """
@@ -317,7 +382,7 @@ def tagify(self) -> "TagChild":
             header = f"❌ Failed to call tool <code>{self.name}</code>"
 
         args = self._arguments_str()
-        content = self._get_value(pretty=True)
+        content = self._get_display_value()
 
         return HTML(
             textwrap.dedent(f"""
@@ -355,7 +420,7 @@ class ContentJson(Content):
     content_type: ContentTypeEnum = "json"
 
     def __str__(self):
-        return json.dumps(self.value, indent=2)
+        return orjson.dumps(self.value, option=orjson.OPT_INDENT_2).decode("utf-8")
 
     def _repr_markdown_(self):
         return f"""```json\n{self.__str__()}\n```"""

chatlas/_google.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 
 import base64
-import json
 from typing import TYPE_CHECKING, Any, Literal, Optional, cast, overload
 
+import orjson
 from pydantic import BaseModel
 
 from ._chat import Chat
@@ -432,7 +432,7 @@ def _as_part_type(self, content: Content) -> "Part":
             if content.error:
                 resp = {"error": content.error}
             else:
-                resp = {"result": str(content.value)}
+                resp = {"result": content.get_model_value()}
             return Part(
                 # TODO: seems function response parts might need role='tool'???
                 # https://github.com/googleapis/python-genai/blame/c8cfef85c/README.md#L344
@@ -470,7 +470,7 @@ def _as_turn(
             text = part.get("text")
             if text:
                 if has_data_model:
-                    contents.append(ContentJson(value=json.loads(text)))
+                    contents.append(ContentJson(value=orjson.loads(text)))
                 else:
                     contents.append(ContentText(text=text))
             function_call = part.get("function_call")

chatlas/_ollama.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
-import json
 import re
 import urllib.request
 from typing import TYPE_CHECKING, Optional
 
+import orjson
+
 from ._chat import Chat
 from ._openai import ChatOpenAI
 from ._turn import Turn
@@ -121,7 +122,7 @@ def ChatOllama(
 
 def ollama_models(base_url: str) -> list[str]:
     res = urllib.request.urlopen(url=f"{base_url}/api/tags")
-    data = json.loads(res.read())
+    data = orjson.loads(res.read())
     return [re.sub(":latest$", "", x["name"]) for x in data["models"]]
 
 

chatlas/_openai.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 
 import base64
-import json
 from typing import TYPE_CHECKING, Any, Literal, Optional, cast, overload
 
+import orjson
 from pydantic import BaseModel
 
 from ._chat import Chat
@@ -433,7 +433,7 @@ def _as_message_param(turns: list[Turn]) -> list["ChatCompletionMessageParam"]:
                                 "id": x.id,
                                 "function": {
                                     "name": x.name,
-                                    "arguments": json.dumps(x.arguments),
+                                    "arguments": orjson.dumps(x.arguments).decode("utf-8"),
                                 },
                                 "type": "function",
                             }
@@ -499,8 +499,8 @@ def _as_message_param(turns: list[Turn]) -> list["ChatCompletionMessageParam"]:
                     elif isinstance(x, ContentToolResult):
                         tool_results.append(
                             ChatCompletionToolMessageParam(
-                                # TODO: a tool could return an image!?!
-                                content=x.get_final_value(),
+                                # Currently, OpenAI only allows for text content in tool results
+                                content=cast(str, x.get_model_value()),
                                 tool_call_id=x.id,
                                 role="tool",
                             )
@@ -529,7 +529,7 @@ def _as_turn(
         contents: list[Content] = []
         if message.content is not None:
             if has_data_model:
-                data = json.loads(message.content)
+                data = orjson.loads(message.content)
                 contents = [ContentJson(value=data)]
             else:
                 contents = [ContentText(text=message.content)]
@@ -544,8 +544,8 @@ def _as_turn(
 
                 args = {}
                 try:
-                    args = json.loads(func.arguments) if func.arguments else {}
-                except json.JSONDecodeError:
+                    args = orjson.loads(func.arguments) if func.arguments else {}
+                except orjson.JSONDecodeError:
                     raise ValueError(
                         f"The model's completion included a tool request ({func.name}) "
                         "with invalid JSON for input arguments: '{func.arguments}'"

pyproject.toml

@@ -8,6 +8,7 @@ dependencies = [
   "requests",
   "pydantic>=2.0",
   "jinja2",
+  "orjson",
   "rich",
 ]
 classifiers = [

tests/conftest.py

@@ -256,3 +256,8 @@ def assert_pdf_local(chat_fun: ChatFun):
     wait=wait_exponential(min=1, max=60),
     reraise=True,
 )
+
+
+@pytest.fixture
+def test_images_dir():
+    return Path(__file__).parent / "images"