chatlas 0.2.0__py3-none-any.whl

chatlas/_content.py

@@ -0,0 +1,242 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from typing import Any, Literal, Optional
+
+ImageContentTypes = Literal[
+    "image/png",
+    "image/jpeg",
+    "image/webp",
+    "image/gif",
+]
+"""
+Allowable content types for images.
+"""
+
+
+class Content:
+    """
+    Base class for all content types that can be appear in a [](`~chatlas.Turn`)
+    """
+
+    def __str__(self):
+        raise NotImplementedError
+
+    def _repr_markdown_(self):
+        raise NotImplementedError
+
+    def __repr__(self, indent: int = 0):
+        raise NotImplementedError
+
+
+@dataclass
+class ContentText(Content):
+    """
+    Text content for a [](`~chatlas.Turn`)
+    """
+
+    text: str
+
+    def __str__(self):
+        return self.text
+
+    def _repr_markdown_(self):
+        return self.text
+
+    def __repr__(self, indent: int = 0):
+        text = self.text[:50] + "..." if len(self.text) > 50 else self.text
+        return " " * indent + f"<ContentText text='{text}'>"
+
+
+class ContentImage(Content):
+    """
+    Base class for image content.
+
+    This class is not meant to be used directly. Instead, use
+    [](`~chatlas.content_image_url`), [](`~chatlas.content_image_file`), or
+    [](`~chatlas.content_image_plot`).
+    """
+
+    pass
+
+
+@dataclass
+class ContentImageRemote(ContentImage):
+    """
+    Image content from a URL.
+
+    This is the return type for [](`~chatlas.content_image_url`).
+    It's not meant to be used directly.
+
+    Parameters
+    ----------
+    url
+        The URL of the image.
+    detail
+        A detail setting for the image. Can be `"auto"`, `"low"`, or `"high"`.
+    """
+
+    url: str
+    detail: Literal["auto", "low", "high"] = "auto"
+
+    def __str__(self):
+        return f"![]({self.url})"
+
+    def _repr_markdown_(self):
+        return self.__str__()
+
+    def __repr__(self, indent: int = 0):
+        return (
+            " " * indent
+            + f"<ContentImageRemote url='{self.url}' detail='{self.detail}'>"
+        )
+
+
+@dataclass
+class ContentImageInline(ContentImage):
+    """
+    Inline image content.
+
+    This is the return type for [](`~chatlas.content_image_file`) and
+    [](`~chatlas.content_image_plot`).
+    It's not meant to be used directly.
+
+    Parameters
+    ----------
+    content_type
+        The content type of the image.
+    data
+        The base64-encoded image data.
+    """
+
+    content_type: ImageContentTypes
+    data: Optional[str] = None
+
+    def __str__(self):
+        return f"![](data:{self.content_type};base64,{self.data})"
+
+    def _repr_markdown_(self):
+        return self.__str__()
+
+    def __repr__(self, indent: int = 0):
+        n_bytes = len(self.data) if self.data else 0
+        return (
+            " " * indent
+            + f"<ContentImageInline content_type='{self.content_type}' size={n_bytes}>"
+        )
+
+
+@dataclass
+class ContentToolRequest(Content):
+    """
+    A request to call 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
+    requested by the model assistant.
+
+    Parameters
+    ----------
+    id
+        A unique identifier for this request.
+    name
+        The name of the tool/function to call.
+    arguments
+        The arguments to pass to the tool/function.
+    """
+
+    id: str
+    name: str
+    arguments: object
+
+    def __str__(self):
+        args_str = self._arguments_str()
+        func_call = f"{self.name}({args_str})"
+        comment = f"# tool request ({self.id})"
+        return f"\n```python\n{comment}\n{func_call}\n```\n"
+
+    def _repr_markdown_(self):
+        return self.__str__()
+
+    def __repr__(self, indent: int = 0):
+        args_str = self._arguments_str()
+        return (
+            " " * indent
+            + f"<ContentToolRequest name='{self.name}' arguments='{args_str}' id='{self.id}'>"
+        )
+
+    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)
+
+
+@dataclass
+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`)).
+
+    Parameters
+    ----------
+    id
+        The unique identifier of the tool request.
+    value
+        The value returned by the tool/function.
+    error
+        An error message if the tool/function call failed.
+    """
+
+    id: str
+    value: Any = None
+    error: Optional[str] = None
+
+    def __str__(self):
+        comment = f"# tool result ({self.id})"
+        val = self.get_final_value()
+        return f"""\n```python\n{comment}\n"{val}"\n```\n"""
+
+    def _repr_markdown_(self):
+        return self.__str__()
+
+    def __repr__(self, indent: int = 0):
+        res = " " * indent
+        res += f"<ContentToolResult value='{self.value}' id='{self.id}'"
+        if self.error:
+            res += f" error='{self.error}'"
+        return res + ">"
+
+    def get_final_value(self) -> str:
+        if self.error:
+            return f"Tool calling failed with error: '{self.error}'"
+        return str(self.value)
+
+
+@dataclass
+class ContentJson(Content):
+    """
+    JSON content
+
+    This content type primarily exists to signal structured data extraction
+    (i.e., data extracted via [](`~chatlas.Chat`)'s `.extract_data()` method)
+
+    Parameters
+    ----------
+    value
+        The JSON data extracted
+    """
+
+    value: dict[str, Any]
+
+    def __str__(self):
+        return json.dumps(self.value, indent=2)
+
+    def _repr_markdown_(self):
+        return f"""\n```json\n{self.__str__()}\n```\n"""
+
+    def __repr__(self, indent: int = 0):
+        return " " * indent + f"<ContentJson value={self.value}>"

chatlas/_content_image.py

@@ -0,0 +1,272 @@
+from __future__ import annotations
+
+import base64
+import io
+import os
+import re
+import warnings
+from typing import Literal, Union, cast
+
+from ._content import ContentImageInline, ContentImageRemote, ImageContentTypes
+from ._utils import MISSING, MISSING_TYPE
+
+__all__ = (
+    "content_image_url",
+    "content_image_file",
+    "content_image_plot",
+)
+
+
+def content_image_url(
+    url: str, detail: Literal["auto", "low", "high"] = "auto"
+) -> Union[ContentImageInline, ContentImageRemote]:
+    """
+    Encode image content from a URL for chat input.
+
+    This function is used to prepare image URLs for input to the chatbot. It can
+    handle both regular URLs and data URLs.
+
+    Parameters
+    ----------
+    url
+        The URL of the image to include in the chat input. Can be a data: URL or a
+        regular URL.
+    detail
+        The detail setting for this image. Can be `"auto"`, `"low"`, or `"high"`.
+
+    Returns
+    -------
+    [](`~chatlas.types.Content`)
+        Content suitable for a [](`~chatlas.Turn`) object.
+
+    Examples
+    --------
+    ```python
+    from chatlas import ChatOpenAI, content_image_url
+
+    chat = ChatOpenAI()
+    chat.chat(
+        "What do you see in this image?",
+        content_image_url("https://www.python.org/static/img/python-logo.png"),
+    )
+    ```
+
+    Raises
+    ------
+    ValueError
+        If the URL is not valid or the detail setting is invalid.
+    """
+    if detail not in ["auto", "low", "high"]:
+        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:]
+        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)
+    else:
+        return ContentImageRemote(url=url, detail=detail)
+
+
+def content_image_file(
+    path: str,
+    content_type: Literal["auto", ImageContentTypes] = "auto",
+    resize: Union[Literal["low", "high", "none"], str, MISSING_TYPE] = MISSING,
+) -> ContentImageInline:
+    """
+    Encode image content from a file for chat input.
+
+    This function is used to prepare image files for input to the chatbot. It
+    can handle various image formats and provides options for resizing.
+
+    Parameters
+    ----------
+    path
+        The path to the image file to include in the chat input.
+    content_type
+        The content type of the image (e.g., `"image/png"`). If `"auto"`, the content
+        type is inferred from the file extension.
+    resize
+        Resizing option for the image. Can be:
+            - `"low"`: Resize to fit within 512x512
+            - `"high"`: Resize to fit within 2000x768 or 768x2000
+            - `"none"`: No resizing
+            - Custom string (e.g., `"200x200"`, `"300x200>!"`, etc.)
+
+    Returns
+    -------
+    [](`~chatlas.types.Content`)
+        Content suitable for a [](`~chatlas.Turn`) object.
+
+    Examples
+    --------
+    ```python
+    from chatlas import ChatOpenAI, content_image_file
+
+    chat = ChatOpenAI()
+    chat.chat(
+        "What do you see in this image?",
+        content_image_file("path/to/image.png"),
+    )
+    ```
+
+    Raises
+    ------
+    FileNotFoundError
+        If the specified file does not exist.
+    ValueError
+        If the file extension is unsupported or the resize option is invalid.
+    """
+
+    if not os.path.isfile(path):
+        raise FileNotFoundError(f"{path} must be an existing file.")
+
+    if content_type == "auto":
+        ext = os.path.splitext(path)[1].lower()
+        if ext == ".png":
+            content_type = "image/png"
+        elif ext in [".jpeg", ".jpg"]:
+            content_type = "image/jpeg"
+        elif ext == ".webp":
+            content_type = "image/webp"
+        elif ext == ".gif":
+            content_type = "image/gif"
+        else:
+            raise ValueError(f"Unsupported image file extension: {ext}")
+
+    if resize == "none":
+        with open(path, "rb") as image_file:
+            base64_data = base64.b64encode(image_file.read()).decode("utf-8")
+    else:
+        try:
+            from PIL import Image
+        except ImportError:
+            raise ImportError(
+                "Image resizing requires the `Pillow` package. "
+                "Install it with `pip install Pillow`."
+            )
+
+        img = Image.open(path)
+
+        if isinstance(resize, MISSING_TYPE):
+            warnings.warn(
+                "The `resize` parameter is missing. Defaulting to `resize='low'`. "
+                "As a result, the image has likely lost quality before the model received it. "
+                "Set `resize='low'` to suppress this warning, `resize='high'` for higher quality, or "
+                "`resize='none'` to disable resizing.",
+                category=MissingResizeWarning,
+                stacklevel=2,
+            )
+            resize = "low"
+
+        if resize == "low":
+            img.thumbnail((512, 512))
+        elif resize == "high":
+            if img.width > img.height:
+                img.thumbnail((2000, 768))
+            else:
+                img.thumbnail((768, 2000))
+        else:
+            match = re.match(r"(\d+)x(\d+)(>?)(!?)", resize)
+            if match:
+                width, height = map(int, match.group(1, 2))
+                only_shrink = ">" in match.group(3)
+                ignore_aspect = "!" in match.group(4)
+
+                if only_shrink and (img.width <= width and img.height <= height):
+                    pass  # No resize needed
+                elif ignore_aspect:
+                    img = img.resize((width, height))
+                else:
+                    img.thumbnail((width, height))
+            else:
+                raise ValueError(f"Invalid resize value: {resize}")
+
+        buffer = io.BytesIO()
+        img.save(buffer, format=img.format)
+        base64_data = base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+    return ContentImageInline(content_type, base64_data)
+
+
+def content_image_plot(
+    width: int = 768, height: int = 768, dpi: int = 72
+) -> ContentImageInline:
+    """
+    Encode the current matplotlib plot as an image for chat input.
+
+    This function captures the current matplotlib plot, resizes it to the specified
+    dimensions, and prepares it for chat input.
+
+    Parameters
+    ----------
+    width
+        The desired width of the output image in pixels.
+    height
+        The desired height of the output image in pixels.
+    dpi
+        The DPI (dots per inch) of the output image.
+
+    Returns
+    -------
+    [](`~chatlas.types.Content`)
+        Content suitable for a [](`~chatlas.Turn`) object.
+
+    Raises
+    ------
+    ValueError
+        If width or height is not a positive integer.
+
+    Examples
+    --------
+
+    ```python
+    from chatlas import ChatOpenAI, content_image_plot
+    import matplotlib.pyplot as plt
+
+    plt.scatter(faithful["eruptions"], faithful["waiting"])
+    chat = ChatOpenAI()
+    chat.chat(
+        "Describe this plot in one paragraph, as suitable for inclusion in "
+        "alt-text. You should briefly describe the plot type, the axes, and "
+        "2-5 major visual patterns.",
+        content_image_plot(),
+    )
+    ```
+    """
+
+    try:
+        import matplotlib.pyplot as plt
+    except ImportError:
+        raise ImportError(
+            "`content_image_plot()` requires the `matplotlib` package. "
+            "Install it with `pip install matplotlib`."
+        )
+
+    if not plt.get_fignums():
+        raise RuntimeError(
+            "No matplotlib figure to save. Please create one before calling `content_image_plot()`."
+        )
+
+    fig = plt.gcf()
+    size = fig.get_size_inches()
+
+    try:
+        fig.set_size_inches(width / dpi, height / dpi)
+
+        buf = io.BytesIO()
+        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)
+    finally:
+        fig.set_size_inches(*size)
+
+
+class MissingResizeWarning(RuntimeWarning):
+    pass

chatlas/_display.py

@@ -0,0 +1,139 @@
+import logging
+from abc import ABC, abstractmethod
+from typing import Any
+from uuid import uuid4
+
+from rich.live import Live
+from rich.logging import RichHandler
+
+from ._logging import logger
+from ._typing_extensions import TypedDict
+
+
+class MarkdownDisplay(ABC):
+    @abstractmethod
+    def update(self, content: str):
+        pass
+
+    @abstractmethod
+    def __enter__(self) -> "MarkdownDisplay":
+        pass
+
+    @abstractmethod
+    def __exit__(self, exc_type, exc_value, traceback):
+        pass
+
+
+class MockMarkdownDisplay(MarkdownDisplay):
+    def update(self, content: str):
+        pass
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        pass
+
+
+class LiveMarkdownDisplay(MarkdownDisplay):
+    """
+    Stream chunks of markdown into a rich-based live updating console.
+    """
+
+    def __init__(self, echo_options: "EchoOptions"):
+        from rich.console import Console
+
+        self.content: str = ""
+        self.live = Live(
+            auto_refresh=False,
+            vertical_overflow="visible",
+            console=Console(
+                **echo_options["rich_console"],
+            ),
+        )
+        self._markdown_options = echo_options["rich_markdown"]
+
+    def update(self, content: str):
+        from rich.markdown import Markdown
+
+        self.content += content
+        self.live.update(
+            Markdown(
+                self.content,
+                **self._markdown_options,
+            ),
+            refresh=True,
+        )
+
+    def __enter__(self):
+        self.live.__enter__()
+        # Live() isn't smart enough to know to automatically display logs when
+        # when they get handled while it Live() is active.
+        # However, if the logging handler is a RichHandler, it can be told
+        # about the live console so it can add logs to the top of the Live console.
+        handlers = [*logging.getLogger().handlers, *logger.handlers]
+        for h in handlers:
+            if isinstance(h, RichHandler):
+                h.console = self.live.console
+
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.content = ""
+        return self.live.__exit__(exc_type, exc_value, traceback)
+
+
+class IPyMarkdownDisplay(MarkdownDisplay):
+    """
+    Stream chunks of markdown into an IPython notebook.
+    """
+
+    def __init__(self, echo_options: "EchoOptions"):
+        self.content: str = ""
+        self._css_styles = echo_options["css_styles"]
+
+    def update(self, content: str):
+        from IPython.display import Markdown, update_display
+
+        self.content += content
+        update_display(
+            Markdown(self.content),
+            display_id=self._ipy_display_id,
+        )
+
+    def _init_display(self) -> str:
+        try:
+            from IPython.display import HTML, Markdown, display
+        except ImportError:
+            raise ImportError(
+                "The IPython package is required for displaying content in a Jupyter notebook. "
+                "Install it with `pip install ipython`."
+            )
+
+        if self._css_styles:
+            id_ = uuid4().hex
+            css = "".join(f"{k}: {v}; " for k, v in self._css_styles.items())
+            display(HTML(f"<style>#{id_} + .chatlas-markdown {{ {css} }}</style>"))
+            display(HTML(f"<div id='{id_}' class='chatlas-markdown'>"))
+        else:
+            # Unfortunately, there doesn't seem to be a proper way to wrap
+            # Markdown() in a div?
+            display(HTML("<div class='chatlas-markdown'>"))
+
+        handle = display(Markdown(""), display_id=True)
+        if handle is None:
+            raise ValueError("Failed to create display handle")
+        return handle.display_id
+
+    def __enter__(self):
+        self._ipy_display_id = self._init_display()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self._ipy_display_id = None
+
+
+class EchoOptions(TypedDict):
+    rich_markdown: dict[str, Any]
+    rich_console: dict[str, Any]
+    css_styles: dict[str, str]