universal-mcp 0.1.23rc2__py3-none-any.whl → 0.1.24rc3__py3-none-any.whl

universal_mcp/integrations/integration.py

@@ -5,7 +5,6 @@ from loguru import logger
 
 from universal_mcp.exceptions import KeyNotFoundError, NotAuthorizedError
 from universal_mcp.stores import BaseStore, MemoryStore
-from universal_mcp.utils.agentr import AgentrClient
 
 
 def sanitize_api_key_name(name: str) -> str:
@@ -19,75 +18,124 @@ def sanitize_api_key_name(name: str) -> str:
 class Integration:
     """Abstract base class for handling application integrations and authentication.
 
-    This class defines the interface for different types of integrations that handle
-    authentication and authorization with external services.
+    This class defines a common interface for various authentication and
+    authorization strategies an application might use to connect with
+    external services. Subclasses implement specific mechanisms like
+    API key handling, OAuth 2.0 flows, or delegation to platforms like AgentR.
 
-    Args:
-        name: The name identifier for this integration
-        store: Optional Store instance for persisting credentials and other data
+    Each integration is associated with a name and can use a `BaseStore`
+    instance for persisting credentials or other relevant data.
 
     Attributes:
-        name: The name identifier for this integration
-        store: Store instance for persisting credentials and other data
+        name (str): The unique name identifying this integration instance
+                    (e.g., "my_app_api_key", "github_oauth").
+        store (BaseStore): The storage backend (e.g., `MemoryStore`,
+                       `KeyringStore`) used for persisting credentials.
+                       Defaults to `MemoryStore` if not provided.
     """
 
     def __init__(self, name: str, store: BaseStore | None = None):
+        """Initializes the Integration.
+
+        Args:
+            name (str): The unique name/identifier for this integration instance.
+            store (BaseStore | None, optional): A store instance for
+                persisting credentials. Defaults to `MemoryStore()`.
+        """
         self.name = name
         self.store = store or MemoryStore()
+        self.type = ""
 
     def authorize(self) -> str | dict[str, Any]:
-        """Authorize the integration.
+        """Initiates or provides details for the authorization process.
+
+        The exact behavior and return type of this method depend on the
+        specific integration subclass. It might return an authorization URL
+        for the user to visit, parameters needed to construct such a URL,
+        or instructions on how to manually provide credentials.
 
         Returns:
-            Union[str, Dict[str, Any]]: Authorization URL or parameters needed for authorization.
+            str | dict[str, Any]: Typically, an authorization URL (str) or a
+                                  dictionary containing parameters needed for
+                                  the authorization flow.
 
         Raises:
-            ValueError: If required configuration is missing.
+            ValueError: If essential configuration for authorization is missing.
+            NotImplementedError: If the subclass does not implement this method.
         """
-        pass
+        raise NotImplementedError("Subclasses must implement the authorize method.")
 
     def get_credentials(self) -> dict[str, Any]:
-        """Get credentials for the integration.
+        """Retrieves the stored credentials for this integration.
+
+        Fetches credentials associated with `self.name` from the `self.store`.
 
         Returns:
-            Dict[str, Any]: Credentials for the integration.
+            dict[str, Any]: A dictionary containing the credentials. The structure
+                            of this dictionary is specific to the integration type.
 
         Raises:
-            NotAuthorizedError: If credentials are not found or invalid.
+            NotAuthorizedError: If credentials are not found in the store
+                                or are otherwise invalid/inaccessible.
+            KeyNotFoundError: If the key (self.name) is not found in the store.
         """
-        credentials = self.store.get(self.name)
-        return credentials
+        try:
+            credentials = self.store.get(self.name)
+            if credentials is None:  # Explicitly check for None if store can return it
+                raise NotAuthorizedError(f"No credentials found for {self.name}")
+            return credentials
+        except KeyNotFoundError as e:
+            raise NotAuthorizedError(f"Credentials not found for {self.name}: {e}") from e
 
     def set_credentials(self, credentials: dict[str, Any]) -> None:
-        """Set credentials for the integration.
+        """Stores the provided credentials for this integration.
+
+        Saves the given credentials dictionary into `self.store` associated
+        with `self.name`.
 
         Args:
