chatlas 0.8.0__py3-none-any.whl → 0.9.0__py3-none-any.whl

chatlas/_provider.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 from abc import ABC, abstractmethod
 from typing import (
-    Any,
     AsyncIterable,
     Generic,
     Iterable,
@@ -17,6 +16,7 @@ from pydantic import BaseModel
 from ._content import Content
 from ._tools import Tool
 from ._turn import Turn
+from ._typing_extensions import TypedDict
 
 ChatCompletionT = TypeVar("ChatCompletionT")
 ChatCompletionChunkT = TypeVar("ChatCompletionChunkT")
@@ -24,8 +24,52 @@ ChatCompletionChunkT = TypeVar("ChatCompletionChunkT")
 ChatCompletionDictT = TypeVar("ChatCompletionDictT")
 
 
+class AnyTypeDict(TypedDict, total=False):
+    pass
+
+
+SubmitInputArgsT = TypeVar("SubmitInputArgsT", bound=AnyTypeDict)
+"""
+A TypedDict representing the provider specific arguments that can specified when
+submitting input to a model provider.
+"""
+
+
+class StandardModelParams(TypedDict, total=False):
+    """
+    A TypedDict representing the standard model parameters that can be set
+    when using a [](`~chatlas.Chat`) instance.
+    """
+
+    temperature: float
+    top_p: float
+    top_k: int
+    frequency_penalty: float
+    presence_penalty: float
+    seed: int
+    max_tokens: int
+    log_probs: bool
+    stop_sequences: list[str]
+
+
+StandardModelParamNames = Literal[
+    "temperature",
+    "top_p",
+    "top_k",
+    "frequency_penalty",
+    "presence_penalty",
+    "seed",
+    "max_tokens",
+    "log_probs",
+    "stop_sequences",
+]
+
+
 class Provider(
-    ABC, Generic[ChatCompletionT, ChatCompletionChunkT, ChatCompletionDictT]
+    ABC,
+    Generic[
+        ChatCompletionT, ChatCompletionChunkT, ChatCompletionDictT, SubmitInputArgsT
+    ],
 ):
     """
     A model provider interface for a [](`~chatlas.Chat`).
@@ -40,6 +84,24 @@ class Provider(
     directly.
     """
 
+    def __init__(self, *, name: str, model: str):
+        self._name = name
+        self._model = model
+
+    @property
+    def name(self):
+        """
+        Get the name of the provider
+        """
+        return self._name
+
+    @property
+    def model(self):
+        """
+        Get the model used by the provider
+        """
+        return self._model
+
     @overload
     @abstractmethod
     def chat_perform(
@@ -49,7 +111,7 @@ class Provider(
         turns: list[Turn],
         tools: dict[str, Tool],
         data_model: Optional[type[BaseModel]],
-        kwargs: Any,
+        kwargs: SubmitInputArgsT,
     ) -> ChatCompletionT: ...
 
     @overload
@@ -61,7 +123,7 @@ class Provider(
         turns: list[Turn],
         tools: dict[str, Tool],
         data_model: Optional[type[BaseModel]],
-        kwargs: Any,
+        kwargs: SubmitInputArgsT,
     ) -> Iterable[ChatCompletionChunkT]: ...
 
     @abstractmethod
@@ -72,7 +134,7 @@ class Provider(
         turns: list[Turn],
         tools: dict[str, Tool],
         data_model: Optional[type[BaseModel]],
-        kwargs: Any,
+        kwargs: SubmitInputArgsT,
     ) -> Iterable[ChatCompletionChunkT] | ChatCompletionT: ...
 
     @overload
@@ -84,7 +146,7 @@ class Provider(
         turns: list[Turn],
         tools: dict[str, Tool],
         data_model: Optional[type[BaseModel]],
-        kwargs: Any,
+        kwargs: SubmitInputArgsT,
     ) -> ChatCompletionT: ...
 
     @overload
@@ -96,7 +158,7 @@ class Provider(
         turns: list[Turn],
         tools: dict[str, Tool],
         data_model: Optional[type[BaseModel]],
-        kwargs: Any,
+        kwargs: SubmitInputArgsT,
     ) -> AsyncIterable[ChatCompletionChunkT]: ...
 
     @abstractmethod
@@ -107,7 +169,7 @@ class Provider(
         turns: list[Turn],
         tools: dict[str, Tool],
         data_model: Optional[type[BaseModel]],
-        kwargs: Any,
+        kwargs: SubmitInputArgsT,
     ) -> AsyncIterable[ChatCompletionChunkT] | ChatCompletionT: ...
 
     @abstractmethod
@@ -149,3 +211,11 @@ class Provider(
         tools: dict[str, Tool],
         data_model: Optional[type[BaseModel]],
     ) -> int: ...
+
+    @abstractmethod
+    def translate_model_params(
+        self, params: StandardModelParams
+    ) -> SubmitInputArgsT: ...
+
+    @abstractmethod
+    def supported_model_params(self) -> set[StandardModelParamNames]: ...

chatlas/_snowflake.py

@@ -20,10 +20,10 @@ from ._content import (
     ContentToolResult,
 )
 from ._logging import log_model_default
-from ._provider import Provider
+from ._provider import Provider, StandardModelParamNames, StandardModelParams
 from ._tokens import tokens_log
 from ._tools import Tool, basemodel_to_param_schema
-from ._turn import Turn, normalize_turns
+from ._turn import Turn
 from ._utils import drop_none
 
 if TYPE_CHECKING:
@@ -61,7 +61,6 @@ 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,
@@ -111,13 +110,6 @@ def ChatSnowflake(
         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
@@ -157,14 +149,13 @@ def ChatSnowflake(
             private_key_file_pwd=private_key_file_pwd,
             kwargs=kwargs,
         ),
-        turns=normalize_turns(
-            turns or [],
-            system_prompt,
-        ),
+        system_prompt=system_prompt,
     )
 
 
-class SnowflakeProvider(Provider["Completion", "CompletionChunk", "CompletionChunk"]):
+class SnowflakeProvider(
+    Provider["Completion", "CompletionChunk", "CompletionChunk", "CompleteRequest"]
+):
     def __init__(
         self,
         *,
@@ -175,6 +166,7 @@ class SnowflakeProvider(Provider["Completion", "CompletionChunk", "CompletionChu
         password: Optional[str],
         private_key_file: Optional[str],
         private_key_file_pwd: Optional[str],
+        name: str = "Snowflake",
         kwargs: Optional[dict[str, "str | int"]],
     ):
         try:
@@ -185,6 +177,7 @@ class SnowflakeProvider(Provider["Completion", "CompletionChunk", "CompletionChu
                 "`ChatSnowflake()` requires the `snowflake-ml-python` package. "
                 "Please install it via `pip install snowflake-ml-python`."
             )
+        super().__init__(name=name, model=model)
 
         configs: dict[str, str | int] = drop_none(
             {
@@ -198,8 +191,6 @@ class SnowflakeProvider(Provider["Completion", "CompletionChunk", "CompletionChu
             }
         )
 
-        self._model = model
-
         session = Session.builder.configs(configs).create()
         self._cortex_service = Root(session).cortex_inference_service
 
@@ -314,7 +305,7 @@ class SnowflakeProvider(Provider["Completion", "CompletionChunk", "CompletionChu
         from snowflake.core.cortex.inference_service import CompleteRequest
 
         req = CompleteRequest(
-            model=self._model,
+            model=self.model,
             messages=self._as_request_messages(turns),
             stream=stream,
         )
@@ -599,6 +590,26 @@ class SnowflakeProvider(Provider["Completion", "CompletionChunk", "CompletionChu
 
         return models.Tool(tool_spec=spec)
 
+    def translate_model_params(self, params: StandardModelParams) -> "CompleteRequest":
+        res: "CompleteRequest" = {}
+        if "temperature" in params:
+            res["temperature"] = params["temperature"]
+
+        if "top_p" in params:
+            res["top_p"] = params["top_p"]
+
+        if "max_tokens" in params:
+            res["max_tokens"] = params["max_tokens"]
+
+        return res
+
+    def supported_model_params(self) -> set[StandardModelParamNames]:
+        return {
+            "temperature",
+            "top_p",
+            "max_tokens",
+        }
+
 
 # Yield parsed event data from the Snowflake SSEClient
 # (this is only needed for the streaming case).

chatlas/_tokens.py

@@ -1,9 +1,13 @@
 from __future__ import annotations
 
 import copy
+import importlib.resources as resources
+import warnings
 from threading import Lock
 from typing import TYPE_CHECKING
 
+import orjson
+
 from ._logging import logger
 from ._typing_extensions import TypedDict
 
@@ -17,8 +21,10 @@ class TokenUsage(TypedDict):
     """
 
     name: str
+    model: str
     input: int
     output: int
+    cost: float | None
 
 
 class ThreadSafeTokenCounter:
@@ -26,7 +32,9 @@ class ThreadSafeTokenCounter:
         self._lock = Lock()
         self._tokens: dict[str, TokenUsage] = {}
 
-    def log_tokens(self, name: str, input_tokens: int, output_tokens: int) -> None:
+    def log_tokens(
+        self, name: str, model: str, input_tokens: int, output_tokens: int
+    ) -> None:
         logger.info(
             f"Provider '{name}' generated a response of {output_tokens} tokens "
             f"from an input of {input_tokens} tokens."
@@ -36,12 +44,21 @@ class ThreadSafeTokenCounter:
             if name not in self._tokens:
                 self._tokens[name] = {
                     "name": name,
+                    "model": model,
                     "input": input_tokens,
                     "output": output_tokens,
+                    "cost": compute_price(name, model, input_tokens, output_tokens),
                 }
             else:
                 self._tokens[name]["input"] += input_tokens
                 self._tokens[name]["output"] += output_tokens
+                price = compute_price(name, model, input_tokens, output_tokens)
+                if price is not None:
+                    cost = self._tokens[name]["cost"]
+                    if cost is None:
+                        self._tokens[name]["cost"] = price
+                    else:
+                        self._tokens[name]["cost"] = cost + price
 
     def get_usage(self) -> list[TokenUsage] | None:
         with self._lock:
@@ -59,8 +76,7 @@ def tokens_log(provider: "Provider", tokens: tuple[int, int]) -> None:
     """
     Log token usage for a provider in a thread-safe manner.
     """
-    name = provider.__class__.__name__.replace("Provider", "")
-    _token_counter.log_tokens(name, tokens[0], tokens[1])
+    _token_counter.log_tokens(provider.name, provider.model, tokens[0], tokens[1])
 
 
 def tokens_reset() -> None:
@@ -71,17 +87,89 @@ def tokens_reset() -> None:
     _token_counter = ThreadSafeTokenCounter()
 
 
+class TokenPrice(TypedDict):
+    """
+    Defines the necessary information to look up pricing for a given turn.
+    """
+
+    provider: str
+    """The provider name (e.g., "OpenAI", "Anthropic", etc.)"""
+    model: str
+    """The model name (e.g., "gpt-3.5-turbo", "claude-2", etc.)"""
+    cached_input: float
+    """The cost per user token in USD per million tokens for cached input"""
+    input: float
+    """The cost per user token in USD per million tokens"""
+    output: float
+    """The cost per assistant token in USD per million tokens"""
+
+
+# Load in pricing pulled from ellmer
+f = resources.files("chatlas").joinpath("data/prices.json").read_text(encoding="utf-8")
+pricing_list: list[TokenPrice] = orjson.loads(f)
+
+
+def get_token_pricing(name: str, model: str) -> TokenPrice | None:
+    """
+    Get token pricing information given a provider name and model
+
+    Note
+    ----
+    Only a subset of providers and models and currently supported.
+    The pricing information derives from ellmer.
+
+    Returns
+    -------
+    TokenPrice | None
+    """
+    result = next(
+        (
+            item
+            for item in pricing_list
+            if item["provider"] == name and item["model"] == model
+        ),
+        None,
+    )
+    if result is None:
+        warnings.warn(
+            f"Token pricing for the provider '{name}' and model '{model}' you selected is not available. "
+            "Please check the provider's documentation."
+        )
+
+    return result
+
+
+def compute_price(
+    name: str, model: str, input_tokens: int, output_tokens: int
+) -> float | None:
+    """
+    Compute the cost of a turn.
+
+    Returns
+    -------
+    float | None
+        The cost of the turn in USD, or None if the cost could not be calculated.
+    """
+    price = get_token_pricing(name, model)
+    if price is None:
+        return None
+    input_price = input_tokens * (price["input"] / 1e6)
+    output_price = output_tokens * (price["output"] / 1e6)
+    return input_price + output_price
+
+
 def token_usage() -> list[TokenUsage] | None:
     """
     Report on token usage in the current session
 
     Call this function to find out the cumulative number of tokens that you
-    have sent and received in the current session.
+    have sent and received in the current session. The price will be shown if known
 
     Returns
     -------
     list[TokenUsage] | None
-        A list of dictionaries with the following keys: "name", "input", and "output".
+        A list of dictionaries with the following keys: "name", "input", "output", and "cost".
+        If no cost data is available for the name/model combination chosen, then "cost" will be None.
         If no tokens have been logged, then None is returned.
     """
     return _token_counter.get_usage()

chatlas/_tools.py

@@ -2,11 +2,17 @@ from __future__ import annotations
 
 import inspect
 import warnings
-from typing import TYPE_CHECKING, Any, Awaitable, Callable, Optional
+from typing import TYPE_CHECKING, Any, AsyncGenerator, Awaitable, Callable, Optional
 
+import openai
 from pydantic import BaseModel, Field, create_model
 
 from . import _utils
+from ._content import (
+    ContentToolResult,
+    ContentToolResultImage,
+    ContentToolResultResource,
+)
 
 __all__ = (
     "Tool",
@@ -14,6 +20,8 @@ __all__ = (
 )
 
 if TYPE_CHECKING:
+    from mcp import ClientSession as MCPClientSession
+    from mcp import Tool as MCPTool
     from openai.types.chat import ChatCompletionToolParam
 
 
@@ -28,26 +36,168 @@ class Tool:
     ----------
     func
         The function to be invoked when the tool is called.
-    model
-        A Pydantic model that describes the input parameters for the function.
-        If not provided, the model will be inferred from the function's type hints.
-        The primary reason why you might want to provide a model in
-        Note that the name and docstring of the model takes precedence over the
-        name and docstring of the function.
+    name
+        The name of the tool.
+    description
+        A description of what the tool does.
+    parameters
+        A dictionary describing the input parameters and their types.
     """
 
     func: Callable[..., Any] | Callable[..., Awaitable[Any]]
 
     def __init__(
         self,
-        func: Callable[..., Any] | Callable[..., Awaitable[Any]],
         *,
-        model: Optional[type[BaseModel]] = None,
+        func: Callable[..., Any] | Callable[..., Awaitable[Any]],
+        name: str,
+        description: str,
+        parameters: dict[str, Any],
     ):
+        self.name = name
         self.func = func
         self._is_async = _utils.is_async_callable(func)
-        self.schema = func_to_schema(func, model)
-        self.name = self.schema["function"]["name"]
+        self.schema: "ChatCompletionToolParam" = {
+            "type": "function",
+            "function": {
+                "name": name,
+                "description": description,
+                "parameters": parameters,
+            },
+        }
+
+    @classmethod
+    def from_func(
+        cls: type["Tool"],
+        func: Callable[..., Any] | Callable[..., Awaitable[Any]],
+        *,
+        model: Optional[type[BaseModel]] = None,
+    ) -> "Tool":
+        """
+        Create a Tool from a Python function
+
+        Parameters
+        ----------
+        func
+            The function to wrap as a tool.
+        model
+            A Pydantic model that describes the input parameters for the function.
+            If not provided, the model will be inferred from the function's type hints.
+            The primary reason why you might want to provide a model in
+            Note that the name and docstring of the model takes precedence over the
+            name and docstring of the function.
+
+        Returns
+        -------
+        Tool
+            A new Tool instance wrapping the provided function.
+
+        Raises
+        ------
+        ValueError
+            If there is a mismatch between model fields and function parameters.
+        """
+
+        if model is None:
+            model = func_to_basemodel(func)
+
+        # Throw if there is a mismatch between the model and the function parameters
+        params = inspect.signature(func).parameters
+        fields = model.model_fields
+        diff = set(params) ^ set(fields)
+        if diff:
+            raise ValueError(
+                f"`model` fields must match tool function parameters exactly. "
+                f"Fields found in one but not the other: {diff}"
+            )
+
+        params = basemodel_to_param_schema(model)
+
+        return cls(
+            func=func,
+            name=model.__name__ or func.__name__,
+            description=model.__doc__ or func.__doc__ or "",
+            parameters=params,
+        )
+
+    @classmethod
+    def from_mcp(
+        cls: type["Tool"],
+        session: "MCPClientSession",
+        mcp_tool: "MCPTool",
+    ) -> "Tool":
+        """
+        Create a Tool from an MCP tool
+
+        Parameters
+        ----------
+        session
+            The MCP client session to use for calling the tool.
+        mcp_tool
+            The MCP tool to wrap.
+
+        Returns
+        -------
+        Tool
+            A new Tool instance wrapping the MCP tool.
+        """
+
+        async def _call(**args: Any) -> AsyncGenerator[ContentToolResult, None]:
+            result = await session.call_tool(mcp_tool.name, args)
+
+            # Raise an error if the tool call resulted in an error. It doesn't seem to be
+            # very well defined how to get at the error message, but it appears that it gets
+            # stored in the `text` attribute of the content. Also, empirically, the error
+            # message seems to include `Error executing tool {tool_name}: ...`, so
+            if result.isError:
+                err_msg = getattr(
+                    result.content[0],
+                    "text",
+                    f"Error executing tool {mcp_tool.name}.",
+                )
+                raise RuntimeError(err_msg)
+
+            for content in result.content:
+                if content.type == "text":
+                    yield ContentToolResult(value=content.text)
+                elif content.type == "image":
+                    if content.mimeType not in (
+                        "image/png",
+                        "image/jpeg",
+                        "image/webp",
+                        "image/gif",
+                    ):
+                        raise ValueError(
+                            f"Unsupported image MIME type: {content.mimeType}"
+                        )
+
+                    yield ContentToolResultImage(
+                        value=content.data,
+                        mime_type=content.mimeType,
+                    )
+                elif content.type == "resource":
+                    from mcp.types import TextResourceContents
+
+                    resource = content.resource
+                    if isinstance(resource, TextResourceContents):
+                        blob = resource.text.encode("utf-8")
+                    else:
+                        blob = resource.blob.encode("utf-8")
+
+                    yield ContentToolResultResource(
+                        value=blob, mime_type=content.resource.mimeType
+                    )
+                else:
+                    raise RuntimeError(f"Unexpected content type: {content.type}")
+
+        params = mcp_tool_input_schema_to_param_schema(mcp_tool.inputSchema)
+
+        return cls(
+            func=_utils.wrap_async(_call),
+            name=mcp_tool.name,
+            description=mcp_tool.description or "",
+            parameters=params,
+        )
 
 
 class ToolRejectError(Exception):
@@ -160,14 +310,6 @@ def func_to_basemodel(func: Callable) -> type[BaseModel]:
 
 
 def basemodel_to_param_schema(model: type[BaseModel]) -> dict[str, object]:
-    try:
-        import openai
-    except ImportError:
-        raise ImportError(
-            "The openai package is required for this functionality. "
-            "Please install it with `pip install openai`."
-        )
-
     # Lean on openai's ability to translate BaseModel.model_json_schema()
     # to a valid tool schema (this wouldn't be impossible to do ourselves,
     # but it's fair amount of logic to substitute `$refs`, etc.)
@@ -177,10 +319,27 @@ def basemodel_to_param_schema(model: type[BaseModel]) -> dict[str, object]:
     if "parameters" not in fn:
         raise ValueError("Expected `parameters` in function definition.")
 
-    params = fn["parameters"]
+    params = rm_param_titles(fn["parameters"])
+
+    return params
+
+
+def mcp_tool_input_schema_to_param_schema(
+    input_schema: dict[str, Any],
+) -> dict[str, object]:
+    params = rm_param_titles(input_schema)
+
+    if "additionalProperties" not in params:
+        params["additionalProperties"] = False
+
+    return params
+
 
-    # For some reason, openai (or pydantic?) wants to include a title
-    # at the model and field level. I don't think we actually need or want this.
+def rm_param_titles(
+    params: dict[str, object],
+) -> dict[str, object]:
+    # For some reason, pydantic wants to include a title at the model and field
+    # level. I don't think we actually need or want this.
     if "title" in params:
         del params["title"]