remdb 0.3.7__py3-none-any.whl

rem/agentic/mcp/tool_wrapper.py

@@ -0,0 +1,107 @@
+"""
+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) -> 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
+
+    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
+
+    # If we need to inject user_id, 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__ = tool_func.__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)
+    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) -> Tool:
+    """
+    Build a Tool instance from an MCP resource URI.
+
+    This is a placeholder for now. A real implementation would create a
+    tool that reads the content of the resource URI.
+
+    Args:
+        uri: The resource URI (e.g., "rem://resources/some-id").
+        usage: The description of how to use the tool.
+
+    Returns:
+        A Pydantic AI Tool instance.
+    """
+    # Placeholder function that would read the resource
+    def read_resource():
+        """Reads content from a resource URI."""
+        return f"Content of {uri}"
+
+    read_resource.__name__ = f"read_{uri.replace('://', '_').replace('/', '_')}"
+    read_resource.__doc__ = usage
+
+    logger.info(f"Built resource tool: {read_resource.__name__} (uri: {uri})")
+    return Tool(read_resource)

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,151 @@
+"""
+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
+        from opentelemetry.sdk.trace.export import BatchSpanProcessor
+        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
+
+        # 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":
+            exporter = GRPCExporter(
+                endpoint=settings.otel.collector_endpoint,
+                timeout=settings.otel.export_timeout,
+            )
+        else:  # http
+            exporter = HTTPExporter(
+                endpoint=f"{settings.otel.collector_endpoint}/v1/traces",
+                timeout=settings.otel.export_timeout,
+            )
+
+        # 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
+        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}")