mcp-mesh 0.5.7__py3-none-any.whl → 0.6.1__py3-none-any.whl

_mcp_mesh/__init__.py

@@ -31,7 +31,7 @@ from .engine.decorator_registry import (
     get_decorator_stats,
 )
 
-__version__ = "0.5.7"
+__version__ = "0.6.1"
 
 # Store reference to runtime processor if initialized
 _runtime_processor = None

_mcp_mesh/engine/base_injector.py

@@ -0,0 +1,171 @@
+"""
+Base injector class with shared wrapper creation logic.
+
+Provides common functionality for DependencyInjector and MeshLlmAgentInjector.
+"""
+
+import functools
+import inspect
+import logging
+import weakref
+from collections.abc import Callable
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+class BaseInjector:
+    """
+    Base class for injection systems.
+
+    Provides shared functionality for creating and managing function wrappers
+    that support dynamic injection with two-phase updates.
+
+    Two-Phase Injection Pattern:
+    1. Phase 1 (decorator time): Create wrapper with initial state (None)
+    2. Phase 2 (runtime): Update wrapper with actual instances via update method
+
+    Subclasses must implement:
+    - Wrapper logic (what to inject and how)
+    - Update method signature
+    """
+
+    def __init__(self):
+        """Initialize base injector with function registry."""
+        self._function_registry: weakref.WeakValueDictionary = (
+            weakref.WeakValueDictionary()
+        )
+        logger.debug(f"🔧 {self.__class__.__name__} initialized")
+
+    def _register_wrapper(self, function_id: str, wrapper: Callable) -> None:
+        """
+        Register a wrapper in the function registry.
+
+        Args:
+            function_id: Unique function identifier
+            wrapper: Wrapper function to register
+        """
+        self._function_registry[function_id] = wrapper
+        logger.debug(f"🔧 Registered wrapper for {function_id} at {hex(id(wrapper))}")
+
+    def _create_async_wrapper(
+        self,
+        func: Callable,
+        function_id: str,
+        injection_logic: Callable[[Callable, tuple, dict], tuple],
+        metadata: dict[str, Any],
+    ) -> Callable:
+        """
+        Create async wrapper with injection logic.
+
+        Args:
+            func: Original async function to wrap
+            function_id: Unique function identifier
+            injection_logic: Callable that takes (func, args, kwargs) and returns (args, kwargs)
+                            This function should modify kwargs to inject dependencies
+            metadata: Additional metadata to store on wrapper
+
+        Returns:
+            Async wrapper function
+        """
+
+        @functools.wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            # Apply injection logic to modify kwargs
+            args, kwargs = injection_logic(func, args, kwargs)
+
+            # Execute original function
+            return await func(*args, **kwargs)
+
+        # Store metadata on wrapper
+        async_wrapper._mesh_original_func = func
+        async_wrapper._mesh_function_id = function_id
+
+        # Store additional metadata
+        for key, value in metadata.items():
+            setattr(async_wrapper, key, value)
+
+        return async_wrapper
+
+    def _create_sync_wrapper(
+        self,
+        func: Callable,
+        function_id: str,
+        injection_logic: Callable[[Callable, tuple, dict], tuple],
+        metadata: dict[str, Any],
+    ) -> Callable:
+        """
+        Create sync wrapper with injection logic.
+
+        Args:
+            func: Original sync function to wrap
+            function_id: Unique function identifier
+            injection_logic: Callable that takes (func, args, kwargs) and returns (args, kwargs)
+                            This function should modify kwargs to inject dependencies
+            metadata: Additional metadata to store on wrapper
+
+        Returns:
+            Sync wrapper function
+        """
+
+        @functools.wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            # Apply injection logic to modify kwargs
+            args, kwargs = injection_logic(func, args, kwargs)
+
+            # Execute original function
+            return func(*args, **kwargs)
+
+        # Store metadata on wrapper
+        sync_wrapper._mesh_original_func = func
+        sync_wrapper._mesh_function_id = function_id
+
+        # Store additional metadata
+        for key, value in metadata.items():
+            setattr(sync_wrapper, key, value)
+
+        return sync_wrapper
+
+    def create_wrapper_with_injection(
+        self,
+        func: Callable,
+        function_id: str,
+        injection_logic: Callable[[Callable, tuple, dict], tuple],
+        metadata: dict[str, Any],
+        register: bool = True,
+    ) -> Callable:
+        """
+        Create wrapper (async or sync) based on function type.
+
+        This is the main entry point for creating wrappers. It automatically
+        detects if the function is async or sync and creates the appropriate wrapper.
+
+        Args:
+            func: Function to wrap
+            function_id: Unique function identifier
+            injection_logic: Callable that takes (func, args, kwargs) and returns (args, kwargs)
+            metadata: Additional metadata to store on wrapper
+            register: Whether to register wrapper in function_registry (default: True)
+
+        Returns:
+            Wrapped function with injection capability
+        """
+        # Detect async vs sync
+        is_async = inspect.iscoroutinefunction(func)
+
+        if is_async:
+            wrapper = self._create_async_wrapper(
+                func, function_id, injection_logic, metadata
+            )
+            logger.debug(f"✅ Created async wrapper for {function_id}")
+        else:
+            wrapper = self._create_sync_wrapper(
+                func, function_id, injection_logic, metadata
+            )
+            logger.debug(f"✅ Created sync wrapper for {function_id}")
+
+        # Register wrapper if requested
+        if register:
+            self._register_wrapper(function_id, wrapper)
+
+        return wrapper

