chatlas 0.6.1__py3-none-any.whl → 0.7.0__py3-none-any.whl

chatlas/_content.py

@@ -1,11 +1,15 @@
 from __future__ import annotations
 
-import json
+import textwrap
 from pprint import pformat
-from typing import Any, Literal, Optional, Union
+from typing import TYPE_CHECKING, Any, Literal, Optional, Union
 
+import orjson
 from pydantic import BaseModel, ConfigDict
 
+if TYPE_CHECKING:
+    from htmltools import TagChild
+
 ImageContentTypes = Literal[
     "image/png",
     "image/jpeg",
@@ -174,7 +178,7 @@ class ContentToolRequest(Content):
     def __str__(self):
         args_str = self._arguments_str()
         func_call = f"{self.name}({args_str})"
-        comment = f"# tool request ({self.id})"
+        comment = f"# 🔧 tool request ({self.id})"
         return f"```python\n{comment}\n{func_call}\n```\n"
 
     def _repr_markdown_(self):
@@ -192,49 +196,104 @@ class ContentToolRequest(Content):
             return ", ".join(f"{k}={v}" for k, v in self.arguments.items())
         return str(self.arguments)
 
+    def tagify(self) -> "TagChild":
+        "Returns an HTML string suitable for passing to htmltools/shiny's `Chat()` component."
+        try:
+            from htmltools import HTML, TagList, head_content, tags
+        except ImportError:
+            raise ImportError(
+                ".tagify() is only intended to be called by htmltools/shiny, ",
+                "but htmltools is not installed. ",
+            )
+
+        html = f"<p></p><span class='chatlas-tool-request'>🔧 Running tool: <code>{self.name}</code></span>"
+
+        return TagList(
+            HTML(html),
+            head_content(tags.style(TOOL_CSS)),
+        )
+
 
 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
     ----------
-    id
-        The unique identifier of the tool request.
     value
-        The value returned by the tool/function.
-    name
-        The name of the tool/function that was called.
+        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 error message if the tool/function call failed.
+        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.
     """
 
-    id: str
-    value: Any = None
-    name: Optional[str] = None
-    error: Optional[str] = None
+    # public
+    value: Any
+    model_format: Literal["auto", "json", "str", "as_is"] = "auto"
+    error: Optional[Exception] = None
+    extra: Any = None
 
+    # "private"
+    request: Optional[ContentToolRequest] = None
     content_type: ContentTypeEnum = "tool_result"
 
-    def _get_value(self, pretty: bool = False) -> str:
-        if self.error:
-            return f"Tool calling 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)
+    @property
+    def id(self):
+        if not self.request:
+            raise ValueError("id is only available after the tool has been called")
+        return self.request.id
+
+    @property
+    def name(self):
+        if not self.request:
+            raise ValueError("name is only available after the tool has been called")
+        return self.request.name
+
+    @property
+    def arguments(self):
+        if not self.request:
+            raise ValueError(
+                "arguments is only available after the tool has been called"
+            )
+        return self.request.arguments
 
     # Primarily used for `echo="all"`...
     def __str__(self):
-        comment = f"# tool result ({self.id})"
-        value = self._get_value(pretty=True)
+        prefix = "✅ tool result" if not self.error else "❌ tool error"
+        comment = f"# {prefix} ({self.id})"
+        value = self._get_display_value()
         return f"""```python\n{comment}\n{value}\n```"""
 
     # ... and for displaying in the notebook
@@ -248,9 +307,99 @@ class ContentToolResult(Content):
             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":
+        """
+        A method for rendering this object via htmltools/shiny.
+        """
+        try:
+            from htmltools import HTML
+        except ImportError:
+            raise ImportError(
+                ".tagify() is only intended to be called by htmltools/shiny, ",
+                "but htmltools is not installed. ",
+            )
+
+        if not self.error:
+            header = f"View result from <code>{self.name}</code>"
+        else:
+            header = f"❌ Failed to call tool <code>{self.name}</code>"
+
+        args = self._arguments_str()
+        content = self._get_display_value()
+
+        return HTML(
+            textwrap.dedent(f"""
+              <details class="chatlas-tool-result">
+                  <summary>{header}</summary>
+                  <div class="chatlas-tool-result-content">
+                      Result: <p><code>{content}</code></p>
+                      Arguments: <p><code>{args}</code></p>
+                  </div>
+              </details>
+            """)
+        )
+
+    def _arguments_str(self) -> str:
+        if isinstance(self.arguments, dict):
+            return ", ".join(f"{k}={v}" for k, v in self.arguments.items())
+        return str(self.arguments)
 
 
 class ContentJson(Content):
@@ -271,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```"""
@@ -345,3 +494,52 @@ def create_content(data: dict[str, Any]) -> ContentUnion:
         return ContentPDF.model_validate(data)
     else:
         raise ValueError(f"Unknown content type: {ct}")
+
+
+TOOL_CSS = """
+/* Get dot to appear inline, even when in a paragraph following the request */
+.chatlas-tool-request + p:has(.markdown-stream-dot) {
+  display: inline;
+}
+
+/* Hide request when anything other than a dot follows it */
+.chatlas-tool-request:not(:has(+ p .markdown-stream-dot)) {
+  display: none;
+}
+
+.chatlas-tool-request, .chatlas-tool-result {
+  font-weight: 300;
+  font-size: 0.9rem;
+}
+
+.chatlas-tool-result {
+  display: inline-block;
+  width: 100%;
+  margin-bottom: 1rem;
+}
+
+.chatlas-tool-result summary {
+  list-style: none;
+  cursor: pointer;
+}
+
+.chatlas-tool-result summary::after {
+  content: url("data:image/svg+xml,%3Csvg xmlns='http://www.w3.org/2000/svg' width='16' height='16' fill='currentColor' class='bi bi-caret-right-fill' viewBox='0 0 16 16'%3E%3Cpath d='m12.14 8.753-5.482 4.796c-.646.566-1.658.106-1.658-.753V3.204a1 1 0 0 1 1.659-.753l5.48 4.796a1 1 0 0 1 0 1.506z'/%3E%3C/svg%3E");
+  font-size: 1.15rem;
+  margin-left: 0.25rem;
+  vertical-align: middle;
+}
+
+.chatlas-tool-result[open] summary::after {
+  content: url("data:image/svg+xml,%3Csvg xmlns='http://www.w3.org/2000/svg' width='16' height='16' fill='currentColor' class='bi bi-caret-down-fill' viewBox='0 0 16 16'%3E%3Cpath d='M7.247 11.14 2.451 5.658C1.885 5.013 2.345 4 3.204 4h9.592a1 1 0 0 1 .753 1.659l-4.796 5.48a1 1 0 0 1-1.506 0z'/%3E%3C/svg%3E");
+}
+
+.chatlas-tool-result-content {
+  border: 1px solid var(--bs-border-color, #0066cc);
+  width: 100%;
+  padding: 1rem;
+  border-radius: var(--bs-border-radius, 0.2rem);
+  margin-top: 1rem;
+  margin-bottom: 1rem;
+}
+"""

chatlas/_databricks.py

@@ -0,0 +1,145 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+from ._chat import Chat
+from ._logging import log_model_default
+from ._openai import OpenAIProvider
+from ._turn import Turn, normalize_turns
+
+if TYPE_CHECKING:
+    from databricks.sdk import WorkspaceClient
+
+    from ._openai import ChatCompletion
+    from .types.openai import SubmitInputArgs
+
+
+def ChatDatabricks(
+    *,
+    system_prompt: Optional[str] = None,
+    model: Optional[str] = None,
+    turns: Optional[list[Turn]] = None,
+    workspace_client: Optional["WorkspaceClient"] = None,
+) -> Chat["SubmitInputArgs", ChatCompletion]:
+    """
+    Chat with a model hosted on Databricks.
+
+    Databricks provides out-of-the-box access to a number of [foundation
+    models](https://docs.databricks.com/en/machine-learning/model-serving/score-foundation-models.html)
+    and can also serve as a gateway for external models hosted by a third party.
+
+    Prerequisites
+    --------------
+
+    ::: {.callout-note}
+    ## Python requirements
+
+    `ChatDatabricks` requires the `databricks-sdk` package: `pip install
+    "chatlas[databricks]"`.
+    :::
+
+    ::: {.callout-note}
+    ## Authentication
+
+    `chatlas` delegates to the `databricks-sdk` package for authentication with
+    Databricks. As such, you can use any of the authentication methods discussed
+    here:
+
+    https://docs.databricks.com/aws/en/dev-tools/sdk-python#authentication
+
+    Note that Python-specific article points to this language-agnostic "unified"
+    approach to authentication:
+
+    https://docs.databricks.com/aws/en/dev-tools/auth/unified-auth
+
+    There, you'll find all the options listed, but a simple approach that
+    generally works well is to set the following environment variables:
+
+    * `DATABRICKS_HOST`: The Databricks host URL for either the Databricks
+      workspace endpoint or the Databricks accounts endpoint.
+    * `DATABRICKS_TOKEN`: The Databricks personal access token.
+    :::
+
+    Parameters
+    ----------
+    system_prompt
+        A system prompt to set the behavior of the assistant.
+    model
+        The model to use for the chat. The default, None, will pick a reasonable
+        default, and warn you about it. We strongly recommend explicitly
+        choosing a model for all but the most casual use.
+    turns
+        A list of turns to start the chat with (i.e., continuing a previous
+        conversation). If not provided, the conversation begins from scratch. Do
+        not provide non-`None` values for both `turns` and `system_prompt`. Each
+        message in the list should be a dictionary with at least `role` (usually
+        `system`, `user`, or `assistant`, but `tool` is also possible). Normally
+        there is also a `content` field, which is a string.
+    workspace_client
+        A `databricks.sdk.WorkspaceClient()` to use for the connection. If not
+        provided, a new client will be created.
+
+    Returns
+    -------
+    Chat
+        A chat object that retains the state of the conversation.
+    """
+    if model is None:
+        model = log_model_default("databricks-dbrx-instruct")
+
+    return Chat(
+        provider=DatabricksProvider(
+            model=model,
+            workspace_client=workspace_client,
+        ),
+        turns=normalize_turns(
+            turns or [],
+            system_prompt,
+        ),
+    )
+
+
+class DatabricksProvider(OpenAIProvider):
+    def __init__(
+        self,
+        *,
+        model: str,
+        workspace_client: Optional["WorkspaceClient"] = None,
+    ):
+        try:
+            from databricks.sdk import WorkspaceClient
+        except ImportError:
+            raise ImportError(
+                "`ChatDatabricks()` requires the `databricks-sdk` package. "
+                "Install it with `pip install databricks-sdk[openai]`."
+            )
+
+        try:
+            import httpx
+            from openai import AsyncOpenAI
+        except ImportError:
+            raise ImportError(
+                "`ChatDatabricks()` requires the `openai` package. "
+                "Install it with `pip install openai`."
+            )
+
+        self._model = model
+        self._seed = None
+
+        if workspace_client is None:
+            workspace_client = WorkspaceClient()
+
+        client = workspace_client.serving_endpoints.get_open_ai_client()
+
+        self._client = client
+
+        # The databricks sdk does currently expose an async client, but we can
+        # effectively mirror what .get_open_ai_client() does internally.
+        # Note also there is a open PR to add async support that does essentially
+        # the same thing:
+        # https://github.com/databricks/databricks-sdk-py/pull/851
+        self._async_client = AsyncOpenAI(
+            base_url=client.base_url,
+            api_key="no-token",  # A placeholder to pass validations, this will not be used
+            http_client=httpx.AsyncClient(auth=client._client.auth),
+        )

chatlas/_display.py

@@ -12,8 +12,14 @@ from ._typing_extensions import TypedDict
 
 
 class MarkdownDisplay(ABC):
+    """Base class for displaying markdown content in different environments."""
+
     @abstractmethod
-    def update(self, content: str):
+    def echo(self, content: str):
+        """
+        Display the provided markdown string. This will append the content
+        to the current display.
+        """
         pass
 
     @abstractmethod
@@ -26,7 +32,7 @@ class MarkdownDisplay(ABC):
 
 
 class MockMarkdownDisplay(MarkdownDisplay):
-    def update(self, content: str):
+    def echo(self, content: str):
         pass
 
     def __enter__(self):
@@ -41,7 +47,7 @@ class LiveMarkdownDisplay(MarkdownDisplay):
     Stream chunks of markdown into a rich-based live updating console.
     """
 
-    def __init__(self, echo_options: "EchoOptions"):
+    def __init__(self, echo_options: "EchoDisplayOptions"):
         from rich.console import Console
 
         self.content: str = ""
@@ -63,7 +69,7 @@ class LiveMarkdownDisplay(MarkdownDisplay):
 
         self._markdown_options = echo_options["rich_markdown"]
 
-    def update(self, content: str):
+    def echo(self, content: str):
         from rich.markdown import Markdown
 
         self.content += content
@@ -98,11 +104,11 @@ class IPyMarkdownDisplay(MarkdownDisplay):
     Stream chunks of markdown into an IPython notebook.
     """
 
-    def __init__(self, echo_options: "EchoOptions"):
+    def __init__(self, echo_options: "EchoDisplayOptions"):
         self.content: str = ""
         self._css_styles = echo_options["css_styles"]
 
-    def update(self, content: str):
+    def echo(self, content: str):
         from IPython.display import Markdown, update_display
 
         self.content += content
@@ -143,7 +149,7 @@ class IPyMarkdownDisplay(MarkdownDisplay):
         self._ipy_display_id = None
 
 
-class EchoOptions(TypedDict):
+class EchoDisplayOptions(TypedDict):
     rich_markdown: dict[str, Any]
     rich_console: dict[str, Any]
     css_styles: dict[str, str]

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 @@ class GoogleProvider(
             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 @@ class GoogleProvider(
             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")
@@ -492,9 +492,13 @@ class GoogleProvider(
                 if name:
                     contents.append(
                         ContentToolResult(
-                            id=function_response.get("id") or name,
                             value=function_response.get("response"),
-                            name=name,
+                            request=ContentToolRequest(
+                                id=function_response.get("id") or name,
+                                name=name,
+                                # TODO: how to get the arguments?
+                                arguments={},
+                            ),
                         )
                     )
 

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
@@ -325,7 +325,8 @@ class OpenAIProvider(Provider[ChatCompletion, ChatCompletionChunk, ChatCompletio
                 del kwargs_full["tools"]
 
         if stream and "stream_options" not in kwargs_full:
-            kwargs_full["stream_options"] = {"include_usage": True}
+            if self.__class__.__name__ != "DatabricksProvider":
+                kwargs_full["stream_options"] = {"include_usage": True}
 
         return kwargs_full
 
@@ -432,7 +433,7 @@ class OpenAIProvider(Provider[ChatCompletion, ChatCompletionChunk, ChatCompletio
                                 "id": x.id,
                                 "function": {
                                     "name": x.name,
-                                    "arguments": json.dumps(x.arguments),
+                                    "arguments": orjson.dumps(x.arguments).decode("utf-8"),
                                 },
                                 "type": "function",
                             }
@@ -498,8 +499,8 @@ class OpenAIProvider(Provider[ChatCompletion, ChatCompletionChunk, ChatCompletio
                     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",
                             )
@@ -528,7 +529,7 @@ class OpenAIProvider(Provider[ChatCompletion, ChatCompletionChunk, ChatCompletio
         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)]
@@ -543,8 +544,8 @@ class OpenAIProvider(Provider[ChatCompletion, ChatCompletionChunk, ChatCompletio
 
                 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}'"