universal-mcp 0.1.7rc2__py3-none-any.whl → 0.1.8rc2__py3-none-any.whl

universal_mcp/servers/server.py

@@ -1,181 +1,240 @@
 import os
 from abc import ABC, abstractmethod
-from typing import Any
+from typing import Any, Callable
+from urllib.parse import urlparse
 
 import httpx
 from loguru import logger
 from mcp.server.fastmcp import FastMCP
-from mcp.server.fastmcp.exceptions import ToolError
 from mcp.types import TextContent
 
-from universal_mcp.applications import app_from_slug
-from universal_mcp.config import AppConfig, IntegrationConfig, StoreConfig
-from universal_mcp.exceptions import NotAuthorizedError
-from universal_mcp.integrations import AgentRIntegration, ApiKeyIntegration
-from universal_mcp.stores import store_from_config
-
-
-class Server(FastMCP, ABC):
-    """
-    Server is responsible for managing the applications and the store
-    It also acts as a router for the applications, and exposed to the client
+from universal_mcp.applications import Application, app_from_slug
+from universal_mcp.config import AppConfig, ServerConfig, StoreConfig
+from universal_mcp.exceptions import NotAuthorizedError, ToolError
+from universal_mcp.integrations import integration_from_config
+from universal_mcp.stores import BaseStore, store_from_config
+from universal_mcp.tools.tools import ToolManager
+
+
+class BaseServer(FastMCP, ABC):
+    """Base server class with common functionality.
+    
+    This class provides core server functionality including store setup,
+    tool management, and application loading.
+    
+    Args:
+        config: Server configuration
+        **kwargs: Additional keyword arguments passed to FastMCP
     """
 
-    def __init__(
-        self, name: str, description: str, store: StoreConfig | None = None, **kwargs
-    ):
-        super().__init__(name, description, **kwargs)
-        logger.info(f"Initializing server: {name} with store: {store}")
-        self.store = store_from_config(store) if store else None
-        self._setup_store(store)
+    def __init__(self, config: ServerConfig, **kwargs):
+        super().__init__(config.name, config.description, **kwargs)
+        logger.info(f"Initializing server: {config.name} with store: {config.store}")
+        
+        self.config = config  # Store config at base level for consistency
+        self._tool_manager = ToolManager(warn_on_duplicate_tools=True)
+        self.store = self._setup_store(config.store)
         self._load_apps()
 
-    def _setup_store(self, store_config: StoreConfig | None):
-        """
-        Setup the store for the server.
+    def _setup_store(self, store_config: StoreConfig | None) -> BaseStore | None:
+        """Setup and configure the store.
+        
+        Args:
+            store_config: Store configuration
+            
+        Returns:
+            Configured store instance or None if no config provided
         """
-        if store_config is None:
-            return
-        self.store = store_from_config(store_config)
-        self.add_tool(self.store.set)
-        self.add_tool(self.store.delete)
-        # self.add_tool(self.store.get)
+        if not store_config:
+            return None
+            
+        store = store_from_config(store_config)
+        self.add_tool(store.set)
+        self.add_tool(store.delete)
+        return store
 
     @abstractmethod
-    def _load_apps(self):
+    def _load_apps(self) -> None:
+        """Load and register applications."""
         pass
 
-    async def call_tool(self, name: str, arguments: dict[str, Any]):
-        """Call a tool by name with arguments."""
+    def add_tool(self, tool: Callable) -> None:
+        """Add a tool to the server.
+        
+        Args:
+            tool: Tool to add
+        """
+        self._tool_manager.add_tool(tool)
+
+    async def list_tools(self) -> list[dict]:
+        """List all available tools in MCP format.
+        
+        Returns:
+            List of tool definitions
+        """
+        return self._tool_manager.list_tools(format='mcp')
+    
+    def _format_tool_result(self, result: Any) -> list[TextContent]:
+        """Format tool result into TextContent list.
+        
+        Args:
+            result: Raw tool result
+            
+        Returns:
+            List of TextContent objects
+        """
+        if isinstance(result, str):
+            return [TextContent(type="text", text=result)]
+        elif isinstance(result, list) and all(isinstance(item, TextContent) for item in result):
+            return result
+        else:
+            logger.warning(f"Tool returned unexpected type: {type(result)}. Wrapping in TextContent.")
+            return [TextContent(type="text", text=str(result))]
+                    
+    async def call_tool(self, name: str, arguments: dict[str, Any]) -> list[TextContent]:
+        """Call a tool with comprehensive error handling.
+        
+        Args:
+            name: Tool name
+            arguments: Tool arguments
+            
+        Returns:
+            List of TextContent results
+            
+        Raises:
+            ToolError: If tool execution fails
+        """
         logger.info(f"Calling tool: {name} with arguments: {arguments}")
         try:
-            result = await super().call_tool(name, arguments)
-            logger.info(f"Tool {name} completed successfully")
-            return result
+            result = await self._tool_manager.call_tool(name, arguments)
+            logger.info(f"Tool '{name}' completed successfully")
+            return self._format_tool_result(result)
+
         except ToolError as e:
-            raised_error = e.__cause__
-            if isinstance(raised_error, NotAuthorizedError):
-                logger.warning(
-                    f"Not authorized to call tool {name}: {raised_error.message}"
-                )
-                return [TextContent(type="text", text=raised_error.message)]
-            else:
-                logger.error(f"Error calling tool {name}: {str(e)}")
-                raise e
-
-
-class LocalServer(Server):
-    """
-    Local server for development purposes
+            if isinstance(e.__cause__, NotAuthorizedError):
+                message = f"Not authorized to call tool {name}: {e.__cause__.message}"
+                logger.warning(message)
+                return [TextContent(type="text", text=message)]
+            elif isinstance(e.__cause__, httpx.HTTPError):
+                message = f"HTTP error calling tool {name}: {str(e.__cause__)}"
+                logger.error(message)
+                return [TextContent(type="text", text=message)]
+            elif isinstance(e.__cause__, ValueError):
+                message = f"Invalid arguments for tool {name}: {str(e.__cause__)}"
+                logger.error(message)
+                return [TextContent(type="text", text=message)]
+            logger.error(f"Error calling tool {name}: {str(e)}", exc_info=True)
+            raise
+
+
+class LocalServer(BaseServer):
+    """Local development server implementation.
+    
+    Args:
+        config: Server configuration
+        **kwargs: Additional keyword arguments passed to FastMCP
     """
 
-    def __init__(
-        self,
-        apps_list: list[AppConfig] = None,
-        **kwargs,
-    ):
-        if not apps_list:
-            self.apps_list = []
-        else:
-            self.apps_list = apps_list
-        super().__init__(**kwargs)
-
-    def _get_store(self, store_config: StoreConfig | None):
-        logger.info(f"Getting store: {store_config}")
-        # No store override, use the one from the server
-        if store_config is None:
-            return self.store
-        return store_from_config(store_config)
-
-    def _get_integration(self, integration_config: IntegrationConfig | None):
-        if not integration_config:
+    def __init__(self, config: ServerConfig, **kwargs):
+        super().__init__(config, **kwargs)
+
+    def _load_app(self, app_config: AppConfig) -> Application | None:
+        """Load a single application with its integration.
+        
+        Args:
+            app_config: Application configuration
+            
+        Returns:
+            Configured application instance or None if loading fails
+        """
+        try:
+            integration = integration_from_config(
+                app_config.integration,
+                store=self.store
+            ) if app_config.integration  else None
+            return app_from_slug(app_config.name)(integration=integration)
+        except Exception as e:
+            logger.error(f"Failed to load app {app_config.name}: {e}", exc_info=True)
             return None
-        if integration_config.type == "api_key":
-            store = self._get_store(integration_config.store)
-            integration = ApiKeyIntegration(integration_config.name, store=store)
-            if integration_config.credentials:
-                integration.set_credentials(integration_config.credentials)
-            return integration
-        return None
-
-    def _load_app(self, app_config: AppConfig):
-        name = app_config.name
-        integration = self._get_integration(app_config.integration)
-        app = app_from_slug(name)(integration=integration)
-        return app
-
-    def _load_apps(self):
-        logger.info(f"Loading apps: {self.apps_list}")
-        for app_config in self.apps_list:
-            try:
-                app = self._load_app(app_config)
-                if app:
-                    tools = app.list_tools()
-                for tool in tools:
-                    tool_name = tool.__name__
-                    name = app.name + "_" + tool_name
-                    description = tool.__doc__
-                    if (
-                        app_config.actions is None
-                        or len(app_config.actions) == 0
-                        or name in app_config.actions
-                    ):
-                        self.add_tool(tool, name=name, description=description)
-            except Exception as e:
-                logger.error(f"Error loading app {app_config.name}: {e}")
-
-
-class AgentRServer(Server):
-    """
-    AgentR server. Connects to the AgentR API to get the apps and tools. Only supports agentr integrations.
+
+    def _load_apps(self) -> None:
+        """Load all configured applications."""
+        logger.info(f"Loading apps: {self.config.apps}")
+        for app_config in self.config.apps:
+            app = self._load_app(app_config)
+            if app:
+                self._tool_manager.register_tools_from_app(app, app_config.actions)
+
+
+class AgentRServer(BaseServer):
+    """AgentR API-connected server implementation.
+    
+    Args:
+        config: Server configuration
+        api_key: Optional API key for AgentR authentication. If not provided,
+                will attempt to read from AGENTR_API_KEY environment variable.
+        **kwargs: Additional keyword arguments passed to FastMCP
     """
 
