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

mesh/decorators.py

@@ -6,7 +6,8 @@ Provides @mesh.tool and @mesh.agent decorators with clean separation of concerns
 
 import logging
 import uuid
-from collections.abc import Callable
+from collections.abc import Awaitable, Callable
+from functools import wraps
 from typing import Any, TypeVar
 
 # Import from _mcp_mesh for registry and runtime integration
@@ -32,7 +33,7 @@ def _start_uvicorn_immediately(http_host: str, http_port: int):
     This prevents the DNS threading conflicts by ensuring uvicorn takes control
     before the script ends and Python enters shutdown state.
     """
-    logger.info(
+    logger.debug(
         f"🎯 IMMEDIATE UVICORN: _start_uvicorn_immediately() called with host={http_host}, port={http_port}"
     )
 
@@ -42,9 +43,9 @@ def _start_uvicorn_immediately(http_host: str, http_port: int):
         import time
 
         import uvicorn
-        from fastapi import FastAPI
+        from fastapi import FastAPI, Response
 
-        logger.info(
+        logger.debug(
             "📦 IMMEDIATE UVICORN: Successfully imported uvicorn, FastAPI, threading, asyncio"
         )
 
@@ -55,11 +56,11 @@ def _start_uvicorn_immediately(http_host: str, http_port: int):
 
             fastmcp_lifespan = DecoratorRegistry.get_fastmcp_lifespan()
             if fastmcp_lifespan:
-                logger.info(
+                logger.debug(
                     "✅ IMMEDIATE UVICORN: Found stored FastMCP lifespan, will integrate with FastAPI"
                 )
             else:
-                logger.info(
+                logger.debug(
                     "🔍 IMMEDIATE UVICORN: No FastMCP lifespan found, creating basic FastAPI app"
                 )
         except Exception as e:
@@ -68,20 +69,61 @@ def _start_uvicorn_immediately(http_host: str, http_port: int):
         # Create FastAPI app with FastMCP lifespan if available
         if fastmcp_lifespan:
             app = FastAPI(title="MCP Mesh Agent (Starting)", lifespan=fastmcp_lifespan)
-            logger.info(
+            logger.debug(
                 "📦 IMMEDIATE UVICORN: Created FastAPI app with FastMCP lifespan integration"
             )
         else:
             app = FastAPI(title="MCP Mesh Agent (Starting)")
-            logger.info("📦 IMMEDIATE UVICORN: Created minimal FastAPI app")
+            logger.debug("📦 IMMEDIATE UVICORN: Created minimal FastAPI app")
+
+        # Add health endpoint that can be updated by pipeline
+        # Store health check result in a shared location that can be updated
+        health_result = {"status": "starting", "message": "Agent is starting"}
 
-        # Add basic health endpoint
         @app.get("/health")
-        def health():
-            return {
-                "status": "immediate_uvicorn",
-                "message": "MCP Mesh agent started via immediate uvicorn",
-            }
+        @app.head("/health")
+        async def health(response: Response):
+            """Health check endpoint that supports custom health checks."""
+            # Check if a custom health check has been configured
+            # The pipeline will update this via DecoratorRegistry
+            custom_health = DecoratorRegistry.get_health_check_result()
+            health_data = custom_health if custom_health else health_result
+
+            # Set HTTP status code based on health status
+            # K8s expects non-200 status for unhealthy services
+            status = health_data.get("status", "starting")
+            if status == "healthy":
+                response.status_code = 200
+            else:
+                # Return 503 for unhealthy, degraded, starting, or unknown
+                response.status_code = 503
+
+            return health_data
+
+        @app.get("/ready")
+        @app.head("/ready")
+        async def ready(response: Response):
+            """Kubernetes readiness probe - service ready to serve traffic."""
+            custom_health = DecoratorRegistry.get_health_check_result()
+            health_data = custom_health if custom_health else health_result
+
+            status = health_data.get("status", "starting")
+            if status == "healthy":
+                response.status_code = 200
+                return {"ready": True, "status": status}
+            else:
+                response.status_code = 503
+                return {
+                    "ready": False,
+                    "status": status,
+                    "reason": f"Service is {status}",
+                }
+
+        @app.get("/livez")
+        @app.head("/livez")
+        async def livez():
+            """Kubernetes liveness probe - always returns 200 if app is running."""
+            return {"alive": True, "message": "Application is running"}
 
         @app.get("/immediate-status")
         def immediate_status():
@@ -90,17 +132,17 @@ def _start_uvicorn_immediately(http_host: str, http_port: int):
                 "message": "This server was started immediately in decorator",
             }
 
-        logger.info("📦 IMMEDIATE UVICORN: Added health endpoints")
+        logger.debug("📦 IMMEDIATE UVICORN: Added status endpoints")
 
         # Determine port (0 means auto-assign)
         port = http_port if http_port > 0 else 8080
 
-        logger.info(
+        logger.debug(
             f"🚀 IMMEDIATE UVICORN: Starting uvicorn server on {http_host}:{port}"
         )
 
         # Use uvicorn.run() for proper signal handling (enables FastAPI lifespan shutdown)
-        logger.info(
+        logger.debug(
             "⚡ IMMEDIATE UVICORN: Starting server with uvicorn.run() for proper signal handling"
         )
 
@@ -108,7 +150,7 @@ def _start_uvicorn_immediately(http_host: str, http_port: int):
         def run_server():
             """Run uvicorn server in background thread with proper signal handling."""
             try:
-                logger.info(
+                logger.debug(
                     f"🌟 IMMEDIATE UVICORN: Starting server on {http_host}:{port}"
                 )
                 # Use uvicorn.run() instead of Server().run() for proper signal handling
@@ -130,7 +172,7 @@ def _start_uvicorn_immediately(http_host: str, http_port: int):
         thread = threading.Thread(target=run_server, daemon=False)
         thread.start()
 
-        logger.info(
+        logger.debug(
             "🔒 IMMEDIATE UVICORN: Server thread started (daemon=False) - can handle signals"
         )
 
@@ -151,14 +193,14 @@ def _start_uvicorn_immediately(http_host: str, http_port: int):
 
         DecoratorRegistry.store_immediate_uvicorn_server(server_info)
 
-        logger.info(
+        logger.debug(
             "🔄 IMMEDIATE UVICORN: Server reference stored in DecoratorRegistry BEFORE pipeline starts"
         )
 
         # Give server a moment to start
         time.sleep(1)
 
-        logger.info(
+        logger.debug(
             f"✅ IMMEDIATE UVICORN: Uvicorn server running on {http_host}:{port} (daemon thread)"
         )
 
@@ -452,6 +494,8 @@ def agent(
     enable_http: bool = True,
     namespace: str = "default",
     health_interval: int = 5,  # Will be overridden by centralized defaults
+    health_check: Callable[[], Awaitable[Any]] | None = None,
+    health_check_ttl: int = 15,
     auto_run: bool = True,  # Changed to True by default!
     auto_run_interval: int = 10,
     **kwargs: Any,
@@ -476,6 +520,10 @@ def agent(
             Environment variable: MCP_MESH_NAMESPACE (takes precedence)
         health_interval: Health check interval in seconds (default: 30)
             Environment variable: MCP_MESH_HEALTH_INTERVAL (takes precedence)
+        health_check: Optional async function that returns HealthStatus
+            Called before heartbeat and on /health endpoint with TTL caching
+        health_check_ttl: Cache TTL for health check results in seconds (default: 15)
+            Reduces expensive health check calls by caching results
         auto_run: Automatically start service and keep process alive (default: True)
             Environment variable: MCP_MESH_AUTO_RUN (takes precedence)
         auto_run_interval: Keep-alive heartbeat interval in seconds (default: 10)
@@ -551,6 +599,14 @@ def agent(
         if auto_run_interval < 1:
             raise ValueError("auto_run_interval must be at least 1 second")
 
+        if health_check is not None and not callable(health_check):
+            raise ValueError("health_check must be a callable (async function)")
+
+        if not isinstance(health_check_ttl, int):
+            raise ValueError("health_check_ttl must be an integer")
+        if health_check_ttl < 1:
+            raise ValueError("health_check_ttl must be at least 1 second")
+
         # Separate binding host (for uvicorn server) from external host (for registry)
         from _mcp_mesh.shared.host_resolver import HostResolver
 
@@ -622,6 +678,8 @@ def agent(
             "enable_http": final_enable_http,
             "namespace": final_namespace,
             "health_interval": final_health_interval,
+            "health_check": health_check,
+            "health_check_ttl": health_check_ttl,
             "auto_run": final_auto_run,
             "auto_run_interval": final_auto_run_interval,
             "agent_id": agent_id,
@@ -646,7 +704,7 @@ def agent(
 
         # Auto-run functionality: start uvicorn immediately to prevent Python shutdown state
         if final_auto_run:
-            logger.info(
+            logger.debug(
                 f"🚀 AGENT DECORATOR: Auto-run enabled for agent '{name}' - starting uvicorn immediately to prevent shutdown state"
             )
 
@@ -654,7 +712,7 @@ def agent(
             fastmcp_lifespan = None
             try:
                 # Try to create FastMCP server and extract lifespan
-                logger.info(
+                logger.debug(
                     "🔍 AGENT DECORATOR: Creating FastMCP server for lifespan extraction"
                 )
 
@@ -666,7 +724,7 @@ def agent(
                     # Look for 'app' attribute (standard FastMCP pattern)
                     if hasattr(current_module, "app"):
                         fastmcp_server = current_module.app
-                        logger.info(
+                        logger.debug(
                             f"🔍 AGENT DECORATOR: Found FastMCP server: {type(fastmcp_server)}"
                         )
 
@@ -680,7 +738,7 @@ def agent(
                                 )
                                 if hasattr(fastmcp_http_app, "lifespan"):
                                     fastmcp_lifespan = fastmcp_http_app.lifespan
-                                    logger.info(
+                                    logger.debug(
                                         "✅ AGENT DECORATOR: Extracted FastMCP lifespan for FastAPI integration"
                                     )
 
@@ -691,7 +749,7 @@ def agent(
                                     DecoratorRegistry.store_fastmcp_http_app(
                                         fastmcp_http_app
                                     )
-                                    logger.info(
+                                    logger.debug(
                                         "✅ AGENT DECORATOR: Stored FastMCP HTTP app for proper mounting"
                                     )
                                 else:
@@ -707,7 +765,7 @@ def agent(
                                 "⚠️ AGENT DECORATOR: FastMCP server has no http_app method"
                             )
                     else:
-                        logger.info(
+                        logger.debug(
                             "🔍 AGENT DECORATOR: No FastMCP 'app' found in current module - will handle in pipeline"
                         )
                 else:
@@ -720,12 +778,12 @@ def agent(
                     f"⚠️ AGENT DECORATOR: FastMCP lifespan creation failed: {e}"
                 )
 
-            logger.info(
+            logger.debug(
                 f"🎯 AGENT DECORATOR: About to call _start_uvicorn_immediately({binding_host}, {final_http_port})"
             )
             # Start basic uvicorn server immediately to prevent interpreter shutdown
             _start_uvicorn_immediately(binding_host, final_http_port)
-            logger.info(
+            logger.debug(
                 "✅ AGENT DECORATOR: _start_uvicorn_immediately() call completed"
             )
 
@@ -758,8 +816,8 @@ def route(
         async def upload_resume(
             request: Request,
             file: UploadFile = File(...),
-            pdf_agent: McpAgent = None,    # Injected by MCP Mesh
-            user_service: McpAgent = None  # Injected by MCP Mesh
+            pdf_agent: mesh.McpMeshAgent = None,    # Injected by MCP Mesh
+            user_service: mesh.McpMeshAgent = None  # Injected by MCP Mesh
         ):
             result = await pdf_agent.extract_text_from_pdf(file)
             await user_service.update_profile(user_data, result)
@@ -929,3 +987,353 @@ def set_shutdown_context(context: dict[str, Any]):
     """Set context for graceful shutdown (called from pipeline)."""
     # Delegate to the shared graceful shutdown manager
     set_global_shutdown_context(context)
+
+
+def llm(
+    filter: dict[str, Any] | list[dict[str, Any] | str] | str | None = None,
+    *,
+    filter_mode: str = "all",
+    provider: str | dict[str, Any] = "claude",
+    model: str | None = None,
+    api_key: str | None = None,
+    max_iterations: int = 10,
+    system_prompt: str | None = None,
+    system_prompt_file: str | None = None,
+    context_param: str | None = None,
+    **kwargs: Any,
+) -> Callable[[T], T]:
+    """
+    LLM agent decorator with automatic agentic loop.
+
+    This decorator enables LLM agents to automatically access mesh tools via
+    dependency injection. The MeshLlmAgent proxy handles the complete agentic loop:
+    - Tool filtering based on filter parameter
+    - LLM API calls (Claude, OpenAI, etc. via LiteLLM)
+    - Tool execution via MCP proxies
+    - Response parsing to Pydantic models
+
+    Configuration Hierarchy (ENV > Decorator):
+        - MESH_LLM_PROVIDER: Override provider
+        - MESH_LLM_MODEL: Override model
+        - ANTHROPIC_API_KEY: Claude API key
+        - OPENAI_API_KEY: OpenAI API key
+        - MESH_LLM_MAX_ITERATIONS: Override max iterations
+
+    Usage:
+        from pydantic import BaseModel
+        import mesh
+
+        class ChatResponse(BaseModel):
+            answer: str
+            confidence: float
+
+        @mesh.llm(
+            filter={"capability": "document", "tags": ["pdf"]},
+            provider="claude",
+            model="claude-3-5-sonnet-20241022"
+        )
+        @mesh.tool(capability="chat")
+        def chat(message: str, llm: mesh.MeshLlmAgent = None) -> ChatResponse:
+            llm.set_system_prompt("You are a helpful assistant.")
+            return llm(message)
+
+    Args:
+        filter: Tool filter (string, dict, or list of mixed)
+        filter_mode: Filter mode ("all", "best_match", "*")
+        provider: LLM provider (string like "claude" for direct LiteLLM, or dict for mesh delegation)
+                  Mesh delegation format: {"capability": "llm", "tags": ["claude"], "version": ">=1.0.0"}
+                  When dict: Uses mesh DI to resolve provider agent instead of calling LiteLLM directly
+        model: Model name (can be overridden by MESH_LLM_MODEL) - only used with string provider
+        api_key: API key (can be overridden by provider-specific env vars) - only used with string provider
+        max_iterations: Max agentic loop iterations (can be overridden by MESH_LLM_MAX_ITERATIONS)
+        system_prompt: Default system prompt
+        system_prompt_file: Path to Jinja2 template file
+        **kwargs: Additional configuration
+
+    Returns:
+        Decorated function with MeshLlmAgent injection
+
+    Raises:
+        ValueError: If no MeshLlmAgent parameter found
+        UserWarning: If multiple MeshLlmAgent parameters or non-Pydantic return type
+    """
+    import inspect
+    import warnings
+
+    def decorator(func: T) -> T:
+        # Step 1: Resolve configuration with hierarchy (ENV > decorator params)
+        # Phase 1: Detect file:// prefix for template files
+        is_template = False
+        template_path = None
+
+        if system_prompt:
+            # Check for file:// prefix
+            if system_prompt.startswith("file://"):
+                is_template = True
+                template_path = system_prompt[7:]  # Strip "file://" prefix
+            # Auto-detect .jinja2 or .j2 extension without file:// prefix
+            elif system_prompt.endswith(".jinja2") or system_prompt.endswith(".j2"):
+                is_template = True
+                template_path = system_prompt
+
+        # Backward compatibility: system_prompt_file (deprecated)
+        if system_prompt_file:
+            logger.warning(
+                f"⚠️ @mesh.llm: 'system_prompt_file' parameter is deprecated. "
+                f"Use 'system_prompt=\"file://{system_prompt_file}\"' instead."
+            )
+            if not is_template:  # Only use if system_prompt didn't specify a template
+                is_template = True
+                template_path = system_prompt_file
+
+        # Validate context_param usage
+        if context_param and not is_template:
+            logger.warning(
+                f"⚠️ @mesh.llm: 'context_param' specified for function '{func.__name__}' "
+                f"but system_prompt is not a template (no file:// prefix or .jinja2/.j2 extension). "
+                f"Context parameter will be ignored."
+            )
+
+        # Handle provider config: dict (mesh delegation) or string (direct LiteLLM)
+        # If provider is dict, don't allow env var override (explicit mesh delegation)
+        if isinstance(provider, dict):
+            resolved_provider = provider
+        else:
+            resolved_provider = get_config_value(
+                "MESH_LLM_PROVIDER",
+                override=provider,
+                default="claude",
+                rule=ValidationRule.STRING_RULE,
+            )
+
+        resolved_config = {
+            "filter": filter,
+            "filter_mode": get_config_value(
+                "MESH_LLM_FILTER_MODE",
+                override=filter_mode,
+                default="all",
+                rule=ValidationRule.STRING_RULE,
+            ),
+            "provider": resolved_provider,
+            "model": get_config_value(
+                "MESH_LLM_MODEL",
+                override=model,
+                default=None,
+                rule=ValidationRule.STRING_RULE,
+            ),
+            "api_key": api_key,  # Will be resolved from provider-specific env vars later
+            "max_iterations": get_config_value(
+                "MESH_LLM_MAX_ITERATIONS",
+                override=max_iterations,
+                default=10,
+                rule=ValidationRule.NONZERO_RULE,
+            ),
+            "system_prompt": system_prompt,
+            "system_prompt_file": system_prompt_file,
+            # Phase 1: Template metadata
+            "is_template": is_template,
+            "template_path": template_path,
+            "context_param": context_param,
+        }
+        resolved_config.update(kwargs)
+
+        # Step 2: Extract output type from return annotation
+        sig = inspect.signature(func)
+        return_annotation = sig.return_annotation
+
+        output_type = None
+        if return_annotation and return_annotation != inspect.Signature.empty:
+            output_type = return_annotation
+
+            # Warn if not a Pydantic model
+            try:
+                from pydantic import BaseModel
+
+                if not (
+                    inspect.isclass(output_type) and issubclass(output_type, BaseModel)
+                ):
+                    warnings.warn(
+                        f"Function '{func.__name__}' decorated with @mesh.llm should return a Pydantic BaseModel subclass, "
+                        f"got {output_type}. This may cause validation errors at runtime.",
+                        UserWarning,
+                        stacklevel=2,
+                    )
+            except ImportError:
+                pass  # Pydantic not available, skip validation
+
+        # Step 3: Find MeshLlmAgent parameter
+        from mesh.types import MeshLlmAgent
+
+        llm_params = []
+        for param_name, param in sig.parameters.items():
+            if param.annotation == MeshLlmAgent or (
+                hasattr(param.annotation, "__origin__")
+                and param.annotation.__origin__ == MeshLlmAgent
+            ):
+                llm_params.append(param_name)
+
+        if not llm_params:
+            raise ValueError(
+                f"Function '{func.__name__}' decorated with @mesh.llm must have at least one parameter "
+                f"of type 'mesh.MeshLlmAgent'. Example: def {func.__name__}(..., llm: mesh.MeshLlmAgent = None)"
+            )
+
+        if len(llm_params) > 1:
+            warnings.warn(
+                f"Function '{func.__name__}' has multiple MeshLlmAgent parameters: {llm_params}. "
+                f"Only the first parameter '{llm_params[0]}' will be injected. "
+                f"Additional parameters will be ignored.",
+                UserWarning,
+                stacklevel=2,
+            )
+
+        param_name = llm_params[0]
+
+        # Step 4: Generate unique function ID
+        function_id = f"{func.__name__}_{uuid.uuid4().hex[:8]}"
+
+        # Step 5: Register with DecoratorRegistry
+        DecoratorRegistry.register_mesh_llm(
+            func=func,
+            config=resolved_config,
+            output_type=output_type,
+            param_name=param_name,
+            function_id=function_id,
+        )
+
+        logger.debug(
+            f"@mesh.llm registered: {func.__name__} "
+            f"(provider={resolved_config['provider']}, param={param_name}, filter={filter})"
+        )
+
+        # Step 6: Enhance existing wrapper from @mesh.tool (if present)
+        # or create new wrapper
+        #
+        # This approach:
+        # - Reuses the wrapper created by @mesh.tool (if present)
+        # - Avoids creating multiple wrapper layers
+        # - Ensures FastMCP caches the SAME wrapper instance we update later
+        # - Combines both DI injection and LLM injection in the same wrapper
+
+        # Check if there's an existing wrapper from @mesh.tool
+        mesh_tools = DecoratorRegistry.get_mesh_tools()
+        existing_wrapper = None
+
+        if func.__name__ in mesh_tools:
+            existing_wrapper = mesh_tools[func.__name__].function
+            logger.info(
+                f"🔗 Found existing @mesh.tool wrapper for '{func.__name__}' at {hex(id(existing_wrapper))} - enhancing it"
+            )
+
+        # Trigger debounced processing
+        _trigger_debounced_processing()
+
+        if existing_wrapper:
+            # ENHANCE the existing wrapper with LLM attributes
+            logger.info(
+                f"✨ Enhancing existing wrapper with LLM injection for '{func.__name__}'"
+            )
+
+            # Store the original wrapped function if not already stored
+            if not hasattr(existing_wrapper, "__wrapped__"):
+                existing_wrapper.__wrapped__ = func
+
+            # Store the original call behavior to preserve DI injection
+            original_call = existing_wrapper
+
+            # Create enhanced wrapper that does BOTH DI injection and LLM injection
+            @wraps(func)
+            def combined_injection_wrapper(*args, **kwargs):
+                """Wrapper that injects both MeshLlmAgent and DI parameters."""
+                # Inject LLM parameter if not provided or if it's None
+                if param_name not in kwargs or kwargs.get(param_name) is None:
+                    kwargs[param_name] = combined_injection_wrapper._mesh_llm_agent
+                # Then call the original wrapper (which handles DI injection)
+                return original_call(*args, **kwargs)
+
+            # Add LLM metadata attributes to combined wrapper
+            combined_injection_wrapper._mesh_llm_agent = (
+                None  # Will be updated during heartbeat
+            )
+            combined_injection_wrapper._mesh_llm_param_name = param_name
+            combined_injection_wrapper._mesh_llm_function_id = function_id
+            combined_injection_wrapper._mesh_llm_config = resolved_config
+            combined_injection_wrapper._mesh_llm_output_type = output_type
+            combined_injection_wrapper.__wrapped__ = func
+
+            # Create update method for heartbeat that updates the COMBINED wrapper
+            def update_llm_agent(agent):
+                combined_injection_wrapper._mesh_llm_agent = agent
+                logger.info(
+                    f"🔄 Updated MeshLlmAgent on combined wrapper for {func.__name__} (function_id={function_id})"
+                )
+
+            combined_injection_wrapper._mesh_update_llm_agent = update_llm_agent
+
+            # Copy any other mesh attributes from existing wrapper
+            for attr in dir(existing_wrapper):
+                if attr.startswith("_mesh_") and not hasattr(
+                    combined_injection_wrapper, attr
+                ):
+                    try:
+                        setattr(
+                            combined_injection_wrapper,
+                            attr,
+                            getattr(existing_wrapper, attr),
+                        )
+                    except AttributeError:
+                        pass  # Some attributes might not be settable
+
+            # Update DecoratorRegistry with the combined wrapper
+            DecoratorRegistry.update_mesh_llm_function(
+                function_id, combined_injection_wrapper
+            )
+            DecoratorRegistry.update_mesh_tool_function(
+                func.__name__, combined_injection_wrapper
+            )
+
+            logger.info(
+                f"✅ Enhanced wrapper for '{func.__name__}' with combined DI + LLM injection at {hex(id(combined_injection_wrapper))}"
+            )
+
+            # Return the enhanced wrapper
+            return combined_injection_wrapper
+
+        else:
+            # FALLBACK: Create new wrapper if no existing @mesh.tool wrapper found
+            logger.info(
+                f"📝 No existing wrapper found for '{func.__name__}' - creating new LLM wrapper"
+            )
+
+            @wraps(func)
+            def llm_injection_wrapper(*args, **kwargs):
+                """Wrapper that injects MeshLlmAgent parameter."""
+                # Inject llm parameter if not provided or if it's None
+                if param_name not in kwargs or kwargs.get(param_name) is None:
+                    kwargs[param_name] = llm_injection_wrapper._mesh_llm_agent
+                return func(*args, **kwargs)
+
+            # Create update method for heartbeat - updates the wrapper, not func
+            def update_llm_agent(agent):
+                llm_injection_wrapper._mesh_llm_agent = agent
+                logger.info(
+                    f"🔄 Updated MeshLlmAgent for {func.__name__} (function_id={function_id})"
+                )
+
+            # Copy all metadata attributes to the wrapper
+            llm_injection_wrapper._mesh_llm_agent = None
+            llm_injection_wrapper._mesh_llm_param_name = param_name
+            llm_injection_wrapper._mesh_llm_function_id = function_id
+            llm_injection_wrapper._mesh_llm_config = resolved_config
+            llm_injection_wrapper._mesh_llm_output_type = output_type
+            llm_injection_wrapper._mesh_update_llm_agent = update_llm_agent
+
+            # Update DecoratorRegistry with the wrapper
+            DecoratorRegistry.update_mesh_llm_function(
+                function_id, llm_injection_wrapper
+            )
+
+            # Return the new wrapper
+            return llm_injection_wrapper
+
+    return decorator