-            credentials: Dictionary containing credentials for the integration.
+            credentials (dict[str, Any]): A dictionary containing the credentials
+                to be stored. The required structure depends on the integration.
 
         Raises:
-            ValueError: If credentials are invalid or missing required fields.
+            ValueError: If the provided credentials are invalid or missing
+                        required fields for the specific integration type.
         """
         self.store.set(self.name, credentials)
 
+    def __str__(self) -> str:
+        return f"Integration(name={self.name}, type={self.type})"
 
-class ApiKeyIntegration(Integration):
-    """Integration class for API key based authentication.
+    def __repr__(self) -> str:
+        return self.__str__()
 
-    This class implements the Integration interface for services that use API key
-    authentication. It handles storing and retrieving API keys using the provided
-    store.
 
-    Args:
-        name: The name identifier for this integration
-        store: Optional Store instance for persisting credentials and other data
-        **kwargs: Additional keyword arguments passed to parent class
+class ApiKeyIntegration(Integration):
+    """Handles integrations that use a simple API key for authentication.
+
+    This class manages storing and retrieving an API key. The key name is
+    automatically sanitized (e.g., uppercased and suffixed with `_API_KEY`)
+    before being used with the store.
 
     Attributes:
-        name: The name identifier for this integration
-        store: Store instance for persisting credentials and other data
+        name (str): The sanitized name used as the key for storing the API key.
+        store (BaseStore): Store for persisting the API key.
+        type (str): Set to "api_key".
+        _api_key (str | None): Cached API key.
     """
 
     def __init__(self, name: str, store: BaseStore | None = None, **kwargs):
+        """Initializes ApiKeyIntegration.
+
+        The provided `name` is sanitized (e.g., 'mykey' becomes 'MYKEY_API_KEY')
+        to form the actual key used for storage.
+
+        Args:
+            name (str): The base name for the API key (e.g., "TAVILY").
+            store (BaseStore | None, optional): Store for credentials.
+                                               Defaults to `MemoryStore()`.
+            **kwargs: Additional arguments passed to the parent `Integration`.
+        """
         self.type = "api_key"
         sanitized_name = sanitize_api_key_name(name)
         super().__init__(sanitized_name, store, **kwargs)
@@ -95,25 +143,39 @@ class ApiKeyIntegration(Integration):
         self._api_key: str | None = None
 
     @property
-    def api_key(self) -> str | None:
+    def api_key(self) -> str:  # Changed to str, as it raises if None effectively
+        """Retrieves the API key, loading it from the store if necessary.
+
+        If the API key is not already cached in `_api_key`, it attempts
+        to load it from `self.store` using `self.name` as the key.
+
+        Returns:
+            str: The API key.
+
+        Raises:
+            NotAuthorizedError: If the API key is not found in the store.
+                                The original `KeyNotFoundError` is chained.
+        """
         if not self._api_key:
             try:
-                credentials = self.store.get(self.name)
+                credentials = self.store.get(self.name)  # type: ignore
                 self._api_key = credentials
             except KeyNotFoundError as e:
                 action = self.authorize()
                 raise NotAuthorizedError(action) from e
-        return self._api_key
+        return self._api_key  # type: ignore
 
     @api_key.setter
     def api_key(self, value: str | None) -> None:
-        """Set the API key.
+        """Sets and stores the API key.
 
         Args:
-            value: The API key value to set.
+            value (str | None): The API key value. If None, `_api_key` is set
+                                to None, but nothing is stored (or cleared from store).
+                                Consider if None should clear the store.
 
         Raises:
-            ValueError: If the API key is invalid.
+            ValueError: If `value` is provided and is not a string.
         """
         if value is not None and not isinstance(value, str):
             raise ValueError("API key must be a string")
@@ -122,63 +184,63 @@ class ApiKeyIntegration(Integration):
             self.store.set(self.name, value)
 
     def get_credentials(self) -> dict[str, str]:
-        """Get API key credentials.
+        """Retrieves the API key and returns it in a standard dictionary format.
 
         Returns:
-            Dict[str, str]: Dictionary containing the API key.
+            dict[str, str]: A dictionary like `{"api_key": "your_api_key_value"}`.
 
         Raises:
-            NotAuthorizedError: If API key is not found.
+            NotAuthorizedError: If the API key cannot be retrieved.
         """
         return {"api_key": self.api_key}
 
     def set_credentials(self, credentials: dict[str, Any]) -> None:
-        """Set API key credentials.
+        """Sets the API key from a dictionary.
+
+        Expects `credentials` to be a dictionary, typically containing
+        an 'api_key' field, but it stores the entire dictionary as is
+        under `self.name`. For direct API key setting, use the `api_key` property.
 
         Args:
-            credentials: Dictionary containing the API key.
+            credentials (dict[str, Any]): A dictionary containing the API key
+                or related credential information.
 
         Raises:
-            ValueError: If credentials are invalid or missing API key.
+            ValueError: If `credentials` is not a dictionary.
         """
         if not credentials or not isinstance(credentials, dict):
             raise ValueError("Invalid credentials format")
         self.store.set(self.name, credentials)
 
     def authorize(self) -> str:
-        """Get authorization instructions for API key.
+        """Provides instructions for setting the API key.
+
+        Since API key setup is typically manual, this method returns a
+        message guiding the user on how to provide the key.
 
         Returns:
-            str: Instructions for setting up API key.
+            str: A message instructing the user to provide the API key
+                 for `self.name`.
         """
         return f"Please ask the user for api key and set the API Key for {self.name} in the store"
 
 
 class OAuthIntegration(Integration):
-    """Integration class for OAuth based authentication.
-
-    This class implements the Integration interface for services that use OAuth
-    authentication. It handles the OAuth flow including authorization, token exchange,
-    and token refresh.
-
-    Args:
-        name: The name identifier for this integration
-        store: Optional Store instance for persisting credentials and other data
-        client_id: OAuth client ID
-        client_secret: OAuth client secret
-        auth_url: OAuth authorization URL
-        token_url: OAuth token exchange URL
-        scope: OAuth scope string
-        **kwargs: Additional keyword arguments passed to parent class
+    """Manages OAuth 2.0 authentication and authorization flows.
+
+    This class implements the necessary steps for an OAuth 2.0 client,
+    including generating authorization request parameters, handling the
+    redirect callback from the authorization server, exchanging the
+    authorization code for access/refresh tokens, and refreshing tokens.
 
     Attributes:
-        name: The name identifier for this integration
-        store: Store instance for persisting credentials and other data
-        client_id: OAuth client ID
-        client_secret: OAuth client secret
-        auth_url: OAuth authorization URL
-        token_url: OAuth token exchange URL
-        scope: OAuth scope string
+        name (str): Name of the integration.
+        store (BaseStore): Store for OAuth tokens.
+        client_id (str | None): The OAuth 2.0 Client ID.
+        client_secret (str | None): The OAuth 2.0 Client Secret.
+        auth_url (str | None): The authorization server's endpoint URL.
+        token_url (str | None): The token server's endpoint URL.
+        scope (str | None): The requested OAuth scopes, space-separated.
     """
 
     def __init__(
@@ -192,7 +254,21 @@ class OAuthIntegration(Integration):
         scope: str | None = None,
         **kwargs,
     ):
+        """Initializes the OAuthIntegration.
+
+        Args:
+            name (str): The unique name for this integration instance.
+            store (BaseStore | None, optional): Store for credentials.
+                                               Defaults to `MemoryStore()`.
+            client_id (str | None, optional): The OAuth 2.0 Client ID.
+            client_secret (str | None, optional): The OAuth 2.0 Client Secret.
+            auth_url (str | None, optional): The authorization server's endpoint URL.
+            token_url (str | None, optional): The token server's endpoint URL.
+            scope (str | None, optional): The requested OAuth scopes, space-separated.
+            **kwargs: Additional arguments passed to the parent `Integration`.
+        """
         super().__init__(name, store, **kwargs)
+        self.type = "oauth"
         self.client_id = client_id
         self.client_secret = client_secret
         self.auth_url = auth_url
@@ -200,24 +276,30 @@ class OAuthIntegration(Integration):
         self.scope = scope
 
     def get_credentials(self) -> dict[str, Any] | None:
-        """Get OAuth credentials.
+        """Retrieves stored OAuth tokens for this integration.
 
         Returns:
-            Optional[Dict[str, Any]]: Dictionary containing OAuth tokens if found, None otherwise.
+            dict[str, Any] | None: A dictionary containing the OAuth tokens
+                                  (e.g., `access_token`, `refresh_token`) if found,
+                                  otherwise None.
         """
         credentials = self.store.get(self.name)
         if not credentials:
             return None