-    def __init__(
-        self, name: str, description: str, api_key: str | None = None, **kwargs
-    ):
+    def __init__(self, config: ServerConfig, api_key: str | None = None, **kwargs):
         self.api_key = api_key or os.getenv("AGENTR_API_KEY")
         self.base_url = os.getenv("AGENTR_BASE_URL", "https://api.agentr.dev")
+        
         if not self.api_key:
             raise ValueError("API key required - get one at https://agentr.dev")
-        super().__init__(name, description=description, **kwargs)
+        parsed = urlparse(self.base_url)
+        if not all([parsed.scheme, parsed.netloc]):
+            raise ValueError(f"Invalid base URL format: {self.base_url}")
+
+        super().__init__(config, **kwargs)
+        
+    def _fetch_apps(self) -> list[AppConfig]:
+        """Fetch available apps from AgentR API.
+        
+        Returns:
+            List of application configurations
+            
+        Raises:
+            httpx.HTTPError: If API request fails
+        """
+        try:
+            response = httpx.get(
+                f"{self.base_url}/api/apps/",
+                headers={"X-API-KEY": self.api_key},
+                timeout=10,
+            )
+            response.raise_for_status()
+            return [AppConfig.model_validate(app) for app in response.json()]
+        except httpx.HTTPError as e:
+            logger.error(f"Failed to fetch apps from AgentR: {e}", exc_info=True)
+            raise
+
+    def _load_app(self, app_config: AppConfig) -> Application | None:
+        """Load a single application with AgentR integration.
+        
+        Args:
+            app_config: Application configuration
+            
+        Returns:
+            Configured application instance or None if loading fails
+        """
+        try:
+            integration = integration_from_config(
+                app_config.integration,
+                api_key=self.api_key
+            )
+            return app_from_slug(app_config.name)(integration=integration)
+        except Exception as e:
+            logger.error(f"Failed to load app {app_config.name}: {e}", exc_info=True)
+            return None
 
-    def _load_app(self, app_config: AppConfig):
-        name = app_config.name
-        if app_config.integration:
-            integration_name = app_config.integration.name
-            integration = AgentRIntegration(integration_name, api_key=self.api_key)
-        else:
-            integration = None
-        app = app_from_slug(name)(integration=integration)
-        return app
-
-    def _list_apps_with_integrations(self) -> list[AppConfig]:
-        # TODO: get this from the API
-        response = httpx.get(
-            f"{self.base_url}/api/apps/", headers={"X-API-KEY": self.api_key}
-        )
-        response.raise_for_status()
-        apps = response.json()
-
-        logger.info(f"Apps: {apps}")
-        return [AppConfig.model_validate(app) for app in apps]
-
-    def _load_apps(self):
-        apps = self._list_apps_with_integrations()
-        for app_config in apps:
-            try:
+    def _load_apps(self) -> None:
+        """Load all apps available from AgentR."""
+        try:
+            for app_config in self._fetch_apps():
                 app = self._load_app(app_config)
                 if app:
-                    tools = app.list_tools()
-                for tool in tools:
-                    tool_name = tool.__name__
-                    name = app.name + "_" + tool_name
-                    description = tool.__doc__
-                    if (
-                        app_config.actions is None
-                        or len(app_config.actions) == 0
-                        or name in app_config.actions
-                    ):
-                        self.add_tool(tool, name=name, description=description)
-            except Exception as e:
-                logger.error(f"Error loading app {app_config.name}: {e}")
+                    self._tool_manager.register_tools_from_app(app, app_config.actions)
+        except Exception:
+            logger.error("Failed to load apps", exc_info=True)
+            raise

universal_mcp/stores/__init__.py

@@ -1,5 +1,10 @@
 from universal_mcp.config import StoreConfig
-from universal_mcp.stores.store import EnvironmentStore, KeyringStore, MemoryStore
+from universal_mcp.stores.store import (
+    BaseStore,
+    EnvironmentStore,
+    KeyringStore,
+    MemoryStore,
+)
 
 
 def store_from_config(store_config: StoreConfig):
@@ -13,4 +18,4 @@ def store_from_config(store_config: StoreConfig):
         raise ValueError(f"Invalid store type: {store_config.type}")
 
 
-__all__ = [MemoryStore, EnvironmentStore, KeyringStore]
+__all__ = [BaseStore, MemoryStore, EnvironmentStore, KeyringStore]