pygeai 0.1.6__py3-none-any.whl → 0.6.0b15__py3-none-any.whl

pygeai/proxy/managers.py

@@ -0,0 +1,232 @@
+from contextlib import AsyncExitStack
+from typing import Any, List, Dict, Tuple
+import asyncio
+import json
+import sys
+import requests
+from urllib3.exceptions import MaxRetryError
+import mcp.types as types
+import a2a.types as a2a_types
+from pygeai.proxy.servers import MCPServer, ToolServer, ProxiedTool, A2AServer
+from pygeai.proxy.config import ProxySettingsManager
+from pygeai.proxy.clients import ProxyClient, ToolProxyData, ToolProxyJobResult
+from pygeai.core.utils.console import Console
+
+class ServerManager:
+    """
+    Manages multiple MCP servers.
+
+    :param servers_cfg: List[Dict[str, Any]] - List of server configurations
+    :param settings: ProxySettingsManager - Proxy settings manager
+    """
+
+    def __init__(self, servers_cfg: List[Dict[str, Any]], settings: ProxySettingsManager):
+        """
+        Initialize the server manager.
+
+        :param servers_cfg: List[Dict[str, Any]] - List of server configurations
+        :param settings: ProxySettingsManager - Proxy settings manager
+        """
+        self.servers_cfg = servers_cfg
+        self.settings = settings
+        self.servers: Dict[str, ToolServer] = {}
+        self.tools: Dict[str, ProxiedTool] = {}
+        self.exit_stack: AsyncExitStack = AsyncExitStack()
+
+    async def _initialize_servers(self) -> None:
+        """
+        Initialize all servers.
+
+        :return: None
+        :raises: Exception - If server initialization fails
+        """
+        for server_cfg in self.servers_cfg:
+            Console.write_stdout(f"Initializing server {server_cfg['name']} type: {server_cfg['type']}")
+            if server_cfg['type'] == 'mcp':
+                server = MCPServer(server_cfg['name'], server_cfg, self.settings)
+            elif server_cfg['type'] == 'a2a':
+                server = A2AServer(server_cfg['name'], server_cfg, self.settings)
+            else:
+                raise ValueError(f"Invalid server type: {server_cfg['type']}")
+            try:
+                await self.exit_stack.enter_async_context(server.exit_stack)
+                await server.initialize()
+                self.servers[server.name] = server
+                Console.write_stdout(f"\nServer {server.name} initialized successfully", end="\n\n")        
+            except Exception as e:
+                Console.write_exception(f"Failed to initialize server {server.name}:", e)
+                raise
+
+        for server in self.servers.values():
+            Console.write_stdout(f"Listing tools for server {server.name}", end="")
+            if server.public_prefix:
+                Console.write_stdout(f" | access scope:public prefix: {server.public_prefix}")
+            else:
+                Console.write_stdout(" ! access scope:private")
+
+            tools = await server.list_tools()
+            for tool in tools:
+                self.tools[tool.get_full_name()] = tool
+                Console.write_stdout(f"\tTool {tool.get_full_name()} added to server {server.name}")
+
+    async def _initialize_client(self) -> ProxyClient:
+        """
+        Initialize the client.
+
+        :return: ProxyClient - Initialized client instance
+        :raises: ConnectionError - If connection fails
+        :raises: MaxRetryError - If max retries are exceeded
+        """
+        try:
+            alias = self.settings.get_current_alias()
+            client = ProxyClient(self.settings.get_api_key(alias), self.settings.get_base_url(alias), self.settings.get_proxy_id(alias))
+            Console.write_stdout(f"\nRegistering proxy {self.settings.get_proxy_id(alias)} with name {self.settings.get_proxy_name(alias)} and description {self.settings.get_proxy_description(alias)}")
+            result = client.register(proxy_data=ToolProxyData(
+                id=self.settings.get_proxy_id(alias),
+                name=self.settings.get_proxy_name(alias),
+                description=self.settings.get_proxy_description(alias),
+                affinity=self.settings.get_proxy_affinity(alias),
+                tools=list(self.tools.values())
+            ))
+            Console.write_stdout("----------------------------------")
+            Console.write_stdout(f"Proxy registered successfully:")
+            if isinstance(result, dict) and isinstance(result.get("Messages"), list):
+                for message in result["Messages"]:
+                    description = message.get("Description", "")
+                    message_type = message.get("Type", None)
+
+                    if message_type == 1:
+                        Console.write_stderr(description)
+                    elif message_type == 2:
+                        Console.write_stdout(description)
+                Console.write_stdout("----------------------------------")
+            return client
+        except (ConnectionError, MaxRetryError):
+            Console.write_exception(f"Error registering proxy {self.settings.get_proxy_id(alias)}:")
+            raise
+    
+    async def execute_tool(
+        self,
+        server_name: str,
+        tool_name: str,
+        arguments: dict[str, Any],
+        retries: int = 2,
+        delay: float = 1.0,
+    ) -> Any:
+        """
+        Execute a tool with retry mechanism.
+
+        :param server_name: str - Name of the server to execute the tool on
+        :param tool_name: str - Name of the tool to execute
+        :param arguments: dict[str, Any] - Tool arguments
+        :param retries: int - Number of retry attempts
+        :param delay: float - Delay between retries in seconds
+        :return: Any - Tool execution result
+        :raises: RuntimeError - If server is not found or not initialized
+        :raises: Exception - If tool execution fails after all retries
+        """
+        if server_name not in self.servers:
+            raise RuntimeError(f"Server {server_name} not found")
+        
+        if tool_name not in self.tools:
+            raise RuntimeError(f"Tool {tool_name} not found")
+            
+        server = self.servers[server_name]
+        
+        try:
+            result = await server.execute_tool(self.tools[tool_name].name, arguments, retries, delay)
+            return result
+        except (RuntimeError, ConnectionError, TimeoutError) as e:
+            raise Exception(f"Failed to execute tool {tool_name} on server {server_name}: {e}") from e
+
+    def extract_function_call_info(self, raw_json: str) -> Tuple[str, str]:
+        """
+        Extract function call info from raw JSON.
+
+        :param raw_json: str - Raw JSON string
+        :return: Tuple[str, str] - Tuple containing function name and arguments
+        """
+        try:
+            data = json.loads(raw_json)
+            return data['function']['name'], data['function']['arguments']
+        except (json.JSONDecodeError, KeyError) as e:
+            Console.write_stdout(f"Error extracting function call info: {e}")
+            return None, None
+        
+    async def start(self) -> None:
+        """
+        Main proxy session handler.
+
+        :return: None
+        """
+        retry_count = 0
+        MAX_RETRIES = 10
+        while retry_count < MAX_RETRIES:
+            try:
+                await self._initialize_servers()
+                try:
+                    client = await self._initialize_client()
+                except (ConnectionError, TimeoutError, RuntimeError) as e:
+                    Console.write_exception(f"Error during client initialization:", e)
+                    for i in range(15, 0, -1):
+                        Console.write_stdout(f"\rRetrying in {i} seconds...   ",'')
+                        await asyncio.sleep(1)
+                    Console.write_stdout("\rRetrying now...           ")
+                    retry_count += 1
+                    continue
+
+                retry_count = 0
+                Console.write_stdout(f"Waiting for jobs...")
+                while True:
+                    try:
+                        jobs = client.dequeue()
+                        retry_count = 0
+                    except requests.exceptions.RequestException as e:
+                        retry_count += 1
+                        if retry_count >= MAX_RETRIES:
+                            Console.write_stderr(f"Failed to dequeue jobs after {MAX_RETRIES} retries.")
+                            Console.write_exception(f"Exception:", e, "\nExiting...")
+                            return
+                        Console.write_stderr(f"Failed to dequeue jobs (attempt {retry_count}/{MAX_RETRIES}):")
+                        for i in range(15, 0, -1):
+                            Console.write_stdout(f"\rRetrying in {i} seconds...   ",'')
+                            await asyncio.sleep(1)
+                        Console.write_stdout("\rRetrying now...           ")
+                        continue
+                    for job in jobs:
+                        Console.write_stdout(f"----------------------------------Job: {job.id}----------------------------------")
+                        tool_name, arguments = self.extract_function_call_info(job.input)
+                        if tool_name:
+                            Console.write_stdout(f"Executing tool {job.server}/{tool_name} with arguments {arguments}")
+                            try:
+                                result = await self.execute_tool(job.server, tool_name, json.loads(arguments))
+                            except (Exception) as e:
+                                Console.write_exception(f"Error executing tool {tool_name}:", e)
+                                continue
+
+                            if isinstance(result.content, list):
+                                text_parts = []
+                                for item in result.content:
+                                    if isinstance(item, types.TextContent):
+                                        text_parts.append(item.text)
+                                    elif isinstance(item, a2a_types.Part):
+                                        if isinstance(item.root, a2a_types.TextPart):
+                                            text_parts.append(item.root.text)
+                                        else:
+                                            Console.write_stdout(f"Unknown content type {type(item.root)}")
+                                    else:
+                                        Console.write_stdout(f"Unknown content type {type(item)}")
+
+                            if text_parts:
+                                job.output = "\n".join(text_parts)
+                                Console.write_stdout(f"result: {job.output} success: {not result.isError}")
+                                try:
+                                    client.send_result(ToolProxyJobResult(success=result.isError, job=job))
+                                except (ConnectionError, TimeoutError, RuntimeError) as e:
+                                    Console.write_exception("Error sending result:", e)
+                            else:
+                                Console.write_stdout(f"{result}")
+                    await asyncio.sleep(1)
+            finally:
+                Console.write_stdout("Proxy stopped")
+                await self.exit_stack.aclose()

