pi-python-core 0.8.1__py3-none-any.whl

pi_python/mcp.py

@@ -0,0 +1,187 @@
+"""Use the tools of an MCP server as Agent tools.
+
+`mcp_tools(session)` wraps every tool of a connected session from the official `mcp` SDK
+(1.x or 2.x); `connect_stdio(...)` also starts a stdio server and closes it afterwards.
+The conversion follows Pi's MCP adapter: text and images pass through, embedded text and
+image resources are unwrapped, other blocks become short text placeholders, and MCP's
+`isError` marks a failed call. Install with ``pip install 'pi-python-core[mcp]'``.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import inspect
+import json
+import re
+import warnings
+from collections.abc import AsyncIterator, Iterable, Mapping, Sequence
+from contextlib import asynccontextmanager
+from typing import Any
+
+from .errors import ConfigurationError
+from .messages import ImageContent, TextContent
+from .tools import Tool, ToolContext, ToolResult
+
+
+def _get(value: Any, snake: str, camel: str | None = None) -> Any:
+    """Read a field in either SDK's naming (2.x snake_case, 1.x camelCase) or from a dict."""
+    for name in (snake, camel) if camel else (snake,):
+        if isinstance(value, Mapping):
+            if name in value:
+                return value[name]
+        elif hasattr(value, name):
+            return getattr(value, name)
+    return None
+
+
+def _block(block: Any) -> TextContent | ImageContent:
+    kind = _get(block, "type")
+    if kind == "text":
+        return TextContent(_get(block, "text") or "")
+    if kind == "image":
+        return ImageContent(_get(block, "data"), _get(block, "mime_type", "mimeType"))
+    if kind == "audio":
+        return TextContent(f"[audio {_get(block, 'mime_type', 'mimeType')} omitted]")
+    if kind == "resource_link":
+        return TextContent(f"{_get(block, 'name')}: {_get(block, 'uri')}")
+    if kind == "resource":
+        resource = _get(block, "resource")
+        mime = _get(resource, "mime_type", "mimeType")
+        if _get(resource, "text") is not None:
+            return TextContent(_get(resource, "text"))
+        if mime and str(mime).startswith("image/"):
+            return ImageContent(_get(resource, "blob"), mime)
+        return TextContent(
+            f"[binary resource {_get(resource, 'uri')} ({mime or 'unknown type'}) omitted]"
+        )
+    return TextContent(f"[unsupported MCP content {kind}]")
+
+
+def _tool_name(prefix: str | None, name: str, taken: set[str]) -> str:
+    # Model APIs accept letters, digits, "_" and "-", up to 64 characters.
+    full = f"{prefix}_{name}" if prefix else name
+    result = re.sub(r"[^A-Za-z0-9_-]", "_", full)[:64]
+    if result in taken:  # "get.file" and "get_file", or two long names with one prefix
+        result = f"{result[:55]}_{hashlib.sha1(full.encode()).hexdigest()[:8]}"
+        warnings.warn(f"MCP tool {name!r} renamed to {result!r} to keep names unique", stacklevel=3)
+    taken.add(result)
+    return result
+
+
+def _wrap(session: Any, spec: Any, prefix: str | None, taken: set[str]) -> Tool:
+    name = _get(spec, "name")
+    schema = dict(_get(spec, "input_schema", "inputSchema") or {})
+    # Model APIs require an object schema, and some reject one without `properties`.
+    schema.update(type="object", properties=schema.get("properties") or {})
+
+    async def execute(args: dict[str, Any], context: ToolContext) -> ToolResult:
+        async def progress(
+            done: float, total: float | None = None, message: str | None = None
+        ) -> None:
+            await context.emit_update({"progress": done, "total": total, "message": message})
+
+        result = await session.call_tool(name, args, progress_callback=progress)
+        blocks = _get(result, "content")
+        if blocks is None:
+            # MCP 2.x can ask the client for input mid-call; this adapter cannot answer.
+            return ToolResult.text(
+                f"MCP tool {name} asked for input this client cannot provide", is_error=True
+            )
+        content = [_block(b) for b in blocks]
+        structured = _get(result, "structured_content", "structuredContent")
+        is_error = _get(result, "is_error", "isError") is True
+        if not content and structured is not None:
+            content = [TextContent(json.dumps(structured, indent=2, ensure_ascii=False))]
+        if is_error and not content:
+            content = [TextContent(f"MCP tool {name} failed")]
+        return ToolResult(content, structured_content=structured, is_error=is_error)
+
+    description = _get(spec, "description") or _get(spec, "title") or f"MCP tool {name}"
+    tool = Tool(_tool_name(prefix, name, set(taken)), description, schema, execute)
+    taken.add(tool.name)
+    return tool
+
+
+async def mcp_tools(
+    session: Any, *, prefix: str | None = None, names: Iterable[str] | None = None
+) -> list[Tool]:
+    """Wrap the tools of an initialized MCP client session.
+
+    `prefix` namespaces the tool names (``prefix_tool``); `names` keeps only those MCP
+    tools. A tool whose input schema cannot be used is skipped with a warning, so one
+    unusual tool does not make the whole server unusable.
+    """
+    wanted = set(names) if names is not None else None
+    specs: list[Any] = []
+    cursor = None
+    while True:
+        page = await _list_page(session, cursor)
+        specs.extend(_get(page, "tools") or [])
+        cursor = _get(page, "next_cursor", "nextCursor")
+        if not cursor:
+            break
+    tools: list[Tool] = []
+    taken: set[str] = set()
+    for spec in specs:
+        if wanted is not None and _get(spec, "name") not in wanted:
+            continue
+        try:
+            tools.append(_wrap(session, spec, prefix, taken))
+        except (ConfigurationError, RecursionError) as exc:
+            warnings.warn(f"Skipping MCP tool {_get(spec, 'name')}: {exc}", stacklevel=2)
+    return tools
+
+
+async def _list_page(session: Any, cursor: Any) -> Any:
+    if cursor is None:
+        return await session.list_tools()
+    if "params" in inspect.signature(session.list_tools).parameters:  # MCP SDK 2.x
+        from mcp.types import PaginatedRequestParams  # type: ignore[import-not-found,unused-ignore]
+
+        return await session.list_tools(params=PaginatedRequestParams(cursor=cursor))
+    return await session.list_tools(cursor=cursor)
+
+
+@asynccontextmanager
+async def connect_stdio(
+    command: str,
+    args: Sequence[str] = (),
+    *,
+    env: dict[str, str] | None = None,
+    cwd: str | None = None,
+    prefix: str | None = None,
+    names: Iterable[str] | None = None,
+    init_timeout: float = 30.0,
+) -> AsyncIterator[list[Tool]]:
+    """Start a stdio MCP server, yield its tools, and stop it when the block ends.
+
+    ``async with connect_stdio("uvx", ["mcp-server-fetch"]) as tools: ...``
+
+    A server that does not complete the MCP handshake within `init_timeout` seconds
+    raises TimeoutError instead of hanging.
+    """
+    try:
+        from mcp import ClientSession, StdioServerParameters  # type: ignore[import-not-found,unused-ignore]
+        from mcp.client.stdio import stdio_client  # type: ignore[import-not-found,unused-ignore]
+    except ImportError as exc:
+        raise ImportError(
+            "connect_stdio needs the MCP SDK: pip install 'pi-python-core[mcp]'"
+        ) from exc
+    params = StdioServerParameters(command=command, args=list(args), env=env, cwd=cwd)
+    try:
+        async with stdio_client(params) as (read, write):
+            async with ClientSession(read, write) as session:
+                async with asyncio.timeout(init_timeout):
+                    await session.initialize()
+                    tools = await mcp_tools(session, prefix=prefix, names=names)
+                yield tools
+    except BaseExceptionGroup as group:
+        # The SDK's task groups wrap a single failure ("Connection closed", a timeout,
+        # or an error from the caller's block), sometimes twice; raise that failure itself.
+        error: BaseException = group
+        while isinstance(error, BaseExceptionGroup) and len(error.exceptions) == 1:
+            error = error.exceptions[0]
+        if error is group:
+            raise
+        raise error from None

