mirror of
https://github.com/open-webui/open-webui.git
synced 2026-05-04 19:29:27 -05:00
137 lines
4.7 KiB
Python
137 lines
4.7 KiB
Python
import asyncio
|
|
from typing import Optional
|
|
from contextlib import AsyncExitStack
|
|
|
|
import anyio
|
|
|
|
from mcp import ClientSession
|
|
from mcp.client.auth import OAuthClientProvider, TokenStorage
|
|
from mcp.client.streamable_http import streamablehttp_client
|
|
from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata, OAuthToken
|
|
import httpx
|
|
from open_webui.env import AIOHTTP_CLIENT_SESSION_TOOL_SERVER_SSL
|
|
|
|
|
|
def create_insecure_httpx_client(headers=None, timeout=None, auth=None):
|
|
"""Create an httpx AsyncClient with SSL verification disabled.
|
|
|
|
Note: verify=False must be passed at construction time because httpx
|
|
configures the SSL context during __init__. Setting client.verify = False
|
|
after construction does not affect the underlying transport's SSL context.
|
|
"""
|
|
kwargs = {
|
|
'follow_redirects': True,
|
|
'verify': False,
|
|
}
|
|
if timeout is not None:
|
|
kwargs['timeout'] = timeout
|
|
if headers is not None:
|
|
kwargs['headers'] = headers
|
|
if auth is not None:
|
|
kwargs['auth'] = auth
|
|
return httpx.AsyncClient(**kwargs)
|
|
|
|
|
|
class MCPClient:
|
|
def __init__(self):
|
|
self.session: Optional[ClientSession] = None
|
|
self.exit_stack = None
|
|
|
|
async def connect(self, url: str, headers: Optional[dict] = None):
|
|
async with AsyncExitStack() as exit_stack:
|
|
try:
|
|
if AIOHTTP_CLIENT_SESSION_TOOL_SERVER_SSL:
|
|
self._streams_context = streamablehttp_client(url, headers=headers)
|
|
else:
|
|
self._streams_context = streamablehttp_client(
|
|
url,
|
|
headers=headers,
|
|
httpx_client_factory=create_insecure_httpx_client,
|
|
)
|
|
|
|
transport = await exit_stack.enter_async_context(self._streams_context)
|
|
read_stream, write_stream, _ = transport
|
|
|
|
self._session_context = ClientSession(read_stream, write_stream) # pylint: disable=W0201
|
|
|
|
self.session = await exit_stack.enter_async_context(self._session_context)
|
|
with anyio.fail_after(10):
|
|
await self.session.initialize()
|
|
self.exit_stack = exit_stack.pop_all()
|
|
except Exception as e:
|
|
await asyncio.shield(self.disconnect())
|
|
raise e
|
|
|
|
async def list_tool_specs(self) -> Optional[dict]:
|
|
if not self.session:
|
|
raise RuntimeError('MCP client is not connected.')
|
|
|
|
result = await self.session.list_tools()
|
|
tools = result.tools
|
|
|
|
tool_specs = []
|
|
for tool in tools:
|
|
name = tool.name
|
|
description = tool.description
|
|
|
|
inputSchema = tool.inputSchema
|
|
|
|
# TODO: handle outputSchema if needed
|
|
outputSchema = getattr(tool, 'outputSchema', None)
|
|
|
|
tool_specs.append({'name': name, 'description': description, 'parameters': inputSchema})
|
|
|
|
return tool_specs
|
|
|
|
async def call_tool(self, function_name: str, function_args: dict) -> Optional[dict]:
|
|
if not self.session:
|
|
raise RuntimeError('MCP client is not connected.')
|
|
|
|
result = await self.session.call_tool(function_name, function_args)
|
|
if not result:
|
|
raise Exception('No result returned from MCP tool call.')
|
|
|
|
result_dict = result.model_dump(mode='json')
|
|
result_content = result_dict.get('content', {})
|
|
|
|
if result.isError:
|
|
raise Exception(result_content)
|
|
else:
|
|
return result_content
|
|
|
|
async def list_resources(self, cursor: Optional[str] = None) -> Optional[dict]:
|
|
if not self.session:
|
|
raise RuntimeError('MCP client is not connected.')
|
|
|
|
result = await self.session.list_resources(cursor=cursor)
|
|
if not result:
|
|
raise Exception('No result returned from MCP list_resources call.')
|
|
|
|
result_dict = result.model_dump()
|
|
resources = result_dict.get('resources', [])
|
|
|
|
return resources
|
|
|
|
async def read_resource(self, uri: str) -> Optional[dict]:
|
|
if not self.session:
|
|
raise RuntimeError('MCP client is not connected.')
|
|
|
|
result = await self.session.read_resource(uri)
|
|
if not result:
|
|
raise Exception('No result returned from MCP read_resource call.')
|
|
result_dict = result.model_dump()
|
|
|
|
return result_dict
|
|
|
|
async def disconnect(self):
|
|
# Clean up and close the session
|
|
await self.exit_stack.aclose()
|
|
|
|
async def __aenter__(self):
|
|
await self.exit_stack.__aenter__()
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc_value, traceback):
|
|
await self.exit_stack.__aexit__(exc_type, exc_value, traceback)
|
|
await self.disconnect()
|