mirascope 1.22.4__py3-none-any.whl → 2.1.1__py3-none-any.whl

mirascope/llm/mcp/mcp_client.py

@@ -0,0 +1,130 @@
+import contextlib
+from collections.abc import AsyncIterator
+from datetime import timedelta
+from typing import cast
+
+from mcp import ClientSession
+from mcp.client.sse import sse_client as mcp_sse_client
+from mcp.client.stdio import StdioServerParameters, stdio_client as mcp_stdio_client
+from mcp.client.streamable_http import (
+    streamable_http_client as mcp_streamable_http_client,
+)
+from mcp.types import CallToolResult, Tool as MCPTool
+
+from ..tools import AsyncTool
+from ..tools.tool_schema import ToolParameterSchema
+from ..types import Jsonable
+
+
+class MCPClient:
+    """Mirascope wrapper around a MCP ClientSession.
+
+    It provides a way to get MCP results that are pre-converted into Mirascope-friendly
+    types.
+
+    The underlying MCP ClientSession may be accessed by .session if needed.
+    """
+
+    def __init__(self, session: ClientSession) -> None:
+        self._session = session
+
+    @property
+    def session(self) -> ClientSession:
+        """Access the underlying MCP ClientSession if needed."""
+        return self._session
+
+    def _convert_mcp_tool_to_async_tool(self, mcp_tool: MCPTool) -> AsyncTool:
+        """Convert an MCP Tool to a Mirascope AsyncTool.
+
+        Args:
+            mcp_tool: The MCP tool to convert.
+
+        Returns:
+            An `AsyncTool` that wraps the MCP tool.
+        """
+
+        # Create an async function that calls the MCP tool
+        async def tool_fn(**kwargs: object) -> Jsonable:
+            tool_result: CallToolResult = await self._session.call_tool(
+                mcp_tool.name, kwargs
+            )
+            # Convert ContentBlock objects to JSON-serializable dicts
+            # Cast to Jsonable since model_dump() returns dict[str, Any]
+            return cast(
+                Jsonable, [content.model_dump() for content in tool_result.content]
+            )
+
+        # Convert MCP tool's inputSchema to Mirascope's ToolParameterSchema
+        input_schema = mcp_tool.inputSchema
+        parameters = ToolParameterSchema(
+            properties=input_schema.get("properties", {}),
+            required=input_schema.get("required", []),
+            additionalProperties=input_schema.get("additionalProperties", False),
+        )
+        if "$defs" in input_schema:
+            parameters.defs = input_schema["$defs"]
+
+        # Create the AsyncTool instance
+        return AsyncTool(
+            fn=tool_fn,
+            name=mcp_tool.name,
+            description=mcp_tool.description or mcp_tool.name,
+            parameters=parameters,
+            strict=False,
+        )
+
+    async def list_tools(self) -> list[AsyncTool]:
+        """List all tools available on the MCP server.
+
+        Returns:
+            A list of dynamically created `AsyncTool`s.
+        """
+        result = await self._session.list_tools()
+        return [self._convert_mcp_tool_to_async_tool(tool) for tool in result.tools]
+
+
+@contextlib.asynccontextmanager
+async def streamable_http_client(
+    url: str,
+) -> AsyncIterator[MCPClient]:  # pragma: no cover
+    """Create a Mirascope MCPClient using StreamableHTTP."""
+    # NOTE: If updating this function, unskip and manually run the TestTransportModes
+    # tests in test_mcp_client.py. (Skipped because they are flaky)
+    async with (
+        mcp_streamable_http_client(url) as (read, write, _),
+        ClientSession(read, write) as session,
+    ):
+        await session.initialize()
+        yield MCPClient(session)
+
+
+@contextlib.asynccontextmanager
+async def stdio_client(
+    server_parameters: StdioServerParameters,
+    name: str | None = None,
+) -> AsyncIterator[MCPClient]:
+    """Create a Mirascope MCPClient using stdio."""
+    async with (
+        mcp_stdio_client(server_parameters) as (read, write),
+        ClientSession(read, write) as session,
+    ):
+        await session.initialize()
+        yield MCPClient(session)
+
+
+@contextlib.asynccontextmanager
+async def sse_client(
+    url: str,
+    read_timeout_seconds: timedelta | None = None,
+) -> AsyncIterator[MCPClient]:  # pragma: no cover
+    """Create a Mirascope MCPClient using sse."""
+    # NOTE: If updating this function, unskip and manually run the TestTransportModes
+    # tests in test_mcp_client.py. (Skipped because they are flaky)
+    async with (
+        mcp_sse_client(url) as (read, write),
+        ClientSession(
+            read, write, read_timeout_seconds=read_timeout_seconds
+        ) as session,
+    ):
+        await session.initialize()
+        yield MCPClient(session)

mirascope/llm/messages/__init__.py

@@ -0,0 +1,35 @@
+"""The messages module for LLM interactions.
+
+This module defines the message types used in LLM interactions. Messages are represented
+as a unified `Message` class with different roles (system, user, assistant) and flexible
+content arrays that can include text, images, audio, documents, and tool interactions.
+"""
+
+from ._utils import is_messages, promote_to_messages
+from .message import (
+    AssistantContent,
+    AssistantMessage,
+    Message,
+    SystemContent,
+    SystemMessage,
+    UserContent,
+    UserMessage,
+    assistant,
+    system,
+    user,
+)
+
+__all__ = [
+    "AssistantContent",
+    "AssistantMessage",
+    "Message",
+    "SystemContent",
+    "SystemMessage",
+    "UserContent",
+    "UserMessage",
+    "assistant",
+    "is_messages",
+    "promote_to_messages",
+    "system",
+    "user",
+]

mirascope/llm/messages/_utils.py

@@ -0,0 +1,34 @@
+"""Utility functions for message handling."""
+
+from collections.abc import Sequence
+from typing_extensions import TypeIs
+
+from .message import (
+    AssistantMessage,
+    Message,
+    SystemMessage,
+    UserContent,
+    UserMessage,
+    user,
+)
+
+
+def is_messages(
+    content: UserContent | Sequence[Message],
+) -> TypeIs[Sequence[Message]]:
+    if isinstance(content, list):
+        if not content:
+            raise ValueError("Empty array may not be used as message content")
+        return isinstance(content[0], SystemMessage | UserMessage | AssistantMessage)
+    return False
+
+
+def promote_to_messages(content: UserContent | Sequence[Message]) -> Sequence[Message]:
+    """Promote a prompt result to a list of messages.
+
+    If the result is already a list of Messages, returns it as-is.
+    If the result is str/UserContentPart/Sequence of content parts, wraps it in a user message.
+    """
+    if is_messages(content):
+        return content
+    return [user(content)]

mirascope/llm/messages/message.py

@@ -0,0 +1,190 @@
+"""The `Message` class and its utility constructors."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Literal, TypeAlias
+
+from ..content import AssistantContentPart, Text, UserContentPart
+from ..types import Jsonable
+
+if TYPE_CHECKING:
+    from ..providers import ModelId, ProviderId
+
+
+@dataclass(kw_only=True)
+class SystemMessage:
+    """A system message that sets context and instructions for the conversation."""
+
+    role: Literal["system"] = "system"
+    """The role of this message. Always "system"."""
+
+    content: Text
+    """The content of this `SystemMesssage`."""
+
+
+@dataclass(kw_only=True)
+class UserMessage:
+    """A user message containing input from the user."""
+
+    role: Literal["user"] = "user"
+    """The role of this message. Always "user"."""
+
+    content: Sequence[UserContentPart]
+    """The content of the user message."""
+
+    name: str | None = None
+    """A name identifying the creator of this message."""
+
+
+@dataclass(kw_only=True)
+class AssistantMessage:
+    """An assistant message containing the model's response."""
+
+    role: Literal["assistant"] = "assistant"
+    """The role of this message. Always "assistant"."""
+
+    content: Sequence[AssistantContentPart]
+    """The content of the assistant message."""
+
+    name: str | None = None
+    """A name identifying the creator of this message."""
+
+    provider_id: ProviderId | None
+    """The LLM provider that generated this assistant message, if available."""
+
+    model_id: ModelId | None
+    """The model identifier of the LLM that generated this assistant message, if available."""
+
+    provider_model_name: str | None
+    """The provider-specific model identifier (e.g. "gpt-5:responses"), if available."""
+
+    raw_message: Jsonable | None
+    """The provider-specific raw representation of this assistant message, if available.
+
+    If raw_content is truthy, then it may be used for provider-specific behavior when
+    resuming an LLM interaction that included this assistant message. For example, we can
+    reuse the provider-specific raw encoding rather than re-encoding the message from it's
+    Mirascope content representation. This may also take advantage of server-side provider
+    context, e.g. identifiers of reasoning context tokens that the provider generated.
+    
+    If present, the content should be encoded as JSON-serializable data, and in a format
+    that matches representation the provider expects representing the Mirascope data.
+    This may involve e.g.  converting Pydantic `BaseModel`s into plain dicts via `model_dump`.
+
+    Raw content is not required, as the Mirascope content can also be used to generate
+    a valid input to the provider (potentially without taking advantage of provider-specific
+    reasoning caches, etc). In that case raw content should be left empty.
+    """
+
+
+Message: TypeAlias = SystemMessage | UserMessage | AssistantMessage
+"""A message in an LLM interaction.
+
+Messages have a role (system, user, or assistant) and content that is a sequence
+of content parts. The content can include text, images, audio, documents, and
+tool interactions.
+
+For most use cases, prefer the convenience functions `system()`, `user()`, and
+`assistant()` instead of directly creating `Message` objects.
+
+Example:
+
+    ```python
+    from mirascope import llm
+
+    messages = [
+        llm.messages.system("You are a helpful assistant."),
+        llm.messages.user("Hello, how are you?"),
+    ]
+    ```
+"""
+
+UserContent: TypeAlias = str | UserContentPart | Sequence[str | UserContentPart]
+"""Type alias for content that can fit into a `UserMessage`."""
+
+AssistantContent: TypeAlias = (
+    str | AssistantContentPart | Sequence[str | AssistantContentPart]
+)
+"""Type alias for content that can fit into an `AssistantMessage`."""
+
+SystemContent: TypeAlias = str | Text
+"""Type alias for content that can fit into a `SystemMessage`."""
+
+
+def system(content: SystemContent) -> SystemMessage:
+    """Creates a system message.
+
+    Args:
+        content: The content of the message, which must be a `str` or `Text` content.
+
+    Returns:
+        A `SystemMessage`.
+    """
+    promoted_content = Text(text=content) if isinstance(content, str) else content
+    return SystemMessage(content=promoted_content)
+
+
+def user(
+    content: UserContent,
+    *,
+    name: str | None = None,
+) -> UserMessage:
+    """Creates a user message.
+
+    Args:
+        content: The content of the message, which can be `str` or any `UserContent`,
+            or a sequence of such user content pieces.
+        name: Optional name to identify a specific user in multi-party conversations.
+
+    Returns:
+        A `UserMessage`.
+    """
+    if isinstance(content, str) or not isinstance(content, Sequence):
+        content = [content]
+    promoted_content = [
+        Text(text=part) if isinstance(part, str) else part for part in content
+    ]
+    return UserMessage(content=promoted_content, name=name)
+
+
+def assistant(
+    content: AssistantContent,
+    *,
+    model_id: ModelId | None,
+    provider_id: ProviderId | None,
+    provider_model_name: str | None = None,
+    raw_message: Jsonable | None = None,
+    name: str | None = None,
+) -> AssistantMessage:
+    """Creates an assistant message.
+
+    Args:
+        content: The content of the message, which can be `str` or any `AssistantContent`,
+            or a sequence of assistant content pieces.
+        model_id: Optional id of the model that produced this message.
+        provider_id: Optional identifier of the provider that produced this message.
+        provider_model_name: Optional provider-specific model name. May include
+            provider-specific additional info (like api mode in "gpt-5:responses").
+        raw_message: Optional Jsonable object that contains the provider-specific
+            "raw" data representation of the content for this assistant message.
+        name: Optional name to identify a specific assistant in multi-party conversations.
+
+    Returns:
+        An `AssistantMessage`.
+    """
+
+    if isinstance(content, str) or not isinstance(content, Sequence):
+        content = [content]
+    promoted_content = [
+        Text(text=part) if isinstance(part, str) else part for part in content
+    ]
+    return AssistantMessage(
+        content=promoted_content,
+        provider_id=provider_id,
+        model_id=model_id,
+        provider_model_name=provider_model_name,
+        raw_message=raw_message,
+        name=name,
+    )

mirascope/llm/models/__init__.py

@@ -0,0 +1,21 @@
+"""The `llm.models` module for implementing the `Model` interface and utilities.
+
+This module provides a unified interface for interacting with different LLM models
+through the `Model` class. The `llm.model()` context manager allows you to override
+the model at runtime, and `llm.use_model()` retrieves the model from context or
+creates a default one.
+"""
+
+from .models import Model, model, model_from_context, use_model
+from .params import Params
+from .thinking_config import ThinkingConfig, ThinkingLevel
+
+__all__ = [
+    "Model",
+    "Params",
+    "ThinkingConfig",
+    "ThinkingLevel",
+    "model",
+    "model_from_context",
+    "use_model",
+]