pygeai/proxy/servers.py

@@ -0,0 +1,304 @@
+"""Server module for managing MCP server connections and tool execution."""
+import asyncio
+import os
+import shutil
+import httpx
+from types import SimpleNamespace
+from uuid import uuid4
+from contextlib import AsyncExitStack
+from abc import ABC, abstractmethod
+from typing import Any, Dict
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.stdio import stdio_client
+from mcp.client.sse import sse_client
+from pygeai.proxy.tool import ProxiedTool
+from pygeai.proxy.config import ProxySettingsManager
+from a2a.client import A2AClient, A2ACardResolver
+from a2a.types import (
+    AgentCard, SendMessageRequest, MessageSendParams,
+    Message, Task, SendMessageSuccessResponse
+)
+
+from pygeai.core.utils.console import Console
+
+class ToolServer(ABC):
+    """
+    Interface for tool servers like MCP and A2A.
+
+    Subclasses must implement methods to initialize the server,
+    list available tools, and execute a tool.
+    """
+   
+    def __init__(self, sever_name: str, config: Dict[str, Any], settings: ProxySettingsManager):
+        """
+        Initialize the server.
+        
+        :param sever_name: str - Name of the server
+        :param config: Dict[str, Any] - Server configuration
+        :param settings: ProxySettingsManager - Proxy settings manager
+        """
+        self.config = config
+        self.settings = settings
+        self.name: str = sever_name
+        self.public_prefix: str | None = None
+        self.exit_stack: AsyncExitStack = AsyncExitStack()
+
+    @abstractmethod
+    async def initialize(self) -> None:
+        """
+        Initialize the server connection.
+
+        :return: None
+        :raises: ValueError - If the command is invalid
+        :raises: RuntimeError - If server initialization fails
+        :raises: ConnectionError - If connection to server fails
+        """
+        pass
+
+    @abstractmethod
+    async def list_tools(self) -> list[ProxiedTool]:
+        """
+        List available tools from the server.
+
+        :return: list[Tool] - List of available tools
+        :raises: RuntimeError - If server is not initialized
+        """
+        pass
+
+    @abstractmethod
+    async def execute_tool(
+        self,
+        tool_name: str,
+        arguments: dict[str, Any],
+        retries: int = 2,
+        delay: float = 1.0,
+    ) -> Any:
+        """
+        Execute a tool with retry mechanism.
+
+        :param tool_name: str - Name of the tool to execute
+        :param arguments: dict[str, Any] - Tool arguments
+        :param retries: int - Number of retry attempts
+        :param delay: float - Delay between retries in seconds
+        :return: Any - Tool execution result
+        :raises: RuntimeError - If server is not initialized
+        :raises: ConnectionError - If connection to server fails
+        :raises: ValueError - If tool execution fails
+        """
+        pass
+
+class A2AServer(ToolServer):
+    """
+    Manages A2A server connections and tool execution.
+    """
+    def __init__(self, sever_name: str, config: Dict[str, Any], settings: ProxySettingsManager):
+        super().__init__(sever_name, config, settings)
+        self.client: A2AClient | None = None
+        self.card_url: str = self.config["url"]
+        self.agent_card: AgentCard | None = None
+        self.httpx_client: httpx.AsyncClient | None = None
+
+    async def initialize(self) -> None:
+        """
+        Initialize the A2A client from the agent card and convert agent skills into tools
+        compatible with OpenAI function call format.
+        """
+        try:
+            self.httpx_client = httpx.AsyncClient(timeout=60.0)
+            self.public_prefix = self.config.get("public_prefix")
+            headers = self.config.get("headers")
+            if headers:
+                self.httpx_client.headers.update(headers)
+            resolver = A2ACardResolver(httpx_client=self.httpx_client, base_url=self.card_url)
+            self.agent_card = await resolver.get_agent_card()
+            self.client = A2AClient(httpx_client=self.httpx_client, agent_card=self.agent_card)
+
+        except httpx.HTTPError as e:
+            Console.write_exception(f"HTTP error initializing A2A server {self.name}:", e)
+            raise ConnectionError(f"Failed to connect to A2A server: {e}") from e
+        except ValueError as e:
+            Console.write_exception(f"Invalid configuration for A2A server {self.name}:" ,e)
+            raise ValueError(f"Invalid A2A server configuration: {e}") from e
+        except RuntimeError as e:
+            Console.write_exception(f"Runtime error initializing A2A server {self.name}:", e)
+            raise RuntimeError(f"A2A server initialization failed: {e}") from e
+        
+    async def list_tools(self) -> list[ProxiedTool]:
+        if not self.client:
+            raise RuntimeError(f"Server {self.name} not initialized")
+        input_schema = {
+            "type": "object",
+            "properties": {
+                "input-text": {
+                    "type": "string",
+                    "description": "The input text to send to the agent for execution."
+                }
+            },
+            "required": ["input-text"]
+        }
+        tools = []
+        for skill in self.agent_card.skills:
+            tools.append(ProxiedTool(self.name, skill.id, skill.description, self.public_prefix, input_schema))
+        return tools
+
+    async def execute_tool(
+        self,
+        tool_name: str,
+        arguments: dict[str, Any],
+        retries: int = 2,
+        delay: float = 1.0,
+    ) -> Any:
+        Console.write_stdout(f"Executing {tool_name}...")
+        message_id = uuid4().hex
+        send_message_payload: dict[str, Any] = {
+            'message': {
+                'role': 'user',
+                'parts': [
+                    {
+                        'kind': 'text',
+                        'text': (
+                            'Use your skill:' + tool_name +
+                            ' with this input:' + arguments["input-text"]
+                        )
+                    }
+                ],
+                'messageId': message_id,
+            },
+        }
+        request = SendMessageRequest(
+            id=message_id,
+            params=MessageSendParams(**send_message_payload)
+        )
+
+        response = await self.client.send_message(request)
+        result = []
+        if isinstance(response.root, SendMessageSuccessResponse):
+            if isinstance(response.root.result, Message):
+                result = response.root.result.parts
+            elif isinstance(response.root, Task):
+                Console.write_stderr(f"Task response: {response.root.result}")
+            else:
+                Console.write_stderr(f"Unknown response type: {type(response.root)}")
+                raise ValueError(f"Unknown response type: {type(response.root)}")
+        
+        return SimpleNamespace(content=result, isError=False)
+
+class MCPServer(ToolServer):
+    """
+    Manages MCP server connections and tool execution.
+
+    :param sever_name: str - Name of the server
+    :param config: Dict[str, Any] - Server configuration
+    :param settings: ProxySettingsManager - Proxy settings manager
+    """
+    def __init__(self, sever_name: str, config: Dict[str, Any], settings: ProxySettingsManager):
+        super().__init__(sever_name, config, settings)
+        self.public_prefix = config.get("public_prefix")
+        self.stdio_context: Any | None = None
+        self.session: ClientSession | None = None
+
+    async def initialize(self) -> None:
+        self.public_prefix = self.config.get("public_prefix")
+        transport = self.config.get("transport") or (
+            "sse" if ("uri" in self.config  or "url" in self.config) else "stdio"
+        )
+        try:
+            if transport == "stdio":
+                command = (
+                    shutil.which("npx")
+                    if self.config["command"] == "npx"
+                    else self.config["command"]
+                )
+                if command is None:
+                    raise ValueError("The command must be una cadena válida")
+
+                server_params = StdioServerParameters(
+                    command=command,
+                    args=self.config["args"],
+                    env=(
+                        {**os.environ, **self.config["env"]}
+                        if self.config.get("env")
+                        else None
+                    ),
+                )
+            
+                stdio_transport = await self.exit_stack.enter_async_context(
+                    stdio_client(server_params)
+                )
+                read, write = stdio_transport
+            elif transport == "sse":
+                uri = self.config.get("uri", self.config["url"])
+                if not uri:
+                    raise ValueError("Missing 'uri' for sse transport")
+                try:
+                    sse_transport = await self.exit_stack.enter_async_context(
+                        sse_client(
+                            url=uri,
+                            headers=self.config.get("headers")
+                        )
+                    )
+                    read, write = sse_transport
+                except httpx.HTTPStatusError as e:
+                    Console.write_exception(f"HTTP error initializing MCP server {self.name}:", e)
+            else:
+                raise ValueError(f"Unsupported transport: {transport}")
+            
+            session = await self.exit_stack.enter_async_context(
+                ClientSession(read, write)
+            )
+            await session.initialize()
+            self.session = session
+        except (RuntimeError, ConnectionError, ValueError) as e:
+            Console.write_exception(f"Error initializing server {self.name}:", e)
+            raise
+
+    async def list_tools(self) -> list[ProxiedTool]:
+        if not self.session:
+            raise RuntimeError(f"Server {self.name} not initialized")
+
+        tools_response = await self.session.list_tools()
+        tools = []
+
+        for item in tools_response:
+            if isinstance(item, tuple) and item[0] == "tools":
+                for tool in item[1]:
+                    tools.append(
+                        ProxiedTool(
+                            self.name,
+                            tool.name,
+                            tool.description,
+                            self.public_prefix,
+                            tool.inputSchema
+                        )
+                    )
+
+        return tools
+
+    async def execute_tool(
+        self,
+        tool_name: str,
+        arguments: dict[str, Any],
+        retries: int = 2,
+        delay: float = 1.0,
+    ) -> Any:
+        if not self.session:
+            raise RuntimeError(f"Server {self.name} not initialized")
+
+        attempt = 0
+        while attempt < retries:
+            try:
+                Console.write_stdout(f"Executing {tool_name}...")
+                result = await self.session.call_tool(tool_name, arguments)
+                return result
+
+            except (RuntimeError, ConnectionError, ValueError) as e:
+                attempt += 1
+                Console.write_exception(
+                    "Error executing tool:", e, f". Attempt {attempt} of {retries}.\n"
+                )
+                if attempt < retries:
+                    Console.write_stdout(f"Retrying in {delay} seconds...")
+                    await asyncio.sleep(delay)
+                else:
+                    Console.write_stdout("Max retries reached. Failing.")
+                    raise

