refactor: modernize type hints and improve async context handling
- Update type hints to use Python 3.10 syntax (dict, list, X | None) - Replace requests with httpx for HTTP client consistency - Improve async context management using AsyncExitStack - Simplify server cleanup method
This commit is contained in:
parent
a0216c3e50
commit
466e1e8eb7
1 changed files with 30 additions and 48 deletions
|
|
@ -3,9 +3,10 @@ import json
|
||||||
import logging
|
import logging
|
||||||
import os
|
import os
|
||||||
import shutil
|
import shutil
|
||||||
from typing import Any, Dict, List, Optional
|
from contextlib import AsyncExitStack
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
import requests
|
import httpx
|
||||||
from dotenv import load_dotenv
|
from dotenv import load_dotenv
|
||||||
from mcp import ClientSession, StdioServerParameters
|
from mcp import ClientSession, StdioServerParameters
|
||||||
from mcp.client.stdio import stdio_client
|
from mcp.client.stdio import stdio_client
|
||||||
|
|
@ -30,7 +31,7 @@ class Configuration:
|
||||||
load_dotenv()
|
load_dotenv()
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def load_config(file_path: str) -> Dict[str, Any]:
|
def load_config(file_path: str) -> dict[str, Any]:
|
||||||
"""Load server configuration from JSON file.
|
"""Load server configuration from JSON file.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
|
|
@ -64,12 +65,13 @@ class Configuration:
|
||||||
class Server:
|
class Server:
|
||||||
"""Manages MCP server connections and tool execution."""
|
"""Manages MCP server connections and tool execution."""
|
||||||
|
|
||||||
def __init__(self, name: str, config: Dict[str, Any]) -> None:
|
def __init__(self, name: str, config: dict[str, Any]) -> None:
|
||||||
self.name: str = name
|
self.name: str = name
|
||||||
self.config: Dict[str, Any] = config
|
self.config: dict[str, Any] = config
|
||||||
self.stdio_context: Optional[Any] = None
|
self.stdio_context: Any | None = None
|
||||||
self.session: Optional[ClientSession] = None
|
self.session: ClientSession | None = None
|
||||||
self._cleanup_lock: asyncio.Lock = asyncio.Lock()
|
self._cleanup_lock: asyncio.Lock = asyncio.Lock()
|
||||||
|
self.exit_stack: AsyncExitStack = AsyncExitStack()
|
||||||
|
|
||||||
async def initialize(self) -> None:
|
async def initialize(self) -> None:
|
||||||
"""Initialize the server connection."""
|
"""Initialize the server connection."""
|
||||||
|
|
@ -89,17 +91,16 @@ class Server:
|
||||||
else None,
|
else None,
|
||||||
)
|
)
|
||||||
try:
|
try:
|
||||||
self.stdio_context = stdio_client(server_params)
|
stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
|
||||||
read, write = await self.stdio_context.__aenter__()
|
read, write = stdio_transport
|
||||||
self.session = ClientSession(read, write)
|
self.session = await self.exit_stack.enter_async_context(ClientSession(read, write))
|
||||||
await self.session.__aenter__()
|
|
||||||
await self.session.initialize()
|
await self.session.initialize()
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logging.error(f"Error initializing server {self.name}: {e}")
|
logging.error(f"Error initializing server {self.name}: {e}")
|
||||||
await self.cleanup()
|
await self.cleanup()
|
||||||
raise
|
raise
|
||||||
|
|
||||||
async def list_tools(self) -> List[Any]:
|
async def list_tools(self) -> list[Any]:
|
||||||
"""List available tools from the server.
|
"""List available tools from the server.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
|
|
@ -124,7 +125,7 @@ class Server:
|
||||||
async def execute_tool(
|
async def execute_tool(
|
||||||
self,
|
self,
|
||||||
tool_name: str,
|
tool_name: str,
|
||||||
arguments: Dict[str, Any],
|
arguments: dict[str, Any],
|
||||||
retries: int = 2,
|
retries: int = 2,
|
||||||
delay: float = 1.0,
|
delay: float = 1.0,
|
||||||
) -> Any:
|
) -> Any:
|
||||||
|
|
@ -170,28 +171,8 @@ class Server:
|
||||||
"""Clean up server resources."""
|
"""Clean up server resources."""
|
||||||
async with self._cleanup_lock:
|
async with self._cleanup_lock:
|
||||||
try:
|
try:
|
||||||
if self.session:
|
await self.exit_stack.aclose()
|
||||||
try:
|
|
||||||
await self.session.__aexit__(None, None, None)
|
|
||||||
except Exception as e:
|
|
||||||
logging.warning(
|
|
||||||
f"Warning during session cleanup for {self.name}: {e}"
|
|
||||||
)
|
|
||||||
finally:
|
|
||||||
self.session = None
|
self.session = None
|
||||||
|
|
||||||
if self.stdio_context:
|
|
||||||
try:
|
|
||||||
await self.stdio_context.__aexit__(None, None, None)
|
|
||||||
except (RuntimeError, asyncio.CancelledError) as e:
|
|
||||||
logging.info(
|
|
||||||
f"Note: Normal shutdown message for {self.name}: {e}"
|
|
||||||
)
|
|
||||||
except Exception as e:
|
|
||||||
logging.warning(
|
|
||||||
f"Warning during stdio cleanup for {self.name}: {e}"
|
|
||||||
)
|
|
||||||
finally:
|
|
||||||
self.stdio_context = None
|
self.stdio_context = None
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logging.error(f"Error during cleanup of server {self.name}: {e}")
|
logging.error(f"Error during cleanup of server {self.name}: {e}")
|
||||||
|
|
@ -201,11 +182,11 @@ class Tool:
|
||||||
"""Represents a tool with its properties and formatting."""
|
"""Represents a tool with its properties and formatting."""
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self, name: str, description: str, input_schema: Dict[str, Any]
|
self, name: str, description: str, input_schema: dict[str, Any]
|
||||||
) -> None:
|
) -> None:
|
||||||
self.name: str = name
|
self.name: str = name
|
||||||
self.description: str = description
|
self.description: str = description
|
||||||
self.input_schema: Dict[str, Any] = input_schema
|
self.input_schema: dict[str, Any] = input_schema
|
||||||
|
|
||||||
def format_for_llm(self) -> str:
|
def format_for_llm(self) -> str:
|
||||||
"""Format tool information for LLM.
|
"""Format tool information for LLM.
|
||||||
|
|
@ -237,7 +218,7 @@ class LLMClient:
|
||||||
def __init__(self, api_key: str) -> None:
|
def __init__(self, api_key: str) -> None:
|
||||||
self.api_key: str = api_key
|
self.api_key: str = api_key
|
||||||
|
|
||||||
def get_response(self, messages: List[Dict[str, str]]) -> str:
|
def get_response(self, messages: list[dict[str, str]]) -> str:
|
||||||
"""Get a response from the LLM.
|
"""Get a response from the LLM.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
|
|
@ -247,7 +228,7 @@ class LLMClient:
|
||||||
The LLM's response as a string.
|
The LLM's response as a string.
|
||||||
|
|
||||||
Raises:
|
Raises:
|
||||||
RequestException: If the request to the LLM fails.
|
httpx.RequestError: If the request to the LLM fails.
|
||||||
"""
|
"""
|
||||||
url = "https://api.groq.com/openai/v1/chat/completions"
|
url = "https://api.groq.com/openai/v1/chat/completions"
|
||||||
|
|
||||||
|
|
@ -266,16 +247,17 @@ class LLMClient:
|
||||||
}
|
}
|
||||||
|
|
||||||
try:
|
try:
|
||||||
response = requests.post(url, headers=headers, json=payload)
|
with httpx.Client() as client:
|
||||||
|
response = client.post(url, headers=headers, json=payload)
|
||||||
response.raise_for_status()
|
response.raise_for_status()
|
||||||
data = response.json()
|
data = response.json()
|
||||||
return data["choices"][0]["message"]["content"]
|
return data["choices"][0]["message"]["content"]
|
||||||
|
|
||||||
except requests.exceptions.RequestException as e:
|
except httpx.RequestError as e:
|
||||||
error_message = f"Error getting LLM response: {str(e)}"
|
error_message = f"Error getting LLM response: {str(e)}"
|
||||||
logging.error(error_message)
|
logging.error(error_message)
|
||||||
|
|
||||||
if e.response is not None:
|
if hasattr(e, 'response'):
|
||||||
status_code = e.response.status_code
|
status_code = e.response.status_code
|
||||||
logging.error(f"Status code: {status_code}")
|
logging.error(f"Status code: {status_code}")
|
||||||
logging.error(f"Response details: {e.response.text}")
|
logging.error(f"Response details: {e.response.text}")
|
||||||
|
|
@ -289,8 +271,8 @@ class LLMClient:
|
||||||
class ChatSession:
|
class ChatSession:
|
||||||
"""Orchestrates the interaction between user, LLM, and tools."""
|
"""Orchestrates the interaction between user, LLM, and tools."""
|
||||||
|
|
||||||
def __init__(self, servers: List[Server], llm_client: LLMClient) -> None:
|
def __init__(self, servers: list[Server], llm_client: LLMClient) -> None:
|
||||||
self.servers: List[Server] = servers
|
self.servers: list[Server] = servers
|
||||||
self.llm_client: LLMClient = llm_client
|
self.llm_client: LLMClient = llm_client
|
||||||
|
|
||||||
async def cleanup_servers(self) -> None:
|
async def cleanup_servers(self) -> None:
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue