aixtools 0.1.4__py3-none-any.whl → 0.1.6__py3-none-any.whl

aixtools/mcp/client.py

@@ -0,0 +1,274 @@
+"""MCP server utilities with caching and robust error handling."""
+
+import asyncio
+from typing import Any
+
+import anyio
+from cachebox import TTLCache
+from mcp import types as mcp_types
+from mcp.shared.exceptions import McpError
+from pydantic_ai import RunContext, exceptions
+from pydantic_ai.mcp import MCPServerStreamableHTTP, ToolResult
+from pydantic_ai.toolsets.abstract import ToolsetTool
+
+from aixtools.context import SessionIdTuple
+from aixtools.logging.logging_config import get_logger
+
+MCP_TOOL_CACHE_TTL = 300  # 5 minutes
+DEFAULT_MCP_CONNECTION_TIMEOUT = 30
+CACHE_KEY = "TOOL_LIST"
+
+logger = get_logger(__name__)
+
+
+def get_mcp_headers(session_id_tuple: SessionIdTuple) -> dict[str, str] | None:
+    """
+    Generate headers for MCP server requests.
+
+    This function creates a dictionary of headers to be used in requests to
+    the MCP servers. If a `user_id` or `session_id` is provided, they are
+    included in the headers.
+
+    Args:
+        session_id_tuple (SessionIdTuple): user_id and session_id tuple
+    Returns:
+        dict[str, str] | None: A dictionary of headers for MCP server requests,
+                               or None if neither user_id nor session_id is
+                               provided. When None is returned, default headers
+                               from the client or transport will be used.
+    """
+    headers = None
+    user_id, session_id = session_id_tuple
+    if session_id or user_id:
+        headers = {}
+        if session_id:
+            headers["session-id"] = session_id
+        if user_id:
+            headers["user-id"] = user_id
+    return headers
+
+
+def get_configured_mcp_servers(
+    session_id_tuple: SessionIdTuple, mcp_urls: list[str], timeout: int = DEFAULT_MCP_CONNECTION_TIMEOUT
+):
+    """
+    Retrieve the configured MCP server instances with optional caching.
+
+    Context values `user_id` and `session_id` are included in the headers for each server request.
+
+    Each server is wrapped in a try-except block to isolate them from each other.
+    If one server fails, it won't affect the others.
+
+    Args:
+        session_id_tuple (SessionIdTuple): A tuple containing (user_id, session_id).
+        mcp_urls: (list[str], optional): A list of MCP server URLs to use.
+        timeout (int, optional): Timeout in seconds for MCP server connections. Defaults to 30 seconds.
+    Returns:
+        list[MCPServerStreamableHTTP]: A list of configured MCP server instances. If
+                              neither user_id nor session_id is provided, the
+                              server instances will use default headers defined
+                              by the underlying HTTP implementation.
+    """
+    headers = get_mcp_headers(session_id_tuple)
+
+    return [CachedMCPServerStreamableHTTP(url=url, headers=headers, timeout=timeout) for url in mcp_urls]
+
+
+class CachedMCPServerStreamableHTTP(MCPServerStreamableHTTP):
+    """StreamableHTTP MCP server with cachebox-based TTL caching and robust error handling.
+
+    This class addresses the cancellation propagation issue by:
+    1. Using complete task isolation to prevent CancelledError propagation
+    2. Implementing comprehensive error handling for all MCP operations
+    3. Using fallback mechanisms when servers become unavailable
+    4. Overriding pydantic_ai methods to fix variable scoping bug
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._tools_cache = TTLCache(maxsize=1, ttl=MCP_TOOL_CACHE_TTL)
+        self._tools_list = None
+        self._isolation_lock = asyncio.Lock()  # Lock for critical operations
+
+    async def _run_direct_or_isolated(self, func, fallback, timeout: float | None):
+        """Run a coroutine in complete isolation to prevent cancellation propagation.
+
+        Args:
+            func: Function that returns a coroutine to run
+            fallback: Function that takes an exception and returns a fallback value
+            timeout: Timeout in seconds. If None, then direct run is performed
+
+        Returns:
+            The result of the coroutine on success, or fallback value on any exception
+        """
+        try:
+            if timeout is None:
+                return await func()
+
+            task = asyncio.create_task(func())
+
+            # Use asyncio.wait to prevent cancellation propagation
+            done, pending = await asyncio.wait([task], timeout=timeout)
+
+            if pending:
+                # Cancel pending tasks safely
+                for t in pending:
+                    t.cancel()
+                    try:
+                        await t
+                    except (asyncio.CancelledError, Exception):  # pylint: disable=broad-except
+                        pass
+                raise TimeoutError(f"Task timed out after {timeout} seconds")
+
+            # Get result from completed task
+            completed_task = done.pop()
+            if exc := completed_task.exception():
+                raise exc
+            return completed_task.result()
+
+        except exceptions.ModelRetry as exc:
+            logger.warning("MCP %s: %s ModelRetry: %s", self.url, func.__name__, exc)
+            raise
+        except TimeoutError as exc:
+            logger.warning("MCP %s: %s timed out: %s", self.url, func.__name__, exc)
+            return fallback(exc)
+        except asyncio.CancelledError as exc:
+            logger.warning("MCP %s: %s was cancelled", self.url, func.__name__)
+            return fallback(exc)
+        except anyio.ClosedResourceError as exc:
+            logger.warning("MCP %s: %s closed resource.", self.url, func.__name__)
+            return fallback(exc)
+        except Exception as exc:  # pylint: disable=broad-except
+            if str(exc) == "Attempted to exit cancel scope in a different task than it was entered in":
+                logger.warning("MCP %s: %s enter/exit cancel scope task mismatch.", self.url, func.__name__)
+            else:
+                logger.warning("MCP %s: %s exception %s: %s", self.url, func.__name__, type(exc), exc)
+            return fallback(exc)
+
+    async def __aenter__(self):
+        """Enter the context of the cached MCP server with complete cancellation isolation."""
+        async with self._isolation_lock:
+
+            async def direct_init():
+                return await super(CachedMCPServerStreamableHTTP, self).__aenter__()  # pylint: disable=super-with-arguments
+
+            def fallback(_exc):
+                self._client = None
+                return self
+
+            return await self._run_direct_or_isolated(direct_init, fallback, timeout=None)
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Exit the context of the cached MCP server with complete cancellation isolation."""
+        async with self._isolation_lock:
+            # If we're being cancelled, just clean up
+            if exc_type is asyncio.CancelledError:
+                logger.warning("MCP %s: __aexit__ called with cancellation - cleaning up", self.url)
+                self._client = None
+                return True
+
+            # If client is already None, skip cleanup
+            if not self._client:
+                logger.warning("MCP %s: is uninitialized -> skipping cleanup", self.url)
+                return True
+
+            async def direct_cleanup():
+                return await super(CachedMCPServerStreamableHTTP, self).__aexit__(exc_type, exc_val, exc_tb)  # pylint: disable=super-with-arguments
+
+            def fallback(_exc):
+                self._client = None
+                return True  # Suppress exceptions to prevent propagation
+
+            return await self._run_direct_or_isolated(direct_cleanup, fallback, timeout=None)
+
+    async def list_tools(self) -> list[mcp_types.Tool]:
+        """Override to fix variable scoping bug and add caching with cancellation isolation."""
+        # If client is not initialized, return empty list
+        if not self._client:
+            logger.warning("MCP %s: is uninitialized -> no tools", self.url)
+            return []
+
+        # First, check if we have a valid cached result
+        if CACHE_KEY in self._tools_cache:
+            logger.info("Using cached tools for %s", self.url)
+            return self._tools_cache[CACHE_KEY]
+
+        # Create isolated task to prevent cancellation propagation
+        async def isolated_list_tools():
+            """Isolated list_tools with variable scoping bug fix."""
+            result = None  # Initialize to prevent UnboundLocalError
+            async with self:  # Ensure server is running
+                result = await self._client.list_tools()
+            if result:
+                self._tools_list = result.tools or []
+                self._tools_cache[CACHE_KEY] = self._tools_list
+                logger.info("MCP %s: list_tools returned %d tools", self.url, len(self._tools_list))
+            else:
+                logger.warning("MCP %s: list_tools returned no result", self.url)
+            return self._tools_list or []
+
+        def fallback(_exc):
+            return self._tools_list or []
+
+        return await self._run_direct_or_isolated(isolated_list_tools, fallback, timeout=5.0)
+
+    async def call_tool(
+        self,
+        name: str,
+        tool_args: dict[str, Any],
+        ctx: RunContext[Any],
+        tool: ToolsetTool[Any],
+    ) -> ToolResult:
+        """Call tool with complete isolation from cancellation using patched pydantic_ai."""
+        logger.info("MCP %s: call_tool '%s' started.", self.url, name)
+
+        # Early returns for uninitialized servers
+        if not self._client:
+            logger.warning("MCP %s: is uninitialized -> cannot call tool", self.url)
+            return f"There was an error with calling tool '{name}': MCP connection is uninitialized."
+
+        # Create isolated task to prevent cancellation propagation
+        async def isolated_call_tool():
+            """Isolated call_tool using patched pydantic_ai methods."""
+            return await super(CachedMCPServerStreamableHTTP, self).call_tool(name, tool_args, ctx, tool)  # pylint: disable=super-with-arguments
+
+        def fallback(exc):
+            return f"Exception {type(exc)} when calling tool '{name}': {exc}. Consider alternative approaches."
+
+        result = await self._run_direct_or_isolated(isolated_call_tool, fallback, timeout=3600.0)
+        logger.info("MCP %s: call_tool '%s' completed.", self.url, name)
+        return result
+
+    async def direct_call_tool(
+        self, name: str, args: dict[str, Any], metadata: dict[str, Any] | None = None
+    ) -> ToolResult:
+        """Override to fix variable scoping bug in direct_call_tool."""
+        result = None  # Initialize to prevent UnboundLocalError
+        async with self:  # Ensure server is running
+            try:
+                result = await self._client.send_request(
+                    mcp_types.ClientRequest(
+                        mcp_types.CallToolRequest(
+                            method="tools/call",
+                            params=mcp_types.CallToolRequestParams(
+                                name=name,
+                                arguments=args,
+                                _meta=mcp_types.RequestParams.Meta(**metadata) if metadata else None,
+                            ),
+                        )
+                    ),
+                    mcp_types.CallToolResult,
+                )
+            except McpError as e:
+                raise exceptions.ModelRetry(e.error.message)
+
+        if not result:
+            raise exceptions.ModelRetry("No result from MCP server")
+
+        content = [await self._map_tool_result_part(part) for part in result.content]
+
+        if result.isError:
+            text = "\n".join(str(part) for part in content)
+            raise exceptions.ModelRetry(text)
+
+        return content[0] if len(content) == 1 else content

