chatlas 0.4.0__py3-none-any.whl → 0.6.0__py3-none-any.whl

chatlas/__init__.py

@@ -1,7 +1,9 @@
 from . import types
 from ._anthropic import ChatAnthropic, ChatBedrockAnthropic
+from ._auto import ChatAuto
 from ._chat import Chat
 from ._content_image import content_image_file, content_image_plot, content_image_url
+from ._content_pdf import content_pdf_file, content_pdf_url
 from ._github import ChatGithub
 from ._google import ChatGoogle, ChatVertex
 from ._groq import ChatGroq
@@ -10,12 +12,19 @@ from ._ollama import ChatOllama
 from ._openai import ChatAzureOpenAI, ChatOpenAI
 from ._perplexity import ChatPerplexity
 from ._provider import Provider
+from ._snowflake import ChatSnowflake
 from ._tokens import token_usage
 from ._tools import Tool
 from ._turn import Turn
 
+try:
+    from ._version import version as __version__
+except ImportError:  # pragma: no cover
+    __version__ = "0.0.0"  # stub value for docs
+
 __all__ = (
     "ChatAnthropic",
+    "ChatAuto",
     "ChatBedrockAnthropic",
     "ChatGithub",
     "ChatGoogle",
@@ -24,11 +33,14 @@ __all__ = (
     "ChatOpenAI",
     "ChatAzureOpenAI",
     "ChatPerplexity",
+    "ChatSnowflake",
     "ChatVertex",
     "Chat",
     "content_image_file",
     "content_image_plot",
     "content_image_url",
+    "content_pdf_file",
+    "content_pdf_url",
     "interpolate",
     "interpolate_file",
     "Provider",

chatlas/_anthropic.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import base64
 import json
 import warnings
 from typing import TYPE_CHECKING, Any, Literal, Optional, Union, cast, overload
@@ -12,6 +13,7 @@ from ._content import (
     ContentImageInline,
     ContentImageRemote,
     ContentJson,
+    ContentPDF,
     ContentText,
     ContentToolRequest,
     ContentToolResult,
@@ -31,6 +33,7 @@ if TYPE_CHECKING:
         ToolParam,
         ToolUseBlock,
     )
+    from anthropic.types.document_block_param import DocumentBlockParam
     from anthropic.types.image_block_param import ImageBlockParam
     from anthropic.types.model_param import ModelParam
     from anthropic.types.text_block_param import TextBlockParam
@@ -45,6 +48,7 @@ if TYPE_CHECKING:
         ImageBlockParam,
         ToolUseBlockParam,
         ToolResultBlockParam,
+        DocumentBlockParam,
     ]
 else:
     Message = object
@@ -72,7 +76,7 @@ def ChatAnthropic(
     ::: {.callout-note}
     ## API key
 
-    Note that a Claude Prop membership does not give you the ability to call
+    Note that a Claude Pro membership does not give you the ability to call
     models via the API. You will need to go to the [developer
     console](https://console.anthropic.com/account/keys) to sign up (and pay
     for) a developer account that will give you an API key that you can use with
@@ -82,7 +86,7 @@ def ChatAnthropic(
     ::: {.callout-note}
     ## Python requirements
 
-    `ChatAnthropic` requires the `anthropic` package (e.g., `pip install anthropic`).
+    `ChatAnthropic` requires the `anthropic` package: `pip install "chatlas[anthropic]"`.
     :::
 
     Examples
@@ -164,7 +168,7 @@ def ChatAnthropic(
     """
 
     if model is None:
-        model = log_model_default("claude-3-5-sonnet-latest")
+        model = log_model_default("claude-3-7-sonnet-latest")
 
     return Chat(
         provider=AnthropicProvider(
@@ -450,12 +454,21 @@ class AnthropicProvider(Provider[Message, RawMessageStreamEvent, Message]):
             return {"text": content.text, "type": "text"}
         elif isinstance(content, ContentJson):
             return {"text": "<structured data/>", "type": "text"}
+        elif isinstance(content, ContentPDF):
+            return {
+                "type": "document",
+                "source": {
+                    "type": "base64",
+                    "media_type": "application/pdf",
+                    "data": base64.b64encode(content.data).decode("utf-8"),
+                },
+            }
         elif isinstance(content, ContentImageInline):
             return {
                 "type": "image",
                 "source": {
                     "type": "base64",
-                    "media_type": content.content_type,
+                    "media_type": content.image_content_type,
                     "data": content.data or "",
                 },
             }
@@ -504,7 +517,7 @@ class AnthropicProvider(Provider[Message, RawMessageStreamEvent, Message]):
         contents = []
         for content in completion.content:
             if content.type == "text":
-                contents.append(ContentText(content.text))
+                contents.append(ContentText(text=content.text))
             elif content.type == "tool_use":
                 if has_data_model and content.name == "_structured_tool_call":
                     if not isinstance(content.input, dict):
@@ -515,11 +528,11 @@ class AnthropicProvider(Provider[Message, RawMessageStreamEvent, Message]):
                         raise ValueError(
                             "Expected data extraction tool to return a 'data' field."
                         )
-                    contents.append(ContentJson(content.input["data"]))
+                    contents.append(ContentJson(value=content.input["data"]))
                 else:
                     contents.append(
                         ContentToolRequest(
-                            content.id,
+                            id=content.id,
                             name=content.name,
                             arguments=content.input,
                         )
@@ -572,8 +585,8 @@ def ChatBedrockAnthropic(
     ::: {.callout-note}
     ## Python requirements
 
-    `ChatBedrockAnthropic`, requires the `anthropic` package with the `bedrock` extras
-    (e.g., `pip install anthropic[bedrock]`).
+    `ChatBedrockAnthropic`, requires the `anthropic` package with the `bedrock` extras:
+    `pip install "chatlas[bedrock-anthropic]"`
     :::
 
     Examples

chatlas/_auto.py

@@ -0,0 +1,183 @@
+from __future__ import annotations
+
+import json
+import os
+from typing import Callable, Literal, Optional
+
+from ._anthropic import ChatAnthropic, ChatBedrockAnthropic
+from ._chat import Chat
+from ._github import ChatGithub
+from ._google import ChatGoogle, ChatVertex
+from ._groq import ChatGroq
+from ._ollama import ChatOllama
+from ._openai import ChatAzureOpenAI, ChatOpenAI
+from ._perplexity import ChatPerplexity
+from ._snowflake import ChatSnowflake
+from ._turn import Turn
+
+AutoProviders = Literal[
+    "anthropic",
+    "bedrock-anthropic",
+    "github",
+    "google",
+    "groq",
+    "ollama",
+    "openai",
+    "azure-openai",
+    "perplexity",
+    "snowflake",
+    "vertex",
+]
+
+_provider_chat_model_map: dict[AutoProviders, Callable[..., Chat]] = {
+    "anthropic": ChatAnthropic,
+    "bedrock-anthropic": ChatBedrockAnthropic,
+    "github": ChatGithub,
+    "google": ChatGoogle,
+    "groq": ChatGroq,
+    "ollama": ChatOllama,
+    "openai": ChatOpenAI,
+    "azure-openai": ChatAzureOpenAI,
+    "perplexity": ChatPerplexity,
+    "snowflake": ChatSnowflake,
+    "vertex": ChatVertex,
+}
+
+
+def ChatAuto(
+    system_prompt: Optional[str] = None,
+    turns: Optional[list[Turn]] = None,
+    *,
+    provider: Optional[AutoProviders] = None,
+    model: Optional[str] = None,
+    **kwargs,
+) -> Chat:
+    """
+    Use environment variables (env vars) to configure the Chat provider and model.
+
+    Creates a `:class:~chatlas.Chat` instance based on the specified provider.
+    The provider may be specified through the `provider` parameter and/or the
+    `CHATLAS_CHAT_PROVIDER` env var. If both are set, the env var takes
+    precedence. Similarly, the provider's model may be specified through the
+    `model` parameter and/or the `CHATLAS_CHAT_MODEL` env var. Also, additional
+    configuration may be provided through the `kwargs` parameter and/or the
+    `CHATLAS_CHAT_ARGS` env var (as a JSON string). In this case, when both are
+    set, they are merged, with the env var arguments taking precedence.
+
+    As a result, `ChatAuto()` provides a convenient way to set a default
+    provider and model in your Python code, while allowing you to override
+    these settings through env vars (i.e., without modifying your code).
+
+    Prerequisites
+    -------------
+
+    ::: {.callout-note}
+    ## API key
+
+    Follow the instructions for the specific provider to obtain an API key.
+    :::
+
+    ::: {.callout-note}
+    ## Python requirements
+
+    Follow the instructions for the specific provider to install the required
+    Python packages.
+    :::
+
+
+    Examples
+    --------
+    First, set the environment variables for the provider, arguments, and API key:
+
+    ```bash
+    export CHATLAS_CHAT_PROVIDER=anthropic
+    export CHATLAS_CHAT_MODEL=claude-3-haiku-20240229
+    export CHATLAS_CHAT_ARGS='{"kwargs": {"max_retries": 3}}'
+    export ANTHROPIC_API_KEY=your_api_key
+    ```
+
+    Then, you can use the `ChatAuto` function to create a Chat instance:
+
+    ```python
+    from chatlas import ChatAuto
+
+    chat = ChatAuto()
+    chat.chat("What is the capital of France?")
+    ```
+
+    Parameters
+    ----------
+    provider
+        The name of the default chat provider to use. Providers are strings
+        formatted in kebab-case, e.g. to use `ChatBedrockAnthropic` set
+        `provider="bedrock-anthropic"`.
+
+        This value can also be provided via the `CHATLAS_CHAT_PROVIDER`
+        environment variable, which takes precedence over `provider`
+        when set.
+    model
+        The name of the default model to use. This value can also be provided
+        via the `CHATLAS_CHAT_MODEL` environment variable, which takes
+        precedence over `model` when set.
+    system_prompt
+        A system prompt to set the behavior of the assistant.
+    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.
+    **kwargs
+        Additional keyword arguments to pass to the Chat constructor. See the
+        documentation for each provider for more details on the available
+        options.
+
+        These arguments can also be provided via the `CHATLAS_CHAT_ARGS`
+        environment variable as a JSON string. When provided, the options
+        in the `CHATLAS_CHAT_ARGS` envvar take precedence over the options
+        passed to `kwargs`.
+
+        Note that `system_prompt` and `turns` in `kwargs` or in
+        `CHATLAS_CHAT_ARGS` are ignored.
+
+    Returns
+    -------
+    Chat
+        A chat instance using the specified provider.
+
+    Raises
+    ------
+    ValueError
+        If no valid provider is specified either through parameters or
+        environment variables.
+    """
+    the_provider = os.environ.get("CHATLAS_CHAT_PROVIDER", provider)
+
+    if the_provider is None:
+        raise ValueError(
+            "Provider name is required as parameter or `CHATLAS_CHAT_PROVIDER` must be set."
+        )
+    if the_provider not in _provider_chat_model_map:
+        raise ValueError(
+            f"Provider name '{the_provider}' is not a known chatlas provider: "
+            f"{', '.join(_provider_chat_model_map.keys())}"
+        )
+
+    # `system_prompt` and `turns` always come from `ChatAuto()`
+    base_args = {"system_prompt": system_prompt, "turns": turns}
+
+    if env_model := os.environ.get("CHATLAS_CHAT_MODEL"):
+        model = env_model
+
+    if model:
+        base_args["model"] = model
+
+    env_kwargs = {}
+    if env_kwargs_str := os.environ.get("CHATLAS_CHAT_ARGS"):
+        env_kwargs = json.loads(env_kwargs_str)
+
+    kwargs = {**kwargs, **env_kwargs, **base_args}
+    kwargs = {k: v for k, v in kwargs.items() if v is not None}
+
+    return _provider_chat_model_map[the_provider](**kwargs)

chatlas/_chat.py

@@ -408,7 +408,9 @@ 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`.
+            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`.
         kwargs
             Additional keyword arguments to pass to the method used for requesting
             the response.
@@ -1246,7 +1248,7 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
         id_: str,
     ) -> ContentToolResult:
         if func is None:
-            return ContentToolResult(id_, value=None, error="Unknown tool")
+            return ContentToolResult(id=id_, value=None, error="Unknown tool")
 
         name = func.__name__
 
@@ -1256,10 +1258,10 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
             else:
                 result = func(arguments)
 
-            return ContentToolResult(id_, value=result, error=None, name=name)
+            return ContentToolResult(id=id_, value=result, error=None, name=name)
         except Exception as e:
             log_tool_error(name, str(arguments), e)
-            return ContentToolResult(id_, value=None, error=str(e), name=name)
+            return ContentToolResult(id=id_, value=None, error=str(e), name=name)
 
     @staticmethod
     async def _invoke_tool_async(
@@ -1268,7 +1270,7 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
         id_: str,
     ) -> ContentToolResult:
         if func is None:
-            return ContentToolResult(id_, value=None, error="Unknown tool")
+            return ContentToolResult(id=id_, value=None, error="Unknown tool")
 
         name = func.__name__
 
@@ -1278,10 +1280,10 @@ class Chat(Generic[SubmitInputArgsT, CompletionT]):
             else:
                 result = await func(arguments)
 
-            return ContentToolResult(id_, value=result, error=None, name=name)
+            return ContentToolResult(id=id_, value=result, error=None, name=name)
         except Exception as e:
             log_tool_error(func.__name__, str(arguments), e)
-            return ContentToolResult(id_, value=None, error=str(e), name=name)
+            return ContentToolResult(id=id_, value=None, error=str(e), name=name)
 
     def _markdown_display(
         self, echo: Literal["text", "all", "none"]

chatlas/_content.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
 import json
-from dataclasses import dataclass
 from pprint import pformat
-from typing import Any, Literal, Optional
+from typing import Any, Literal, Optional, Union
+
+from pydantic import BaseModel, ConfigDict
 
 ImageContentTypes = Literal[
     "image/png",
@@ -15,12 +16,28 @@ ImageContentTypes = Literal[
 Allowable content types for images.
 """
 
+ContentTypeEnum = Literal[
+    "text",
+    "image_remote",
+    "image_inline",
+    "tool_request",
+    "tool_result",
+    "json",
+    "pdf",
+]
+"""
+A discriminated union of all content types.
+"""
+
 
-class Content:
+class Content(BaseModel):
     """
     Base class for all content types that can be appear in a [](`~chatlas.Turn`)
     """
 
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    content_type: ContentTypeEnum
+
     def __str__(self):
         raise NotImplementedError
 
@@ -31,13 +48,13 @@ class Content:
         raise NotImplementedError
 
 
-@dataclass
 class ContentText(Content):
     """
     Text content for a [](`~chatlas.Turn`)
     """
 
     text: str
+    content_type: ContentTypeEnum = "text"
 
     def __str__(self):
         return self.text
@@ -62,7 +79,6 @@ class ContentImage(Content):
     pass
 
 
-@dataclass
 class ContentImageRemote(ContentImage):
     """
     Image content from a URL.
@@ -81,6 +97,8 @@ class ContentImageRemote(ContentImage):
     url: str
     detail: Literal["auto", "low", "high"] = "auto"
 
+    content_type: ContentTypeEnum = "image_remote"
+
     def __str__(self):
         return f"![]({self.url})"
 
@@ -94,7 +112,6 @@ class ContentImageRemote(ContentImage):
         )
 
 
-@dataclass
 class ContentImageInline(ContentImage):
     """
     Inline image content.
@@ -105,17 +122,19 @@ class ContentImageInline(ContentImage):
 
     Parameters
     ----------
-    content_type
+    image_content_type
         The content type of the image.
     data
         The base64-encoded image data.
     """
 
-    content_type: ImageContentTypes
+    image_content_type: ImageContentTypes
     data: Optional[str] = None
 
+    content_type: ContentTypeEnum = "image_inline"
+
     def __str__(self):
-        return f"![](data:{self.content_type};base64,{self.data})"
+        return f"![](data:{self.image_content_type};base64,{self.data})"
 
     def _repr_markdown_(self):
         return self.__str__()
@@ -124,11 +143,10 @@ class ContentImageInline(ContentImage):
         n_bytes = len(self.data) if self.data else 0
         return (
             " " * indent
-            + f"<ContentImageInline content_type='{self.content_type}' size={n_bytes}>"
+            + f"<ContentImageInline content_type='{self.image_content_type}' size={n_bytes}>"
         )
 
 
-@dataclass
 class ContentToolRequest(Content):
     """
     A request to call a tool/function
@@ -151,6 +169,8 @@ class ContentToolRequest(Content):
     name: str
     arguments: object
 
+    content_type: ContentTypeEnum = "tool_request"
+
     def __str__(self):
         args_str = self._arguments_str()
         func_call = f"{self.name}({args_str})"
@@ -173,7 +193,6 @@ class ContentToolRequest(Content):
         return str(self.arguments)
 
 
-@dataclass
 class ContentToolResult(Content):
     """
     The result of calling a tool/function
@@ -199,21 +218,26 @@ class ContentToolResult(Content):
     name: Optional[str] = None
     error: Optional[str] = None
 
-    def _get_value_and_language(self) -> tuple[str, str]:
+    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}'", ""
+            return f"Tool calling failed with error: '{self.error}'"
+        if not pretty:
+            return str(self.value)
         try:
-            json_val = json.loads(self.value)
-            return pformat(json_val, indent=2, sort_dicts=False), "python"
-        except:  # noqa: E722
-            return str(self.value), ""
+            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):
         comment = f"# tool result ({self.id})"
-        value, language = self._get_value_and_language()
-
-        return f"""```{language}\n{comment}\n{value}\n```"""
+        value = self._get_value(pretty=True)
+        return f"""```python\n{comment}\n{value}\n```"""
 
+    # ... and for displaying in the notebook
     def _repr_markdown_(self):
         return self.__str__()
 
@@ -224,12 +248,11 @@ class ContentToolResult(Content):
             res += f" error='{self.error}'"
         return res + ">"
 
+    # The actual value to send to the model
     def get_final_value(self) -> str:
-        value, _language = self._get_value_and_language()
-        return value
+        return self._get_value()
 
 
-@dataclass
 class ContentJson(Content):
     """
     JSON content
@@ -245,6 +268,8 @@ class ContentJson(Content):
 
     value: dict[str, Any]
 
+    content_type: ContentTypeEnum = "json"
+
     def __str__(self):
         return json.dumps(self.value, indent=2)
 
@@ -253,3 +278,70 @@ class ContentJson(Content):
 
     def __repr__(self, indent: int = 0):
         return " " * indent + f"<ContentJson value={self.value}>"
+
+
+class ContentPDF(Content):
+    """
+    PDF content
+
+    This content type primarily exists to signal PDF data extraction
+    (i.e., data extracted via [](`~chatlas.Chat`)'s `.extract_data()` method)
+
+    Parameters
+    ----------
+    value
+        The PDF data extracted
+    """
+
+    data: bytes
+
+    content_type: ContentTypeEnum = "pdf"
+
+    def __str__(self):
+        return "<PDF document>"
+
+    def _repr_markdown_(self):
+        return self.__str__()
+
+    def __repr__(self, indent: int = 0):
+        return " " * indent + f"<ContentPDF size={len(self.data)}>"
+
+
+ContentUnion = Union[
+    ContentText,
+    ContentImageRemote,
+    ContentImageInline,
+    ContentToolRequest,
+    ContentToolResult,
+    ContentJson,
+    ContentPDF,
+]
+
+
+def create_content(data: dict[str, Any]) -> ContentUnion:
+    """
+    Factory function to create the appropriate Content subclass based on the data.
+
+    This is useful when deserializing content from JSON.
+    """
+    if not isinstance(data, dict):
+        raise ValueError("Content data must be a dictionary")
+
+    ct = data.get("content_type")
+
+    if ct == "text":
+        return ContentText.model_validate(data)
+    elif ct == "image_remote":
+        return ContentImageRemote.model_validate(data)
+    elif ct == "image_inline":
+        return ContentImageInline.model_validate(data)
+    elif ct == "tool_request":
+        return ContentToolRequest.model_validate(data)
+    elif ct == "tool_result":
+        return ContentToolResult.model_validate(data)
+    elif ct == "json":
+        return ContentJson.model_validate(data)
+    elif ct == "pdf":
+        return ContentPDF.model_validate(data)
+    else:
+        raise ValueError(f"Unknown content type: {ct}")

chatlas/_content_image.py

@@ -8,6 +8,7 @@ import warnings
 from typing import Literal, Union, cast
 
 from ._content import ContentImageInline, ContentImageRemote, ImageContentTypes
+from ._content_pdf import parse_data_url
 from ._utils import MISSING, MISSING_TYPE
 
 __all__ = (
@@ -60,15 +61,11 @@ def content_image_url(
         raise ValueError("detail must be 'auto', 'low', or 'high'")
 
     if url.startswith("data:"):
-        parts = url[5:].split(";", 1)
-        if len(parts) != 2 or not parts[1].startswith("base64,"):
-            raise ValueError("url is not a valid data URL.")
-        content_type = parts[0]
-        base64_data = parts[1][7:]
+        content_type, base64_data = parse_data_url(url)
         if content_type not in ["image/png", "image/jpeg", "image/webp", "image/gif"]:
             raise ValueError(f"Unsupported image content type: {content_type}")
         content_type = cast(ImageContentTypes, content_type)
-        return ContentImageInline(content_type, base64_data)
+        return ContentImageInline(image_content_type=content_type, data=base64_data)
     else:
         return ContentImageRemote(url=url, detail=detail)
 
@@ -191,7 +188,7 @@ def content_image_file(
         img.save(buffer, format=img.format)
         base64_data = base64.b64encode(buffer.getvalue()).decode("utf-8")
 
-    return ContentImageInline(content_type, base64_data)
+    return ContentImageInline(image_content_type=content_type, data=base64_data)
 
 
 def content_image_plot(
@@ -263,7 +260,7 @@ def content_image_plot(
         fig.savefig(buf, format="png", dpi=dpi, bbox_inches="tight")
         buf.seek(0)
         base64_data = base64.b64encode(buf.getvalue()).decode("utf-8")
-        return ContentImageInline("image/png", base64_data)
+        return ContentImageInline(image_content_type="image/png", data=base64_data)
     finally:
         fig.set_size_inches(*size)