openai-agents 0.0.19__py3-none-any.whl → 0.2.0__py3-none-any.whl

agents/mcp/server.py

@@ -2,21 +2,27 @@ from __future__ import annotations
 
 import abc
 import asyncio
+import inspect
 from contextlib import AbstractAsyncContextManager, AsyncExitStack
 from datetime import timedelta
 from pathlib import Path
-from typing import Any, Literal
+from typing import TYPE_CHECKING, Any, Literal, cast
 
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from mcp import ClientSession, StdioServerParameters, Tool as MCPTool, stdio_client
 from mcp.client.sse import sse_client
 from mcp.client.streamable_http import GetSessionIdCallback, streamablehttp_client
 from mcp.shared.message import SessionMessage
-from mcp.types import CallToolResult, InitializeResult
+from mcp.types import CallToolResult, GetPromptResult, InitializeResult, ListPromptsResult
 from typing_extensions import NotRequired, TypedDict
 
 from ..exceptions import UserError
 from ..logger import logger
+from ..run_context import RunContextWrapper
+from .util import ToolFilter, ToolFilterCallable, ToolFilterContext, ToolFilterStatic
+
+if TYPE_CHECKING:
+    from ..agent import AgentBase
 
 
 class MCPServer(abc.ABC):
@@ -44,7 +50,11 @@ class MCPServer(abc.ABC):
         pass
 
     @abc.abstractmethod
-    async def list_tools(self) -> list[MCPTool]:
+    async def list_tools(
+        self,
+        run_context: RunContextWrapper[Any] | None = None,
+        agent: AgentBase | None = None,
+    ) -> list[MCPTool]:
         """List the tools available on the server."""
         pass
 
@@ -53,11 +63,30 @@ class MCPServer(abc.ABC):
         """Invoke a tool on the server."""
         pass
 
+    @abc.abstractmethod
+    async def list_prompts(
+        self,
+    ) -> ListPromptsResult:
+        """List the prompts available on the server."""
+        pass
+
+    @abc.abstractmethod
+    async def get_prompt(
+        self, name: str, arguments: dict[str, Any] | None = None
+    ) -> GetPromptResult:
+        """Get a specific prompt from the server."""
+        pass
+
 
 class _MCPServerWithClientSession(MCPServer, abc.ABC):
     """Base class for MCP servers that use a `ClientSession` to communicate with the server."""
 
-    def __init__(self, cache_tools_list: bool, client_session_timeout_seconds: float | None):
+    def __init__(
+        self,
+        cache_tools_list: bool,
+        client_session_timeout_seconds: float | None,
+        tool_filter: ToolFilter = None,
+    ):
         """
         Args:
             cache_tools_list: Whether to cache the tools list. If `True`, the tools list will be
@@ -68,6 +97,7 @@ class _MCPServerWithClientSession(MCPServer, abc.ABC):
             (by avoiding a round-trip to the server every time).
 
             client_session_timeout_seconds: the read timeout passed to the MCP ClientSession.
+            tool_filter: The tool filter to use for filtering tools.
         """
         self.session: ClientSession | None = None
         self.exit_stack: AsyncExitStack = AsyncExitStack()
@@ -81,6 +111,86 @@ class _MCPServerWithClientSession(MCPServer, abc.ABC):
         self._cache_dirty = True
         self._tools_list: list[MCPTool] | None = None
 
+        self.tool_filter = tool_filter
+
+    async def _apply_tool_filter(
+        self,
+        tools: list[MCPTool],
+        run_context: RunContextWrapper[Any],
+        agent: AgentBase,
+    ) -> list[MCPTool]:
+        """Apply the tool filter to the list of tools."""
+        if self.tool_filter is None:
+            return tools
+
+        # Handle static tool filter
+        if isinstance(self.tool_filter, dict):
+            return self._apply_static_tool_filter(tools, self.tool_filter)
+
+        # Handle callable tool filter (dynamic filter)
+        else:
+            return await self._apply_dynamic_tool_filter(tools, run_context, agent)
+
+    def _apply_static_tool_filter(
+        self, tools: list[MCPTool], static_filter: ToolFilterStatic
+    ) -> list[MCPTool]:
+        """Apply static tool filtering based on allowlist and blocklist."""
+        filtered_tools = tools
+
+        # Apply allowed_tool_names filter (whitelist)
+        if "allowed_tool_names" in static_filter:
+            allowed_names = static_filter["allowed_tool_names"]
+            filtered_tools = [t for t in filtered_tools if t.name in allowed_names]
+
+        # Apply blocked_tool_names filter (blacklist)
+        if "blocked_tool_names" in static_filter:
+            blocked_names = static_filter["blocked_tool_names"]
+            filtered_tools = [t for t in filtered_tools if t.name not in blocked_names]
+
+        return filtered_tools
+
+    async def _apply_dynamic_tool_filter(
+        self,
+        tools: list[MCPTool],
+        run_context: RunContextWrapper[Any],
+        agent: AgentBase,
+    ) -> list[MCPTool]:
+        """Apply dynamic tool filtering using a callable filter function."""
+
+        # Ensure we have a callable filter and cast to help mypy
+        if not callable(self.tool_filter):
+            raise ValueError("Tool filter must be callable for dynamic filtering")
+        tool_filter_func = cast(ToolFilterCallable, self.tool_filter)
+
+        # Create filter context
+        filter_context = ToolFilterContext(
+            run_context=run_context,
+            agent=agent,
+            server_name=self.name,
+        )
+
+        filtered_tools = []
+        for tool in tools:
+            try:
+                # Call the filter function with context
+                result = tool_filter_func(filter_context, tool)
+
+                if inspect.isawaitable(result):
+                    should_include = await result
+                else:
+                    should_include = result
+
+                if should_include:
+                    filtered_tools.append(tool)
+            except Exception as e:
+                logger.error(
+                    f"Error applying tool filter to tool '{tool.name}' on server '{self.name}': {e}"
+                )
+                # On error, exclude the tool for safety
+                continue
+
+        return filtered_tools
+
     @abc.abstractmethod
     def create_streams(
         self,
@@ -131,21 +241,32 @@ class _MCPServerWithClientSession(MCPServer, abc.ABC):
             await self.cleanup()
             raise
 
-    async def list_tools(self) -> list[MCPTool]:
+    async def list_tools(
+        self,
+        run_context: RunContextWrapper[Any] | None = None,
+        agent: AgentBase | None = None,
+    ) -> list[MCPTool]:
         """List the tools available on the server."""
         if not self.session:
             raise UserError("Server not initialized. Make sure you call `connect()` first.")
 
         # Return from cache if caching is enabled, we have tools, and the cache is not dirty
         if self.cache_tools_list and not self._cache_dirty and self._tools_list:
-            return self._tools_list
-
-        # Reset the cache dirty to False
-        self._cache_dirty = False
-
-        # Fetch the tools from the server
-        self._tools_list = (await self.session.list_tools()).tools
-        return self._tools_list
+            tools = self._tools_list
+        else:
+            # Reset the cache dirty to False
+            self._cache_dirty = False
+            # Fetch the tools from the server
+            self._tools_list = (await self.session.list_tools()).tools
+            tools = self._tools_list
+
+        # Filter tools based on tool_filter
+        filtered_tools = tools
+        if self.tool_filter is not None:
+            if run_context is None or agent is None:
+                raise UserError("run_context and agent are required for dynamic tool filtering")
+            filtered_tools = await self._apply_tool_filter(filtered_tools, run_context, agent)
+        return filtered_tools
 
     async def call_tool(self, tool_name: str, arguments: dict[str, Any] | None) -> CallToolResult:
         """Invoke a tool on the server."""
@@ -154,6 +275,24 @@ class _MCPServerWithClientSession(MCPServer, abc.ABC):
 
         return await self.session.call_tool(tool_name, arguments)
 
+    async def list_prompts(
+        self,
+    ) -> ListPromptsResult:
+        """List the prompts available on the server."""
+        if not self.session:
+            raise UserError("Server not initialized. Make sure you call `connect()` first.")
+
+        return await self.session.list_prompts()
+
+    async def get_prompt(
+        self, name: str, arguments: dict[str, Any] | None = None
+    ) -> GetPromptResult:
+        """Get a specific prompt from the server."""
+        if not self.session:
+            raise UserError("Server not initialized. Make sure you call `connect()` first.")
+
+        return await self.session.get_prompt(name, arguments)
+
     async def cleanup(self):
         """Cleanup the server."""
         async with self._cleanup_lock:
@@ -206,6 +345,7 @@ class MCPServerStdio(_MCPServerWithClientSession):
         cache_tools_list: bool = False,
         name: str | None = None,
         client_session_timeout_seconds: float | None = 5,
+        tool_filter: ToolFilter = None,
     ):
         """Create a new MCP server based on the stdio transport.
 
@@ -223,8 +363,13 @@ class MCPServerStdio(_MCPServerWithClientSession):
             name: A readable name for the server. If not provided, we'll create one from the
                 command.
             client_session_timeout_seconds: the read timeout passed to the MCP ClientSession.
+            tool_filter: The tool filter to use for filtering tools.
         """
-        super().__init__(cache_tools_list, client_session_timeout_seconds)
+        super().__init__(
+            cache_tools_list,
+            client_session_timeout_seconds,
+            tool_filter,
+        )
 
         self.params = StdioServerParameters(
             command=params["command"],
@@ -283,6 +428,7 @@ class MCPServerSse(_MCPServerWithClientSession):
         cache_tools_list: bool = False,
         name: str | None = None,
         client_session_timeout_seconds: float | None = 5,
+        tool_filter: ToolFilter = None,
     ):
         """Create a new MCP server based on the HTTP with SSE transport.
 
@@ -302,8 +448,13 @@ class MCPServerSse(_MCPServerWithClientSession):
                 URL.
 
             client_session_timeout_seconds: the read timeout passed to the MCP ClientSession.
+            tool_filter: The tool filter to use for filtering tools.
         """
-        super().__init__(cache_tools_list, client_session_timeout_seconds)
+        super().__init__(
+            cache_tools_list,
+            client_session_timeout_seconds,
+            tool_filter,
+        )
 
         self.params = params
         self._name = name or f"sse: {self.params['url']}"
@@ -362,6 +513,7 @@ class MCPServerStreamableHttp(_MCPServerWithClientSession):
         cache_tools_list: bool = False,
         name: str | None = None,
         client_session_timeout_seconds: float | None = 5,
+        tool_filter: ToolFilter = None,
     ):
         """Create a new MCP server based on the Streamable HTTP transport.
 
@@ -382,8 +534,13 @@ class MCPServerStreamableHttp(_MCPServerWithClientSession):
                 URL.
 
             client_session_timeout_seconds: the read timeout passed to the MCP ClientSession.
+            tool_filter: The tool filter to use for filtering tools.
         """
-        super().__init__(cache_tools_list, client_session_timeout_seconds)
+        super().__init__(
+            cache_tools_list,
+            client_session_timeout_seconds,
+            tool_filter,
+        )
 
         self.params = params
         self._name = name or f"streamable_http: {self.params['url']}"

agents/mcp/util.py

@@ -1,34 +1,113 @@
 import functools
 import json
-from typing import TYPE_CHECKING, Any
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Callable, Optional, Union
 
-from agents.strict_schema import ensure_strict_json_schema
+from typing_extensions import NotRequired, TypedDict
 
 from .. import _debug
 from ..exceptions import AgentsException, ModelBehaviorError, UserError
 from ..logger import logger
 from ..run_context import RunContextWrapper
+from ..strict_schema import ensure_strict_json_schema
 from ..tool import FunctionTool, Tool
 from ..tracing import FunctionSpanData, get_current_span, mcp_tools_span
+from ..util._types import MaybeAwaitable
 
 if TYPE_CHECKING:
     from mcp.types import Tool as MCPTool
 
+    from ..agent import AgentBase
     from .server import MCPServer
 
 
+@dataclass
+class ToolFilterContext:
+    """Context information available to tool filter functions."""
+
+    run_context: RunContextWrapper[Any]
+    """The current run context."""
+
+    agent: "AgentBase"
+    """The agent that is requesting the tool list."""
+
+    server_name: str
+    """The name of the MCP server."""
+
+
+ToolFilterCallable = Callable[["ToolFilterContext", "MCPTool"], MaybeAwaitable[bool]]
+"""A function that determines whether a tool should be available.
+
+Args:
+    context: The context information including run context, agent, and server name.
+    tool: The MCP tool to filter.
+
+Returns:
+    Whether the tool should be available (True) or filtered out (False).
+"""
+
+
+class ToolFilterStatic(TypedDict):
+    """Static tool filter configuration using allowlists and blocklists."""
+
+    allowed_tool_names: NotRequired[list[str]]
+    """Optional list of tool names to allow (whitelist).
+    If set, only these tools will be available."""
+
+    blocked_tool_names: NotRequired[list[str]]
+    """Optional list of tool names to exclude (blacklist).
+    If set, these tools will be filtered out."""
+
+
+ToolFilter = Union[ToolFilterCallable, ToolFilterStatic, None]
+"""A tool filter that can be either a function, static configuration, or None (no filtering)."""
+
+
+def create_static_tool_filter(
+    allowed_tool_names: Optional[list[str]] = None,
+    blocked_tool_names: Optional[list[str]] = None,
+) -> Optional[ToolFilterStatic]:
+    """Create a static tool filter from allowlist and blocklist parameters.
+
+    This is a convenience function for creating a ToolFilterStatic.
+
+    Args:
+        allowed_tool_names: Optional list of tool names to allow (whitelist).
+        blocked_tool_names: Optional list of tool names to exclude (blacklist).
+
+    Returns:
+        A ToolFilterStatic if any filtering is specified, None otherwise.
+    """
+    if allowed_tool_names is None and blocked_tool_names is None:
+        return None
+
+    filter_dict: ToolFilterStatic = {}
+    if allowed_tool_names is not None:
+        filter_dict["allowed_tool_names"] = allowed_tool_names
+    if blocked_tool_names is not None:
+        filter_dict["blocked_tool_names"] = blocked_tool_names
+
+    return filter_dict
+
+
 class MCPUtil:
     """Set of utilities for interop between MCP and Agents SDK tools."""
 
     @classmethod
     async def get_all_function_tools(
-        cls, servers: list["MCPServer"], convert_schemas_to_strict: bool
+        cls,
+        servers: list["MCPServer"],
+        convert_schemas_to_strict: bool,
+        run_context: RunContextWrapper[Any],
+        agent: "AgentBase",
     ) -> list[Tool]:
         """Get all function tools from a list of MCP servers."""
         tools = []
         tool_names: set[str] = set()
         for server in servers:
-            server_tools = await cls.get_function_tools(server, convert_schemas_to_strict)
+            server_tools = await cls.get_function_tools(
+                server, convert_schemas_to_strict, run_context, agent
+            )
             server_tool_names = {tool.name for tool in server_tools}
             if len(server_tool_names & tool_names) > 0:
                 raise UserError(
@@ -42,12 +121,16 @@ class MCPUtil:
 
     @classmethod
     async def get_function_tools(
-        cls, server: "MCPServer", convert_schemas_to_strict: bool
+        cls,
+        server: "MCPServer",
+        convert_schemas_to_strict: bool,
+        run_context: RunContextWrapper[Any],
+        agent: "AgentBase",
     ) -> list[Tool]:
         """Get all function tools from a single MCP server."""
 
         with mcp_tools_span(server=server.name) as span:
-            tools = await server.list_tools()
+            tools = await server.list_tools(run_context, agent)
             span.span_data.result = [tool.name for tool in tools]
 
         return [cls.to_function_tool(tool, server, convert_schemas_to_strict) for tool in tools]

agents/memory/__init__.py

@@ -0,0 +1,3 @@
+from .session import Session, SQLiteSession
+
+__all__ = ["Session", "SQLiteSession"]