chatlas 0.9.2__py3-none-any.whl → 0.11.0__py3-none-any.whl

chatlas/_provider_ollama.py

@@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Optional
 import orjson
 
 from ._chat import Chat
-from ._provider_openai import OpenAIProvider
+from ._provider_openai import ModelInfo, OpenAIProvider
 from ._utils import MISSING_TYPE, is_testing
 
 if TYPE_CHECKING:
@@ -90,18 +90,19 @@ def ChatOllama(
         raise RuntimeError("Can't find locally running ollama.")
 
     if model is None:
-        models = ollama_models(base_url)
+        models = ollama_model_info(base_url)
+        model_ids = [m["id"] for m in models]
         raise ValueError(
-            f"Must specify model. Locally installed models: {', '.join(models)}"
+            f"Must specify model. Locally installed models: {', '.join(model_ids)}"
         )
     if isinstance(seed, MISSING_TYPE):
         seed = 1014 if is_testing() else None
 
     return Chat(
-        provider=OpenAIProvider(
+        provider=OllamaProvider(
             api_key="ollama",  # ignored
             model=model,
-            base_url=f"{base_url}/v1",
+            base_url=base_url,
             seed=seed,
             name="Ollama",
             kwargs=kwargs,
@@ -110,10 +111,40 @@ def ChatOllama(
     )
 
 
-def ollama_models(base_url: str) -> list[str]:
-    res = urllib.request.urlopen(url=f"{base_url}/api/tags")
-    data = orjson.loads(res.read())
-    return [re.sub(":latest$", "", x["name"]) for x in data["models"]]
+class OllamaProvider(OpenAIProvider):
+    def __init__(self, *, api_key, model, base_url, seed, name, kwargs):
+        super().__init__(
+            api_key=api_key,
+            model=model,
+            base_url=f"{base_url}/v1",
+            seed=seed,
+            name=name,
+            kwargs=kwargs,
+        )
+        self.base_url = base_url
+
+    def list_models(self):
+        return ollama_model_info(self.base_url)
+
+
+def ollama_model_info(base_url: str) -> list[ModelInfo]:
+    response = urllib.request.urlopen(url=f"{base_url}/api/tags")
+    data = orjson.loads(response.read())
+    models = data.get("models", [])
+    if not models:
+        return []
+
+    res: list[ModelInfo] = []
+    for model in models:
+        # TODO: add capabilities
+        info: ModelInfo = {
+            "id": re.sub(":latest$", "", model["name"]),
+            "created_at": model["modified_at"],
+            "size": model["size"],
+        }
+        res.append(info)
+
+    return res
 
 
 def has_ollama(base_url):

chatlas/_provider_openai.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import base64
+from datetime import datetime
 from typing import TYPE_CHECKING, Any, Literal, Optional, cast, overload
 
 import orjson
@@ -23,8 +24,8 @@ from ._content import (
 )
 from ._logging import log_model_default
 from ._merge import merge_dicts
-from ._provider import Provider, StandardModelParamNames, StandardModelParams
-from ._tokens import tokens_log
+from ._provider import ModelInfo, Provider, StandardModelParamNames, StandardModelParams
+from ._tokens import get_token_pricing, tokens_log
 from ._tools import Tool, basemodel_to_param_schema
 from ._turn import Turn, user_turn
 from ._utils import MISSING, MISSING_TYPE, is_testing, split_http_client_kwargs
@@ -200,6 +201,32 @@ class OpenAIProvider(
         self._client = OpenAI(**sync_kwargs)  # type: ignore
         self._async_client = AsyncOpenAI(**async_kwargs)
 
+    def list_models(self):
+        models = self._client.models.list()
+
+        res: list[ModelInfo] = []
+        for m in models:
+            pricing = get_token_pricing(self.name, m.id) or {}
+            info: ModelInfo = {
+                "id": m.id,
+                "owned_by": m.owned_by,
+                "input": pricing.get("input"),
+                "output": pricing.get("output"),
+                "cached_input": pricing.get("cached_input"),
+            }
+            # DeepSeek compatibility
+            if m.created is not None:
+                info["created_at"] = datetime.fromtimestamp(m.created).date()
+            res.append(info)
+
+        # More recent models first
+        res.sort(
+            key=lambda x: x.get("created_at", 0),
+            reverse=True,
+        )
+
+        return res
+
     @overload
     def chat_perform(
         self,
@@ -310,8 +337,7 @@ class OpenAIProvider(
                 del kwargs_full["tools"]
 
         if stream and "stream_options" not in kwargs_full:
-            if self.__class__.__name__ != "DatabricksProvider":
-                kwargs_full["stream_options"] = {"include_usage": True}
+            kwargs_full["stream_options"] = {"include_usage": True}
 
         return kwargs_full
 
@@ -411,7 +437,9 @@ class OpenAIProvider(
                     if isinstance(x, ContentText):
                         content_parts.append({"type": "text", "text": x.text})
                     elif isinstance(x, ContentJson):
-                        content_parts.append({"type": "text", "text": ""})
+                        content_parts.append(
+                            {"type": "text", "text": "<structured data/>"}
+                        )
                     elif isinstance(x, ContentToolRequest):
                         tool_calls.append(
                             {
@@ -450,7 +478,7 @@ class OpenAIProvider(
                     if isinstance(x, ContentText):
                         contents.append({"type": "text", "text": x.text})
                     elif isinstance(x, ContentJson):
-                        contents.append({"type": "text", "text": ""})
+                        contents.append({"type": "text", "text": "<structured data/>"})
                     elif isinstance(x, ContentPDF):
                         contents.append(
                             {
@@ -522,7 +550,10 @@ class OpenAIProvider(
         contents: list[Content] = []
         if message.content is not None:
             if has_data_model:
-                data = orjson.loads(message.content)
+                data = message.content
+                # Some providers (e.g., Cloudflare) may already provide a dict
+                if not isinstance(data, dict):
+                    data = orjson.loads(data)
                 contents = [ContentJson(value=data)]
             else:
                 contents = [ContentText(text=message.content)]

chatlas/_provider_openrouter.py

@@ -0,0 +1,149 @@
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING, Optional
+
+from ._chat import Chat
+from ._logging import log_model_default
+from ._provider_openai import OpenAIProvider
+from ._utils import MISSING, MISSING_TYPE, is_testing
+
+if TYPE_CHECKING:
+    from ._provider_openai import ChatCompletion
+    from .types.openai import ChatClientArgs, SubmitInputArgs
+
+
+def ChatOpenRouter(
+    *,
+    system_prompt: Optional[str] = None,
+    model: Optional[str] = None,
+    api_key: Optional[str] = None,
+    base_url: str = "https://openrouter.ai/api/v1",
+    seed: Optional[int] | MISSING_TYPE = MISSING,
+    kwargs: Optional["ChatClientArgs"] = None,
+) -> Chat["SubmitInputArgs", ChatCompletion]:
+    """
+    Chat with one of the many models hosted on OpenRouter.
+
+    OpenRouter provides access to a wide variety of language models from different providers
+    through a unified API. Support for features depends on the underlying model that you use.
+
+    Prerequisites
+    -------------
+
+    ::: {.callout-note}
+    ## API key
+
+    Sign up at <https://openrouter.ai> to get an API key.
+    :::
+
+    Examples
+    --------
+
+    ```python
+    import os
+    from chatlas import ChatOpenRouter
+
+    chat = ChatOpenRouter(api_key=os.getenv("OPENROUTER_API_KEY"))
+    chat.chat("What is the capital of France?")
+    ```
+
+    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. See <https://openrouter.ai/models>
+        for available models.
+    api_key
+        The API key to use for authentication. You generally should not supply
+        this directly, but instead set the `OPENROUTER_API_KEY` environment variable.
+    base_url
+        The base URL to the endpoint; the default uses OpenRouter's API.
+    seed
+        Optional integer seed that the model uses to try and make output more
+        reproducible.
+    kwargs
+        Additional arguments to pass to the `openai.OpenAI()` client constructor.
+
+    Returns
+    -------
+    Chat
+        A chat object that retains the state of the conversation.
+
+    Note
+    ----
+    This function is a lightweight wrapper around [](`~chatlas.ChatOpenAI`) with
+    the defaults tweaked for OpenRouter.
+
+    Note
+    ----
+    Pasting an API key into a chat constructor (e.g., `ChatOpenRouter(api_key="...")`)
+    is the simplest way to get started, and is fine for interactive use, but is
+    problematic for code that may be shared with others.
+
+    Instead, consider using environment variables or a configuration file to manage
+    your credentials. One popular way to manage credentials is to use a `.env` file
+    to store your credentials, and then use the `python-dotenv` package to load them
+    into your environment.
+
+    ```shell
+    pip install python-dotenv
+    ```
+
+    ```shell
+    # .env
+    OPENROUTER_API_KEY=...
+    ```
+
+    ```python
+    from chatlas import ChatOpenRouter
+    from dotenv import load_dotenv
+
+    load_dotenv()
+    chat = ChatOpenRouter()
+    chat.console()
+    ```
+
+    Another, more general, solution is to load your environment variables into the shell
+    before starting Python (maybe in a `.bashrc`, `.zshrc`, etc. file):
+
+    ```shell
+    export OPENROUTER_API_KEY=...
+    ```
+    """
+    if model is None:
+        model = log_model_default("gpt-4.1")
+
+    if api_key is None:
+        api_key = os.getenv("OPENROUTER_API_KEY")
+
+    if isinstance(seed, MISSING_TYPE):
+        seed = 1014 if is_testing() else None
+
+    kwargs2 = add_default_headers(kwargs or {})
+
+    return Chat(
+        provider=OpenAIProvider(
+            api_key=api_key,
+            model=model,
+            base_url=base_url,
+            seed=seed,
+            name="OpenRouter",
+            kwargs=kwargs2,
+        ),
+        system_prompt=system_prompt,
+    )
+
+
+def add_default_headers(kwargs: "ChatClientArgs") -> "ChatClientArgs":
+    headers = kwargs.get("default_headers", None)
+    # https://openrouter.ai/docs/api-keys
+    default_headers = {
+        "HTTP-Referer": "https://posit-dev.github.io/chatlas",
+        "X-Title": "chatlas",
+        **(headers or {}),
+    }
+    return {"default_headers": default_headers, **kwargs}

chatlas/_provider_perplexity.py

@@ -126,7 +126,7 @@ def ChatPerplexity(
         seed = 1014 if is_testing() else None
 
     return Chat(
-        provider=OpenAIProvider(
+        provider=PerplexityProvider(
             api_key=api_key,
             model=model,
             base_url=base_url,
@@ -136,3 +136,11 @@ def ChatPerplexity(
         ),
         system_prompt=system_prompt,
     )
+
+
+class PerplexityProvider(OpenAIProvider):
+    def list_models(self):
+        raise NotImplementedError(
+            ".list_models() is not yet implemented for Perplexity."
+            " To view available models online, see https://docs.perplexity.ai/getting-started/models"
+        )

chatlas/_provider_portkey.py

@@ -0,0 +1,131 @@
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING, Optional
+
+from ._chat import Chat
+from ._logging import log_model_default
+from ._provider_openai import OpenAIProvider
+from ._utils import drop_none
+
+if TYPE_CHECKING:
+    from ._provider_openai import ChatCompletion
+    from .types.openai import ChatClientArgs, SubmitInputArgs
+
+
+def ChatPortkey(
+    *,
+    system_prompt: Optional[str] = None,
+    model: Optional[str] = None,
+    api_key: Optional[str] = None,
+    virtual_key: Optional[str] = None,
+    base_url: str = "https://api.portkey.ai/v1",
+    kwargs: Optional["ChatClientArgs"] = None,
+) -> Chat["SubmitInputArgs", ChatCompletion]:
+    """
+    Chat with a model hosted on PortkeyAI
+
+    [PortkeyAI](https://portkey.ai/docs/product/ai-gateway/universal-api)
+    provides an interface (AI Gateway) to connect through its Universal API to a
+    variety of LLMs providers with a single endpoint.
+
+    Prerequisites
+    -------------
+
+    ::: {.callout-note}
+    ## Portkey credentials
+
+    Follow the instructions at <https://portkey.ai/docs/introduction/make-your-first-request>
+    to get started making requests to PortkeyAI. You will need to set the
+    `PORTKEY_API_KEY` environment variable to your Portkey API key, and optionally
+    the `PORTKEY_VIRTUAL_KEY` environment variable to your virtual key.
+    :::
+
+    Examples
+    --------
+    ```python
+    import os
+    from chatlas import ChatPortkey
+
+    chat = ChatPortkey(api_key=os.getenv("PORTKEY_API_KEY"))
+    chat.chat("What is the capital of France?")
+    ```
+
+    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.
+    api_key
+        The API key to use for authentication. You generally should not supply
+        this directly, but instead set the `PORTKEY_API_KEY` environment variable.
+    virtual_key
+        An (optional) virtual identifier, storing the LLM provider's API key. See
+        [documentation](https://portkey.ai/docs/product/ai-gateway/virtual-keys).
+        You generally should not supply this directly, but instead set the
+        `PORTKEY_VIRTUAL_KEY` environment variable.
+    base_url
+        The base URL for the Portkey API. The default is suitable for most users.
+    kwargs
+        Additional arguments to pass to the OpenAIProvider, such as headers or
+        other client configuration options.
+
+    Returns
+    -------
+    Chat
+        A chat object that retains the state of the conversation.
+
+    Notes
+    -----
+    This function is a lightweight wrapper around [](`~chatlas.ChatOpenAI`) with
+    the defaults tweaked for PortkeyAI.
+
+    """
+    if model is None:
+        model = log_model_default("gpt-4.1")
+    if api_key is None:
+        api_key = os.getenv("PORTKEY_API_KEY")
+
+    kwargs2 = add_default_headers(
+        kwargs or {},
+        api_key=api_key,
+        virtual_key=virtual_key,
+    )
+
+    return Chat(
+        provider=OpenAIProvider(
+            api_key=api_key,
+            model=model,
+            base_url=base_url,
+            name="Portkey",
+            kwargs=kwargs2,
+        ),
+        system_prompt=system_prompt,
+    )
+
+
+def add_default_headers(
+    kwargs: "ChatClientArgs",
+    api_key: Optional[str] = None,
+    virtual_key: Optional[str] = None,
+) -> "ChatClientArgs":
+    headers = kwargs.get("default_headers", None)
+    default_headers = drop_none(
+        {
+            "x-portkey-api-key": api_key,
+            "x-portkey-virtual-key": virtual_key,
+            **(headers or {}),
+        }
+    )
+    return {"default_headers": default_headers, **kwargs}
+
+
+class PortkeyProvider(OpenAIProvider):
+    def list_models(self):
+        raise NotImplementedError(
+            ".list_models() is not yet implemented for Portkey. "
+            "To view model availability online, see https://portkey.ai/docs/product/model-catalog"
+        )

chatlas/_provider_snowflake.py

@@ -194,6 +194,12 @@ class SnowflakeProvider(
         session = Session.builder.configs(configs).create()
         self._cortex_service = Root(session).cortex_inference_service
 
+    def list_models(self):
+        raise NotImplementedError(
+            ".list_models() is not yet implemented for Snowflake. "
+            "To view model availability online, see https://docs.snowflake.com/user-guide/snowflake-cortex/aisql#availability"
+        )
+
     @overload
     def chat_perform(
         self,

chatlas/_tokens.py

@@ -8,7 +8,7 @@ from typing import TYPE_CHECKING
 import orjson
 
 from ._logging import logger
-from ._typing_extensions import TypedDict
+from ._typing_extensions import NotRequired, TypedDict
 
 if TYPE_CHECKING:
     from ._provider import Provider
@@ -109,11 +109,11 @@ class TokenPrice(TypedDict):
     """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
+    cached_input: NotRequired[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
+    output: NotRequired[float]
     """The cost per assistant token in USD per million tokens"""
 
 
@@ -160,8 +160,8 @@ def compute_cost(
     if price is None:
         return None
     input_price = input_tokens * (price["input"] / 1e6)
-    output_price = output_tokens * (price["output"] / 1e6)
-    cached_price = cached_tokens * (price["cached_input"] / 1e6)
+    output_price = output_tokens * (price.get("output", 0) / 1e6)
+    cached_price = cached_tokens * (price.get("cached_input", 0) / 1e6)
     return input_price + output_price + cached_price
 
 

chatlas/_tools.py

@@ -22,6 +22,7 @@ __all__ = (
 if TYPE_CHECKING:
     from mcp import ClientSession as MCPClientSession
     from mcp import Tool as MCPTool
+    from mcp.types import ToolAnnotations
     from openai.types.chat import ChatCompletionToolParam
 
 
@@ -42,6 +43,9 @@ class Tool:
         A description of what the tool does.
     parameters
         A dictionary describing the input parameters and their types.
+    annotations
+        Additional properties that describe the tool and its behavior. Should be
+        a `from mcp.types import ToolAnnotations` instance.
     """
 
     func: Callable[..., Any] | Callable[..., Awaitable[Any]]
@@ -53,9 +57,11 @@ class Tool:
         name: str,
         description: str,
         parameters: dict[str, Any],
+        annotations: "Optional[ToolAnnotations]" = None,
     ):
         self.name = name
         self.func = func
+        self.annotations = annotations
         self._is_async = _utils.is_async_callable(func)
         self.schema: "ChatCompletionToolParam" = {
             "type": "function",
@@ -72,6 +78,7 @@ class Tool:
         func: Callable[..., Any] | Callable[..., Awaitable[Any]],
         *,
         model: Optional[type[BaseModel]] = None,
+        annotations: "Optional[ToolAnnotations]" = None,
     ) -> "Tool":
         """
         Create a Tool from a Python function
@@ -86,6 +93,9 @@ class Tool:
             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.
+        annotations
+            Additional properties that describe the tool and its behavior. Should be
+            a `from mcp.types import ToolAnnotations` instance.
 
         Returns
         -------
@@ -118,6 +128,7 @@ class Tool:
             name=model.__name__ or func.__name__,
             description=model.__doc__ or func.__doc__ or "",
             parameters=params,
+            annotations=annotations,
         )
 
     @classmethod
@@ -197,6 +208,7 @@ class Tool:
             name=mcp_tool.name,
             description=mcp_tool.description or "",
             parameters=params,
+            annotations=mcp_tool.annotations,
         )
 
 

chatlas/_typing_extensions.py

@@ -14,13 +14,13 @@ else:
 # they should both come from the same typing module.
 # https://peps.python.org/pep-0655/#usage-in-python-3-11
 if sys.version_info >= (3, 11):
-    from typing import Required, TypedDict
+    from typing import NotRequired, Required, TypedDict
 else:
-    from typing_extensions import Required, TypedDict
+    from typing_extensions import NotRequired, Required, TypedDict
 
 
 # The only purpose of the following line is so that pyright will put all of the
 # conditional imports into the .pyi file when generating type stubs. Without this line,
 # pyright will not include the above imports in the generated .pyi file, and it will
 # result in a lot of red squiggles in user code.
-_: "ParamSpec | TypeGuard | is_typeddict | Required | TypedDict"  # type: ignore
+_: "ParamSpec | TypeGuard | is_typeddict | NotRequired | Required | TypedDict"  # type: ignore

chatlas/_version.py

@@ -1,7 +1,14 @@
 # file generated by setuptools-scm
 # don't change, don't track in version control
 
-__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
 
 TYPE_CHECKING = False
 if TYPE_CHECKING:
@@ -9,13 +16,19 @@ if TYPE_CHECKING:
     from typing import Union
 
     VERSION_TUPLE = Tuple[Union[int, str], ...]
+    COMMIT_ID = Union[str, None]
 else:
     VERSION_TUPLE = object
+    COMMIT_ID = object
 
 version: str
 __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID
 
-__version__ = version = '0.9.2'
-__version_tuple__ = version_tuple = (0, 9, 2)
+__version__ = version = '0.11.0'
+__version_tuple__ = version_tuple = (0, 11, 0)
+
+__commit_id__ = commit_id = None