chatlas 0.7.0__py3-none-any.whl → 0.8.0__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": {},
@@ -631,31 +635,18 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
     def stream(
         self,
         *args: Content | str,
-    ) -> 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],
+        echo: EchoOptions = "none",
+        kwargs: Optional[SubmitInputArgsT] = None,
     ) -> 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(
@@ -712,31 +703,18 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
     async def stream_async(
         self,
         *args: Content | str,
-    ) -> 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],
+        echo: EchoOptions = "none",
+        kwargs: Optional[SubmitInputArgsT] = None,
     ) -> 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(
@@ -987,6 +965,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 +1442,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 +1493,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 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)
+
+        await self._on_tool_result_callbacks.invoke_async(result)
+        return result
 
     def _markdown_display(self, echo: EchoOptions) -> ChatMarkdownDisplay:
         """
@@ -1545,6 +1600,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(
@@ -111,17 +111,11 @@ class DatabricksProvider(OpenAIProvider):
         except ImportError:
             raise ImportError(
                 "`ChatDatabricks()` requires the `databricks-sdk` package. "
-                "Install it with `pip install databricks-sdk[openai]`."
+                "Install it with `pip install databricks-sdk`."
             )
 
-        try:
-            import httpx
-            from openai import AsyncOpenAI
-        except ImportError:
-            raise ImportError(
-                "`ChatDatabricks()` requires the `openai` package. "
-                "Install it with `pip install openai`."
-            )
+        import httpx
+        from openai import AsyncOpenAI
 
         self._model = model
         self._seed = None

chatlas/_github.py

@@ -40,12 +40,6 @@ def ChatGithub(
     You may need to apply for and be accepted into a beta access program.
     :::
 
-    ::: {.callout-note}
-    ## Python requirements
-
-    `ChatGithub` requires the `openai` package: `pip install "chatlas[github]"`.
-    :::
-
 
     Examples
     --------

chatlas/_google.py

@@ -291,7 +291,8 @@ class GoogleProvider(
                 GoogleTool(
                     function_declarations=[
                         FunctionDeclaration.from_callable(
-                            client=self._client, callable=tool.func
+                            client=self._client._api_client,
+                            callable=tool.func,
                         )
                         for tool in tools.values()
                     ]

chatlas/_groq.py

@@ -38,12 +38,6 @@ def ChatGroq(
     Sign up at <https://groq.com> to get an API key.
     :::
 
-    ::: {.callout-note}
-    ## Python requirements
-
-    `ChatGroq` requires the `openai` package: `pip install "chatlas[groq]"`.
-    :::
-
     Examples
     --------
 

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

chatlas/_ollama.py

@@ -49,12 +49,6 @@ def ChatOllama(
     (e.g. `ollama pull llama3.2`).
     :::
 
-    ::: {.callout-note}
-    ## Python requirements
-
-    `ChatOllama` requires the `openai` package: `pip install "chatlas[ollama]"`.
-    :::
-
 
     Examples
     --------

chatlas/_openai.py

@@ -78,12 +78,6 @@ def ChatOpenAI(
     account that will give you an API key that you can use with this package.
     :::
 
-    ::: {.callout-note}
-    ## Python requirements
-
-    `ChatOpenAI` requires the `openai` package: `pip install "chatlas[openai]"`.
-    :::
-
     Examples
     --------
     ```python
@@ -194,13 +188,7 @@ class OpenAIProvider(Provider[ChatCompletion, ChatCompletionChunk, ChatCompletio
         seed: Optional[int] = None,
         kwargs: Optional["ChatClientArgs"] = None,
     ):
-        try:
-            from openai import AsyncOpenAI, OpenAI
-        except ImportError:
-            raise ImportError(
-                "`ChatOpenAI()` requires the `openai` package. "
-                "Install it with `pip install openai`."
-            )
+        from openai import AsyncOpenAI, OpenAI
 
         self._model = model
         self._seed = seed
@@ -433,7 +421,9 @@ class OpenAIProvider(Provider[ChatCompletion, ChatCompletionChunk, ChatCompletio
                                 "id": x.id,
                                 "function": {
                                     "name": x.name,
-                                    "arguments": orjson.dumps(x.arguments).decode("utf-8"),
+                                    "arguments": orjson.dumps(x.arguments).decode(
+                                        "utf-8"
+                                    ),
                                 },
                                 "type": "function",
                             }
@@ -602,16 +592,6 @@ def ChatAzureOpenAI(
     hosts a number of open source models as well as proprietary models
     from OpenAI.
 
-    Prerequisites
-    -------------
-
-    ::: {.callout-note}
-    ## Python requirements
-
-    `ChatAzureOpenAI` requires the `openai` package:
-    `pip install "chatlas[azure-openai]"`.
-    :::
-
     Examples
     --------
     ```python
@@ -693,13 +673,7 @@ class OpenAIAzureProvider(OpenAIProvider):
         seed: int | None = None,
         kwargs: Optional["ChatAzureClientArgs"] = None,
     ):
-        try:
-            from openai import AsyncAzureOpenAI, AzureOpenAI
-        except ImportError:
-            raise ImportError(
-                "`ChatAzureOpenAI()` requires the `openai` package. "
-                "Install it with `pip install openai`."
-            )
+        from openai import AsyncAzureOpenAI, AzureOpenAI
 
         self._model = deployment_id
         self._seed = seed

chatlas/_perplexity.py

@@ -40,12 +40,6 @@ def ChatPerplexity(
     Sign up at <https://www.perplexity.ai> to get an API key.
     :::
 
-    ::: {.callout-note}
-    ## Python requirements
-
-    `ChatPerplexity` requires the `openai` package: `pip install "chatlas[perplexity]"`.
-    :::
-
 
     Examples
     --------