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

tools/ai/providers/openai_provider.py

@@ -0,0 +1,204 @@
+"""OpenAI AI provider implementation."""
+
+import json
+from typing import Any, Dict, List
+
+import mcp.types as mcp_types
+import openai
+
+from tools.common import (
+    ConversationStep,
+    OperationStatus,
+    Result,
+    ToolCall,
+    get_logger,
+)
+
+from .base_provider import BaseAIProvider
+
+logger = get_logger("OpenAIProvider")
+
+
+class OpenAIProvider(BaseAIProvider):
+    """OpenAI AI provider implementation."""
+
+    def __init__(self, api_key: str):
+        """Initialize OpenAI provider."""
+        super().__init__(api_key)
+        self.client = openai.AsyncOpenAI(api_key=api_key)
+
+    @property
+    def provider_name(self) -> str:
+        """Get the provider name."""
+        return "openai"
+
+    @property
+    def default_model(self) -> str:
+        """Get the default model for this provider."""
+        return "gpt-4o"
+
+    async def generate_step(
+        self,
+        messages: List[Dict[str, Any]],
+        available_tools: Dict[str, tuple[mcp_types.Tool, str]],
+        step_number: int,
+    ) -> Result[ConversationStep]:
+        """Generate a single conversation step with potential tool calls."""
+        try:
+            # Build system prompt with tools context
+            tools_context = self.build_tools_context(available_tools)
+            system_prompt = f"""You are an AI assistant that can control multiple creative applications through MCP servers.
+
+{tools_context}
+
+You can use the available tools to interact with the connected servers. When you need to perform actions, use the appropriate tools. For conversational responses, respond normally with helpful information."""
+
+            # Format messages for OpenAI (includes system message in messages array)
+            formatted_messages = [
+                {"role": "system", "content": system_prompt}
+            ] + self.format_messages_for_api(messages)
+
+            # Build tools for OpenAI's function calling
+            openai_tools = []
+            for tool_name, (mcp_tool, server_name) in available_tools.items():
+                openai_tool = self.format_tool_for_api(mcp_tool, server_name)
+                openai_tools.append(openai_tool)
+
+            logger.debug(f"Making OpenAI API call with {len(openai_tools)} tools")
+
+            # Make API call with explicit parameters
+            if openai_tools:
+                response = await self.client.chat.completions.create(
+                    model=self.default_model,
+                    messages=formatted_messages,  # type: ignore
+                    max_tokens=4000,
+                    tools=openai_tools,  # type: ignore
+                    tool_choice="auto",
+                )
+            else:
+                response = await self.client.chat.completions.create(
+                    model=self.default_model,
+                    messages=formatted_messages,  # type: ignore
+                    max_tokens=4000,
+                )
+
+            # Create conversation step
+            step = ConversationStep(step_number=step_number)
+            message = response.choices[0].message
+
+            # Parse the response
+            if message.content:
+                step.text = message.content
+
+            if message.tool_calls:
+                # Parse OpenAI tool calls
+                tool_calls = self.parse_tool_calls(message)
+                for tool_call in tool_calls:
+                    step.add_tool_call(tool_call)
+
+            # Set finish reason if no tool calls and no text
+            if not step.tool_calls and not step.text:
+                step.finish_reason = "stop"
+
+            return Result(status=OperationStatus.SUCCESS, data=step)
+
+        except Exception as e:
+            logger.error("OpenAI API call failed", error=e)
+            return Result(
+                status=OperationStatus.FAILED,
+                error=f"OpenAI API error: {e}",
+                error_code="OPENAI_API_ERROR",
+            )
+
+    def build_tools_context(
+        self, available_tools: Dict[str, tuple[mcp_types.Tool, str]]
+    ) -> str:
+        """Build a context description of available tools."""
+        if not available_tools:
+            return "No connected servers or tools available."
+
+        tools_desc = []
+        tools_by_server: Dict[str, List[mcp_types.Tool]] = {}
+
+        # Group tools by server
+        for tool_name, (mcp_tool, server_name) in available_tools.items():
+            if server_name not in tools_by_server:
+                tools_by_server[server_name] = []
+            tools_by_server[server_name].append(mcp_tool)
+
+        # Build description
+        for server_name, server_tools in tools_by_server.items():
+            tools_desc.append(f"**{server_name} Server**:")
+            for tool in server_tools:
+                description = tool.description or "No description available"
+                tools_desc.append(f"  - {tool.name}: {description}")
+
+        return "Connected Servers and Available Tools:\n" + "\n".join(tools_desc)
+
+    def format_tool_for_api(
+        self, mcp_tool: mcp_types.Tool, server_name: str
+    ) -> Dict[str, Any]:
+        """Convert MCP tool to OpenAI tool format."""
+        return {
+            "type": "function",
+            "function": {
+                "name": mcp_tool.name,
+                "description": mcp_tool.description
+                or f"Call {mcp_tool.name} on {server_name} server",
+                "parameters": mcp_tool.inputSchema,
+            },
+        }
+
+    def parse_tool_calls(self, message: Any) -> List[ToolCall]:
+        """Parse tool calls from OpenAI response."""
+        tool_calls = []
+
+        if hasattr(message, "tool_calls") and message.tool_calls:
+            for tool_call in message.tool_calls:
+                try:
+                    arguments = json.loads(
+                        tool_call.function.arguments
+                        if hasattr(tool_call.function, "arguments")
+                        else "{}"
+                    )
+                except json.JSONDecodeError:
+                    arguments = {}
+
+                parsed_call = ToolCall(
+                    id=tool_call.id if hasattr(tool_call, "id") else "",
+                    tool_name=tool_call.function.name
+                    if hasattr(tool_call.function, "name")
+                    else "",
+                    arguments=arguments,
+                )
+                tool_calls.append(parsed_call)
+
+        return tool_calls
+
+    def format_messages_for_api(
+        self, messages: List[Dict[str, Any]]
+    ) -> List[Dict[str, Any]]:
+        """Format messages for OpenAI's API."""
+        # OpenAI accepts all message types in the messages array
+        return messages
+
+    def format_tool_results_for_api(
+        self, tool_calls: List[ToolCall], tool_results: List[Any]
+    ) -> List[Dict[str, Any]]:
+        """Format tool results for OpenAI's API."""
+        # OpenAI expects tool results as separate tool messages
+        tool_messages = []
+
+        for tool_call, result in zip(tool_calls, tool_results):
+            tool_message = {
+                "role": "tool",
+                "tool_call_id": tool_call.id,
+                "content": json.dumps(result.result)
+                if hasattr(result, "result") and result.result
+                else result.error
+                if hasattr(result, "error")
+                else "No result",
+            }
+            tool_messages.append(tool_message)
+
+        return tool_messages

tools/ai/tool_executor.py

@@ -0,0 +1,257 @@
+"""Tool executor for handling MCP tool calls with connection pooling."""
+
+import asyncio
+import json
+import time
+from typing import Any, Dict, List, Optional, Tuple
+
+import mcp.types as mcp_types
+
+from tools.common import (
+    ToolCall,
+    ToolExecutionError,
+    ToolResult,
+    ToolTimeoutError,
+    get_logger,
+    run_concurrent_operations,
+    with_correlation_id,
+)
+from tools.common.async_utils import ConnectionPool
+
+logger = get_logger("ToolExecutor")
+
+
+class ToolExecutor:
+    """Handles execution of MCP tool calls with connection pooling and concurrency."""
+
+    def __init__(self, max_concurrent: int = 5, tool_timeout: float = 30.0):
+        """Initialize the tool executor."""
+        self.max_concurrent = max_concurrent
+        self.tool_timeout = tool_timeout
+        self.available_tools: Dict[str, Tuple[mcp_types.Tool, str]] = {}
+        self.connection_pool: Optional[ConnectionPool] = None
+
+    async def update_tools(
+        self,
+        available_tools: Dict[str, Tuple[mcp_types.Tool, str]],
+        connection_pool: ConnectionPool,
+    ):
+        """Update the available tools and connection pool."""
+        self.available_tools = available_tools
+        self.connection_pool = connection_pool
+        logger.info(f"Updated tool executor with {len(available_tools)} tools")
+
+    @with_correlation_id
+    async def execute_tool(self, tool_call: ToolCall) -> ToolResult:
+        """Execute a single tool call."""
+        start_time = time.time()
+
+        # Check if tool exists
+        if tool_call.tool_name not in self.available_tools:
+            return ToolResult(
+                id=tool_call.id,
+                tool_name=tool_call.tool_name,
+                arguments=tool_call.arguments,
+                error=f"Tool {tool_call.tool_name} not found",
+                error_code="TOOL_NOT_FOUND",
+                duration_ms=(time.time() - start_time) * 1000,
+            )
+
+        mcp_tool, server_name = self.available_tools[tool_call.tool_name]
+
+        if not self.connection_pool:
+            return ToolResult(
+                id=tool_call.id,
+                tool_name=tool_call.tool_name,
+                arguments=tool_call.arguments,
+                error="Connection pool not available",
+                error_code="NO_CONNECTION_POOL",
+                duration_ms=(time.time() - start_time) * 1000,
+            )
+
+        try:
+            # Execute with timeout
+            result = await asyncio.wait_for(
+                self._execute_tool_with_connection(tool_call, server_name),
+                timeout=self.tool_timeout,
+            )
+
+            duration_ms = (time.time() - start_time) * 1000
+            result.duration_ms = duration_ms
+            result.server_name = server_name
+
+            logger.debug(
+                f"Tool {tool_call.tool_name} executed successfully",
+                tool_name=tool_call.tool_name,
+                server_name=server_name,
+                duration_ms=duration_ms,
+            )
+
+            return result
+
+        except asyncio.TimeoutError:
+            error = ToolTimeoutError(
+                f"Tool {tool_call.tool_name} timed out after {self.tool_timeout}s",
+                tool_name=tool_call.tool_name,
+                server_name=server_name,
+            )
+            logger.error("Tool execution timed out", error=error)
+            return ToolResult(
+                id=tool_call.id,
+                tool_name=tool_call.tool_name,
+                arguments=tool_call.arguments,
+                error=str(error),
+                error_code=error.error_code,
+                server_name=server_name,
+                duration_ms=(time.time() - start_time) * 1000,
+            )
+
+        except Exception as e:
+            execution_error = ToolExecutionError(
+                f"Error executing tool {tool_call.tool_name}: {e}",
+                tool_name=tool_call.tool_name,
+                server_name=server_name,
+                cause=e,
+            )
+            logger.error("Tool execution failed", error=execution_error)
+            return ToolResult(
+                id=tool_call.id,
+                tool_name=tool_call.tool_name,
+                arguments=tool_call.arguments,
+                error=str(execution_error),
+                error_code=execution_error.error_code,
+                server_name=server_name,
+                duration_ms=(time.time() - start_time) * 1000,
+            )
+
+    async def _execute_tool_with_connection(
+        self, tool_call: ToolCall, server_name: str
+    ) -> ToolResult:
+        """Execute tool using connection pool."""
+        if not self.connection_pool:
+            raise ToolExecutionError(
+                "Connection pool not available",
+                tool_name=tool_call.tool_name,
+                server_name=server_name,
+            )
+        async with self.connection_pool.get_connection(server_name) as client:
+            # Call the tool
+            mcp_result = await client.call_tool(
+                tool_call.tool_name, tool_call.arguments
+            )
+
+            # Convert MCP result to our format
+            result = self._convert_mcp_result(mcp_result, tool_call)
+            return result
+
+    def _convert_mcp_result(self, mcp_result: Any, tool_call: ToolCall) -> ToolResult:
+        """Convert MCP tool call result to our ToolResult format."""
+        result = ToolResult(
+            id=tool_call.id,
+            tool_name=tool_call.tool_name,
+            arguments=tool_call.arguments,
+        )
+
+        # Handle different response formats
+        if hasattr(mcp_result, "content"):
+            content = mcp_result.content
+        elif isinstance(mcp_result, list):
+            content = mcp_result
+        else:
+            content = [mcp_result] if mcp_result else []
+
+        if content and len(content) > 0:
+            # Handle different content types safely
+            first_content = content[0]
+            if hasattr(first_content, "text"):
+                try:
+                    # Try to parse as JSON first
+                    result.result = json.loads(first_content.text)
+                except json.JSONDecodeError:
+                    # If not JSON, store as text
+                    result.result = first_content.text
+            else:
+                # Handle other content types
+                result.result = {
+                    "type": type(first_content).__name__,
+                    "content": str(first_content),
+                }
+        else:
+            result.error = "No result returned"
+            result.error_code = "EMPTY_RESULT"
+
+        return result
+
+    async def execute_tools_concurrently(
+        self, tool_calls: List[ToolCall]
+    ) -> List[ToolResult]:
+        """Execute multiple tool calls concurrently."""
+        if not tool_calls:
+            return []
+
+        logger.info(f"Executing {len(tool_calls)} tools concurrently")
+
+        # Create operations for concurrent execution
+        async def execute_single_tool(tool_call: ToolCall) -> ToolResult:
+            return await self.execute_tool(tool_call)
+
+        operations = [
+            lambda tc=tool_call: execute_single_tool(tc) for tool_call in tool_calls
+        ]
+        operation_names = [f"tool_{tc.tool_name}" for tc in tool_calls]
+
+        # Execute concurrently
+        results = await run_concurrent_operations(
+            operations,
+            max_concurrent=self.max_concurrent,
+            operation_names=operation_names,
+        )
+
+        # Extract the actual ToolResult objects from Result wrappers
+        tool_results = []
+        for result in results:
+            if result.is_success:
+                tool_results.append(result.data)
+            else:
+                # Create error ToolResult for failed operations
+                tool_results.append(
+                    ToolResult(
+                        id="",
+                        tool_name="unknown",
+                        arguments={},
+                        error=result.error,
+                        error_code=result.error_code,
+                    )
+                )
+
+        successful = sum(1 for tr in tool_results if not tr.error)
+        logger.info(
+            f"Tool execution completed: {successful}/{len(tool_calls)} successful"
+        )
+
+        return tool_results
+
+    def get_available_tools(self) -> Dict[str, str]:
+        """Get available tools mapped to their server names."""
+        return {
+            tool_name: server_name
+            for tool_name, (_, server_name) in self.available_tools.items()
+        }
+
+    def get_tool_info(self, tool_name: str) -> Optional[Tuple[mcp_types.Tool, str]]:
+        """Get information about a specific tool."""
+        return self.available_tools.get(tool_name)
+
+    def validate_tool_call(self, tool_call: ToolCall) -> Tuple[bool, Optional[str]]:
+        """Validate a tool call before execution."""
+        if tool_call.tool_name not in self.available_tools:
+            return False, f"Tool {tool_call.tool_name} not found"
+
+        if not tool_call.id:
+            return False, "Tool call ID is required"
+
+        if not isinstance(tool_call.arguments, dict):
+            return False, "Tool arguments must be a dictionary"
+
+        # Could add more validation here (e.g., schema validation)
+        return True, None

tools/common/__init__.py

@@ -0,0 +1,99 @@
+"""Common utilities and types for the tools package."""
+
+# Import shared types from common module
+try:
+    from common import (
+        HealthStatus,
+        OperationStatus,
+        ServerInfo,
+        ServerState,
+        ToolCall,
+        ToolCallState,
+        ToolResult,
+    )
+except ImportError:
+    # Fallback for development/testing
+    import sys
+    from pathlib import Path
+
+    sys.path.insert(0, str(Path(__file__).parent.parent.parent))
+    from common import (
+        HealthStatus,
+        OperationStatus,
+        ServerInfo,
+        ServerState,
+        ToolCall,
+        ToolCallState,
+        ToolResult,
+    )
+
+from .async_utils import (
+    ConnectionPool,
+    RetryManager,
+    get_connection_pool,
+    run_concurrent_operations,
+    shutdown_connection_pool,
+)
+from .errors import (
+    AIProviderError,
+    ConfigurationError,
+    ConversationError,
+    LightfastMCPError,
+    ServerConnectionError,
+    ServerError,
+    ServerStartupError,
+    ToolExecutionError,
+    ToolNotFoundError,
+    ToolTimeoutError,
+    ValidationError,
+)
+from .logging import (
+    correlation_id,
+    get_logger,
+    operation_context,
+    with_correlation_id,
+    with_operation_context,
+)
+from .types import (
+    ConversationResult,
+    ConversationStep,
+    Result,
+)
+
+__all__ = [
+    # Types
+    "OperationStatus",
+    "ServerState",
+    "Result",
+    "ServerInfo",
+    "ConversationResult",
+    "ConversationStep",
+    "ToolCall",
+    "ToolCallState",
+    "ToolResult",
+    "HealthStatus",
+    # Errors
+    "LightfastMCPError",
+    "ConfigurationError",
+    "ConversationError",
+    "ServerError",
+    "ServerStartupError",
+    "ServerConnectionError",
+    "ToolExecutionError",
+    "ToolNotFoundError",
+    "ToolTimeoutError",
+    "AIProviderError",
+    "ValidationError",
+    # Logging
+    "get_logger",
+    "with_correlation_id",
+    "with_operation_context",
+    "correlation_id",
+    "operation_context",
+    # Async utilities
+    "ConnectionPool",
+    "RetryManager",
+    "get_connection_pool",
+    "run_concurrent_operations",
+    "shutdown_connection_pool",
+]