pygeai/proxy/tool.py

@@ -0,0 +1,69 @@
+from typing import Any
+import json
+
+
+class ProxiedTool:
+    """
+    Represents a tool exposed by an ToolServer (Multi-Channel Proxy).
+
+    :param server_name: str - Name of the server exposing the tool
+    :param name: str - Name of the tool
+    :param description: str - Description of the tool's functionality
+    :param public_prefix: str - Public prefix of the tool
+    :param input_schema: dict[str, Any] - JSON schema defining the expected input for the tool
+    :return: ProxiedTool - Instance of the ProxiedTool class
+    :raises: ValueError - If any required parameter is invalid or missing
+    """
+
+    def __init__(
+        self,
+        server_name: str,
+        name: str,
+        description: str,
+        public_prefix: str,
+        input_schema: dict[str, Any]
+    ) -> None:
+        self.server_name: str = server_name
+        self.name: str = name
+        self.openai_compatible_name: str = self._get_openai_compatible_name(name)
+        self.description: str = description
+        self.input_schema: dict[str, Any] = input_schema
+        self.public_prefix: str = public_prefix
+
+    def get_full_name(self) -> str:
+        """Get the full name of the tool."""
+        return f"{self.server_name}__{self.openai_compatible_name}"
+
+    def is_public(self) -> bool:
+        """Check if the tool is public."""
+        return self.public_prefix is not None
+
+    def get_public_name(self) -> str:
+        """Get the public name of the tool."""
+        if not self.public_prefix:
+            return ""
+        if self.server_name in self.public_prefix:
+            return f"{self.public_prefix}.{self.name}"
+        else:
+            return f"{self.public_prefix}.{self.get_full_name()}"
+
+    def format_for_llm(self) -> str:
+        """Format tool information for LLM.
+
+        Returns:
+            A formatted string describing the tool.
+        """
+        return json.dumps({
+            'type': 'function',
+            'function': {
+                'name': self.get_full_name(),
+                'description': self.description or '',
+                'parameters': self.input_schema
+            }
+        })
+    
+    def _get_openai_compatible_name(self, name: str) -> str:
+        compatible_name = ''.join(
+            c if c.isalnum() or c in '-_' else '_' for c in name.lower()
+        )
+        return compatible_name