chatlas 0.3.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

chatlas/_live_render.py

@@ -0,0 +1,116 @@
+# A 'patched' version of LiveRender that adds the 'crop_above' vertical overflow method.
+# Derives from https://github.com/Textualize/rich/pull/3637
+import sys
+from typing import Optional, Tuple
+
+if sys.version_info >= (3, 8):
+    from typing import Literal
+else:
+    from typing_extensions import Literal  # pragma: no cover
+
+from rich._loop import loop_last
+from rich.console import Console, ConsoleOptions, RenderableType, RenderResult
+from rich.control import Control
+from rich.segment import ControlType, Segment
+from rich.style import StyleType
+from rich.text import Text
+
+VerticalOverflowMethod = Literal["crop", "crop_above", "ellipsis", "visible"]
+
+
+class LiveRender:
+    """Creates a renderable that may be updated.
+
+    Args:
+        renderable (RenderableType): Any renderable object.
+        style (StyleType, optional): An optional style to apply to the renderable. Defaults to "".
+    """
+
+    def __init__(
+        self,
+        renderable: RenderableType,
+        style: StyleType = "",
+        vertical_overflow: VerticalOverflowMethod = "ellipsis",
+    ) -> None:
+        self.renderable = renderable
+        self.style = style
+        self.vertical_overflow = vertical_overflow
+        self._shape: Optional[Tuple[int, int]] = None
+
+    def set_renderable(self, renderable: RenderableType) -> None:
+        """Set a new renderable.
+
+        Args:
+            renderable (RenderableType): Any renderable object, including str.
+        """
+        self.renderable = renderable
+
+    def position_cursor(self) -> Control:
+        """Get control codes to move cursor to beginning of live render.
+
+        Returns:
+            Control: A control instance that may be printed.
+        """
+        if self._shape is not None:
+            _, height = self._shape
+            return Control(
+                ControlType.CARRIAGE_RETURN,
+                (ControlType.ERASE_IN_LINE, 2),
+                *(
+                    (
+                        (ControlType.CURSOR_UP, 1),
+                        (ControlType.ERASE_IN_LINE, 2),
+                    )
+                    * (height - 1)
+                ),
+            )
+        return Control()
+
+    def restore_cursor(self) -> Control:
+        """Get control codes to clear the render and restore the cursor to its previous position.
+
+        Returns:
+            Control: A Control instance that may be printed.
+        """
+        if self._shape is not None:
+            _, height = self._shape
+            return Control(
+                ControlType.CARRIAGE_RETURN,
+                *((ControlType.CURSOR_UP, 1), (ControlType.ERASE_IN_LINE, 2)) * height,
+            )
+        return Control()
+
+    def __rich_console__(
+        self, console: Console, options: ConsoleOptions
+    ) -> RenderResult:
+        renderable = self.renderable
+        style = console.get_style(self.style)
+        lines = console.render_lines(renderable, options, style=style, pad=False)
+        shape = Segment.get_shape(lines)
+
+        _, height = shape
+        if height > options.size.height:
+            if self.vertical_overflow == "crop":
+                lines = lines[: options.size.height]
+                shape = Segment.get_shape(lines)
+            elif self.vertical_overflow == "crop_above":
+                lines = lines[-(options.size.height) :]
+                shape = Segment.get_shape(lines)
+            elif self.vertical_overflow == "ellipsis":
+                lines = lines[: (options.size.height - 1)]
+                overflow_text = Text(
+                    "...",
+                    overflow="crop",
+                    justify="center",
+                    end="",
+                    style="live.ellipsis",
+                )
+                lines.append(list(console.render(overflow_text)))
+                shape = Segment.get_shape(lines)
+        self._shape = shape
+
+        new_line = Segment.line()
+        for last, line in loop_last(lines):
+            yield from line
+            if not last:
+                yield new_line

chatlas/_merge.py

@@ -1,5 +1,5 @@
 # Adapted from https://github.com/langchain-ai/langchain/blob/master/libs/core/langchain_core/utils/_merge.py
-# Also tweaked to more closely match https://github.com/hadley/elmer/blob/main/R/utils-merge.R
+# Also tweaked to more closely match https://github.com/hadley/ellmer/blob/main/R/utils-merge.R
 
 from __future__ import annotations
 

chatlas/_ollama.py

@@ -51,7 +51,7 @@ def ChatOllama(
     ::: {.callout-note}
     ## Python requirements
 
-    `ChatOllama` requires the `openai` package (e.g., `pip install openai`).
+    `ChatOllama` requires the `openai` package: `pip install "chatlas[ollama]"`.
     :::
 
 

chatlas/_openai.py

@@ -79,7 +79,7 @@ def ChatOpenAI(
     ::: {.callout-note}
     ## Python requirements
 
-    `ChatOpenAI` requires the `openai` package (e.g., `pip install openai`).
+    `ChatOpenAI` requires the `openai` package: `pip install "chatlas[openai]"`.
     :::
 
     Examples
@@ -338,7 +338,7 @@ class OpenAIProvider(Provider[ChatCompletion, ChatCompletionChunk, ChatCompletio
             return chunkd
         return merge_dicts(completion, chunkd)
 
-    def stream_turn(self, completion, has_data_model, stream) -> Turn:
+    def stream_turn(self, completion, has_data_model) -> Turn:
         from openai.types.chat import ChatCompletion
 
         delta = completion["choices"][0].pop("delta")  # type: ignore
@@ -346,9 +346,6 @@ class OpenAIProvider(Provider[ChatCompletion, ChatCompletionChunk, ChatCompletio
         completion = ChatCompletion.construct(**completion)
         return self._as_turn(completion, has_data_model)
 
-    async def stream_turn_async(self, completion, has_data_model, stream):
-        return self.stream_turn(completion, has_data_model, stream)
-
     def value_turn(self, completion, has_data_model) -> Turn:
         return self._as_turn(completion, has_data_model)
 
@@ -595,7 +592,8 @@ def ChatAzureOpenAI(
     ::: {.callout-note}
     ## Python requirements
 
-    `ChatAzureOpenAI` requires the `openai` package (e.g., `pip install openai`).
+    `ChatAzureOpenAI` requires the `openai` package:
+    `pip install "chatlas[azure-openai]"`.
     :::
 
     Examples

chatlas/_perplexity.py

@@ -43,7 +43,7 @@ def ChatPerplexity(
     ::: {.callout-note}
     ## Python requirements
 
-    `ChatPerplexity` requires the `openai` package (e.g., `pip install openai`).
+    `ChatPerplexity` requires the `openai` package: `pip install "chatlas[perplexity]"`.
     :::
 
 

chatlas/_provider.py

@@ -125,15 +125,6 @@ class Provider(
         self,
         completion: ChatCompletionDictT,
         has_data_model: bool,
-        stream: Any,
-    ) -> Turn: ...
-
-    @abstractmethod
-    async def stream_turn_async(
-        self,
-        completion: ChatCompletionDictT,
-        has_data_model: bool,
-        stream: Any,
     ) -> Turn: ...
 
     @abstractmethod

chatlas/_snowflake.py

@@ -0,0 +1,321 @@
+from typing import TYPE_CHECKING, Literal, Optional, TypedDict, overload
+
+from pydantic import BaseModel
+
+from ._chat import Chat
+from ._content import Content
+from ._logging import log_model_default
+from ._provider import Provider
+from ._tools import Tool
+from ._turn import Turn, normalize_turns
+from ._utils import drop_none
+
+if TYPE_CHECKING:
+    from snowflake.snowpark import Column
+
+    # Types inferred from the return type of the `snowflake.cortex.complete` function
+    Completion = str | Column
+    CompletionChunk = str
+
+    from .types.snowflake import SubmitInputArgs
+
+
+# The main prompt input type for Snowflake
+# This was copy-pasted from `snowflake.cortex._complete.ConversationMessage`
+class ConversationMessage(TypedDict):
+    role: str
+    content: str
+
+
+def ChatSnowflake(
+    *,
+    system_prompt: Optional[str] = None,
+    model: Optional[str] = None,
+    turns: Optional[list[Turn]] = None,
+    connection_name: Optional[str] = None,
+    account: Optional[str] = None,
+    user: Optional[str] = None,
+    password: Optional[str] = None,
+    private_key_file: Optional[str] = None,
+    private_key_file_pwd: Optional[str] = None,
+    kwargs: Optional[dict[str, "str | int"]] = None,
+) -> Chat["SubmitInputArgs", "Completion"]:
+    """
+    Chat with a Snowflake Cortex LLM
+
+    https://docs.snowflake.com/en/user-guide/snowflake-cortex/llm-functions
+
+    Prerequisites
+    -------------
+
+    ::: {.callout-note}
+    ## Python requirements
+
+    `ChatSnowflake`, requires the `snowflake-ml-python` package:
+    `pip install "chatlas[snowflake]"`.
+    :::
+
+    ::: {.callout-note}
+    ## Snowflake credentials
+
+    Snowflake provides a handful of ways to authenticate, but it's recommended
+    to use [key-pair
+    auth](https://docs.snowflake.com/en/developer-guide/python-connector/python-connector-connect#label-python-connection-toml)
+    to generate a `private_key_file`. It's also recommended to place your
+    credentials in a [`connections.toml`
+    file](https://docs.snowflake.com/en/developer-guide/snowpark/python/creating-session#connect-by-using-the-connections-toml-file).
+
+    This way, once your credentials are in the `connections.toml` file, you can
+    simply call `ChatSnowflake(connection_name="my_connection")` to
+    authenticate. If you don't want to use a `connections.toml` file, you can
+    specify the connection parameters directly (with `account`, `user`,
+    `password`, etc.).
+    :::
+
+
+    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.
+    connection_name
+        The name of the connection (i.e., section) within the connections.toml file.
+        This is useful if you want to keep your credentials in a connections.toml file
+        rather than specifying them directly in the arguments.
+        https://docs.snowflake.com/en/developer-guide/snowpark/python/creating-session#connect-by-using-the-connections-toml-file
+    account
+        Your Snowflake account identifier. Required if `connection_name` is not provided.
+        https://docs.snowflake.com/en/user-guide/admin-account-identifier
+    user
+        Your Snowflake user name. Required if `connection_name` is not provided.
+    password
+        Your Snowflake password. Required if doing password authentication and
+        `connection_name` is not provided.
+    private_key_file
+        The path to your private key file. Required if you are using key pair authentication.
+        https://docs.snowflake.com/en/user-guide/key-pair-auth
+    private_key_file_pwd
+        The password for your private key file. Required if you are using key pair authentication.
+        https://docs.snowflake.com/en/user-guide/key-pair-auth
+    kwargs
+        Additional keyword arguments passed along to the Snowflake connection builder. These can
+        include any parameters supported by the `snowflake-ml-python` package.
+        https://docs.snowflake.com/en/developer-guide/snowpark/python/creating-session#connect-by-specifying-connection-parameters
+    """
+
+    if model is None:
+        model = log_model_default("llama3.1-70b")
+
+    return Chat(
+        provider=SnowflakeProvider(
+            model=model,
+            connection_name=connection_name,
+            account=account,
+            user=user,
+            password=password,
+            private_key_file=private_key_file,
+            private_key_file_pwd=private_key_file_pwd,
+            kwargs=kwargs,
+        ),
+        turns=normalize_turns(
+            turns or [],
+            system_prompt,
+        ),
+    )
+
+
+class SnowflakeProvider(Provider["Completion", "CompletionChunk", "CompletionChunk"]):
+    def __init__(
+        self,
+        *,
+        model: str,
+        connection_name: Optional[str],
+        account: Optional[str],
+        user: Optional[str],
+        password: Optional[str],
+        private_key_file: Optional[str],
+        private_key_file_pwd: Optional[str],
+        kwargs: Optional[dict[str, "str | int"]],
+    ):
+        try:
+            from snowflake.snowpark import Session
+        except ImportError:
+            raise ImportError(
+                "`ChatSnowflake()` requires the `snowflake-ml-python` package. "
+                "Please install it via `pip install snowflake-ml-python`."
+            )
+
+        configs: dict[str, str | int] = drop_none(
+            {
+                "connection_name": connection_name,
+                "account": account,
+                "user": user,
+                "password": password,
+                "private_key_file": private_key_file,
+                "private_key_file_pwd": private_key_file_pwd,
+                **(kwargs or {}),
+            }
+        )
+
+        self._model = model
+        self._session = Session.builder.configs(configs).create()
+
+    @overload
+    def chat_perform(
+        self,
+        *,
+        stream: Literal[False],
+        turns: list[Turn],
+        tools: dict[str, Tool],
+        data_model: Optional[type[BaseModel]] = None,
+        kwargs: Optional["SubmitInputArgs"] = None,
+    ): ...
+
+    @overload
+    def chat_perform(
+        self,
+        *,
+        stream: Literal[True],
+        turns: list[Turn],
+        tools: dict[str, Tool],
+        data_model: Optional[type[BaseModel]] = None,
+        kwargs: Optional["SubmitInputArgs"] = None,
+    ): ...
+
+    def chat_perform(
+        self,
+        *,
+        stream: bool,
+        turns: list[Turn],
+        tools: dict[str, Tool],
+        data_model: Optional[type[BaseModel]] = None,
+        kwargs: Optional["SubmitInputArgs"] = None,
+    ):
+        from snowflake.cortex import complete
+
+        kwargs = self._chat_perform_args(stream, turns, tools, data_model, kwargs)
+        return complete(**kwargs)
+
+    @overload
+    async def chat_perform_async(
+        self,
+        *,
+        stream: Literal[False],
+        turns: list[Turn],
+        tools: dict[str, Tool],
+        data_model: Optional[type[BaseModel]] = None,
+        kwargs: Optional["SubmitInputArgs"] = None,
+    ): ...
+
+    @overload
+    async def chat_perform_async(
+        self,
+        *,
+        stream: Literal[True],
+        turns: list[Turn],
+        tools: dict[str, Tool],
+        data_model: Optional[type[BaseModel]] = None,
+        kwargs: Optional["SubmitInputArgs"] = None,
+    ): ...
+
+    async def chat_perform_async(
+        self,
+        *,
+        stream: bool,
+        turns: list[Turn],
+        tools: dict[str, Tool],
+        data_model: Optional[type[BaseModel]] = None,
+        kwargs: Optional["SubmitInputArgs"] = None,
+    ):
+        raise NotImplementedError(
+            "Snowflake does not currently support async completions."
+        )
+
+    def _chat_perform_args(
+        self,
+        stream: bool,
+        turns: list[Turn],
+        tools: dict[str, Tool],
+        data_model: Optional[type[BaseModel]] = None,
+        kwargs: Optional["SubmitInputArgs"] = None,
+    ):
+        # Cortex doesn't seem to support tools
+        if tools:
+            raise ValueError("Snowflake does not currently support tools.")
+
+        # TODO: implement data_model when this PR makes it into snowflake-ml-python
+        # https://github.com/snowflakedb/snowflake-ml-python/pull/141
+        # https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-llm-rest-api#structured-output-example
+        if data_model:
+            raise NotImplementedError(
+                "The snowflake-ml-python package currently doesn't support structured output. "
+                "Upvote this PR to help prioritize it: "
+                "https://github.com/snowflakedb/snowflake-ml-python/pull/141"
+            )
+
+        kwargs_full: "SubmitInputArgs" = {
+            "stream": stream,
+            "prompt": self._as_prompt_input(turns),
+            "model": self._model,
+            **(kwargs or {}),
+        }
+
+        return kwargs_full
+
+    def stream_text(self, chunk):
+        return chunk
+
+    def stream_merge_chunks(self, completion, chunk):
+        if completion is None:
+            return chunk
+        return completion + chunk
+
+    def stream_turn(self, completion, has_data_model) -> Turn:
+        return self._as_turn(completion, has_data_model)
+
+    def value_turn(self, completion, has_data_model) -> Turn:
+        return self._as_turn(completion, has_data_model)
+
+    def token_count(
+        self,
+        *args: "Content | str",
+        tools: dict[str, Tool],
+        data_model: Optional[type[BaseModel]],
+    ) -> int:
+        raise NotImplementedError(
+            "Snowflake does not currently support token counting."
+        )
+
+    async def token_count_async(
+        self,
+        *args: "Content | str",
+        tools: dict[str, Tool],
+        data_model: Optional[type[BaseModel]],
+    ) -> int:
+        raise NotImplementedError(
+            "Snowflake does not currently support token counting."
+        )
+
+    def _as_prompt_input(self, turns: list[Turn]) -> list["ConversationMessage"]:
+        res: list["ConversationMessage"] = []
+        for turn in turns:
+            res.append(
+                {
+                    "role": turn.role,
+                    "content": turn.text,
+                }
+            )
+        return res
+
+    def _as_turn(self, completion, has_data_model) -> Turn:
+        return Turn("assistant", completion)

chatlas/_utils.py

@@ -61,6 +61,13 @@ def is_async_callable(
     return False
 
 
+T = TypeVar("T")
+
+
+def drop_none(x: dict[str, T | None]) -> dict[str, T]:
+    return {k: v for k, v in x.items() if v is not None}
+
+
 # https://docs.pytest.org/en/latest/example/simple.html#pytest-current-test-environment-variable
 def is_testing():
     return os.environ.get("PYTEST_CURRENT_TEST", None) is not None

chatlas/_version.py

@@ -0,0 +1,21 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '0.5.0'
+__version_tuple__ = version_tuple = (0, 5, 0)

chatlas/types/__init__.py

@@ -1,4 +1,8 @@
-from .._chat import ChatResponse, ChatResponseAsync, SubmitInputArgsT
+from .._chat import (  # noqa: A005
+  ChatResponse,
+  ChatResponseAsync,
+  SubmitInputArgsT,
+)
 from .._content import (
     Content,
     ContentImage,

chatlas/types/anthropic/_submit.py

@@ -8,18 +8,24 @@ from typing import Iterable, Literal, Mapping, Optional, TypedDict, Union
 import anthropic
 import anthropic.types.message_param
 import anthropic.types.text_block_param
+import anthropic.types.thinking_config_disabled_param
+import anthropic.types.thinking_config_enabled_param
+import anthropic.types.tool_bash_20250124_param
 import anthropic.types.tool_choice_any_param
 import anthropic.types.tool_choice_auto_param
+import anthropic.types.tool_choice_none_param
 import anthropic.types.tool_choice_tool_param
 import anthropic.types.tool_param
+import anthropic.types.tool_text_editor_20250124_param
 
 
 class SubmitInputArgs(TypedDict, total=False):
     max_tokens: int
     messages: Iterable[anthropic.types.message_param.MessageParam]
     model: Union[
-        str,
         Literal[
+            "claude-3-7-sonnet-latest",
+            "claude-3-7-sonnet-20250219",
             "claude-3-5-haiku-latest",
             "claude-3-5-haiku-20241022",
             "claude-3-5-sonnet-latest",
@@ -32,6 +38,7 @@ class SubmitInputArgs(TypedDict, total=False):
             "claude-2.1",
             "claude-2.0",
         ],
+        str,
     ]
     stop_sequences: Union[list[str], anthropic.NotGiven]
     stream: Union[Literal[False], Literal[True], anthropic.NotGiven]
@@ -41,13 +48,28 @@ class SubmitInputArgs(TypedDict, total=False):
         anthropic.NotGiven,
     ]
     temperature: float | anthropic.NotGiven
+    thinking: Union[
+        anthropic.types.thinking_config_enabled_param.ThinkingConfigEnabledParam,
+        anthropic.types.thinking_config_disabled_param.ThinkingConfigDisabledParam,
+        anthropic.NotGiven,
+    ]
     tool_choice: Union[
         anthropic.types.tool_choice_auto_param.ToolChoiceAutoParam,
         anthropic.types.tool_choice_any_param.ToolChoiceAnyParam,
         anthropic.types.tool_choice_tool_param.ToolChoiceToolParam,
+        anthropic.types.tool_choice_none_param.ToolChoiceNoneParam,
+        anthropic.NotGiven,
+    ]
+    tools: Union[
+        Iterable[
+            Union[
+                anthropic.types.tool_param.ToolParam,
+                anthropic.types.tool_bash_20250124_param.ToolBash20250124Param,
+                anthropic.types.tool_text_editor_20250124_param.ToolTextEditor20250124Param,
+            ]
+        ],
         anthropic.NotGiven,
     ]
-    tools: Union[Iterable[anthropic.types.tool_param.ToolParam], anthropic.NotGiven]
     top_k: int | anthropic.NotGiven
     top_p: float | anthropic.NotGiven
     extra_headers: Optional[Mapping[str, Union[str, anthropic.Omit]]]