_mcp_mesh/engine/decorator_registry.py

@@ -32,6 +32,18 @@ class DecoratedFunction:
             self.metadata["function_name"] = self.function.__name__
 
 
+@dataclass
+class LLMAgentMetadata:
+    """Metadata for a function decorated with @mesh.llm."""
+
+    function: Callable
+    config: dict[str, Any]  # LLM configuration (provider, model, filter, etc.)
+    output_type: Optional[type]  # Pydantic model type from return annotation
+    param_name: str  # Name of MeshLlmAgent parameter
+    function_id: str  # Unique function ID for registry
+    registered_at: datetime
+
+
 class DecoratorRegistry:
     """
     Central registry for ALL MCP Mesh decorators.
@@ -52,19 +64,23 @@ class DecoratorRegistry:
     _mesh_tools: dict[str, DecoratedFunction] = {}  # Future use
     _mesh_resources: dict[str, DecoratedFunction] = {}  # Future use
     _mesh_workflows: dict[str, DecoratedFunction] = {}  # Future use
+    _mesh_llm_agents: dict[str, "LLMAgentMetadata"] = {}  # LLM agents with agentic loop
 
     # Registry for new decorator types (extensibility)
     _custom_decorators: dict[str, dict[str, DecoratedFunction]] = {}
-    
+
     # Immediate uvicorn server storage (for preventing shutdown state)
     _immediate_uvicorn_server: Optional[dict[str, Any]] = None
-    
+
     # FastMCP lifespan storage (for proper integration with FastAPI)
     _fastmcp_lifespan: Optional[Any] = None
-    
+
     # FastMCP HTTP app storage (the same app instance whose lifespan was extracted)
     _fastmcp_http_app: Optional[Any] = None
 
+    # FastMCP server info storage (for schema extraction during heartbeat)
+    _fastmcp_server_info: Optional[dict[str, Any]] = None
+
     @classmethod
     def register_mesh_agent(cls, func: Callable, metadata: dict[str, Any]) -> None:
         """
@@ -101,11 +117,13 @@ class DecoratorRegistry:
         if func_name in cls._mesh_tools:
             old_func = cls._mesh_tools[func_name].function
             cls._mesh_tools[func_name].function = new_func
