tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-17 08:25:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: extract transport context creation into separate method (#19794)

Browse Source
This commit is contained in:
Alexsander Hamir
2026-01-26 09:57:53 -08:00
committed by GitHub
parent b920be3ee7
commit de3bbc216d
1 changed files with 47 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/experimental_mcp_client/client.py
+47 -38
View File
@@ -4,7 +4,7 @@ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers.
import asyncio
import base64
from typing import Awaitable, Callable, Dict, List, Optional, TypeVar, Union
from typing import Any, Awaitable, Callable, Dict, List, Optional, TypeVar, Union
import httpx
from mcp import ClientSession, ReadResourceResult, Resource, StdioServerParameters
@@ -74,52 +74,61 @@ class MCPClient:
if auth_value:
self.update_auth_value(auth_value)
def _create_transport_context(
self,
) -> tuple[Any, Optional[httpx.AsyncClient]]:
"""Create the appropriate transport context based on transport type."""
http_client: Optional[httpx.AsyncClient] = None
if self.transport_type == MCPTransport.stdio:
if not self.stdio_config:
raise ValueError("stdio_config is required for stdio transport")
server_params = StdioServerParameters(
command=self.stdio_config.get("command", ""),
args=self.stdio_config.get("args", []),
env=self.stdio_config.get("env", {}),
)
transport_ctx = stdio_client(server_params)
elif self.transport_type == MCPTransport.sse:
headers = self._get_auth_headers()
httpx_client_factory = self._create_httpx_client_factory()
transport_ctx = sse_client(
url=self.server_url,
timeout=self.timeout,
headers=headers,
httpx_client_factory=httpx_client_factory,
)
else:
headers = self._get_auth_headers()
httpx_client_factory = self._create_httpx_client_factory()
verbose_logger.debug(
"litellm headers for streamable_http_client: %s", headers
)
http_client = httpx_client_factory(
headers=headers,
timeout=httpx.Timeout(self.timeout),
)
transport_ctx = streamable_http_client(
url=self.server_url,
http_client=http_client,
)
if transport_ctx is None:
raise RuntimeError("Failed to create transport context")
return transport_ctx, http_client
async def run_with_session(
self, operation: Callable[[ClientSession], Awaitable[TSessionResult]]
) -> TSessionResult:
"""Open a session, run the provided coroutine, and clean up."""
transport_ctx = None
http_client: Optional[httpx.AsyncClient] = None
transport = None
session_ctx = None
try:
if self.transport_type == MCPTransport.stdio:
if not self.stdio_config:
raise ValueError("stdio_config is required for stdio transport")
server_params = StdioServerParameters(
command=self.stdio_config.get("command", ""),
args=self.stdio_config.get("args", []),
env=self.stdio_config.get("env", {}),
)
transport_ctx = stdio_client(server_params)
elif self.transport_type == MCPTransport.sse:
headers = self._get_auth_headers()
httpx_client_factory = self._create_httpx_client_factory()
transport_ctx = sse_client(
url=self.server_url,
timeout=self.timeout,
headers=headers,
httpx_client_factory=httpx_client_factory,
)
else:
headers = self._get_auth_headers()
httpx_client_factory = self._create_httpx_client_factory()
verbose_logger.debug(
"litellm headers for streamable_http_client: %s", headers
)
http_client = httpx_client_factory(
headers=headers,
timeout=httpx.Timeout(self.timeout),
)
transport_ctx = streamable_http_client(
url=self.server_url,
http_client=http_client,
)
if transport_ctx is None:
raise RuntimeError("Failed to create transport context")
transport_ctx, http_client = self._create_transport_context()
# Enter transport context
transport = await transport_ctx.__aenter__()