-        return credentials
+        return credentials  # type: ignore
 
     def set_credentials(self, credentials: dict[str, Any]) -> None:
-        """Set OAuth credentials.
+        """Stores OAuth tokens for this integration.
+
+        Validates that essential fields like 'access_token' are present.
 
         Args:
-            credentials: Dictionary containing OAuth tokens.
+            credentials (dict[str, Any]): A dictionary containing OAuth tokens.
+                Must include at least 'access_token'.
 
         Raises:
-            ValueError: If credentials are invalid or missing required tokens.
+            ValueError: If `credentials` is not a dictionary or if 'access_token'
+                        is missing.
         """
         if not credentials or not isinstance(credentials, dict):
             raise ValueError("Invalid credentials format")
@@ -226,13 +308,19 @@ class OAuthIntegration(Integration):
         self.store.set(self.name, credentials)
 
     def authorize(self) -> dict[str, Any]:
-        """Get OAuth authorization parameters.
+        """Constructs parameters required for the OAuth authorization request.
+
+        These parameters are typically used to build the URL to which the
+        user must be redirected to grant authorization.
 
         Returns:
-            Dict[str, Any]: Dictionary containing OAuth authorization parameters.
+            dict[str, Any]: A dictionary containing the authorization endpoint URL
+                            (`url`), query parameters (`params`), client secret,
+                            and token URL.
 
         Raises:
-            ValueError: If required OAuth configuration is missing.
+            ValueError: If essential OAuth configuration like `client_id`,
+                        `client_secret`, `auth_url`, or `token_url` is missing.
         """
         if not all([self.client_id, self.client_secret, self.auth_url, self.token_url]):
             raise ValueError("Missing required OAuth configuration")
@@ -251,19 +339,24 @@ class OAuthIntegration(Integration):
         }
 
     def handle_callback(self, code: str) -> dict[str, Any]:
-        """Handle OAuth callback and exchange code for tokens.
+        """Handles the OAuth callback by exchanging the authorization code for tokens.
+
+        This method is called after the user authorizes the application and the
+        authorization server redirects back with an authorization code.
 
         Args:
-            code: Authorization code from OAuth callback.
+            code (str): The authorization code received from the OAuth server.
 
         Returns:
-            Dict[str, Any]: Dictionary containing OAuth tokens.
+            dict[str, Any]: A dictionary containing the access token, refresh token
+                            (if any), and other token response data. These are also
+                            stored via `set_credentials`.
 
         Raises:
-            ValueError: If required OAuth configuration is missing.
-            httpx.HTTPError: If token exchange request fails.
+            ValueError: If essential OAuth configuration is missing.
+            httpx.HTTPStatusError: If the token exchange request to `token_url` fails.
         """
-        if not all([self.client_id, self.client_secret, self.token_url]):
+        if not all([self.client_id, self.client_secret, self.token_url]):  # type: ignore
             raise ValueError("Missing required OAuth configuration")
 
         token_params = {
@@ -273,24 +366,26 @@ class OAuthIntegration(Integration):
             "grant_type": "authorization_code",
         }
 
-        response = httpx.post(self.token_url, data=token_params)
+        response = httpx.post(self.token_url, data=token_params)  # type: ignore
         response.raise_for_status()
         credentials = response.json()
         self.store.set(self.name, credentials)
         return credentials
 
     def refresh_token(self) -> dict[str, Any]:
-        """Refresh OAuth access token using refresh token.
+        """Refreshes an expired access token using a stored refresh token.
 
         Returns:
-            Dict[str, Any]: Dictionary containing new OAuth tokens.
+            dict[str, Any]: A dictionary containing the new access token,
+                            refresh token, and other token response data.
+                            These are also stored.
 
         Raises:
-            ValueError: If required OAuth configuration is missing.
-            httpx.HTTPError: If token refresh request fails.
-            KeyError: If refresh token is not found in current credentials.
+            ValueError: If essential OAuth configuration is missing.
+            KeyError: If a refresh token is not found in the stored credentials.
+            httpx.HTTPStatusError: If the token refresh request fails.
         """
