mseep-lightfast-mcp 0.0.1__py3-none-any.whl

tools/ai/conversation_client.py

@@ -0,0 +1,399 @@
+"""Conversation client for managing AI conversations across multiple MCP servers."""
+
+import os
+import uuid
+from typing import Any, Dict, List, Optional
+
+import mcp.types as mcp_types
+
+from tools.common import (
+    AIProviderError,
+    ConversationResult,
+    ConversationStep,
+    OperationStatus,
+    Result,
+    ToolCall,
+    ToolResult,
+    get_connection_pool,
+    get_logger,
+    with_correlation_id,
+    with_operation_context,
+)
+from tools.common.async_utils import ConnectionPool
+
+from .conversation_session import ConversationSession
+from .providers.base_provider import BaseAIProvider
+from .providers.claude_provider import ClaudeProvider
+from .providers.openai_provider import OpenAIProvider
+from .tool_executor import ToolExecutor
+
+logger = get_logger("ConversationClient")
+
+
+class ConversationClient:
+    """Manages AI conversations across multiple MCP servers."""
+
+    def __init__(
+        self,
+        servers: Dict[str, Dict[str, Any]],
+        ai_provider: str = "claude",
+        api_key: Optional[str] = None,
+        max_steps: int = 5,
+        max_concurrent_tools: int = 5,
+    ):
+        """Initialize the conversation client."""
+        self.servers = servers
+        self.ai_provider_name = ai_provider.lower()
+        self.api_key = api_key or self._get_api_key()
+        self.max_steps = max_steps
+        self.max_concurrent_tools = max_concurrent_tools
+
+        # Initialize components
+        self.ai_provider = self._create_ai_provider()
+        self.tool_executor = ToolExecutor(max_concurrent=max_concurrent_tools)
+        self.connection_pool: Optional[ConnectionPool] = None
+
+        # Server and tool tracking
+        self.connected_servers: Dict[str, Dict[str, Any]] = {}
+        self.available_tools: Dict[str, tuple[mcp_types.Tool, str]] = {}
+
+        # Active sessions
+        self.active_sessions: Dict[str, ConversationSession] = {}
+
+    def _get_api_key(self) -> str:
+        """Get API key from environment variables."""
+        if self.ai_provider_name == "claude":
+            key = os.getenv("ANTHROPIC_API_KEY")
+            if not key:
+                raise AIProviderError(
+                    "ANTHROPIC_API_KEY environment variable required for Claude",
+                    provider="claude",
+                    error_code="MISSING_API_KEY",
+                )
+        elif self.ai_provider_name == "openai":
+            key = os.getenv("OPENAI_API_KEY")
+            if not key:
+                raise AIProviderError(
+                    "OPENAI_API_KEY environment variable required for OpenAI",
+                    provider="openai",
+                    error_code="MISSING_API_KEY",
+                )
+        else:
+            raise AIProviderError(
+                f"Unsupported AI provider: {self.ai_provider_name}",
+                provider=self.ai_provider_name,
+                error_code="UNSUPPORTED_PROVIDER",
+            )
+        return key
+
+    def _create_ai_provider(self) -> BaseAIProvider:
+        """Create the appropriate AI provider."""
+        if self.ai_provider_name == "claude":
+            return ClaudeProvider(api_key=self.api_key)
+        elif self.ai_provider_name == "openai":
+            return OpenAIProvider(api_key=self.api_key)
+        else:
+            raise AIProviderError(
+                f"Unsupported AI provider: {self.ai_provider_name}",
+                provider=self.ai_provider_name,
+                error_code="UNSUPPORTED_PROVIDER",
+            )
+
+    @with_correlation_id
+    @with_operation_context(operation="connect_to_servers")
+    async def connect_to_servers(self) -> Result[Dict[str, bool]]:
+        """Connect to all configured servers."""
+        self.connection_pool = await get_connection_pool()
+        connection_results = {}
+
+        for server_name, server_config in self.servers.items():
+            try:
+                logger.info(f"Connecting to {server_name}")
+
+                # Register server with connection pool
+                if self.connection_pool is not None:
+                    await self.connection_pool.register_server(
+                        server_name, server_config
+                    )
+
+                    # Test connection by getting tools
+                    async with self.connection_pool.get_connection(
+                        server_name
+                    ) as client:
+                        tools_result = await client.list_tools()
+
+                        # Handle different response formats
+                        if hasattr(tools_result, "tools"):
+                            mcp_tools = tools_result.tools
+                        elif isinstance(tools_result, list):
+                            mcp_tools = tools_result
+                        else:
+                            mcp_tools = []
+
+                        # Store tools
+                        for mcp_tool in mcp_tools:
+                            self.available_tools[mcp_tool.name] = (
+                                mcp_tool,
+                                server_name,
+                            )
+                            logger.debug(
+                                f"Added tool {mcp_tool.name} from {server_name}"
+                            )
+
+                    self.connected_servers[server_name] = server_config
+                    connection_results[server_name] = True
+                    logger.info(f"Successfully connected to {server_name}")
+                else:
+                    logger.error("Connection pool is None")
+                    connection_results[server_name] = False
+
+            except Exception as e:
+                logger.error(f"Failed to connect to {server_name}", error=e)
+                connection_results[server_name] = False
+
+        # Update tool executor with available tools
+        if self.connection_pool is not None:
+            await self.tool_executor.update_tools(
+                self.available_tools, self.connection_pool
+            )
+
+        successful_connections = sum(
+            1 for success in connection_results.values() if success
+        )
+        logger.info(
+            f"Connected to {successful_connections}/{len(self.servers)} servers"
+        )
+
+        return Result(status=OperationStatus.SUCCESS, data=connection_results)
+
+    @with_correlation_id
+    @with_operation_context(operation="start_conversation")
+    async def start_conversation(
+        self,
+        initial_message: Optional[str] = None,
+        max_steps: Optional[int] = None,
+        session_id: Optional[str] = None,
+    ) -> Result[ConversationSession]:
+        """Start a new conversation session."""
+        if session_id is None:
+            session_id = str(uuid.uuid4())
+
+        if session_id in self.active_sessions:
+            return Result(
+                status=OperationStatus.FAILED,
+                error=f"Session {session_id} already exists",
+                error_code="SESSION_EXISTS",
+            )
+
+        session = ConversationSession(
+            session_id=session_id,
+            max_steps=max_steps or self.max_steps,
+            ai_provider=self.ai_provider,
+            tool_executor=self.tool_executor,
+            available_tools=self.available_tools,
+        )
+
+        self.active_sessions[session_id] = session
+
+        # If initial message provided, process it
+        if initial_message:
+            result = await session.process_message(initial_message)
+            if not result.is_success:
+                # Clean up session on failure
+                del self.active_sessions[session_id]
+                return Result(
+                    status=OperationStatus.FAILED,
+                    error=f"Failed to process initial message: {result.error}",
+                    error_code="INITIAL_MESSAGE_FAILED",
+                )
+
+        logger.info(f"Started conversation session {session_id}")
+        return Result(status=OperationStatus.SUCCESS, data=session)
+
+    @with_correlation_id
+    @with_operation_context(operation="chat")
+    async def chat(
+        self,
+        message: str,
+        session_id: Optional[str] = None,
+        max_steps: Optional[int] = None,
+    ) -> Result[ConversationResult]:
+        """Send a message and get a complete conversation result."""
+        # Create session if none provided
+        if session_id is None:
+            session_result = await self.start_conversation(max_steps=max_steps)
+            if not session_result.is_success:
+                return Result(
+                    status=OperationStatus.FAILED,
+                    error=f"Failed to create session: {session_result.error}",
+                    error_code="SESSION_CREATION_FAILED",
+                )
+            session = session_result.data
+            session_id = session.session_id
+        else:
+            session = self.active_sessions.get(session_id)
+            if not session:
+                return Result(
+                    status=OperationStatus.FAILED,
+                    error=f"Session {session_id} not found",
+                    error_code="SESSION_NOT_FOUND",
+                )
+
+        # Process the message
+        result = await session.process_message(message)
+        if not result.is_success:
+            return Result(
+                status=OperationStatus.FAILED,
+                error=f"Failed to process message: {result.error}",
+                error_code="MESSAGE_PROCESSING_FAILED",
+            )
+
+        # Create conversation result
+        conversation_result = ConversationResult(
+            session_id=session_id,
+            steps=session.steps.copy(),
+            total_duration_ms=sum(step.duration_ms or 0 for step in session.steps),
+        )
+
+        return Result(status=OperationStatus.SUCCESS, data=conversation_result)
+
+    @with_correlation_id
+    async def continue_conversation(
+        self, session_id: str, message: str
+    ) -> Result[ConversationResult]:
+        """Continue an existing conversation."""
+        return await self.chat(message, session_id=session_id)
+
+    async def get_conversation_history(
+        self, session_id: str
+    ) -> Result[List[ConversationStep]]:
+        """Get the conversation history for a session."""
+        session = self.active_sessions.get(session_id)
+        if not session:
+            return Result(
+                status=OperationStatus.FAILED,
+                error=f"Session {session_id} not found",
+                error_code="SESSION_NOT_FOUND",
+            )
+
+        return Result(status=OperationStatus.SUCCESS, data=session.steps.copy())
+
+    async def execute_tools(
+        self, tool_calls: List[ToolCall]
+    ) -> Result[List[ToolResult]]:
+        """Execute a list of tool calls."""
+        if not tool_calls:
+            return Result(status=OperationStatus.SUCCESS, data=[])
+
+        results = await self.tool_executor.execute_tools_concurrently(tool_calls)
+
+        return Result(status=OperationStatus.SUCCESS, data=results)
+
+    def get_connected_servers(self) -> List[str]:
+        """Get list of connected server names."""
+        return list(self.connected_servers.keys())
+
+    def get_available_tools(self) -> Dict[str, List[str]]:
+        """Get all available tools organized by server."""
+        tools_by_server: Dict[str, List[str]] = {}
+        for tool_name, (mcp_tool, server_name) in self.available_tools.items():
+            if server_name not in tools_by_server:
+                tools_by_server[server_name] = []
+            tools_by_server[server_name].append(tool_name)
+        return tools_by_server
+
+    def find_tool_server(self, tool_name: str) -> Optional[str]:
+        """Find which server has a specific tool."""
+        if tool_name in self.available_tools:
+            return self.available_tools[tool_name][1]
+        return None
+
+    def get_server_status(self) -> Dict[str, Dict[str, Any]]:
+        """Get status information for all servers."""
+        status = {}
+        for server_name in self.connected_servers:
+            server_tools = [
+                tool
+                for tool, (_, srv) in self.available_tools.items()
+                if srv == server_name
+            ]
+            status[server_name] = {
+                "connected": True,
+                "tools_count": len(server_tools),
+                "tools": server_tools,
+            }
+        return status
+
+    def get_active_sessions(self) -> Dict[str, ConversationSession]:
+        """Get all active conversation sessions."""
+        return self.active_sessions.copy()
+
+    async def close_session(self, session_id: str) -> Result[None]:
+        """Close a conversation session."""
+        if session_id not in self.active_sessions:
+            return Result(
+                status=OperationStatus.FAILED,
+                error=f"Session {session_id} not found",
+                error_code="SESSION_NOT_FOUND",
+            )
+
+        session = self.active_sessions[session_id]
+        await session.close()
+        del self.active_sessions[session_id]
+
+        logger.info(f"Closed conversation session {session_id}")
+        return Result(status=OperationStatus.SUCCESS)
+
+    async def disconnect_from_servers(self) -> Result[None]:
+        """Disconnect from all servers and clean up resources."""
+        logger.info("Disconnecting from all servers...")
+
+        # Close all active sessions
+        for session_id in list(self.active_sessions.keys()):
+            await self.close_session(session_id)
+
+        # Clear server and tool data
+        self.connected_servers.clear()
+        self.available_tools.clear()
+
+        # Shutdown connection pool if we have one
+        if self.connection_pool:
+            from tools.common.async_utils import shutdown_connection_pool
+
+            await shutdown_connection_pool()
+            self.connection_pool = None
+
+        logger.info("Disconnected from all servers")
+        return Result(status=OperationStatus.SUCCESS)
+
+
+async def create_conversation_client(
+    servers: Dict[str, Dict[str, Any]],
+    ai_provider: str = "claude",
+    api_key: Optional[str] = None,
+    max_steps: int = 5,
+) -> Result[ConversationClient]:
+    """Create and connect a conversation client from configuration."""
+    try:
+        client = ConversationClient(
+            servers=servers,
+            ai_provider=ai_provider,
+            api_key=api_key,
+            max_steps=max_steps,
+        )
+
+        connection_result = await client.connect_to_servers()
+        if not connection_result.is_success:
+            return Result(
+                status=OperationStatus.FAILED,
+                error=f"Failed to connect to servers: {connection_result.error}",
+                error_code="CONNECTION_FAILED",
+            )
+
+        return Result(status=OperationStatus.SUCCESS, data=client)
+
+    except Exception as e:
+        logger.error("Failed to create conversation client", error=e)
+        return Result(
+            status=OperationStatus.FAILED, error=str(e), error_code=type(e).__name__
+        )