pi_python/messages.py

@@ -0,0 +1,405 @@
+"""Versioned, JSON-only messages. All boundaries copy caller-owned values."""
+
+from __future__ import annotations
+
+import json
+import math
+import time
+from dataclasses import asdict, dataclass, field
+from typing import Any, Literal, TypeAlias
+
+from .errors import MessageValidationError, UnsupportedCapabilityError
+import base64
+
+JsonValue: TypeAlias = "None | bool | int | float | str | list[JsonValue] | dict[str, JsonValue]"
+SCHEMA_VERSION = 3
+
+
+def now() -> float:
+    return time.time()
+
+
+def validate_json(value: Any) -> None:
+    if value is None or type(value) in (str, bool, int):
+        return
+    if type(value) is float and math.isfinite(value):
+        return
+    if type(value) is list:
+        for item in value:
+            validate_json(item)
+        return
+    if type(value) is dict and all(type(k) is str for k in value):
+        for item in value.values():
+            validate_json(item)
+        return
+    raise MessageValidationError(f"Not a finite JSON value: {type(value).__name__}")
+
+
+@dataclass
+class TextContent:
+    text: str
+    text_signature: str | None = None
+    type: Literal["text"] = field(default="text", init=False)
+
+
+@dataclass
+class ImageContent:
+    data: str
+    mime_type: str
+    type: Literal["image"] = field(default="image", init=False)
+
+
+@dataclass
+class ThinkingContent:
+    thinking: str
+    thinking_signature: str | None = None
+    redacted: bool = False
+    type: Literal["thinking"] = field(default="thinking", init=False)
+
+
+@dataclass
+class ToolCall:
+    id: str
+    name: str
+    arguments: dict[str, Any]
+    thought_signature: str | None = None
+    namespace: str | None = None
+    type: Literal["tool_call"] = field(default="tool_call", init=False)
+
+
+@dataclass
+class ToolDeclaration:
+    name: str
+    description: str
+    input_schema: dict[str, Any]
+
+
+@dataclass
+class SystemMessage:
+    content: str | list[TextContent] = ""
+    sections: dict[str, str | None] = field(default_factory=dict)
+    tools_added: list[ToolDeclaration] = field(default_factory=list)
+    tools_removed: list[str] = field(default_factory=list)
+    timestamp: float = field(default_factory=now)
+    role: Literal["system"] = field(default="system", init=False)
+
+
+@dataclass
+class UserMessage:
+    content: str | list[TextContent | ImageContent]
+    timestamp: float = field(default_factory=now)
+    role: Literal["user"] = field(default="user", init=False)
+
+
+@dataclass
+class AssistantMessage:
+    content: list[TextContent | ThinkingContent | ToolCall]
+    stop_reason: str = "stop"
+    provider: str = "mock"
+    model: str = "mock"
+    usage: dict[str, Any] = field(default_factory=dict)
+    timestamp: float = field(default_factory=now)
+    error: str | None = None
+    api: str | None = None
+    provider_thinking_level: str | None = None
+    response_model: str | None = None
+    response_id: str | None = None
+    thinking_level: str | None = None
+    diagnostics: list[dict[str, Any]] | None = None
+    raw_stop_reason: str | None = None
+    end_turn: bool | None = None
+    deferred: dict[str, Any] | None = None
+    role: Literal["assistant"] = field(default="assistant", init=False)
+
+    @classmethod
+    def text(cls, text: str, **kwargs: Any) -> AssistantMessage:
+        return cls([TextContent(text)], **kwargs)
+
+    @property
+    def tool_calls(self) -> list[ToolCall]:
+        return [b for b in self.content if isinstance(b, ToolCall)]
+
+
+@dataclass
+class ToolResultMessage:
+    call_id: str
+    name: str
+    content: list[TextContent | ImageContent]
+    is_error: bool = False
+    timestamp: float = field(default_factory=now)
+    details: Any = None
+    usage: dict[str, Any] | None = None
+    nested_calls: dict[str, Any] | None = None
+    role: Literal["tool_result"] = field(default="tool_result", init=False)
+
+
+@dataclass
+class CustomMessage:
+    custom_type: str
+    data: Any
+    timestamp: float = field(default_factory=now)
+    role: Literal["custom"] = field(default="custom", init=False)
+
+
+Message: TypeAlias = (
+    SystemMessage | UserMessage | AssistantMessage | ToolResultMessage | CustomMessage
+)
+
+
+def message_to_dict(message: Message) -> dict[str, Any]:
+    if not isinstance(
+        message, (SystemMessage, UserMessage, AssistantMessage, ToolResultMessage, CustomMessage)
+    ):
+        raise MessageValidationError("Expected a supported Message")
+    optional_fields = {
+        "text_signature",
+        "thought_signature",
+        "namespace",
+        "api",
+        "provider_thinking_level",
+        "response_model",
+        "response_id",
+        "thinking_level",
+        "diagnostics",
+        "raw_stop_reason",
+        "end_turn",
+        "deferred",
+        "details",
+        "usage",
+        "nested_calls",
+    }
+    # dict_factory visits dataclass fields, not arbitrary user dictionaries.
+    data = asdict(
+        message,
+        dict_factory=lambda pairs: {
+            k: v for k, v in pairs if not (k in optional_fields and v is None)
+        },
+    )
+    validate_json(data)
+    # Parsing validates semantic fields as well as JSON representability.
+    message_from_dict(data)
+    return data
+
+
+def _text(value: Any, label: str) -> str:
+    if not isinstance(value, str):
+        raise MessageValidationError(f"{label} must be a string")
+    return value
+
+
+def _blocks(values: Any, tools: bool = False, images: bool = False) -> list:
+    if not isinstance(values, list):
+        raise MessageValidationError("content must be a list")
+    result: list = []
+    for b in values:
+        if not isinstance(b, dict):
+            raise MessageValidationError("content block must be an object")
+        if b.get("type") == "text":
+            if set(b) - {"type", "text", "text_signature"} or "text" not in b:
+                raise MessageValidationError("Invalid text block fields")
+            signature = b.get("text_signature")
+            if signature is not None:
+                _text(signature, "text_signature")
+            result.append(TextContent(_text(b["text"], "text"), signature))
+        elif b.get("type") == "image" and images:
+            if set(b) != {"type", "data", "mime_type"}:
+                raise MessageValidationError("Invalid image block")
+            data, mime = _text(b["data"], "image data"), _text(b["mime_type"], "mime_type")
+            if mime not in {"image/png", "image/jpeg", "image/gif", "image/webp"}:
+                raise UnsupportedCapabilityError("Unsupported image MIME type")
+            try:
+                base64.b64decode(data, validate=True)
+            except ValueError as exc:
+                raise MessageValidationError("Invalid image base64") from exc
+            result.append(ImageContent(data, mime))
+        elif b.get("type") == "thinking" and tools:
+            if set(b) - {"type", "thinking", "thinking_signature", "redacted"}:
+                raise MessageValidationError("Invalid thinking block")
+            signature = b.get("thinking_signature")
+            if signature is not None:
+                _text(signature, "thinking_signature")
+            if type(b.get("redacted", False)) is not bool:
+                raise MessageValidationError("redacted must be boolean")
+            result.append(
+                ThinkingContent(
+                    _text(b["thinking"], "thinking"), signature, b.get("redacted", False)
+                )
+            )
+        elif b.get("type") == "tool_call" and tools:
+            if (
+                set(b) - {"type", "id", "name", "arguments", "thought_signature", "namespace"}
+                or not {"id", "name", "arguments"} <= b.keys()
+            ):
+                raise MessageValidationError("Invalid tool call fields")
+            if not b["id"] or not b["name"] or not isinstance(b["arguments"], dict):
+                raise MessageValidationError("Invalid tool call ID, name or arguments")
+            for key in ("thought_signature", "namespace"):
+                if b.get(key) is not None:
+                    _text(b[key], key)
+            result.append(
+                ToolCall(
+                    _text(b["id"], "id"),
+                    _text(b["name"], "name"),
+                    b["arguments"],
+                    b.get("thought_signature"),
+                    b.get("namespace"),
+                )
+            )
+        else:
+            raise UnsupportedCapabilityError(f"Unsupported content: {b.get('type')}")
+    return result
+
+
+def message_from_dict(data: dict[str, Any]) -> Message:
+    validate_json(data)
+    if not isinstance(data, dict):
+        raise MessageValidationError("Message must be an object")
+    d = json.loads(json.dumps(data, allow_nan=False))
+    role = d.pop("role", None)
+    try:
+        if "timestamp" in d and (type(d["timestamp"]) not in (float, int)):
+            raise MessageValidationError("timestamp must be a number")
+        if role == "user":
+            if not isinstance(d["content"], str):
+                d["content"] = _blocks(d["content"], images=True)
+            return UserMessage(**d)
+        if role == "system":
+            if isinstance(d.get("content", ""), str):
+                _text(d.get("content", ""), "content")
+            else:
+                d["content"] = _blocks(d["content"])
+            sections = d.get("sections", {})
+            if not isinstance(sections, dict) or any(
+                v is not None and not isinstance(v, str) for v in sections.values()
+            ):
+                raise MessageValidationError("Invalid sections")
+            d["tools_added"] = [ToolDeclaration(**t) for t in d.get("tools_added", [])]
+            for t in d["tools_added"]:
+                _text(t.name, "tool name")
+                _text(t.description, "description")
+                if not isinstance(t.input_schema, dict):
+                    raise MessageValidationError("Invalid tool schema")
+            if not isinstance(d.get("tools_removed", []), list):
+                raise MessageValidationError("Invalid tool removals")
+            for name in d.get("tools_removed", []):
+                _text(name, "tool removal")
+            return SystemMessage(**d)
+        if role == "assistant":
+            d["content"] = _blocks(d["content"], tools=True)
+            if d.get("stop_reason", "stop") not in {
+                "stop",
+                "tool_use",
+                "length",
+                "error",
+                "aborted",
+                "pending",
+                "deferred",
+            }:
+                raise UnsupportedCapabilityError("Unsupported stop reason")
+            calls = [b.id for b in d["content"] if isinstance(b, ToolCall)]
+            if len(set(calls)) != len(calls):
+                raise MessageValidationError("Duplicate tool call ID")
+            for key in ("provider", "model"):
+                _text(d.get(key, "mock"), key)
+            for key in (
+                "api",
+                "provider_thinking_level",
+                "response_model",
+                "response_id",
+                "thinking_level",
+                "raw_stop_reason",
+            ):
+                if d.get(key) is not None:
+                    _text(d[key], key)
+            if not isinstance(d.get("usage", {}), dict):
+                raise MessageValidationError("Invalid usage")
+            if d.get("error") is not None:
+                _text(d["error"], "error")
+            if d.get("end_turn") is not None and type(d["end_turn"]) is not bool:
+                raise MessageValidationError("end_turn must be boolean")
+            if d.get("diagnostics") is not None and (
+                not isinstance(d["diagnostics"], list)
+                or any(not isinstance(v, dict) for v in d["diagnostics"])
+            ):
+                raise MessageValidationError("diagnostics must be a list of objects")
+            if d.get("deferred") is not None:
+                handle = d["deferred"]
+                if not isinstance(handle, dict) or not all(
+                    isinstance(handle.get(k), str) for k in ("provider", "model_id", "api", "id")
+                ):
+                    raise MessageValidationError("Invalid deferred handle")
+            return AssistantMessage(**d)
+        if role == "tool_result":
+            d["content"] = _blocks(d["content"], images=True)
+            _text(d["call_id"], "call_id")
+            _text(d["name"], "name")
+            if type(d.get("is_error", False)) is not bool:
+                raise MessageValidationError("is_error must be boolean")
+            if d.get("usage") is not None and not isinstance(d["usage"], dict):
+                raise MessageValidationError("Invalid tool usage")
+            if d.get("nested_calls") is not None:
+                nested = d["nested_calls"]
+                if (
+                    not isinstance(nested, dict)
+                    or type(nested.get("complete")) is not bool
+                    or not isinstance(nested.get("calls"), list)
+                ):
+                    raise MessageValidationError("Invalid nested calls")
+                for call in nested["calls"]:
+                    if (
+                        not isinstance(call, dict)
+                        or not all(isinstance(call.get(k), str) for k in ("id", "name"))
+                        or call.get("status") not in {"ok", "error", "unfinished"}
+                    ):
+                        raise MessageValidationError("Invalid nested call")
+            return ToolResultMessage(**d)
+        if role == "custom":
+            _text(d["custom_type"], "custom_type")
+            return CustomMessage(**d)
+        raise UnsupportedCapabilityError(f"Unsupported role: {role}")
+    except (TypeError, KeyError, ValueError) as exc:
+        raise MessageValidationError(str(exc)) from exc
+
+
+def validate_history(messages: list[Message]) -> None:
+    pending: dict[str, str] = {}
+    for message in messages:
+        message_to_dict(message)
+        if isinstance(message, ToolResultMessage):
+            if pending.pop(message.call_id, None) != message.name:
+                raise MessageValidationError("Unmatched or duplicate tool result")
+        else:
+            if pending:
+                raise MessageValidationError("Dangling tool calls")
+            if isinstance(message, AssistantMessage):
+                if message.stop_reason == "pending":
+                    raise MessageValidationError("Pending response cannot be committed to history")
+                if message.tool_calls and message.stop_reason in {"error", "aborted"}:
+                    raise MessageValidationError("Failed assistant cannot declare tool calls")
+                pending = {c.id: c.name for c in message.tool_calls}
+    if pending:
+        raise MessageValidationError("Dangling tool calls")
+
+
+def encode_messages(messages: list[Message]) -> str:
+    validate_history(messages)
+    return json.dumps(
+        {"schema_version": SCHEMA_VERSION, "messages": [message_to_dict(m) for m in messages]},
+        ensure_ascii=False,
+        allow_nan=False,
+    )
+
+
+def decode_messages(value: str) -> list[Message]:
+    try:
+        data = json.loads(value)
+        validate_json(data)
+        if type(data.get("schema_version")) is not int or data["schema_version"] != SCHEMA_VERSION:
+            raise MessageValidationError("Unsupported message schema version")
+        if set(data) != {"schema_version", "messages"} or not isinstance(data["messages"], list):
+            raise MessageValidationError("Invalid message envelope")
+        messages = [message_from_dict(m) for m in data["messages"]]
+        validate_history(messages)
+        return messages
+    except (ValueError, TypeError, AttributeError, KeyError) as exc:
+        raise MessageValidationError(str(exc)) from exc