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

mesh/decorators.py

@@ -7,6 +7,7 @@ Provides @mesh.tool and @mesh.agent decorators with clean separation of concerns
 import logging
 import uuid
 from collections.abc import Callable
+from functools import wraps
 from typing import Any, TypeVar
 
 # Import from _mcp_mesh for registry and runtime integration
@@ -758,8 +759,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 +930,344 @@ 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 = "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 ("claude", "openai", "custom")
+        model: Model name (can be overridden by MESH_LLM_MODEL)
+        api_key: API key (can be overridden by provider-specific env vars)
+        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."
+            )
+
+        resolved_config = {
+            "filter": filter,
+            "filter_mode": get_config_value(
+                "MESH_LLM_FILTER_MODE",
+                override=filter_mode,
+                default="all",
+                rule=ValidationRule.STRING_RULE,
+            ),
+            "provider": get_config_value(
+                "MESH_LLM_PROVIDER",
+                override=provider,
+                default="claude",
+                rule=ValidationRule.STRING_RULE,
+            ),
+            "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

mesh/types.py

@@ -180,113 +180,95 @@ class McpMeshAgent(Protocol):
             }
 
 
-class McpAgent(Protocol):
+class MeshLlmAgent(Protocol):
     """
-    DEPRECATED: Use McpMeshAgent instead.
-
-    This type has been unified with McpMeshAgent. All features previously exclusive
-    to McpAgent are now available in McpMeshAgent using FastMCP's superior client.
-
-    Migration:
-        # Old way (deprecated)
-        def process_files(file_service: McpAgent) -> str:
-            pass
-
-        # New way (recommended)
-        def process_files(file_service: McpMeshAgent) -> str:
-            pass
-
-    McpMeshAgent now provides all MCP protocol features including streaming,
-    session management, and CallToolResult objects via FastMCP client.
+    LLM agent proxy with automatic agentic loop.
+
+    This protocol defines the interface for LLM agents that are automatically injected
+    by the @mesh.llm decorator. The proxy handles the entire agentic loop internally:
+    - Tool formatting for provider (Claude, OpenAI, etc.)
+    - LLM API calls
+    - Tool execution via MCP proxies
+    - Response parsing to Pydantic models
+
+    The MeshLlmAgent is injected by the mesh framework and configured via the
+    @mesh.llm decorator. Users only need to call the proxy with their message.
+
+    Usage Example:
+        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: MeshLlmAgent = None) -> ChatResponse:
+            # Optional: Override system prompt
+            llm.set_system_prompt("You are a helpful document assistant.")
+
+            # Execute automatic agentic loop
+            return llm(message)
+
+    Configuration Hierarchy:
+        - Decorator parameters provide defaults
+        - Environment variables override decorator settings:
+          * 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
+
+    The proxy is automatically injected with:
+        - Filtered tools from registry (based on @mesh.llm filter)
+        - Provider configuration (provider, model, api_key)
+        - Output type (inferred from function return annotation)
+        - System prompt (from decorator or file)
     """
 
-    # Basic compatibility with McpMeshAgent
-    def __call__(self, arguments: Optional[dict[str, Any]] = None) -> Any:
-        """Call the bound remote function (McpMeshAgent compatibility)."""
-        ...
-
-    def invoke(self, arguments: Optional[dict[str, Any]] = None) -> Any:
-        """Explicitly invoke the bound remote function (McpMeshAgent compatibility)."""
-        ...
-
-    # Vanilla MCP Protocol Methods (100% compatibility)
-    async def list_tools(self) -> list:
-        """List available tools from remote agent (vanilla MCP method)."""
-        ...
-
-    async def list_resources(self) -> list:
-        """List available resources from remote agent (vanilla MCP method)."""
-        ...
-
-    async def read_resource(self, uri: str) -> Any:
-        """Read resource contents from remote agent (vanilla MCP method)."""
-        ...
-
-    async def list_prompts(self) -> list:
-        """List available prompts from remote agent (vanilla MCP method)."""
-        ...
-
-    async def get_prompt(self, name: str, arguments: Optional[dict] = None) -> Any:
-        """Get prompt template from remote agent (vanilla MCP method)."""
-        ...
-
-    # Streaming Support - THE BREAKTHROUGH METHOD!
-    async def call_tool_streaming(
-        self, name: str, arguments: dict | None = None
-    ) -> AsyncIterator[dict]:
+    def set_system_prompt(self, prompt: str) -> None:
         """
-        Call a tool with streaming response using FastMCP's text/event-stream.
-
-        This enables multihop streaming (A→B→C chains) by leveraging FastMCP's
-        built-in streaming support with Accept: text/event-stream header.
+        Override the system prompt at runtime.
 
         Args:
-            name: Tool name to call
-            arguments: Tool arguments
-
-        Yields:
-            Streaming response chunks as dictionaries
-        """
-        ...
+            prompt: System prompt to use for LLM calls
 
-    # Phase 6: Explicit Session Management
-    async def create_session(self) -> str:
-        """
-        Create a new session and return session ID.
-
-        For Phase 6 explicit session management. In Phase 8, this will be
-        automated based on @mesh.tool(session_required=True) annotations.
-
-        Returns:
-            New session ID string
+        Example:
+            llm.set_system_prompt("You are an expert document analyst.")
         """
         ...
 
-    async def call_with_session(self, session_id: str, **kwargs) -> Any:
+    def __call__(self, message: str, **kwargs) -> Any:
         """
-        Call tool with explicit session ID for stateful operations.
+        Execute automatic agentic loop and return typed response.
 
-        This ensures all calls with the same session_id route to the same
-        agent instance for session affinity.
+        This method handles the complete agentic loop:
+        1. Format tools for provider (via LiteLLM)
+        2. Call LLM API with tools
+        3. If tool_use: execute via MCP proxies, loop back to LLM
+        4. If final response: parse into output type (Pydantic model)
+        5. Return typed response
 
         Args:
-            session_id: Session ID to include in request headers
-            **kwargs: Tool arguments to pass
+            message: User message to send to LLM
+            **kwargs: Additional context passed to LLM (provider-specific)
 
         Returns:
-            Tool response
-        """
-        ...
+            Pydantic model instance (type inferred from function return annotation)
 
-    async def close_session(self, session_id: str) -> bool:
-        """
-        Close session and cleanup session state.
-
-        Args:
-            session_id: Session ID to close
+        Raises:
+            MaxIterationsError: If max_iterations exceeded without final response
+            ValidationError: If LLM response doesn't match output type schema
+            ToolExecutionError: If tool execution fails during agentic loop
 
-        Returns:
-            True if session was closed successfully
+        Example:
+            response = llm("Analyze this document: /path/to/file.pdf")
+            # Returns ChatResponse(answer="...", confidence=0.95)
         """
         ...
 
@@ -299,13 +281,15 @@ class McpAgent(Protocol):
             handler: Any,
         ) -> core_schema.CoreSchema:
             """
-            Custom Pydantic core schema for McpAgent.
+            Custom Pydantic core schema for MeshLlmAgent.
+
+            This makes MeshLlmAgent parameters appear as optional/nullable in MCP schemas,
+            preventing serialization errors while maintaining type safety for dependency injection.
 
-            Similar to McpMeshAgent, this makes McpAgent parameters appear as
-            optional/nullable in MCP schemas, preventing serialization errors
-            while maintaining type safety for dependency injection.
+            The MeshLlmAgentInjector will replace None values with actual proxy objects
+            at runtime, so MCP callers never need to provide these parameters.
             """
-            # Treat McpAgent as an optional Any type for MCP serialization
+            # Treat MeshLlmAgent as an optional Any type for MCP serialization
             return core_schema.with_default_schema(
                 core_schema.nullable_schema(core_schema.any_schema()),
                 default=None,
@@ -320,3 +304,70 @@ class McpAgent(Protocol):
                 "schema": {"type": "nullable", "schema": {"type": "any"}},
                 "default": None,
             }
+
+
+# Import BaseModel for MeshContextModel
+try:
+    from pydantic import BaseModel
+
+    class MeshContextModel(BaseModel):
+        """
+        Base model for LLM prompt template contexts.
+
+        Use this to create type-safe, validated context models for
+        Jinja2 prompt templates in @mesh.llm decorated functions.
+
+        The MeshContextModel provides:
+        - Type safety via Pydantic validation
+        - Field descriptions for LLM schema generation
+        - Strict mode (extra fields forbidden)
+        - Automatic .model_dump() for template rendering
+
+        Example:
+            from mesh import MeshContextModel
+            from pydantic import Field
+
+            class ChatContext(MeshContextModel):
+                user_name: str = Field(description="Name of the user")
+                domain: str = Field(description="Chat domain: support, sales, etc.")
+                expertise_level: str = Field(
+                    default="beginner",
+                    description="User expertise: beginner, intermediate, expert"
+                )
+
+            @mesh.llm(
+                system_prompt="file://prompts/chat.jinja2",
+                context_param="ctx"
+            )
+            @mesh.tool(capability="chat")
+            def chat(message: str, ctx: ChatContext, llm: MeshLlmAgent = None):
+                return llm(message)  # Template auto-rendered with ctx!
+
+        Field Descriptions in LLM Chains:
+            When a specialist LLM agent has MeshContextModel parameters, the Field
+            descriptions are extracted and included in the tool schema sent to
+            calling LLM agents. This helps orchestrator LLMs construct context
+            objects correctly.
+
+            Without descriptions:
+                {"domain": "string"}  # LLM doesn't know what this means
+
+            With descriptions:
+                {"domain": {"type": "string", "description": "Chat domain: support, sales"}}
+                # LLM understands what to provide!
+
+        Template Rendering:
+            When used with @mesh.llm(system_prompt="file://..."), the context is
+            automatically converted to a dict via .model_dump() and passed to the
+            Jinja2 template renderer.
+        """
+
+        class Config:
+            extra = "forbid"  # Strict mode - reject unexpected fields
+
+except ImportError:
+    # Fallback if Pydantic not available (should not happen in practice)
+    class MeshContextModel:  # type: ignore
+        """Placeholder when Pydantic unavailable."""
+
+        pass