remdb 0.3.242__py3-none-any.whl

rem/agentic/mcp/tool_wrapper.py

@@ -0,0 +1,273 @@
+"""
+MCP Tool Wrappers for Pydantic AI.
+
+This module provides functions to convert MCP tool functions and resources
+into a format compatible with the Pydantic AI library.
+"""
+
+from typing import Any, Callable
+
+from loguru import logger
+from pydantic_ai.tools import Tool
+
+
+def create_pydantic_tool(func: Callable[..., Any]) -> Tool:
+    """
+    Create a Pydantic AI Tool from a given function.
+
+    This uses the Tool constructor, which inspects the
+    function's signature and docstring to create the tool schema.
+
+    Args:
+        func: The function to wrap as a tool.
+
+    Returns:
+        A Pydantic AI Tool instance.
+    """
+    logger.debug(f"Creating Pydantic tool from function: {func.__name__}")
+    return Tool(func)
+
+
+def create_mcp_tool_wrapper(
+    tool_name: str,
+    mcp_tool: Any,
+    user_id: str | None = None,
+    description_suffix: str | None = None,
+) -> Tool:
+    """
+    Create a Pydantic AI Tool from a FastMCP FunctionTool.
+
+    FastMCP tools are FunctionTool objects that wrap the actual async function.
+    We pass the function directly to Pydantic AI's Tool class, which will
+    inspect its signature properly. User ID injection happens in the wrapper.
+
+    Args:
+        tool_name: Name of the MCP tool
+        mcp_tool: The FastMCP FunctionTool object
+        user_id: Optional user_id to inject into tool calls
+        description_suffix: Optional text to append to the tool's docstring.
+            Used to add schema-specific context (e.g., default table for search_rem).
+
+    Returns:
+        A Pydantic AI Tool instance
+    """
+    # Extract the actual function from FastMCP FunctionTool
+    tool_func = mcp_tool.fn
+
+    # Check if function accepts user_id parameter
+    import inspect
+    sig = inspect.signature(tool_func)
+    has_user_id = "user_id" in sig.parameters
+
+    # Build the docstring with optional suffix
+    base_doc = tool_func.__doc__ or ""
+    final_doc = base_doc + description_suffix if description_suffix else base_doc
+
+    # If we need to inject user_id or modify docstring, create a wrapper
+    # Otherwise, use the function directly for better signature preservation
+    if user_id and has_user_id:
+        async def wrapped_tool(**kwargs) -> Any:
+            """Wrapper that injects user_id."""
+            if "user_id" not in kwargs:
+                kwargs["user_id"] = user_id
+                logger.debug(f"Injecting user_id={user_id} into tool {tool_name}")
+
+            # Filter kwargs to only include parameters that the function accepts
+            valid_params = set(sig.parameters.keys())
+            filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params}
+
+            return await tool_func(**filtered_kwargs)
+
+        # Copy signature from original function for Pydantic AI inspection
+        wrapped_tool.__name__ = tool_name
+        wrapped_tool.__doc__ = final_doc
+        wrapped_tool.__annotations__ = tool_func.__annotations__
+        wrapped_tool.__signature__ = sig  # Important: preserve full signature
+
+        logger.debug(f"Creating MCP tool wrapper with user_id injection: {tool_name}")
+        return Tool(wrapped_tool)
+    elif description_suffix:
+        # Need to wrap just for docstring modification
+        async def wrapped_tool(**kwargs) -> Any:
+            """Wrapper for docstring modification."""
+            valid_params = set(sig.parameters.keys())
+            filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params}
+            return await tool_func(**filtered_kwargs)
+
+        wrapped_tool.__name__ = tool_name
+        wrapped_tool.__doc__ = final_doc
+        wrapped_tool.__annotations__ = tool_func.__annotations__
+        wrapped_tool.__signature__ = sig
+
+        logger.debug(f"Creating MCP tool wrapper with description suffix: {tool_name}")
+        return Tool(wrapped_tool)
+    else:
+        # No injection needed - use original function directly
+        logger.debug(f"Creating MCP tool wrapper (no injection): {tool_name}")
+        return Tool(tool_func)
+
+
+def create_resource_tool(uri: str, usage: str = "", mcp_server: Any = None) -> Tool:
+    """
+    Build a Tool instance from an MCP resource URI.
+
+    Creates a tool that fetches the resource content when called.
+    Resources declared in agent YAML become callable tools - this eliminates
+    the artificial MCP distinction between tools and resources.
+
+    Supports both:
+    - Concrete URIs: "rem://agents" -> tool with no parameters
+    - Template URIs: "patient-profile://field/{field_key}" -> tool with field_key parameter
+
+    Args:
+        uri: The resource URI (concrete or template with {variable} placeholders).
+        usage: The description of what this resource provides.
+        mcp_server: Optional FastMCP server instance to resolve resources from.
+            If provided, resources are resolved from this server's registry.
+            If not provided, falls back to REM's built-in load_resource().
+
+    Returns:
+        A Pydantic AI Tool instance that fetches the resource.
+
+    Example:
+        # Concrete URI -> no-param tool
+        tool = create_resource_tool("rem://agents", "List all agent schemas")
+
+        # Template URI -> parameterized tool
+        tool = create_resource_tool("patient-profile://field/{field_key}", "Get field definition", mcp_server=mcp)
+        # Agent calls: get_patient_profile_field(field_key="safety.suicidality")
+    """
+    import json
+    import re
+
+    # Extract template variables from URI (e.g., {field_key}, {domain_name})
+    template_vars = re.findall(r'\{([^}]+)\}', uri)
+
+    # Parse URI to create function name (strip template vars for cleaner name)
+    clean_uri = re.sub(r'\{[^}]+\}', '', uri)
+    parts = clean_uri.replace("://", "_").replace("-", "_").replace("/", "_").replace(".", "_")
+    parts = re.sub(r'_+', '_', parts).strip('_')  # Clean up multiple underscores
+    func_name = f"get_{parts}"
+
+    # For parameterized URIs, append _by_{params} to avoid naming conflicts
+    # e.g., rem://agents/{name} -> get_rem_agents_by_name (distinct from get_rem_agents)
+    if template_vars:
+        param_suffix = "_by_" + "_".join(template_vars)
+        func_name = f"{func_name}{param_suffix}"
+
+    # Build description including parameter info
+    description = usage or f"Fetch {uri} resource"
+    if template_vars:
+        param_desc = ", ".join(template_vars)
+        description = f"{description}\n\nParameters: {param_desc}"
+
+    # Capture mcp_server reference at tool creation time (for closure)
+    # This ensures the correct server is used even if called later
+    _captured_mcp_server = mcp_server
+    _captured_uri = uri  # Also capture URI for consistent logging
+
+    if template_vars:
+        # Template URI -> create parameterized tool
+        async def wrapper(**kwargs: Any) -> str:
+            """Fetch MCP resource with substituted parameters."""
+            import asyncio
+            import inspect
+
+            logger.debug(f"Resource tool invoked: uri={_captured_uri}, kwargs={kwargs}, mcp_server={'set' if _captured_mcp_server else 'None'}")
+
+            # Try to resolve from MCP server's resource templates first
+            if _captured_mcp_server is not None:
+                try:
+                    # Get resource templates from MCP server
+                    templates = await _captured_mcp_server.get_resource_templates()
+                    logger.debug(f"MCP server templates: {list(templates.keys())}")
+                    if _captured_uri in templates:
+                        template = templates[_captured_uri]
+                        logger.debug(f"Found template for {_captured_uri}, calling fn with kwargs={kwargs}")
+                        # Call the template's underlying function directly
+                        # The fn expects the template variables as kwargs
+                        fn_result = template.fn(**kwargs)
+                        # Handle both sync and async functions
+                        if inspect.iscoroutine(fn_result):
+                            fn_result = await fn_result
+                        if isinstance(fn_result, str):
+                            return fn_result
+                        return json.dumps(fn_result, indent=2)
+                    else:
+                        logger.warning(f"Template {_captured_uri} not found in MCP server templates: {list(templates.keys())}")
+                except Exception as e:
+                    logger.warning(f"Failed to resolve resource {_captured_uri} from MCP server: {e}", exc_info=True)
+            else:
+                logger.warning(f"No MCP server provided for resource tool {_captured_uri}, using fallback")
+
+            # Fallback: substitute template variables and use load_resource
+            resolved_uri = _captured_uri
+            for var in template_vars:
+                if var in kwargs:
+                    resolved_uri = resolved_uri.replace(f"{{{var}}}", str(kwargs[var]))
+                else:
+                    return json.dumps({"error": f"Missing required parameter: {var}"})
+
+            logger.debug(f"Using fallback load_resource for resolved URI: {resolved_uri}")
+            from rem.api.mcp_router.resources import load_resource
+            result = await load_resource(resolved_uri)
+            if isinstance(result, str):
+                return result
+            return json.dumps(result, indent=2)
+
+        # Build parameter annotations for Pydantic AI
+        wrapper.__name__ = func_name
+        wrapper.__doc__ = description
+        # Add type hints for parameters
+        wrapper.__annotations__ = {var: str for var in template_vars}
+        wrapper.__annotations__['return'] = str
+
+        logger.info(f"Built parameterized resource tool: {func_name} (uri: {uri}, params: {template_vars}, mcp_server={'provided' if mcp_server else 'None'})")
+    else:
+        # Concrete URI -> no-param tool
+        async def wrapper(**kwargs: Any) -> str:
+            """Fetch MCP resource and return contents."""
+            import asyncio
+            import inspect
+
+            if kwargs:
+                logger.warning(f"Resource tool {func_name} called with unexpected kwargs: {list(kwargs.keys())}")
+
+            logger.debug(f"Concrete resource tool invoked: uri={_captured_uri}, mcp_server={'set' if _captured_mcp_server else 'None'}")
+
+            # Try to resolve from MCP server's resources first
+            if _captured_mcp_server is not None:
+                try:
+                    resources = await _captured_mcp_server.get_resources()
+                    logger.debug(f"MCP server resources: {list(resources.keys())}")
+                    if _captured_uri in resources:
+                        resource = resources[_captured_uri]
+                        logger.debug(f"Found resource for {_captured_uri}")
+                        # Call the resource's underlying function
+                        fn_result = resource.fn()
+                        if inspect.iscoroutine(fn_result):
+                            fn_result = await fn_result
+                        if isinstance(fn_result, str):
+                            return fn_result
+                        return json.dumps(fn_result, indent=2)
+                    else:
+                        logger.warning(f"Resource {_captured_uri} not found in MCP server resources: {list(resources.keys())}")
+                except Exception as e:
+                    logger.warning(f"Failed to resolve resource {_captured_uri} from MCP server: {e}", exc_info=True)
+            else:
+                logger.warning(f"No MCP server provided for resource tool {_captured_uri}, using fallback")
+
+            # Fallback to load_resource
+            logger.debug(f"Using fallback load_resource for URI: {_captured_uri}")
+            from rem.api.mcp_router.resources import load_resource
+            result = await load_resource(_captured_uri)
+            if isinstance(result, str):
+                return result
+            return json.dumps(result, indent=2)
+
+        wrapper.__name__ = func_name
+        wrapper.__doc__ = description
+
+        logger.info(f"Built resource tool: {func_name} (uri: {uri}, mcp_server={'provided' if mcp_server else 'None'})")
+
+    return Tool(wrapper)

rem/agentic/otel/__init__.py

@@ -0,0 +1,5 @@
+"""OpenTelemetry instrumentation for REM agents."""
+
+from .setup import setup_instrumentation, set_agent_resource_attributes
+
+__all__ = ["setup_instrumentation", "set_agent_resource_attributes"]

rem/agentic/otel/setup.py

@@ -0,0 +1,240 @@
+"""
+OpenTelemetry instrumentation setup for REM agents.
+
+Provides:
+- OTLP exporter configuration
+- Phoenix integration (OpenInference conventions)
+- Resource attributes for agent metadata
+- Idempotent setup (safe to call multiple times)
+"""
+
+from typing import Any
+
+from loguru import logger
+
+from ...settings import settings
+
+
+# Global flag to track if instrumentation is initialized
+_instrumentation_initialized = False
+
+
+def setup_instrumentation() -> None:
+    """
+    Initialize OpenTelemetry instrumentation for REM agents.
+
+    Idempotent - safe to call multiple times, only initializes once.
+
+    Configures:
+    - OTLP exporter (HTTP or gRPC)
+    - Phoenix integration if enabled
+    - Pydantic AI instrumentation (automatic via agent.instrument=True)
+    - Resource attributes (service name, environment, etc.)
+
+    Environment variables:
+        OTEL__ENABLED - Enable instrumentation (default: false)
+        OTEL__SERVICE_NAME - Service name (default: rem-api)
+        OTEL__COLLECTOR_ENDPOINT - OTLP endpoint (default: http://localhost:4318)
+        OTEL__PROTOCOL - Protocol (http or grpc, default: http)
+        PHOENIX__ENABLED - Enable Phoenix (default: false)
+        PHOENIX__COLLECTOR_ENDPOINT - Phoenix endpoint (default: http://localhost:6006/v1/traces)
+    """
+    global _instrumentation_initialized
+
+    if _instrumentation_initialized:
+        logger.debug("OTEL instrumentation already initialized, skipping")
+        return
+
+    if not settings.otel.enabled:
+        logger.debug("OTEL instrumentation disabled (OTEL__ENABLED=false)")
+        return
+
+    logger.info("Initializing OpenTelemetry instrumentation...")
+
+    try:
+        from opentelemetry import trace
+        from opentelemetry.sdk.trace import TracerProvider, ReadableSpan
+        from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanExporter, SpanExportResult
+        from opentelemetry.sdk.resources import Resource, SERVICE_NAME, DEPLOYMENT_ENVIRONMENT
+        from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter as HTTPExporter
+        from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter as GRPCExporter
+
+        class SanitizingSpanExporter(SpanExporter):
+            """
+            Wrapper exporter that sanitizes span attributes before export.
+
+            Removes None values that cause OTLP encoding failures like:
+            - llm.input_messages.3.message.content: None
+            """
+
+            def __init__(self, wrapped_exporter: SpanExporter):
+                self._wrapped = wrapped_exporter
+
+            def _sanitize_value(self, value):
+                """Recursively sanitize a value, replacing None with empty string."""
+                if value is None:
+                    return ""  # Replace None with empty string
+                if isinstance(value, dict):
+                    return {k: self._sanitize_value(v) for k, v in value.items()}
+                if isinstance(value, (list, tuple)):
+                    return [self._sanitize_value(v) for v in value]
+                return value
+
+            def export(self, spans: tuple[ReadableSpan, ...]) -> SpanExportResult:
+                # Create sanitized copies of spans
+                sanitized_spans = []
+                for span in spans:
+                    if span.attributes:
+                        # Sanitize all attribute values - replace None with empty string
+                        sanitized_attrs = {}
+                        for k, v in span.attributes.items():
+                            sanitized_attrs[k] = self._sanitize_value(v)
+                        sanitized_spans.append(_SanitizedSpan(span, sanitized_attrs))
+                    else:
+                        sanitized_spans.append(span)
+
+                return self._wrapped.export(tuple(sanitized_spans))
+
+            def shutdown(self) -> None:
+                self._wrapped.shutdown()
+
+            def force_flush(self, timeout_millis: int = 30000) -> bool:
+                return self._wrapped.force_flush(timeout_millis)
+
+        class _SanitizedSpan(ReadableSpan):
+            """ReadableSpan wrapper with sanitized attributes."""
+
+            def __init__(self, original: ReadableSpan, sanitized_attributes: dict):
+                self._original = original
+                self._sanitized_attributes = sanitized_attributes
+
+            @property
+            def name(self): return self._original.name
+            @property
+            def context(self): return self._original.context
+            @property
+            def parent(self): return self._original.parent
+            @property
+            def resource(self): return self._original.resource
+            @property
+            def instrumentation_scope(self): return self._original.instrumentation_scope
+            @property
+            def status(self): return self._original.status
+            @property
+            def start_time(self): return self._original.start_time
+            @property
+            def end_time(self): return self._original.end_time
+            @property
+            def links(self): return self._original.links
+            @property
+            def events(self): return self._original.events
+            @property
+            def kind(self): return self._original.kind
+            @property
+            def attributes(self): return self._sanitized_attributes
+            @property
+            def dropped_attributes(self): return self._original.dropped_attributes
+            @property
+            def dropped_events(self): return self._original.dropped_events
+            @property
+            def dropped_links(self): return self._original.dropped_links
+
+            def get_span_context(self): return self._original.get_span_context()
+
+        # Create resource with service metadata
+        resource = Resource(
+            attributes={
+                SERVICE_NAME: settings.otel.service_name,
+                DEPLOYMENT_ENVIRONMENT: settings.environment,
+                "service.team": settings.team,
+            }
+        )
+
+        # Create tracer provider
+        tracer_provider = TracerProvider(resource=resource)
+
+        # Configure OTLP exporter based on protocol
+        if settings.otel.protocol == "grpc":
+            base_exporter = GRPCExporter(
+                endpoint=settings.otel.collector_endpoint,
+                timeout=settings.otel.export_timeout,
+                insecure=settings.otel.insecure,
+            )
+        else:  # http
+            base_exporter = HTTPExporter(
+                endpoint=f"{settings.otel.collector_endpoint}/v1/traces",
+                timeout=settings.otel.export_timeout,
+            )
+
+        # Wrap with sanitizing exporter to handle None values
+        exporter = SanitizingSpanExporter(base_exporter)
+
+        # Add span processor
+        tracer_provider.add_span_processor(BatchSpanProcessor(exporter))
+
+        # Set as global tracer provider
+        trace.set_tracer_provider(tracer_provider)
+
+        logger.info(
+            f"OTLP exporter configured: {settings.otel.collector_endpoint} ({settings.otel.protocol})"
+        )
+
+        # Add OpenInference span processor for Pydantic AI
+        # This adds rich attributes (openinference.span.kind, input/output, etc.) to ALL traces
+        # Phoenix receives these traces via the OTLP collector - no separate "Phoenix integration" needed
+        # Note: The OTEL exporter may log warnings about None values in tool call messages,
+        # but this is a known limitation in openinference-instrumentation-pydantic-ai
+        try:
+            from openinference.instrumentation.pydantic_ai import OpenInferenceSpanProcessor as PydanticAISpanProcessor
+
+            tracer_provider.add_span_processor(PydanticAISpanProcessor())
+            logger.info("Added OpenInference span processor for Pydantic AI")
+
+        except ImportError:
+            logger.warning(
+                "openinference-instrumentation-pydantic-ai not installed - traces will lack OpenInference attributes. "
+                "Install with: pip install openinference-instrumentation-pydantic-ai"
+            )
+
+        _instrumentation_initialized = True
+        logger.info("OpenTelemetry instrumentation initialized successfully")
+
+    except Exception as e:
+        logger.error(f"Failed to initialize OTEL instrumentation: {e}")
+        # Don't raise - allow application to continue without tracing
+
+
+def set_agent_resource_attributes(agent_schema: dict[str, Any] | None = None) -> None:
+    """
+    Set resource attributes for agent execution.
+
+    Called before creating agent to set span attributes with agent metadata.
+
+    Args:
+        agent_schema: Agent schema with metadata (kind, name, version, etc.)
+    """
+    if not settings.otel.enabled or not agent_schema:
+        return
+
+    try:
+        from opentelemetry import trace
+
+        # Get current span and set attributes
+        span = trace.get_current_span()
+        if span.is_recording():
+            json_extra = agent_schema.get("json_schema_extra", {})
+            kind = json_extra.get("kind")
+            name = json_extra.get("name")
+            version = json_extra.get("version", "unknown")
+
+            if kind:
+                span.set_attribute("agent.kind", kind)
+            if name:
+                span.set_attribute("agent.name", name)
+            if version:
+                span.set_attribute("agent.version", version)
+
+            logger.debug(f"Set agent resource attributes: kind={kind}, name={name}, version={version}")
+
+    except Exception as e:
+        logger.warning(f"Failed to set agent resource attributes: {e}")