chatlas 0.10.0__py3-none-any.whl → 0.11.1__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,

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

@@ -121,3 +121,11 @@ def add_default_headers(
         }
     )
     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/_tools.py

@@ -2,7 +2,15 @@ from __future__ import annotations
 
 import inspect
 import warnings
-from typing import TYPE_CHECKING, Any, AsyncGenerator, Awaitable, Callable, Optional
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    AsyncGenerator,
+    Awaitable,
+    Callable,
+    Optional,
+    cast,
+)
 
 import openai
 from pydantic import BaseModel, Field, create_model
@@ -12,6 +20,7 @@ from ._content import (
     ContentToolResult,
     ContentToolResultImage,
     ContentToolResultResource,
+    ToolAnnotations,
 )
 
 __all__ = (
@@ -42,6 +51,8 @@ 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.
     """
 
     func: Callable[..., Any] | Callable[..., Awaitable[Any]]
@@ -53,9 +64,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",
@@ -71,7 +84,9 @@ class Tool:
         cls: type["Tool"],
         func: Callable[..., Any] | Callable[..., Awaitable[Any]],
         *,
+        name: Optional[str] = None,
         model: Optional[type[BaseModel]] = None,
+        annotations: "Optional[ToolAnnotations]" = None,
     ) -> "Tool":
         """
         Create a Tool from a Python function
@@ -80,12 +95,17 @@ class Tool:
         ----------
         func
             The function to wrap as a tool.
+        name
+            The name of the tool. If not provided, the name will be inferred from the
+            function's name.
         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.
+        annotations
+            Additional properties that describe the tool and its behavior.
 
         Returns
         -------
@@ -104,7 +124,8 @@ class Tool:
         # 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)
+        fields_alias = [val.alias if val.alias else key for key, val in fields.items()]
+        diff = set(params) ^ set(fields_alias)
         if diff:
             raise ValueError(
                 f"`model` fields must match tool function parameters exactly. "
@@ -115,9 +136,10 @@ class Tool:
 
         return cls(
             func=func,
-            name=model.__name__ or func.__name__,
+            name=name or model.__name__ or func.__name__,
             description=model.__doc__ or func.__doc__ or "",
             parameters=params,
+            annotations=annotations,
         )
 
     @classmethod
@@ -192,11 +214,17 @@ class Tool:
 
         params = mcp_tool_input_schema_to_param_schema(mcp_tool.inputSchema)
 
+        # Convert MCP ToolAnnotations to our TypedDict format
+        annotations = None
+        if mcp_tool.annotations:
+            annotations = cast(ToolAnnotations, mcp_tool.annotations.model_dump())
+
         return cls(
             func=_utils.wrap_async(_call),
             name=mcp_tool.name,
             description=mcp_tool.description or "",
             parameters=params,
+            annotations=annotations,
         )
 
 

chatlas/_typing_extensions.py

@@ -13,7 +13,7 @@ else:
 # Even though TypedDict is available in Python 3.8, because it's used with NotRequired,
 # 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):
+if sys.version_info >= (3, 12):
     from typing import NotRequired, Required, TypedDict
 else:
     from typing_extensions import NotRequired, Required, TypedDict

chatlas/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.10.0'
-__version_tuple__ = version_tuple = (0, 10, 0)
+__version__ = version = '0.11.1'
+__version_tuple__ = version_tuple = (0, 11, 1)
 
 __commit_id__ = commit_id = None