mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-12-03 09:53:45 +00:00
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:
parent
a605cc2e14
commit
76fdff4a85
4 changed files with 82 additions and 34 deletions
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 = (
|
||||
|
|
|
|||
|
|
@ -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] = []
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue