created a single helper function and updated list_mcp_tools and invoke_mcp_tool. Removed the comments in openai_responses.py

This commit is contained in:
Omar Abdelwahab 2025-11-05 13:12:28 -08:00
parent a605cc2e14
commit 76fdff4a85
4 changed files with 82 additions and 34 deletions

View file

@ -496,8 +496,6 @@ class OpenAIResponseInputToolMCP(BaseModel):
server_label: str
server_url: str
headers: dict[str, Any] | None = None
# OAuth access token for MCP server authentication
# Provide just the token (e.g., "my-secret-token"), the "Bearer " prefix will be added automatically
authorization: str | None = None
require_approval: Literal["always"] | Literal["never"] | ApprovalFilter = "never"

View file

@ -1079,23 +1079,12 @@ class StreamingResponseOrchestrator:
"server_url": mcp_tool.server_url,
"mcp_list_tools_id": list_id,
}
# Prepare headers with authorization from tool config
headers = dict(mcp_tool.headers or {})
if mcp_tool.authorization:
# Check if Authorization header already exists (case-insensitive check)
existing_keys_lower = {k.lower() for k in headers.keys()}
if "authorization" in existing_keys_lower:
raise ValueError(
"Cannot specify Authorization in both 'headers' and 'authorization' fields. "
"Please use only the 'authorization' field."
)
# OAuth access token - add "Bearer " prefix
headers["Authorization"] = f"Bearer {mcp_tool.authorization}"
# List MCP tools with authorization from tool config
async with tracing.span("list_mcp_tools", attributes):
tool_defs = await list_mcp_tools(
endpoint=mcp_tool.server_url,
headers=headers,
headers=mcp_tool.headers,
authorization=mcp_tool.authorization,
)
# Create the MCP list tools message

View file

@ -299,25 +299,14 @@ class ToolExecutor:
"server_url": mcp_tool.server_url,
"tool_name": function_name,
}
# Prepare headers with authorization from tool config
headers = dict(mcp_tool.headers or {})
if mcp_tool.authorization:
# Check if Authorization header already exists (case-insensitive check)
existing_keys_lower = {k.lower() for k in headers.keys()}
if "authorization" in existing_keys_lower:
raise ValueError(
"Cannot specify Authorization in both 'headers' and 'authorization' fields. "
"Please use only the 'authorization' field."
)
# OAuth access token - add "Bearer " prefix
headers["Authorization"] = f"Bearer {mcp_tool.authorization}"
# Invoke MCP tool with authorization from tool config
async with tracing.span("invoke_mcp_tool", attributes):
result = await invoke_mcp_tool(
endpoint=mcp_tool.server_url,
headers=headers,
tool_name=function_name,
kwargs=tool_kwargs,
headers=mcp_tool.headers,
authorization=mcp_tool.authorization,
)
elif function_name == "knowledge_search":
response_file_search_tool = (

View file

@ -27,6 +27,36 @@ from llama_stack.providers.utils.tools.ttl_dict import TTLDict
logger = get_logger(__name__, category="tools")
def prepare_mcp_headers(base_headers: dict[str, str] | None, authorization: str | None) -> dict[str, str]:
"""Prepare headers for MCP requests with authorization handling.
Args:
base_headers: Base headers to use (e.g., from mcp_tool.headers)
authorization: OAuth access token (just the token, not "Bearer <token>")
Returns:
Final headers dict with Authorization header if authorization is provided
Raises:
ValueError: If both base_headers contains Authorization and authorization parameter is provided
"""
headers = dict(base_headers or {})
if authorization:
# Check if Authorization header already exists (case-insensitive check)
existing_keys_lower = {k.lower() for k in headers.keys()}
if "authorization" in existing_keys_lower:
raise ValueError(
"Cannot specify Authorization in both 'headers' and 'authorization' fields. "
"Please use only the 'authorization' field."
)
# OAuth access token - add "Bearer " prefix
headers["Authorization"] = f"Bearer {authorization}"
return headers
protocol_cache = TTLDict(ttl_seconds=3600)
@ -109,9 +139,29 @@ async def client_wrapper(endpoint: str, headers: dict[str, str]) -> AsyncGenerat
raise
async def list_mcp_tools(endpoint: str, headers: dict[str, str]) -> ListToolDefsResponse:
async def list_mcp_tools(
endpoint: str,
headers: dict[str, str] | None = None,
authorization: str | None = None,
) -> ListToolDefsResponse:
"""List tools available from an MCP server.
Args:
endpoint: MCP server endpoint URL
headers: Optional base headers to include
authorization: Optional OAuth access token (just the token, not "Bearer <token>")
Returns:
List of tool definitions from the MCP server
Raises:
ValueError: If both headers contains Authorization and authorization parameter is provided
"""
# Prepare headers with authorization handling
final_headers = prepare_mcp_headers(headers, authorization)
tools = []
async with client_wrapper(endpoint, headers) as session:
async with client_wrapper(endpoint, final_headers) as session:
tools_result = await session.list_tools()
for tool in tools_result.tools:
tools.append(
@ -129,9 +179,31 @@ async def list_mcp_tools(endpoint: str, headers: dict[str, str]) -> ListToolDefs
async def invoke_mcp_tool(
endpoint: str, headers: dict[str, str], tool_name: str, kwargs: dict[str, Any]
endpoint: str,
tool_name: str,
kwargs: dict[str, Any],
headers: dict[str, str] | None = None,
authorization: str | None = None,
) -> ToolInvocationResult:
async with client_wrapper(endpoint, headers) as session:
"""Invoke an MCP tool with the given arguments.
Args:
endpoint: MCP server endpoint URL
tool_name: Name of the tool to invoke
kwargs: Tool invocation arguments
headers: Optional base headers to include
authorization: Optional OAuth access token (just the token, not "Bearer <token>")
Returns:
Tool invocation result with content and error information
Raises:
ValueError: If both headers contains Authorization and authorization parameter is provided
"""
# Prepare headers with authorization handling
final_headers = prepare_mcp_headers(headers, authorization)
async with client_wrapper(endpoint, final_headers) as session:
result = await session.call_tool(tool_name, kwargs)
content: list[InterleavedContentItem] = []