aixtools/server/utils.py

@@ -8,7 +8,7 @@ from functools import wraps
 from fastmcp import Context
 from fastmcp.server import dependencies
 
-from ..context import session_id_var, user_id_var
+from ..context import DEFAULT_SESSION_ID, DEFAULT_USER_ID, session_id_var, user_id_var
 
 
 def get_session_id_tuple(ctx: Context | None = None) -> tuple[str, str]:
@@ -18,9 +18,9 @@ def get_session_id_tuple(ctx: Context | None = None) -> tuple[str, str]:
     Returns: Tuple of (user_id, session_id).
     """
     user_id = get_user_id_from_request(ctx)
-    user_id = user_id or user_id_var.get("default_user")
+    user_id = user_id or user_id_var.get(DEFAULT_USER_ID)
     session_id = get_session_id_from_request(ctx)
-    session_id = session_id or session_id_var.get("default_session")
+    session_id = session_id or session_id_var.get(DEFAULT_SESSION_ID)
     return user_id, session_id
 
 

aixtools/utils/config.py

@@ -9,6 +9,7 @@ from pathlib import Path
 from dotenv import dotenv_values, load_dotenv
 
 from aixtools.utils.config_util import find_env_file, get_project_root, get_variable_env
+from aixtools.utils.utils import str2bool
 
 # Debug mode
 LOG_LEVEL = logging.DEBUG
@@ -116,3 +117,15 @@ BEDROCK_MODEL_NAME = get_variable_env("BEDROCK_MODEL_NAME", allow_empty=True)
 # LogFire
 LOGFIRE_TOKEN = get_variable_env("LOGFIRE_TOKEN", True, "")
 LOGFIRE_TRACES_ENDPOINT = get_variable_env("LOGFIRE_TRACES_ENDPOINT", True, "")
+
+# Google Vertex AI
+GOOGLE_GENAI_USE_VERTEXAI = str2bool(get_variable_env("GOOGLE_GENAI_USE_VERTEXAI", True, True))
+GOOGLE_CLOUD_PROJECT = get_variable_env("GOOGLE_CLOUD_PROJECT", True)
+GOOGLE_CLOUD_LOCATION = get_variable_env("GOOGLE_CLOUD_LOCATION", True)
+
+# vault parameters.
+VAULT_ADDRESS = get_variable_env("VAULT_ADDRESS", default="http://localhost:8200")
+VAULT_TOKEN = get_variable_env("VAULT_TOKEN", default="vault-token")
+VAULT_ENV = get_variable_env("ENV", default="dev")
+VAULT_MOUNT_POINT = get_variable_env("VAULT_MOUNT_POINT", default="secret")
+VAULT_PATH_PREFIX = get_variable_env("VAULT_PATH_PREFIX", default="path")

aixtools/utils/files.py

@@ -0,0 +1,17 @@
+"""File utilities"""
+
+
+def is_text_content(data: bytes, mime_type: str) -> bool:
+    """Check if content is text based on mime type and content analysis."""
+    # Check mime type first
+    if mime_type and (
+        mime_type.startswith("text/") or mime_type in ["application/json", "application/xml", "application/javascript"]
+    ):
+        return True
+
+    # Try to decode as UTF-8 to check if it's text
+    try:
+        data.decode("utf-8")
+        return True
+    except UnicodeDecodeError:
+        return False

aixtools/utils/utils.py

@@ -154,6 +154,13 @@ def timestamp_uuid_tuple() -> tuple[str, str, str]:
     return (now.strftime("%Y-%m-%d"), now.strftime("%H:%M:%S"), str(uuid.uuid4()))
 
 
+def str2bool(v: str | None) -> bool:
+    """Convert a string to a boolean value."""
+    if not v:
+        return False
+    return str(v).lower() in ("yes", "true", "on", "1")
+
+
 async def async_iter(items):
     """Asynchronously iterate over items."""
     for item in items:

aixtools/vault/__init__.py

@@ -0,0 +1,7 @@
+"""
+Provides a Vault client for storing and retrieving user service api keys.
+"""
+
+from .vault import VaultClient
+
+__all__ = ["VaultClient"]

aixtools/vault/vault.py

@@ -0,0 +1,73 @@
+# ruff: noqa: PLR0913
+"""
+Provides a Vault client for storing and retrieving user service api keys.
+"""
+
+import logging
+from typing import Optional
+
+import hvac
+from hvac.exceptions import InvalidPath
+
+from aixtools.utils import config
+
+logger = logging.getLogger(__name__)
+
+
+class VaultAuthError(Exception):
+    """Exception raised for vault authentication errors."""
+
+
+class VaultClient:
+    """Vault client for storing and retrieving user service api keys."""
+
+    def __init__(self):
+        self.client = hvac.Client(url=config.VAULT_ADDRESS, token=config.VAULT_TOKEN)
+
+        if not self.client.is_authenticated():
+            raise VaultAuthError("Vault client authentication failed. Check vault_token.")
+
+    def store_user_service_api_key(self, *, user_id: str, service_name: str, user_api_key: str):
+        """
+        Store user's service api key in the Vault at the specified vault mount
+        point, where the path is <path_prefix>/<env>/<user_id>/<service_name>.
+        """
+        secret_path = None
+        try:
+            secret_path = f"{config.VAULT_PATH_PREFIX}/{config.VAULT_ENV}/{user_id}/{service_name}"
+            print("secret_path", secret_path)
+            secret_dict = {"user-api-key": user_api_key}
+            self.client.secrets.kv.v2.create_or_update_secret(
+                secret_path, secret=secret_dict, mount_point=config.VAULT_MOUNT_POINT
+            )
+
+            logger.info("Secret written to path %s", secret_path)
+        except Exception as e:
+            logger.error("Failed to write secret to path %s: %s", secret_path, str(e))
+            raise VaultAuthError(e) from e
+
+    def read_user_service_api_key(self, *, user_id: str, service_name) -> Optional[str]:
+        """
+        Read user's service api key in from vault at the specified mount point,
+        where the path is <path_prefix>/<env>/<user_id>/<service_name>.
+        """
+        secret_path = None
+
+        try:
+            secret_path = f"{config.VAULT_PATH_PREFIX}/{config.VAULT_ENV}/{user_id}/{service_name}"
+            logger.info("Reading secret from path %s", secret_path)
+            response = self.client.secrets.kv.v2.read_secret_version(
+                secret_path, mount_point=config.VAULT_MOUNT_POINT, raise_on_deleted_version=True
+            )
+            secret_data = response["data"]["data"]
+            user_api_key = secret_data["user-api-key"]
+            logger.info("Secret read from path %s ", secret_path)
+            return user_api_key
+        except InvalidPath:
+            # Secret path does not exist
+            logger.warning("Secret path does not exist %s ", secret_path)
+            return None
+
+        except Exception as e:
+            logger.error("Failed to read secret from path %s: %s", secret_path, str(e))
+            raise VaultAuthError(e) from e