-            print(
+            logger.debug(
                 f"🔄 DecoratorRegistry: Updated '{func_name}' from {hex(id(old_func))} to {hex(id(new_func))}"
             )
         else:
-            print(f"⚠️ DecoratorRegistry: Function '{func_name}' not found for update")
+            logger.debug(
+                f"⚠️ DecoratorRegistry: Function '{func_name}' not found for update"
+            )
 
     @classmethod
     def register_mesh_resource(cls, func: Callable, metadata: dict[str, Any]) -> None:
@@ -131,6 +149,53 @@ class DecoratorRegistry:
 
         cls._mesh_workflows[func.__name__] = decorated_func
 
+    @classmethod
+    def register_mesh_llm(
+        cls,
+        func: Callable,
+        config: dict[str, Any],
+        output_type: Optional[type],
+        param_name: str,
+        function_id: str,
+    ) -> None:
+        """
+        Register a @mesh.llm decorated function.
+
+        Args:
+            func: The decorated function
+            config: LLM configuration (provider, model, filter, etc.)
+            output_type: Pydantic model type from return annotation
+            param_name: Name of MeshLlmAgent parameter
+            function_id: Unique function ID for registry
+        """
+        llm_metadata = LLMAgentMetadata(
+            function=func,
+            config=config.copy(),
+            output_type=output_type,
+            param_name=param_name,
+            function_id=function_id,
+            registered_at=datetime.now(),
+        )
+
+        cls._mesh_llm_agents[function_id] = llm_metadata
+        logger.info(
+            f"🤖 Registered LLM agent: {func.__name__} (function_id={function_id}, param={param_name}, filter={config.get('filter')}, provider={config.get('provider')})"
+        )
+
+    @classmethod
+    def update_mesh_llm_function(cls, function_id: str, new_func: Callable) -> None:
+        """Update the function reference for a registered LLM agent (used for wrapper injection)."""
+        if function_id in cls._mesh_llm_agents:
+            old_func = cls._mesh_llm_agents[function_id].function
+            cls._mesh_llm_agents[function_id].function = new_func
+            logger.info(
+                f"🔄 DecoratorRegistry: Updated LLM function '{function_id}' from {hex(id(old_func))} to {hex(id(new_func))}"
+            )
+        else:
+            logger.warning(
+                f"⚠️ DecoratorRegistry: LLM function '{function_id}' not found for update"
+            )
+
     @classmethod
     def register_custom_decorator(
         cls, decorator_type: str, func: Callable, metadata: dict[str, Any]
@@ -175,6 +240,11 @@ class DecoratorRegistry:
         """Get all @mesh_workflow decorated functions."""
         return cls._mesh_workflows.copy()
 
+    @classmethod
+    def get_mesh_llm_agents(cls) -> dict[str, LLMAgentMetadata]:
+        """Get all @mesh.llm decorated functions."""
+        return cls._mesh_llm_agents.copy()
+
     @classmethod
     def get_all_by_type(cls, decorator_type: str) -> dict[str, DecoratedFunction]:
         """
@@ -256,6 +326,7 @@ class DecoratorRegistry:
         cls._mesh_tools.clear()
         cls._mesh_resources.clear()
         cls._mesh_workflows.clear()
+        cls._mesh_llm_agents.clear()
         cls._custom_decorators.clear()
 
         # Also clear the shared agent ID from mesh.decorators
@@ -290,23 +361,21 @@ class DecoratorRegistry:
     def update_agent_config(cls, updates: dict[str, Any]) -> None:
         """
         Update the cached agent configuration with new values.
-        
+
         This is useful for API services that generate their agent ID
         during pipeline execution and need to store it for telemetry.
-        
+
         Args:
             updates: Dictionary of config values to update
         """
         if cls._cached_agent_config is None:
             # Initialize with current resolved config if not cached yet
             cls._cached_agent_config = cls.get_resolved_agent_config().copy()
-        
+
         # Update with new values
         cls._cached_agent_config.update(updates)
-        
-        logger.debug(
-            f"🔧 Updated cached agent configuration with: {updates}"
-        )
+
+        logger.debug(f"🔧 Updated cached agent configuration with: {updates}")
 
     @classmethod
     def get_resolved_agent_config(cls) -> dict[str, Any]:
@@ -320,7 +389,9 @@ class DecoratorRegistry:
             dict: Pre-resolved configuration with consistent agent_id
         """
         # Step 1: Check if cached configuration already has agent_id (from API pipeline)
-        if cls._cached_agent_config is not None and cls._cached_agent_config.get('agent_id'):
+        if cls._cached_agent_config is not None and cls._cached_agent_config.get(
+            "agent_id"
+        ):
             logger.debug(
                 f"🔧 Using cached agent configuration: agent_id='{cls._cached_agent_config.get('agent_id')}'"
             )
@@ -348,15 +419,16 @@ class DecoratorRegistry:
         # Check if we're in an API context (have mesh_route decorators)
         mesh_routes = cls.get_all_by_type("mesh_route")
         is_api_context = len(mesh_routes) > 0
-        
+
         if is_api_context:
             # Use API service ID generation logic for consistency
             agent_id = cls._generate_api_service_id_fallback()
         else:
             # Use standard MCP agent ID generation
             from mesh.decorators import _get_or_create_agent_id
+
             agent_id = _get_or_create_agent_id()
-        
+
         fallback_config = {
             "name": None,
             "version": get_config_value(
@@ -415,44 +487,44 @@ class DecoratorRegistry:
     def _generate_api_service_id_fallback(cls) -> str:
         """
         Generate API service ID as fallback using same priority logic as API pipeline.
-        
+
         Priority order:
-        1. MCP_MESH_API_NAME environment variable 
+        1. MCP_MESH_API_NAME environment variable
         2. MCP_MESH_AGENT_NAME environment variable (fallback)
         3. Default to "api-{uuid8}"
-        
+
         Returns:
             Generated service ID with UUID suffix matching API service format
         """
         import uuid
-        
+
         from ..shared.config_resolver import ValidationRule, get_config_value
-        
+
         # Check for API-specific environment variable first (same as API pipeline)
         api_name = get_config_value(
             "MCP_MESH_API_NAME",
             default=None,
             rule=ValidationRule.STRING_RULE,
         )
-        
+
         # Fallback to general agent name env var
         if not api_name:
             api_name = get_config_value(
-                "MCP_MESH_AGENT_NAME", 
+                "MCP_MESH_AGENT_NAME",
                 default=None,
                 rule=ValidationRule.STRING_RULE,
             )
-        
+
         # Clean the service name if provided
         if api_name:
             cleaned_name = api_name.lower().replace(" ", "-").replace("_", "-")
             cleaned_name = "-".join(part for part in cleaned_name.split("-") if part)
         else:
             cleaned_name = ""
-        
+
         # Generate UUID suffix
         uuid_suffix = str(uuid.uuid4())[:8]
-        
+
         # Apply same naming logic as API pipeline
         if not cleaned_name:
             # No name provided: default to "api-{uuid8}"
@@ -463,8 +535,10 @@ class DecoratorRegistry:
         else:
             # Name doesn't contain "api": use "{name}-api-{uuid8}"
             service_id = f"{cleaned_name}-api-{uuid_suffix}"
-        
-        logger.debug(f"Generated fallback API service ID: '{service_id}' from env name: '{api_name}'")
+
+        logger.debug(
+            f"Generated fallback API service ID: '{service_id}' from env name: '{api_name}'"
+        )
         return service_id
 
     @classmethod
@@ -521,23 +595,25 @@ class DecoratorRegistry:
     def store_immediate_uvicorn_server(cls, server_info: dict[str, Any]) -> None:
         """
         Store reference to immediate uvicorn server started in decorator.
-        
+
         Args:
             server_info: Dictionary containing server information:
                 - 'app': FastAPI app instance
                 - 'host': Server host
-                - 'port': Server port  
+                - 'port': Server port
                 - 'thread': Thread object
                 - Any other relevant server metadata
         """
         cls._immediate_uvicorn_server = server_info
-        logger.debug(f"🔄 REGISTRY: Stored immediate uvicorn server reference: {server_info.get('host')}:{server_info.get('port')}")
+        logger.debug(
+            f"🔄 REGISTRY: Stored immediate uvicorn server reference: {server_info.get('host')}:{server_info.get('port')}"
+        )
 
     @classmethod
     def get_immediate_uvicorn_server(cls) -> Optional[dict[str, Any]]:
         """
         Get stored immediate uvicorn server reference.
-        
+
         Returns:
             Server info dict if available, None otherwise
         """
@@ -549,11 +625,33 @@ class DecoratorRegistry:
         cls._immediate_uvicorn_server = None
         logger.debug("🔄 REGISTRY: Cleared immediate uvicorn server reference")
 
+    # Health check result storage
+    _health_check_result: dict | None = None
+
+    @classmethod
+    def store_health_check_result(cls, result: dict) -> None:
+        """Store health check result for /health endpoint."""
+        cls._health_check_result = result
+        logger.debug(
+            f"💾 REGISTRY: Stored health check result: {result.get('status', 'unknown')}"
+        )
+
+    @classmethod
+    def get_health_check_result(cls) -> dict | None:
+        """Get stored health check result."""
+        return cls._health_check_result
+
+    @classmethod
+    def clear_health_check_result(cls) -> None:
+        """Clear stored health check result."""
+        cls._health_check_result = None
+        logger.debug("🗑️ REGISTRY: Cleared health check result")
+
     @classmethod
     def store_fastmcp_lifespan(cls, lifespan: Any) -> None:
         """
         Store FastMCP lifespan for integration with FastAPI.
-        
+
         Args:
             lifespan: FastMCP lifespan function
         """
@@ -564,7 +662,7 @@ class DecoratorRegistry:
     def get_fastmcp_lifespan(cls) -> Optional[Any]:
         """
         Get stored FastMCP lifespan.
-        
+
         Returns:
             FastMCP lifespan if available, None otherwise
         """
@@ -580,7 +678,7 @@ class DecoratorRegistry:
     def store_fastmcp_http_app(cls, http_app: Any) -> None:
         """
         Store FastMCP HTTP app (the same instance whose lifespan was extracted).
-        
+
         Args:
             http_app: FastMCP HTTP app instance
         """
@@ -591,7 +689,7 @@ class DecoratorRegistry:
     def get_fastmcp_http_app(cls) -> Optional[Any]:
         """
         Get stored FastMCP HTTP app.
-        
+
         Returns:
             FastMCP HTTP app if available, None otherwise
         """
@@ -603,6 +701,35 @@ class DecoratorRegistry:
         cls._fastmcp_http_app = None
         logger.debug("🔄 REGISTRY: Cleared FastMCP HTTP app reference")
 
+    @classmethod
+    def store_fastmcp_server_info(cls, server_info: dict[str, Any]) -> None:
+        """
+        Store FastMCP server info for schema extraction during heartbeat.
+
+        Args:
+            server_info: Dictionary of server_name -> server metadata (including tools)
+        """
+        cls._fastmcp_server_info = server_info
+        logger.debug(
+            f"🔄 REGISTRY: Stored FastMCP server info for {len(server_info)} servers"
+        )
+
+    @classmethod
+    def get_fastmcp_server_info(cls) -> Optional[dict[str, Any]]:
+        """
+        Get stored FastMCP server info.
+
+        Returns:
+            FastMCP server info if available, None otherwise
+        """
+        return cls._fastmcp_server_info
+
+    @classmethod
+    def clear_fastmcp_server_info(cls) -> None:
+        """Clear stored FastMCP server info reference."""
+        cls._fastmcp_server_info = None
+        logger.debug("🔄 REGISTRY: Cleared FastMCP server info reference")
+
 
 # Convenience functions for external access
 def get_all_mesh_agents() -> dict[str, DecoratedFunction]: