chatlas 0.7.1__py3-none-any.whl → 0.8.1__py3-none-any.whl

chatlas/__init__.py

@@ -16,7 +16,7 @@ from ._perplexity import ChatPerplexity
 from ._provider import Provider
 from ._snowflake import ChatSnowflake
 from ._tokens import token_usage
-from ._tools import Tool
+from ._tools import Tool, ToolRejectError
 from ._turn import Turn
 
 try:
@@ -51,6 +51,7 @@ __all__ = (
     "Provider",
     "token_usage",
     "Tool",
+    "ToolRejectError",
     "Turn",
     "types",
 )

chatlas/_anthropic.py

@@ -451,10 +451,7 @@ class AnthropicProvider(Provider[Message, RawMessageStreamEvent, Message]):
     @staticmethod
     def _as_content_block(content: Content) -> "ContentBlockParam":
         if isinstance(content, ContentText):
-            text = content.text
-            if text == "" or text.isspace():
-                text = "[empty string]"
-            return {"type": "text", "text": text}
+            return {"text": content.text, "type": "text"}
         elif isinstance(content, ContentJson):
             return {"text": "<structured data/>", "type": "text"}
         elif isinstance(content, ContentPDF):

chatlas/_callbacks.py

@@ -0,0 +1,56 @@
+from collections import OrderedDict
+from typing import Any, Callable
+
+from ._utils import is_async_callable
+
+
+class CallbackManager:
+    def __init__(self) -> None:
+        self._callbacks: dict[str, Callable[..., Any]] = OrderedDict()
+        self._id: int = 1
+
+    def add(self, callback: Callable[..., Any]) -> Callable[[], None]:
+        callback_id = self._next_id()
+        self._callbacks[callback_id] = callback
+
+        def _rm_callback() -> None:
+            self._callbacks.pop(callback_id, None)
+
+        return _rm_callback
+
+    def invoke(self, *args: Any, **kwargs: Any) -> None:
+        if not self._callbacks:
+            return
+
+        # Invoke in reverse insertion order
+        for callback_id in reversed(list(self._callbacks.keys())):
+            callback = self._callbacks[callback_id]
+            if is_async_callable(callback):
+                raise RuntimeError(
+                    "Can't use async callbacks with `.chat()`/`.stream()`."
+                    "Async callbacks can only be used with `.chat_async()`/`.stream_async()`."
+                )
+            callback(*args, **kwargs)
+
+    async def invoke_async(self, *args: Any, **kwargs: Any) -> None:
+        if not self._callbacks:
+            return
+
+        # Invoke in reverse insertion order
+        for callback_id in reversed(list(self._callbacks.keys())):
+            callback = self._callbacks[callback_id]
+            if is_async_callable(callback):
+                await callback(*args, **kwargs)
+            else:
+                callback(*args, **kwargs)
+
+    def count(self) -> int:
+        return len(self._callbacks)
+
+    def get_callbacks(self) -> list[Callable[..., Any]]:
+        return list(self._callbacks.values())
+
+    def _next_id(self) -> str:
+        current_id = self._id
+        self._id += 1
+        return str(current_id)

chatlas/_chat.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import copy
 import inspect
 import os
 import sys
@@ -25,6 +26,7 @@ from typing import (
 
 from pydantic import BaseModel
 
+from ._callbacks import CallbackManager
 from ._content import (
     Content,
     ContentJson,
@@ -41,7 +43,7 @@ from ._display import (
 )
 from ._logging import log_tool_error
 from ._provider import Provider
-from ._tools import Tool
+from ._tools import Tool, ToolRejectError
 from ._turn import Turn, user_turn
 from ._typing_extensions import TypedDict
 from ._utils import html_escape, wrap_async
@@ -95,6 +97,8 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
         self.provider = provider
         self._turns: list[Turn] = list(turns or [])
         self._tools: dict[str, Tool] = {}
+        self._on_tool_request_callbacks = CallbackManager()
+        self._on_tool_result_callbacks = CallbackManager()
         self._current_display: Optional[MarkdownDisplay] = None
         self._echo_options: EchoDisplayOptions = {
             "rich_markdown": {},
@@ -414,9 +418,12 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
             Whether to run the app in a background thread. If `None`, the app will
             run in a background thread if the current environment is a notebook.
         echo
-            Whether to echo text content, all content (i.e., tool calls), or no
-            content. Defaults to `"none"` when `stream=True` and `"text"` when
-            `stream=False`.
+            One of the following (defaults to `"none"` when `stream=True` and `"text"` when
+            `stream=False`):
+              - `"text"`: Echo just the text content of the response.
+              - `"output"`: Echo text and tool call content.
+              - `"all"`: Echo both the assistant and user turn.
+              - `"none"`: Do not echo any content.
         content
             Whether to display text content or all content (i.e., tool calls).
         kwargs
@@ -504,8 +511,11 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
         Parameters
         ----------
         echo
-            Whether to echo text content, all content (i.e., tool calls), or no
-            content.
+            One of the following (default is "output"):
+              - `"text"`: Echo just the text content of the response.
+              - `"output"`: Echo text and tool call content.
+              - `"all"`: Echo both the assistant and user turn.
+              - `"none"`: Do not echo any content.
         stream
             Whether to stream the response (i.e., have the response appear in chunks).
         kwargs
@@ -542,8 +552,11 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
         args
             The user input(s) to generate a response from.
         echo
-            Whether to echo text content, all content (i.e., tool calls), or no
-            content.
+            One of the following (default is "output"):
+              - `"text"`: Echo just the text content of the response.
+              - `"output"`: Echo text and tool call content.
+              - `"all"`: Echo both the assistant and user turn.
+              - `"none"`: Do not echo any content.
         stream
             Whether to stream the response (i.e., have the response appear in
             chunks).
@@ -592,8 +605,11 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
         args
             The user input(s) to generate a response from.
         echo
-            Whether to echo text content, all content (i.e., tool calls, images,
-            etc), or no content.
+            One of the following (default is "output"):
+              - `"text"`: Echo just the text content of the response.
+              - `"output"`: Echo text and tool call content.
+              - `"all"`: Echo both the assistant and user turn.
+              - `"none"`: Do not echo any content.
         stream
             Whether to stream the response (i.e., have the response appear in
             chunks).
@@ -631,38 +647,25 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
     def stream(
         self,
         *args: Content | str,
+        content: Literal["text"] = "text",
+        echo: EchoOptions = "none",
+        kwargs: Optional[SubmitInputArgsT] = None,
     ) -> Generator[str, None, None]: ...
 
     @overload
     def stream(
         self,
         *args: Content | str,
-        echo: EchoOptions,
-    ) -> Generator[str, None, None]: ...
-
-    @overload
-    def stream(
-        self,
-        *args: Content | str,
-        echo: EchoOptions,
-        content: Literal["text"],
-        kwargs: Optional[SubmitInputArgsT],
-    ) -> Generator[str, None, None]: ...
-
-    @overload
-    def stream(
-        self,
-        *args: Content | str,
-        echo: EchoOptions,
         content: Literal["all"],
-        kwargs: Optional[SubmitInputArgsT],
+        echo: EchoOptions = "none",
+        kwargs: Optional[SubmitInputArgsT] = None,
     ) -> Generator[str | ContentToolRequest | ContentToolResult, None, None]: ...
 
     def stream(
         self,
         *args: Content | str,
-        echo: EchoOptions = "none",
         content: Literal["text", "all"] = "text",
+        echo: EchoOptions = "none",
         kwargs: Optional[SubmitInputArgsT] = None,
     ) -> Generator[str | ContentToolRequest | ContentToolResult, None, None]:
         """
@@ -672,11 +675,15 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
         ----------
         args
             The user input(s) to generate a response from.
-        echo
-            Whether to echo text content, all content (i.e., tool calls), or no
-            content.
         content
-            Whether to yield just text content, or all content (i.e., tool calls).
+            Whether to yield just text content or include rich content objects
+            (e.g., tool calls) when relevant.
+        echo
+            One of the following (default is "none"):
+              - `"text"`: Echo just the text content of the response.
+              - `"output"`: Echo text and tool call content.
+              - `"all"`: Echo both the assistant and user turn.
+              - `"none"`: Do not echo any content.
         kwargs
             Additional keyword arguments to pass to the method used for requesting
             the response.
@@ -712,38 +719,25 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
     async def stream_async(
         self,
         *args: Content | str,
+        content: Literal["text"] = "text",
+        echo: EchoOptions = "none",
+        kwargs: Optional[SubmitInputArgsT] = None,
     ) -> AsyncGenerator[str, None]: ...
 
     @overload
     async def stream_async(
         self,
         *args: Content | str,
-        echo: EchoOptions,
-    ) -> AsyncGenerator[str, None]: ...
-
-    @overload
-    async def stream_async(
-        self,
-        *args: Content | str,
-        echo: EchoOptions,
-        content: Literal["text"],
-        kwargs: Optional[SubmitInputArgsT],
-    ) -> AsyncGenerator[str, None]: ...
-
-    @overload
-    async def stream_async(
-        self,
-        *args: Content | str,
-        echo: EchoOptions,
         content: Literal["all"],
-        kwargs: Optional[SubmitInputArgsT],
+        echo: EchoOptions = "none",
+        kwargs: Optional[SubmitInputArgsT] = None,
     ) -> AsyncGenerator[str | ContentToolRequest | ContentToolResult, None]: ...
 
     async def stream_async(
         self,
         *args: Content | str,
-        echo: EchoOptions = "none",
         content: Literal["text", "all"] = "text",
+        echo: EchoOptions = "none",
         kwargs: Optional[SubmitInputArgsT] = None,
     ) -> AsyncGenerator[str | ContentToolRequest | ContentToolResult, None]:
         """
@@ -753,11 +747,15 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
         ----------
         args
             The user input(s) to generate a response from.
-        echo
-            Whether to echo text content, all content (i.e., tool calls), or no
-            content.
         content
-            Whether to yield just text content, or all content (i.e., tool calls).
+            Whether to yield just text content or include rich content objects
+            (e.g., tool calls) when relevant.
+        echo
+            One of the following (default is "none"):
+              - `"text"`: Echo just the text content of the response.
+              - `"output"`: Echo text and tool call content.
+              - `"all"`: Echo both the assistant and user turn.
+              - `"none"`: Do not echo any content.
         kwargs
             Additional keyword arguments to pass to the method used for requesting
             the response.
@@ -804,7 +802,11 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
         data_model
             A Pydantic model describing the structure of the data to extract.
         echo
-            Whether to echo text content, all content (i.e., tool calls), or no content.
+            One of the following (default is "none"):
+              - `"text"`: Echo just the text content of the response.
+              - `"output"`: Echo text and tool call content.
+              - `"all"`: Echo both the assistant and user turn.
+              - `"none"`: Do not echo any content.
         stream
             Whether to stream the response (i.e., have the response appear in chunks).
 
@@ -862,7 +864,11 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
         data_model
             A Pydantic model describing the structure of the data to extract.
         echo
-            Whether to echo text content, all content (i.e., tool calls), or no content
+            One of the following (default is "none"):
+              - `"text"`: Echo just the text content of the response.
+              - `"output"`: Echo text and tool call content.
+              - `"all"`: Echo both the assistant and user turn.
+              - `"none"`: Do not echo any content.
         stream
             Whether to stream the response (i.e., have the response appear in chunks).
             Defaults to `True` if `echo` is not "none".
@@ -987,6 +993,53 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
         tool = Tool(func, model=model)
         self._tools[tool.name] = tool
 
+    def on_tool_request(self, callback: Callable[[ContentToolRequest], None]):
+        """
+        Register a callback for a tool request event.
+
+        A tool request event occurs when the assistant requests a tool to be
+        called on its behalf. Before invoking the tool, `on_tool_request`
+        handlers are called with the relevant `ContentToolRequest` object. This
+        is useful if you want to handle tool requests in a custom way, such as
+        requiring logging them or requiring user approval before invoking the
+        tool
+
+        Parameters
+        ----------
+        callback
+            A function to be called when a tool request event occurs.
+            This function must have a single argument, which will be the
+            tool request (i.e., a `ContentToolRequest` object).
+
+        Returns
+        -------
+        A callable that can be used to remove the callback later.
+        """
+        return self._on_tool_request_callbacks.add(callback)
+
+    def on_tool_result(self, callback: Callable[[ContentToolResult], None]):
+        """
+        Register a callback for a tool result event.
+
+        A tool result event occurs when a tool has been invoked and the
+        result is ready to be provided to the assistant. After the tool
+        has been invoked, `on_tool_result` handlers are called with the
+        relevant `ContentToolResult` object. This is useful if you want to
+        handle tool results in a custom way such as logging them.
+
+        Parameters
+        ----------
+        callback
+            A function to be called when a tool result event occurs.
+            This function must have a single argument, which will be the
+            tool result (i.e., a `ContentToolResult` object).
+
+        Returns
+        -------
+        A callable that can be used to remove the callback later.
+        """
+        return self._on_tool_result_callbacks.add(callback)
+
     @property
     def current_display(self) -> Optional[MarkdownDisplay]:
         """
@@ -1417,28 +1470,43 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
             e = RuntimeError(f"Unknown tool: {x.name}")
             return ContentToolResult(value=None, error=e, request=x)
 
-        args = x.arguments
-
+        # First, invoke the request callbacks. If a ToolRejectError is raised,
+        # treat it like a tool failure (i.e., gracefully handle it).
+        result: ContentToolResult | None = None
         try:
-            if isinstance(args, dict):
-                result = func(**args)
-            else:
-                result = func(args)
+            self._on_tool_request_callbacks.invoke(x)
+        except ToolRejectError as e:
+            result = ContentToolResult(value=None, error=e, request=x)
+
+        # Invoke the tool (if it hasn't been rejected).
+        if result is None:
+            try:
+                if isinstance(x.arguments, dict):
+                    res = func(**x.arguments)
+                else:
+                    res = func(x.arguments)
+
+                if isinstance(res, ContentToolResult):
+                    result = res
+                else:
+                    result = ContentToolResult(value=res)
 
-            if not isinstance(result, ContentToolResult):
-                result = ContentToolResult(value=result)
+                result.request = x
+            except Exception as e:
+                result = ContentToolResult(value=None, error=e, request=x)
 
-            result.request = x
-            return result
-        except Exception as e:
+        # If we've captured an error, notify and log it.
+        if result.error:
             warnings.warn(
                 f"Calling tool '{x.name}' led to an error.",
                 ToolFailureWarning,
                 stacklevel=2,
             )
             traceback.print_exc()
-            log_tool_error(x.name, str(args), e)
-            return ContentToolResult(value=None, error=e, request=x)
+            log_tool_error(x.name, str(x.arguments), result.error)
+
+        self._on_tool_result_callbacks.invoke(result)
+        return result
 
     async def _invoke_tool_async(self, x: ContentToolRequest) -> ContentToolResult:
         tool_def = self._tools.get(x.name, None)
@@ -1453,28 +1521,43 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
             e = RuntimeError(f"Unknown tool: {x.name}")
             return ContentToolResult(value=None, error=e, request=x)
 
-        args = x.arguments
-
+        # First, invoke the request callbacks. If a ToolRejectError is raised,
+        # treat it like a tool failure (i.e., gracefully handle it).
+        result: ContentToolResult | None = None
         try:
-            if isinstance(args, dict):
-                result = await func(**args)
-            else:
-                result = await func(args)
+            await self._on_tool_request_callbacks.invoke_async(x)
+        except ToolRejectError as e:
+            result = ContentToolResult(value=None, error=e, request=x)
+
+        # Invoke the tool (if it hasn't been rejected).
+        if result is None:
+            try:
+                if isinstance(x.arguments, dict):
+                    res = await func(**x.arguments)
+                else:
+                    res = await func(x.arguments)
 
-            if not isinstance(result, ContentToolResult):
-                result = ContentToolResult(value=result)
+                if isinstance(res, ContentToolResult):
+                    result = res
+                else:
+                    result = ContentToolResult(value=res)
+
+                result.request = x
+            except Exception as e:
+                result = ContentToolResult(value=None, error=e, request=x)
 
-            result.request = x
-            return result
-        except Exception as e:
+        # If we've captured an error, notify and log it.
+        if result.error:
             warnings.warn(
                 f"Calling tool '{x.name}' led to an error.",
                 ToolFailureWarning,
                 stacklevel=2,
             )
             traceback.print_exc()
-            log_tool_error(x.name, str(args), e)
-            return ContentToolResult(value=None, error=e, request=x)
+            log_tool_error(x.name, str(x.arguments), result.error)
+
+        await self._on_tool_result_callbacks.invoke_async(result)
+        return result
 
     def _markdown_display(self, echo: EchoOptions) -> ChatMarkdownDisplay:
         """
@@ -1545,6 +1628,21 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
             res += "\n" + turn.__repr__(indent=2)
         return res + "\n"
 
+    def __deepcopy__(self, memo):
+        result = self.__class__.__new__(self.__class__)
+
+        # Avoid recursive references
+        memo[id(self)] = result
+
+        # Copy all attributes except the problematic provider attribute
+        for key, value in self.__dict__.items():
+            if key != "provider":
+                setattr(result, key, copy.deepcopy(value, memo))
+            else:
+                setattr(result, key, value)
+
+        return result
+
 
 class ChatResponse:
     """

chatlas/_content.py

@@ -60,6 +60,12 @@ class ContentText(Content):
     text: str
     content_type: ContentTypeEnum = "text"
 
+    def __init__(self, **data: Any):
+        super().__init__(**data)
+
+        if self.text == "" or self.text.isspace():
+            self.text = "[empty string]"
+
     def __str__(self):
         return self.text
 

chatlas/_databricks.py

@@ -85,7 +85,7 @@ def ChatDatabricks(
         A chat object that retains the state of the conversation.
     """
     if model is None:
-        model = log_model_default("databricks-dbrx-instruct")
+        model = log_model_default("databricks-claude-3-7-sonnet")
 
     return Chat(
         provider=DatabricksProvider(

chatlas/_logging.py

@@ -1,6 +1,7 @@
 import logging
 import os
 import warnings
+from typing import Literal
 
 from rich.logging import RichHandler
 
@@ -12,15 +13,38 @@ def _rich_handler() -> RichHandler:
     return handler
 
 
-logger = logging.getLogger("chatlas")
-
-if os.environ.get("CHATLAS_LOG") == "info":
+def setup_logger(x: str, level: Literal["debug", "info"]) -> logging.Logger:
+    logger = logging.getLogger(x)
+    if level == "debug":
+        logger.setLevel(logging.DEBUG)
+    elif level == "info":
+        logger.setLevel(logging.INFO)
     # By adding a RichHandler to chatlas' logger, we can guarantee that they
     # never get dropped, even if the root logger's handlers are not
     # RichHandlers.
-    logger.setLevel(logging.INFO)
-    logger.addHandler(_rich_handler())
+    if not any(isinstance(h, RichHandler) for h in logger.handlers):
+        logger.addHandler(_rich_handler())
     logger.propagate = False
+    return logger
+
+
+logger = logging.getLogger("chatlas")
+log_level = os.environ.get("CHATLAS_LOG")
+if log_level:
+    if log_level != "debug" and log_level != "info":
+        warnings.warn(
+            f"CHATLAS_LOG is set to '{log_level}', but the log level must "
+            "be one of 'debug' or 'info'. Defaulting to 'info'.",
+        )
+        log_level = "info"
+
+    # Manually setup the logger for each dependency we care about. This way, we
+    # can ensure that the logs won't get dropped when a rich display is activate
+    logger = setup_logger("chatlas", log_level)
+    openai_logger = setup_logger("openai", log_level)
+    anthropic_logger = setup_logger("anthropic", log_level)
+    google_logger = setup_logger("google_genai.models", log_level)
+    httpx_logger = setup_logger("httpx", log_level)
 
     # Add a RichHandler to the root logger if there are no handlers. Note that
     # if chatlas is imported before other libraries that set up logging, (like