universal-mcp 0.1.8rc1__py3-none-any.whl → 0.1.8rc3__py3-none-any.whl

universal_mcp/servers/server.py

@@ -1,184 +1,239 @@
 import os
 from abc import ABC, abstractmethod
+from collections.abc import Callable
 from typing import Any
+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
+from universal_mcp.analytics import analytics
+from universal_mcp.applications import Application, app_from_slug
+from universal_mcp.config import AppConfig, ServerConfig, StoreConfig
+from universal_mcp.integrations import AgentRIntegration, integration_from_config
+from universal_mcp.stores import BaseStore, store_from_config
+from universal_mcp.tools.tools import ToolManager
 
 
-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
+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)
-        self._load_apps()
+    def __init__(self, config: ServerConfig, **kwargs):
+        super().__init__(config.name, config.description, **kwargs)
+        logger.info(
+            f"Initializing server: {config.name} ({config.type}) with store: {config.store}"
+        )
 
-    def _setup_store(self, store_config: StoreConfig | None):
-        """
-        Setup the store for the server.
-        """
-        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)
+        self.config = config  # Store config at base level for consistency
+        self._tool_manager = ToolManager(warn_on_duplicate_tools=True)
 
     @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."""
-        logger.info(f"Calling tool: {name} with arguments: {arguments}")
-        try:
-            result = await super().call_tool(name, arguments)
-            logger.info(f"Tool {name} completed successfully")
+    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
-        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
+        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.
 
-class LocalServer(Server):
-    """
-    Local server for development purposes
+        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}")
+        result = await self._tool_manager.call_tool(name, arguments)
+        logger.info(f"Tool '{name}' completed successfully")
+        return self._format_tool_result(result)
+
+
+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)
+        self.store = self._setup_store(config.store)
+        self._load_apps()
+
+    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 not store_config:
             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.
+
+        store = store_from_config(store_config)
+        self.add_tool(store.set)
+        self.add_tool(store.delete)
+        return store
+
+    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
+            )
+            analytics.track_app_loaded(app_config.name)  # Track app loading
+            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_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)
-
-    def _load_app(self, app_config: AppConfig):
-        logger.info(f"Loading app: {app_config}")
-        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},
-            timeout=10,
-        )
-        response.raise_for_status()
-        apps = response.json()
+        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)
+        self.integration = AgentRIntegration(name="agentr", api_key=self.api_key)
+        self._load_apps()
 
-        logger.info(f"Loaded apps: {apps}")
-        return [AppConfig.model_validate(app) for app in apps]
+    def _fetch_apps(self) -> list[AppConfig]:
+        """Fetch available apps from AgentR API.
 
-    def _load_apps(self):
-        apps = self._list_apps_with_integrations()
-        for app_config in apps:
-            try:
+        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 = (
+                AgentRIntegration(
+                    name=app_config.integration.name, api_key=self.api_key
+                )
+                if app_config.integration
+                else None
+            )
+            analytics.track_app_loaded(app_config.name)  # Track app loading
+            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_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]