inspect-ai 0.3.90__py3-none-any.whl → 0.3.91__py3-none-any.whl

inspect_ai/tool/_mcp/_sandbox.py

@@ -0,0 +1,104 @@
+import sys
+from contextlib import asynccontextmanager
+from typing import TextIO
+
+import anyio
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from mcp import JSONRPCRequest, StdioServerParameters
+from mcp.types import JSONRPCMessage, JSONRPCNotification
+
+from inspect_ai.tool._tool_support_helpers import (
+    exec_model_request,
+    exec_notification,
+    exec_scalar_request,
+    tool_container_sandbox,
+)
+
+from ._context import MCPServerContext
+
+
+# Pardon the type: ignore's here. This code is a modified clone of Anthropic code
+# for stdio_client. In their case, they don't provide a type hint for the return
+# value. We suspect that if they did, they'd encounter the same issues we're
+# suppressing. Nevertheless, we're confident that the runtime behavior of the
+# code is what we want, and that the errors are purely in the type domain.
+@asynccontextmanager  # type: ignore
+async def sandbox_client(  # type: ignore
+    server: StdioServerParameters,
+    *,
+    sandbox_name: str | None = None,
+    errlog: TextIO = sys.stderr,
+) -> MCPServerContext:  # type: ignore
+    sandbox_environment = await tool_container_sandbox(
+        "mcp support", sandbox_name=sandbox_name
+    )
+
+    # read_stream is remote process's stdout
+    read_stream: MemoryObjectReceiveStream[JSONRPCMessage | Exception]
+    read_stream_writer: MemoryObjectSendStream[JSONRPCMessage | Exception]
+
+    # write_stream is remote process's stdin
+    write_stream: MemoryObjectSendStream[JSONRPCMessage]
+    write_stream_reader: MemoryObjectReceiveStream[JSONRPCMessage]
+
+    read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
+    write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
+
+    session_id = await exec_scalar_request(
+        sandbox=sandbox_environment,
+        method="mcp_launch_server",
+        params={"server_params": server.model_dump()},
+        result_type=int,
+    )
+
+    async def stdout_reader() -> None:
+        # This is NYI until we support unsolicited messages from the sandbox
+        # back to the client
+        pass
+
+    async def stdin_writer() -> None:
+        try:
+            async with write_stream_reader:
+                # This reads messages until the stream is closed
+                async for message in write_stream_reader:
+                    root = message.root
+                    if isinstance(root, JSONRPCRequest):
+                        await read_stream_writer.send(
+                            await exec_model_request(
+                                sandbox=sandbox_environment,
+                                method="mcp_send_request",
+                                params={
+                                    "session_id": session_id,
+                                    "request": root.model_dump(),
+                                },
+                                result_type=JSONRPCMessage,
+                            )
+                        )
+                    elif isinstance(root, JSONRPCNotification):
+                        await exec_notification(
+                            sandbox=sandbox_environment,
+                            method="mcp_send_notification",
+                            params={
+                                "session_id": session_id,
+                                "notification": root.model_dump(),
+                            },
+                        )
+                    else:
+                        assert False, f"Unexpected message type {message=}"
+
+        except anyio.ClosedResourceError:
+            await anyio.lowlevel.checkpoint()
+
+    async with anyio.create_task_group() as tg:
+        tg.start_soon(stdout_reader)
+        tg.start_soon(stdin_writer)
+
+        try:
+            yield read_stream, write_stream
+        finally:
+            await exec_scalar_request(
+                sandbox=sandbox_environment,
+                method="mcp_kill_server",
+                params={"session_id": session_id},
+                result_type=type(None),
+            )

inspect_ai/tool/_mcp/_types.py

@@ -0,0 +1,31 @@
+import abc
+from logging import getLogger
+from typing import Literal
+
+from .._tool import Tool, ToolSource
+
+logger = getLogger(__name__)
+
+
+class MCPServer(ToolSource):
+    """Model Context Protocol server interface.
+
+    `MCPServer` can be passed in the `tools` argument as a source of tools
+    (use the `mcp_tools()` function to filter the list of tools)
+
+    """
+
+    async def tools(self) -> list[Tool]:
+        """List of all tools provided by this server."""
+        return await self._list_tools()
+
+    @abc.abstractmethod
+    async def _connect(self) -> None: ...
+
+    @abc.abstractmethod
+    async def _close(self) -> None: ...
+
+    @abc.abstractmethod
+    async def _list_tools(
+        self, tools: Literal["all"] | list[str] = "all"
+    ) -> list[Tool]: ...

inspect_ai/tool/_mcp/connection.py

@@ -0,0 +1,60 @@
+import contextlib
+from types import TracebackType
+from typing import AsyncIterator, Sequence
+
+from .._tool import Tool, ToolSource
+from .._tool_def import ToolDef
+from ._types import MCPServer
+from .tools import MCPToolSource
+
+
+@contextlib.asynccontextmanager
+async def mcp_connection(
+    tools: Sequence[Tool | ToolDef | ToolSource] | ToolSource,
+) -> AsyncIterator[None]:
+    """Context manager for running MCP servers required by tools.
+
+    Any `ToolSource` passed in tools will be examined to see
+    if it references an MCPServer, and if so, that server will be
+    connected to upon entering the context and disconnected from
+    upon exiting the context.
+
+    Args:
+       tools: Tools in current context.
+    """
+    # discover mcp servers in tools
+    tools = tools if isinstance(tools, Sequence) else [tools]
+    tool_sources = [tool for tool in tools if isinstance(tool, ToolSource)]
+    mcp_servers: list[MCPServer] = []
+    for tool_source in tool_sources:
+        if isinstance(tool_source, MCPServer):
+            mcp_servers.append(tool_source)
+        elif isinstance(tool_source, MCPToolSource):
+            mcp_servers.append(tool_source._server)
+
+    # enter connection contexts
+    async with contextlib.AsyncExitStack() as exit_stack:
+        for connection in [
+            MCPServerConnection(mcp_server) for mcp_server in mcp_servers
+        ]:
+            await exit_stack.enter_async_context(connection)
+
+        # onward
+        yield
+
+
+class MCPServerConnection:
+    def __init__(self, server: MCPServer) -> None:
+        self._server = server
+
+    async def __aenter__(self) -> "MCPServerConnection":
+        await self._server._connect()
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        await self._server._close()

inspect_ai/tool/_mcp/sampling.py

@@ -0,0 +1,118 @@
+from typing import Any
+
+from mcp.client.session import ClientSession
+from mcp.shared.context import RequestContext
+from mcp.types import (
+    INTERNAL_ERROR,
+    CreateMessageRequestParams,
+    CreateMessageResult,
+    EmbeddedResource,
+    ErrorData,
+    ImageContent,
+    TextContent,
+    TextResourceContents,
+)
+from mcp.types import (
+    StopReason as MCPStopReason,
+)
+
+from inspect_ai._util.content import Content, ContentImage, ContentText
+from inspect_ai._util.error import exception_message
+from inspect_ai._util.url import data_uri_mime_type, data_uri_to_base64
+
+
+async def sampling_fn(
+    context: RequestContext[ClientSession, Any],
+    params: CreateMessageRequestParams,
+) -> CreateMessageResult | ErrorData:
+    from inspect_ai.model._chat_message import (
+        ChatMessage,
+        ChatMessageAssistant,
+        ChatMessageSystem,
+        ChatMessageUser,
+    )
+    from inspect_ai.model._generate_config import GenerateConfig
+    from inspect_ai.model._model import get_model
+
+    try:
+        # build message list
+        messages: list[ChatMessage] = []
+        if params.systemPrompt:
+            messages.append(ChatMessageSystem(content=params.systemPrompt))
+
+        for message in params.messages:
+            if message.role == "assistant":
+                messages.append(
+                    ChatMessageAssistant(content=[as_inspect_content(message.content)])
+                )
+            elif message.role == "user":
+                messages.append(
+                    ChatMessageUser(content=[as_inspect_content(message.content)])
+                )
+
+        # sample w/ requested params
+        output = await get_model().generate(
+            messages,
+            config=GenerateConfig(
+                temperature=params.temperature,
+                max_tokens=params.maxTokens,
+                stop_seqs=params.stopSequences,
+            ),
+        )
+
+        # convert stop reason
+        stop_reason: MCPStopReason = (
+            "maxTokens" if output.stop_reason == "max_tokens" else "endTurn"
+        )
+
+        # return first compatible content
+        if isinstance(output.message.content, str):
+            return CreateMessageResult(
+                role="assistant",
+                content=TextContent(type="text", text=output.message.content),
+                model=output.model,
+                stopReason=stop_reason,
+            )
+        else:
+            for content in output.message.content:
+                if isinstance(content, ContentText | ContentImage):
+                    return CreateMessageResult(
+                        role="assistant",
+                        content=as_mcp_content(content),
+                        model=output.model,
+                        stopReason=stop_reason,
+                    )
+
+            # if we get this far then no valid content was returned
+            return ErrorData(
+                code=INTERNAL_ERROR, message="No text or image content was generated."
+            )
+
+    except Exception as ex:
+        return ErrorData(code=INTERNAL_ERROR, message=exception_message(ex))
+
+
+def as_inspect_content(
+    content: TextContent | ImageContent | EmbeddedResource,
+) -> Content:
+    if isinstance(content, TextContent):
+        return ContentText(text=content.text)
+    elif isinstance(content, ImageContent):
+        return ContentImage(
+            image=f"data:image/{content.mimeType};base64,{content.data}"
+        )
+    elif isinstance(content.resource, TextResourceContents):
+        return ContentText(text=content.resource.text)
+    else:
+        raise ValueError(f"Unexpected content: {content}")
+
+
+def as_mcp_content(content: ContentText | ContentImage) -> TextContent | ImageContent:
+    if isinstance(content, ContentText):
+        return TextContent(type="text", text=content.text)
+    else:
+        return ImageContent(
+            type="image",
+            mimeType=data_uri_mime_type(content.image) or "image/png",
+            data=data_uri_to_base64(content.image),
+        )

inspect_ai/tool/_mcp/server.py

@@ -0,0 +1,112 @@
+from logging import getLogger
+from pathlib import Path
+from typing import Any
+
+from inspect_ai._util.error import pip_dependency_error
+from inspect_ai._util.version import verify_required_version
+
+from ._types import MCPServer
+
+logger = getLogger(__name__)
+
+
+def mcp_server_sse(
+    *,
+    url: str,
+    headers: dict[str, Any] | None = None,
+    timeout: float = 5,
+    sse_read_timeout: float = 60 * 5,
+) -> MCPServer:
+    """MCP Server (SSE).
+
+    SSE interface to MCP server.  Use this for MCP servers available via a URL endpoint.
+
+    Args:
+        url: URL to remote server
+        headers: Headers to send server (typically authorization is included here)
+        timeout: Timeout for HTTP operations
+        sse_read_timeout: How long (in seconds) the client will wait for a new
+            event before disconnecting.
+
+    Returns:
+        McpClient: Client for MCP Server
+    """
+    verfify_mcp_package()
+    from ._mcp import create_server_sse
+
+    return create_server_sse(url, headers, timeout, sse_read_timeout)
+
+
+def mcp_server_stdio(
+    *,
+    command: str,
+    args: list[str] = [],
+    cwd: str | Path | None = None,
+    env: dict[str, str] | None = None,
+) -> MCPServer:
+    """MCP Server (Stdio).
+
+    Stdio interface to MCP server. Use this for MCP servers that run locally.
+
+    Args:
+        command: The executable to run to start the server.
+        args: Command line arguments to pass to the executable.
+        env: The environment to use when spawning the process
+            in addition to the platform specific set of default
+            environment variables (e.g. "HOME", "LOGNAME", "PATH",
+            "SHELL", "TERM", and "USER" for Posix-based systems).
+        cwd: The working directory to use when spawning the process.
+
+    Returns:
+        McpClient: Client for MCP Server
+    """
+    verfify_mcp_package()
+    from ._mcp import create_server_stdio
+
+    return create_server_stdio(command, args, cwd, env)
+
+
+def mcp_server_sandbox(
+    *,
+    command: str,
+    args: list[str] = [],
+    cwd: str | Path | None = None,
+    env: dict[str, str] | None = None,
+    sandbox: str | None = None,
+) -> MCPServer:
+    """MCP Server (Sandbox).
+
+    Interface to MCP server running in an Inspect sandbox.
+
+    Args:
+        command: The executable to run to start the server.
+        args: Command line arguments to pass to the executable.
+        env: The environment to use when spawning the process
+            in addition to the platform specific set of default
+            environment variables (e.g. "HOME", "LOGNAME", "PATH",
+            "SHELL", "TERM", and "USER" for Posix-based systems).
+        cwd: The working directory to use when spawning the process.
+        sandbox: The sandbox to use when spawning the process.
+
+    Returns:
+        McpClient: Client for MCP Server
+    """
+    verfify_mcp_package()
+    from ._mcp import create_server_sandbox
+
+    return create_server_sandbox(command, args, cwd, env, sandbox)
+
+
+def verfify_mcp_package() -> None:
+    FEATURE = "MCP tools"
+    PACKAGE = "mcp"
+    MIN_VERSION = "1.6.0"
+
+    # verify we have the package
+    try:
+        import mcp  # noqa: F401
+    except ImportError:
+        raise pip_dependency_error(FEATURE, [PACKAGE])
+
+    # verify version
+    verify_required_version(FEATURE, PACKAGE, MIN_VERSION)

inspect_ai/tool/_mcp/tools.py

@@ -0,0 +1,34 @@
+from typing import Literal
+
+from .._tool import Tool, ToolSource
+from ._types import MCPServer
+
+
+def mcp_tools(
+    server: MCPServer,
+    *,
+    tools: Literal["all"] | list[str] = "all",
+) -> ToolSource:
+    """Tools from MCP server.
+
+    Args:
+       server: MCP server created with `mcp_server_stdio()` or `mcp_server_sse()`
+       tools: List of tool names (or globs) (defaults to "all")
+          which returns all tools.
+
+    Returns:
+       ToolSource: Source for specified MCP server tools.
+    """
+    return MCPToolSource(server, tools)
+
+
+class MCPToolSource(ToolSource):
+    def __init__(self, server: MCPServer, tools: Literal["all"] | list[str]) -> None:
+        self._server = server
+        self._tools = tools
+        self._cached_tool_list: list[Tool] | None = None
+
+    async def tools(self) -> list[Tool]:
+        if self._cached_tool_list is None:
+            self._cached_tool_list = await self._server._list_tools(self._tools)
+        return self._cached_tool_list

inspect_ai/tool/_tool.py

@@ -105,6 +105,19 @@ class Tool(Protocol):
         ...
 
 
+@runtime_checkable
+class ToolSource(Protocol):
+    """Protocol for dynamically providing a set of tools."""
+
+    async def tools(self) -> list[Tool]:
+        """Retrieve tools from tool source.
+
+        Returns:
+            List of tools
+        """
+        ...
+
+
 P = ParamSpec("P")
 
 

inspect_ai/tool/_tool_def.py

@@ -3,6 +3,7 @@ from typing import (
     Any,
     Callable,
     NamedTuple,
+    Sequence,
 )
 
 from inspect_ai._util.registry import (
@@ -13,7 +14,14 @@ from inspect_ai._util.registry import (
     set_registry_params,
 )
 
-from ._tool import TOOL_MODEL_INPUT, TOOL_PARALLEL, TOOL_PROMPT, TOOL_VIEWER, Tool
+from ._tool import (
+    TOOL_MODEL_INPUT,
+    TOOL_PARALLEL,
+    TOOL_PROMPT,
+    TOOL_VIEWER,
+    Tool,
+    ToolSource,
+)
 from ._tool_call import ToolCallModelInput, ToolCallViewer
 from ._tool_description import (
     ToolDescription,
@@ -157,10 +165,21 @@ def apply_description_overrides(target: ToolParams, overrides: dict[str, str]) -
         target.properties[param].description = value
 
 
-def tool_defs(
-    tools: list[Tool] | list[ToolDef] | list[Tool | ToolDef],
+async def tool_defs(
+    tools: Sequence[Tool | ToolDef | ToolSource] | ToolSource,
 ) -> list[ToolDef]:
-    return [ToolDef(tool) if isinstance(tool, Tool) else tool for tool in tools]
+    if isinstance(tools, ToolSource):
+        tools = await tools.tools()
+
+    tool_defs: list[ToolDef] = []
+    for tool in tools:
+        if isinstance(tool, ToolSource):
+            tool_defs.extend([ToolDef(t) for t in await tool.tools()])
+        elif not isinstance(tool, ToolDef):
+            tool_defs.append(ToolDef(tool))
+        else:
+            tool_defs.append(tool)
+    return tool_defs
 
 
 class ToolDefFields(NamedTuple):
@@ -245,7 +264,5 @@ def validate_tool_parameters(tool_name: str, parameters: dict[str, ToolParam]) -
                 f"{context} provided for parameter '{bound_name}' of function '{tool_name}'."
             )
 
-        if param.type is None and not param.anyOf and not param.enum:
-            raise_not_provided_error("Unsupported type or type annotation")
-        elif not param.description:
+        if not param.description:
             raise_not_provided_error("Description not")