chatlas 0.8.1__py3-none-any.whl → 0.9.1__py3-none-any.whl

chatlas/_interpolate.py

@@ -23,7 +23,7 @@ def interpolate(
     This is a light-weight wrapper around the Jinja2 templating engine, making
     it easier to interpolate dynamic data into a prompt template. Compared to
     f-strings, which expects you to wrap dynamic values in `{ }`, this function
-    expects `{{ }}` instead, making it easier to include Python code and JSON in
+    expects `{{{ }}}` instead, making it easier to include Python code and JSON in
     your prompt.
 
     Parameters
@@ -80,7 +80,7 @@ def interpolate_file(
     This is a light-weight wrapper around the Jinja2 templating engine, making
     it easier to interpolate dynamic data into a static prompt. Compared to
     f-strings, which expects you to wrap dynamic values in `{ }`, this function
-    expects `{{ }}` instead, making it easier to include Python code and JSON in
+    expects `{{{ }}}` instead, making it easier to include Python code and JSON in
     your prompt.
 
     Parameters
@@ -102,8 +102,7 @@ def interpolate_file(
 
     See Also
     --------
-    interpolate
-        Interpolating data into a system prompt
+    * :func:`~chatlas.interpolate` : Interpolating data into a prompt
     """
     if variables is None:
         frame = inspect.currentframe()

chatlas/_mcp_manager.py

@@ -0,0 +1,306 @@
+from __future__ import annotations
+
+import asyncio
+import warnings
+from abc import ABC, abstractmethod
+from contextlib import AsyncExitStack
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Optional, Sequence
+
+from ._tools import Tool
+
+if TYPE_CHECKING:
+    from mcp import ClientSession
+
+
+@dataclass
+class SessionInfo(ABC):
+    # Input parameters
+    name: str
+    include_tools: Sequence[str] = field(default_factory=list)
+    exclude_tools: Sequence[str] = field(default_factory=list)
+    namespace: str | None = None
+
+    # Primary derived attributes
+    session: ClientSession | None = None
+    tools: dict[str, Tool] = field(default_factory=dict)
+
+    # Background task management
+    ready_event: asyncio.Event = field(default_factory=asyncio.Event)
+    shutdown_event: asyncio.Event = field(default_factory=asyncio.Event)
+    exit_stack: AsyncExitStack = field(default_factory=AsyncExitStack)
+    task: asyncio.Task | None = None
+    error: asyncio.CancelledError | Exception | None = None
+
+    @abstractmethod
+    async def open_session(self) -> None: ...
+
+    async def close_session(self) -> None:
+        await self.exit_stack.aclose()
+
+    async def request_tools(self) -> None:
+        if self.session is None:
+            raise ValueError("Session must be opened before requesting tools.")
+
+        if self.include_tools and self.exclude_tools:
+            raise ValueError("Cannot specify both include_tools and exclude_tools.")
+
+        # Request the MCP tools available
+        response = await self.session.list_tools()
+        tool_names = set(x.name for x in response.tools)
+
+        # Warn if tools are mis-specified
+        include = set(self.include_tools or [])
+        missing_include = include.difference(tool_names)
+        if missing_include:
+            warnings.warn(
+                f"Specified include_tools {missing_include} did not match any tools from the MCP server. "
+                f"The tools available are: {tool_names}",
+                stacklevel=2,
+            )
+        exclude = set(self.exclude_tools or [])
+        missing_exclude = exclude.difference(tool_names)
+        if missing_exclude:
+            warnings.warn(
+                f"Specified exclude_tools {missing_exclude} did not match any tools from the MCP server. "
+                f"The tools available are: {tool_names}",
+                stacklevel=2,
+            )
+
+        # Filter the tool names
+        if include:
+            tool_names = include.intersection(tool_names)
+        if exclude:
+            tool_names = tool_names.difference(exclude)
+
+        # Apply namespace and convert to chatlas.Tool instances
+        self_tools: dict[str, Tool] = {}
+        for tool in response.tools:
+            if tool.name not in tool_names:
+                continue
+            if self.namespace:
+                tool.name = f"{self.namespace}.{tool.name}"
+            self_tools[tool.name] = Tool.from_mcp(
+                session=self.session,
+                mcp_tool=tool,
+            )
+
+        # Store the tools
+        self.tools = self_tools
+
+
+@dataclass
+class HTTPSessionInfo(SessionInfo):
+    url: str = ""
+    transport_kwargs: dict[str, Any] = field(default_factory=dict)
+
+    async def open_session(self):
+        mcp = try_import_mcp()
+        from mcp.client.streamable_http import streamablehttp_client
+
+        read, write, _ = await self.exit_stack.enter_async_context(
+            streamablehttp_client(
+                self.url,
+                **self.transport_kwargs,
+            )
+        )
+        session = await self.exit_stack.enter_async_context(
+            mcp.ClientSession(read, write)
+        )
+        server = await session.initialize()
+        self.session = session
+        if not self.name:
+            self.name = server.serverInfo.name or "mcp"
+
+
+@dataclass
+class STDIOSessionInfo(SessionInfo):
+    command: str = ""
+    args: list[str] = field(default_factory=list)
+    transport_kwargs: dict[str, Any] = field(default_factory=dict)
+
+    async def open_session(self):
+        mcp = try_import_mcp()
+        from mcp.client.stdio import stdio_client
+
+        server_params = mcp.StdioServerParameters(
+            command=self.command,
+            args=self.args,
+            **self.transport_kwargs,
+        )
+
+        transport = await self.exit_stack.enter_async_context(
+            stdio_client(server_params)
+        )
+        session = await self.exit_stack.enter_async_context(
+            mcp.ClientSession(*transport)
+        )
+        server = await session.initialize()
+        self.session = session
+        if not self.name:
+            self.name = server.serverInfo.name or "mcp"
+
+
+class MCPSessionManager:
+    """Manages MCP (Model Context Protocol) server connections and tools."""
+
+    def __init__(self):
+        self._mcp_sessions: dict[str, SessionInfo] = {}
+
+    async def register_http_stream_tools(
+        self,
+        *,
+        url: str,
+        name: str | None,
+        include_tools: Sequence[str],
+        exclude_tools: Sequence[str],
+        namespace: str | None,
+        transport_kwargs: dict[str, Any],
+    ):
+        session_info = HTTPSessionInfo(
+            name=name or "",
+            url=url,
+            include_tools=include_tools,
+            exclude_tools=exclude_tools,
+            namespace=namespace,
+            transport_kwargs=transport_kwargs or {},
+        )
+
+        # Launch background task that runs until MCP session is *shutdown*
+        # N.B. this is needed since mcp sessions must be opened and closed in the same task
+        asyncio.create_task(self.open_session(session_info))
+
+        # Wait for a ready event from the task (signals that tools are registered)
+        await session_info.ready_event.wait()
+
+        # An error might have been caught in the background task
+        if session_info.error:
+            raise RuntimeError(
+                f"Failed to register tools from MCP server '{name}' at URL '{url}'"
+            ) from session_info.error
+
+        return session_info
+
+    async def register_stdio_tools(
+        self,
+        *,
+        command: str,
+        args: list[str],
+        name: str | None,
+        include_tools: Sequence[str],
+        exclude_tools: Sequence[str],
+        namespace: str | None,
+        transport_kwargs: dict[str, Any],
+    ):
+        session_info = STDIOSessionInfo(
+            name=name or "",
+            command=command,
+            args=args,
+            include_tools=include_tools,
+            exclude_tools=exclude_tools,
+            namespace=namespace,
+            transport_kwargs=transport_kwargs or {},
+        )
+
+        # Launch a background task to initialize the MCP server
+        # N.B. this is needed since mcp sessions must be opened and closed in the same task
+        asyncio.create_task(self.open_session(session_info))
+
+        # Wait for a ready event from the task (signals that tools are registered)
+        await session_info.ready_event.wait()
+
+        # An error might have been caught in the background task
+        if session_info.error:
+            raise RuntimeError(
+                f"Failed to register tools from MCP server '{name}' with command '{command} {args}'"
+            ) from session_info.error
+
+        return session_info
+
+    async def open_session(self, session_info: "SessionInfo"):
+        session_info.task = asyncio.current_task()
+
+        try:
+            # Open the MCP session
+            await session_info.open_session()
+            # Request the tools
+            await session_info.request_tools()
+            # Make sure session can be added to the manager
+            self.add_session(session_info)
+        except (asyncio.CancelledError, Exception) as err:
+            # Keep the error so we can handle in the main task
+            session_info.error = err
+            # Make sure the session is closed
+            try:
+                await session_info.close_session()
+            except Exception:
+                pass
+            return
+        finally:
+            # Whether successful or not, set ready state to prevent deadlock
+            session_info.ready_event.set()
+
+        # If successful, wait for shutdown signal
+        await session_info.shutdown_event.wait()
+
+        # On shutdown close connection to MCP server
+        # This is why we're using a background task in the 1st place...
+        # we must close in the same task that opened the session
+        await session_info.close_session()
+
+    async def close_sessions(self, names: Optional[Sequence[str]] = None):
+        if names is None:
+            names = list(self._mcp_sessions.keys())
+
+        if isinstance(names, str):
+            names = [names]
+
+        closed_sessions: list[SessionInfo] = []
+        for x in names:
+            session = await self.close_background_session(x)
+            if session is None:
+                continue
+            closed_sessions.append(session)
+
+        return closed_sessions
+
+    async def close_background_session(self, name: str) -> SessionInfo | None:
+        session = self.remove_session(name)
+        if session is None:
+            return None
+
+        # Signal shutdown and wait for the task to finish
+        session.shutdown_event.set()
+        if session.task is not None:
+            await session.task
+
+        return session
+
+    def add_session(self, session_info: SessionInfo) -> None:
+        name = session_info.name
+        if name in self._mcp_sessions:
+            raise ValueError(f"Already connected to an MCP server named: '{name}'.")
+        self._mcp_sessions[name] = session_info
+
+    def remove_session(self, name: str) -> SessionInfo | None:
+        if name not in self._mcp_sessions:
+            warnings.warn(
+                f"Cannot close MCP session named '{name}' since it was not found.",
+                stacklevel=2,
+            )
+            return None
+        session = self._mcp_sessions[name]
+        del self._mcp_sessions[name]
+        return session
+
+
+def try_import_mcp():
+    try:
+        import mcp
+
+        return mcp
+    except ImportError:
+        raise ImportError(
+            "The `mcp` package is required to connect to MCP servers. "
+            "Install it with `pip install mcp`."
+        )

chatlas/_ollama.py

@@ -7,8 +7,8 @@ from typing import TYPE_CHECKING, Optional
 import orjson
 
 from ._chat import Chat
-from ._openai import ChatOpenAI
-from ._turn import Turn
+from ._openai import OpenAIProvider
+from ._utils import MISSING_TYPE, is_testing
 
 if TYPE_CHECKING:
     from ._openai import ChatCompletion
@@ -19,7 +19,6 @@ def ChatOllama(
     model: Optional[str] = None,
     *,
     system_prompt: Optional[str] = None,
-    turns: Optional[list[Turn]] = None,
     base_url: str = "http://localhost:11434",
     seed: Optional[int] = None,
     kwargs: Optional["ChatClientArgs"] = None,
@@ -67,13 +66,6 @@ def ChatOllama(
         models will be printed.
     system_prompt
         A system prompt to set the behavior of the assistant.
-    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.
     base_url
         The base URL to the endpoint; the default uses ollama's API.
     seed
@@ -102,15 +94,19 @@ def ChatOllama(
         raise ValueError(
             f"Must specify model. Locally installed models: {', '.join(models)}"
         )
-
-    return ChatOpenAI(
+    if isinstance(seed, MISSING_TYPE):
+        seed = 1014 if is_testing() else None
+
+    return Chat(
+        provider=OpenAIProvider(
+            api_key="ollama",  # ignored
+            model=model,
+            base_url=f"{base_url}/v1",
+            seed=seed,
+            name="Ollama",
+            kwargs=kwargs,
+        ),
         system_prompt=system_prompt,
-        api_key="ollama",  # ignored
-        turns=turns,
-        base_url=f"{base_url}/v1",
-        model=model,
-        seed=seed,
-        kwargs=kwargs,
     )
 
 

chatlas/_openai.py

@@ -4,6 +4,7 @@ import base64
 from typing import TYPE_CHECKING, Any, Literal, Optional, cast, overload
 
 import orjson
+from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI
 from pydantic import BaseModel
 
 from ._chat import Chat
@@ -17,14 +18,16 @@ from ._content import (
     ContentText,
     ContentToolRequest,
     ContentToolResult,
+    ContentToolResultImage,
+    ContentToolResultResource,
 )
 from ._logging import log_model_default
 from ._merge import merge_dicts
-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, user_turn
-from ._utils import MISSING, MISSING_TYPE, is_testing
+from ._turn import Turn, user_turn
+from ._utils import MISSING, MISSING_TYPE, is_testing, split_http_client_kwargs
 
 if TYPE_CHECKING:
     from openai.types.chat import (
@@ -53,7 +56,6 @@ ChatCompletionDict = dict[str, Any]
 def ChatOpenAI(
     *,
     system_prompt: Optional[str] = None,
-    turns: Optional[list[Turn]] = None,
     model: "Optional[ChatModel | str]" = None,
     api_key: Optional[str] = None,
     base_url: str = "https://api.openai.com/v1",
@@ -92,13 +94,6 @@ def ChatOpenAI(
     ----------
     system_prompt
         A system prompt to set the behavior of the assistant.
-    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.
     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
@@ -161,7 +156,7 @@ def ChatOpenAI(
         seed = 1014 if is_testing() else None
 
     if model is None:
-        model = log_model_default("gpt-4o")
+        model = log_model_default("gpt-4.1")
 
     return Chat(
         provider=OpenAIProvider(
@@ -171,14 +166,13 @@ def ChatOpenAI(
             seed=seed,
             kwargs=kwargs,
         ),
-        turns=normalize_turns(
-            turns or [],
-            system_prompt,
-        ),
+        system_prompt=system_prompt,
     )
 
 
-class OpenAIProvider(Provider[ChatCompletion, ChatCompletionChunk, ChatCompletionDict]):
+class OpenAIProvider(
+    Provider[ChatCompletion, ChatCompletionChunk, ChatCompletionDict, "SubmitInputArgs"]
+):
     def __init__(
         self,
         *,
@@ -186,11 +180,11 @@ class OpenAIProvider(Provider[ChatCompletion, ChatCompletionChunk, ChatCompletio
         model: str,
         base_url: str = "https://api.openai.com/v1",
         seed: Optional[int] = None,
+        name: str = "OpenAI",
         kwargs: Optional["ChatClientArgs"] = None,
     ):
-        from openai import AsyncOpenAI, OpenAI
+        super().__init__(name=name, model=model)
 
-        self._model = model
         self._seed = seed
 
         kwargs_full: "ChatClientArgs" = {
@@ -199,9 +193,12 @@ class OpenAIProvider(Provider[ChatCompletion, ChatCompletionChunk, ChatCompletio
             **(kwargs or {}),
         }
 
+        # Avoid passing the wrong sync/async client to the OpenAI constructor.
+        sync_kwargs, async_kwargs = split_http_client_kwargs(kwargs_full)
+
         # TODO: worth bringing in AsyncOpenAI types?
-        self._client = OpenAI(**kwargs_full)  # type: ignore
-        self._async_client = AsyncOpenAI(**kwargs_full)
+        self._client = OpenAI(**sync_kwargs)  # type: ignore
+        self._async_client = AsyncOpenAI(**async_kwargs)
 
     @overload
     def chat_perform(
@@ -284,7 +281,7 @@ class OpenAIProvider(Provider[ChatCompletion, ChatCompletionChunk, ChatCompletio
         kwargs_full: "SubmitInputArgs" = {
             "stream": stream,
             "messages": self._as_message_param(turns),
-            "model": self._model,
+            "model": self.model,
             **(kwargs or {}),
         }
 
@@ -487,6 +484,12 @@ class OpenAIProvider(Provider[ChatCompletion, ChatCompletionChunk, ChatCompletio
                             }
                         )
                     elif isinstance(x, ContentToolResult):
+                        if isinstance(
+                            x, (ContentToolResultImage, ContentToolResultResource)
+                        ):
+                            raise NotImplementedError(
+                                "OpenAI does not support tool results with images or resources."
+                            )
                         tool_results.append(
                             ChatCompletionToolMessageParam(
                                 # Currently, OpenAI only allows for text content in tool results
@@ -573,6 +576,46 @@ class OpenAIProvider(Provider[ChatCompletion, ChatCompletionChunk, ChatCompletio
             completion=completion,
         )
 
+    def translate_model_params(self, params: StandardModelParams) -> "SubmitInputArgs":
+        res: "SubmitInputArgs" = {}
+        if "temperature" in params:
+            res["temperature"] = params["temperature"]
+
+        if "top_p" in params:
+            res["top_p"] = params["top_p"]
+
+        if "frequency_penalty" in params:
+            res["frequency_penalty"] = params["frequency_penalty"]
+
+        if "presence_penalty" in params:
+            res["presence_penalty"] = params["presence_penalty"]
+
+        if "seed" in params:
+            res["seed"] = params["seed"]
+
+        if "max_tokens" in params:
+            res["max_tokens"] = params["max_tokens"]
+
+        if "log_probs" in params:
+            res["logprobs"] = params["log_probs"]
+
+        if "stop_sequences" in params:
+            res["stop"] = params["stop_sequences"]
+
+        return res
+
+    def supported_model_params(self) -> set[StandardModelParamNames]:
+        return {
+            "temperature",
+            "top_p",
+            "frequency_penalty",
+            "presence_penalty",
+            "seed",
+            "max_tokens",
+            "log_probs",
+            "stop_sequences",
+        }
+
 
 def ChatAzureOpenAI(
     *,
@@ -581,7 +624,6 @@ def ChatAzureOpenAI(
     api_version: str,
     api_key: Optional[str] = None,
     system_prompt: Optional[str] = None,
-    turns: Optional[list[Turn]] = None,
     seed: int | None | MISSING_TYPE = MISSING,
     kwargs: Optional["ChatAzureClientArgs"] = None,
 ) -> Chat["SubmitInputArgs", ChatCompletion]:
@@ -624,13 +666,6 @@ def ChatAzureOpenAI(
         variable.
     system_prompt
         A system prompt to set the behavior of the assistant.
-    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.
     seed
         Optional integer seed that ChatGPT uses to try and make output more
         reproducible.
@@ -655,10 +690,7 @@ def ChatAzureOpenAI(
             seed=seed,
             kwargs=kwargs,
         ),
-        turns=normalize_turns(
-            turns or [],
-            system_prompt,
-        ),
+        system_prompt=system_prompt,
     )
 
 
@@ -667,15 +699,16 @@ class OpenAIAzureProvider(OpenAIProvider):
         self,
         *,
         endpoint: Optional[str] = None,
-        deployment_id: Optional[str] = None,
+        deployment_id: str,
         api_version: Optional[str] = None,
         api_key: Optional[str] = None,
         seed: int | None = None,
+        name: str = "OpenAIAzure",
+        model: Optional[str] = "UnusedValue",
         kwargs: Optional["ChatAzureClientArgs"] = None,
     ):
-        from openai import AsyncAzureOpenAI, AzureOpenAI
+        super().__init__(name=name, model=deployment_id)
 
-        self._model = deployment_id
         self._seed = seed
 
         kwargs_full: "ChatAzureClientArgs" = {
@@ -686,8 +719,10 @@ class OpenAIAzureProvider(OpenAIProvider):
             **(kwargs or {}),
         }
 
-        self._client = AzureOpenAI(**kwargs_full)  # type: ignore
-        self._async_client = AsyncAzureOpenAI(**kwargs_full)  # type: ignore
+        sync_kwargs, async_kwargs = split_http_client_kwargs(kwargs_full)
+
+        self._client = AzureOpenAI(**sync_kwargs)  # type: ignore
+        self._async_client = AsyncAzureOpenAI(**async_kwargs)  # type: ignore
 
 
 class InvalidJSONParameterWarning(RuntimeWarning):

chatlas/_perplexity.py

@@ -5,9 +5,8 @@ from typing import TYPE_CHECKING, Optional
 
 from ._chat import Chat
 from ._logging import log_model_default
-from ._openai import ChatOpenAI
-from ._turn import Turn
-from ._utils import MISSING, MISSING_TYPE
+from ._openai import OpenAIProvider
+from ._utils import MISSING, MISSING_TYPE, is_testing
 
 if TYPE_CHECKING:
     from ._openai import ChatCompletion
@@ -17,7 +16,6 @@ if TYPE_CHECKING:
 def ChatPerplexity(
     *,
     system_prompt: Optional[str] = None,
-    turns: Optional[list[Turn]] = None,
     model: Optional[str] = None,
     api_key: Optional[str] = None,
     base_url: str = "https://api.perplexity.ai/",
@@ -56,13 +54,6 @@ def ChatPerplexity(
     ----------
     system_prompt
         A system prompt to set the behavior of the assistant.
-    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.
     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
@@ -131,12 +122,17 @@ def ChatPerplexity(
     if api_key is None:
         api_key = os.getenv("PERPLEXITY_API_KEY")
 
-    return ChatOpenAI(
+    if isinstance(seed, MISSING_TYPE):
+        seed = 1014 if is_testing() else None
+
+    return Chat(
+        provider=OpenAIProvider(
+            api_key=api_key,
+            model=model,
+            base_url=base_url,
+            seed=seed,
+            name="Perplexity",
+            kwargs=kwargs,
+        ),
         system_prompt=system_prompt,
-        turns=turns,
-        model=model,
-        api_key=api_key,
-        base_url=base_url,
-        seed=seed,
-        kwargs=kwargs,
     )