-        if not all([self.client_id, self.client_secret, self.token_url]):
+        if not all([self.client_id, self.client_secret, self.token_url]):  # type: ignore
             raise ValueError("Missing required OAuth configuration")
 
         credentials = self.get_credentials()
@@ -304,86 +399,8 @@ class OAuthIntegration(Integration):
             "refresh_token": credentials["refresh_token"],
         }
 
-        response = httpx.post(self.token_url, data=token_params)
+        response = httpx.post(self.token_url, data=token_params)  # type: ignore
         response.raise_for_status()
         credentials = response.json()
         self.store.set(self.name, credentials)
         return credentials
-
-
-class AgentRIntegration(Integration):
-    """Integration class for AgentR API authentication and authorization.
-
-    This class handles API key authentication and OAuth authorization flow for AgentR services.
-
-    Args:
-        name (str): Name of the integration
-        api_key (str, optional): AgentR API key. If not provided, will look for AGENTR_API_KEY env var
-        **kwargs: Additional keyword arguments passed to parent Integration class
-
-    Raises:
-        ValueError: If no API key is provided or found in environment variables
-    """
-
-    def __init__(self, name: str, api_key: str | None = None, base_url: str | None = None, **kwargs):
-        super().__init__(name, **kwargs)
-        self.client = AgentrClient(api_key=api_key, base_url=base_url)
-        self._credentials = None
-
-    def set_credentials(self, credentials: dict | None = None):
-        """Set credentials for the integration.
-
-        This method is not implemented for AgentR integration. Instead it redirects to the authorize flow.
-
-        Args:
-            credentials (dict | None, optional): Credentials dict (not used). Defaults to None.
-
-        Returns:
-            str: Authorization URL from authorize() method
-        """
-        return self.authorize()
-
-    @property
-    def credentials(self):
-        """Get credentials for the integration from the AgentR API.
-
-        Makes API request to retrieve stored credentials for this integration.
-
-        Returns:
-            dict: Credentials data from API response
-
-        Raises:
-            NotAuthorizedError: If credentials are not found (404 response)
-            HTTPError: For other API errors
-        """
-        if self._credentials is not None:
-            return self._credentials
-        self._credentials = self.client.get_credentials(self.name)
-        return self._credentials
-
-    def get_credentials(self):
-        """Get credentials for the integration from the AgentR API.
-
-        Makes API request to retrieve stored credentials for this integration.
-
-        Returns:
-            dict: Credentials data from API response
-
-        Raises:
-            NotAuthorizedError: If credentials are not found (404 response)
-            HTTPError: For other API errors
-        """
-        return self.credentials
-
-    def authorize(self):
-        """Get authorization URL for the integration.
-
-        Makes API request to get OAuth authorization URL.
-
-        Returns:
-            str: Message containing authorization URL
-
-        Raises:
-            HTTPError: If API request fails
-        """
-        return self.client.get_authorization_url(self.name)

universal_mcp/servers/__init__.py

@@ -1,15 +1,3 @@
-from universal_mcp.config import ServerConfig
-from universal_mcp.servers.server import AgentRServer, BaseServer, LocalServer, SingleMCPServer
+from universal_mcp.servers.server import BaseServer, LocalServer, SingleMCPServer
 
-
-def server_from_config(config: ServerConfig):
-    if config.type == "agentr":
-        return AgentRServer(config=config, api_key=config.api_key)
-
-    elif config.type == "local":
-        return LocalServer(config=config)
-    else:
-        raise ValueError(f"Unsupported server type: {config.type}")
-
-
-__all__ = [AgentRServer, LocalServer, SingleMCPServer, BaseServer, server_from_config]
+__all__ = ["LocalServer", "SingleMCPServer", "BaseServer"]