mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-26 19:24:27 +00:00
40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
from typing import List, Literal, Union
|
|
|
|
from mcp import ClientSession
|
|
from mcp.types import Tool as MCPTool
|
|
from openai.types.chat import ChatCompletionToolParam
|
|
from openai.types.shared_params.function_definition import FunctionDefinition
|
|
|
|
|
|
def transform_mcp_tool_to_openai_tool(mcp_tool: MCPTool) -> ChatCompletionToolParam:
|
|
"""Convert an MCP tool to an OpenAI tool."""
|
|
return ChatCompletionToolParam(
|
|
type="function",
|
|
function=FunctionDefinition(
|
|
name=mcp_tool.name,
|
|
description=mcp_tool.description or "",
|
|
parameters=mcp_tool.inputSchema,
|
|
strict=False,
|
|
),
|
|
)
|
|
|
|
|
|
async def load_mcp_tools(
|
|
session: ClientSession, format: Literal["mcp", "openai"] = "mcp"
|
|
) -> Union[List[MCPTool], List[ChatCompletionToolParam]]:
|
|
"""
|
|
Load all available MCP tools
|
|
|
|
Args:
|
|
session: The MCP session to use
|
|
format: The format to convert the tools to
|
|
By default, the tools are returned in MCP format.
|
|
|
|
If format is set to "openai", the tools are converted to OpenAI tools.
|
|
"""
|
|
tools = await session.list_tools()
|
|
if format == "openai":
|
|
return [
|
|
transform_mcp_tool_to_openai_tool(mcp_tool=tool) for tool in tools.tools
|
|
]
|
|
return tools.tools
|