gatewaysdk 0.3.2__py3-none-any.whl

gatewaysdk/instrumentation/base.py

@@ -0,0 +1,119 @@
+"""
+Base class for instrumentors.
+
+Instrumentors provide automatic metric collection and tracing for specific libraries.
+They follow a plugin architecture where each instrumentor self-registers with the
+InstrumentorRegistry.
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Dict, Optional
+
+
+@dataclass
+class MetricSchema:
+    """
+    Describes a metric that an instrumentor may log.
+
+    This is used for documentation/discovery, NOT validation.
+    The schema-less design means any metric can be logged,
+    but MetricSchema helps developers understand what metrics
+    are available from each instrumentor.
+    """
+
+    name: str
+    description: str
+    unit: str = ""
+    metric_type: str = "gauge"  # gauge, counter, histogram
+    labels: Dict[str, str] = field(default_factory=dict)  # Label name -> description
+
+
+class Instrumentor(ABC):
+    """
+    Base class for auto-instrumentation plugins.
+
+    Instrumentors provide automatic metric collection and tracing
+    for specific libraries (e.g., OpenAI, Anthropic, custom algorithms).
+
+    Usage:
+        class OpenAIInstrumentor(Instrumentor):
+            @property
+            def name(self) -> str:
+                return "openai"
+
+            def instrument(self) -> None:
+                # Apply monkey patches or hooks
+                pass
+
+            def uninstrument(self) -> None:
+                # Remove patches
+                pass
+
+    Instrumentors are registered via the InstrumentorRegistry:
+        InstrumentorRegistry.register("openai", OpenAIInstrumentor)
+
+    Or via Python entry points in pyproject.toml:
+        [project.entry-points."gatewaysdk.instrumentors"]
+        openai = "gatewaysdk.instrumentation.openai:OpenAIInstrumentor"
+    """
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """
+        Unique identifier for this instrumentor.
+
+        Examples: 'openai', 'anthropic', 'gepa', 'langchain'
+        """
+        pass
+
+    @property
+    def default_labels(self) -> Dict[str, str]:
+        """
+        Labels added to all metrics from this instrumentor.
+
+        Override to add custom default labels. By default,
+        includes the instrumentor name.
+        """
+        return {"instrumentor": self.name}
+
+    def get_metric_schema(self) -> Dict[str, MetricSchema]:
+        """
+        Describe metrics this instrumentor may log.
+
+        Used for documentation/discovery, NOT validation.
+        Returns a dict mapping metric names to their schemas.
+
+        Override this to document the metrics your instrumentor logs.
+        """
+        return {}
+
+    @abstractmethod
+    def instrument(self) -> None:
+        """
+        Apply patches/hooks to the target library.
+
+        This method is called when the instrumentor is activated.
+        It should apply any necessary monkey patches, register
+        callbacks, or otherwise hook into the target library.
+        """
+        pass
+
+    def uninstrument(self) -> None:
+        """
+        Remove patches/hooks from the target library.
+
+        Override this method to provide cleanup functionality.
+        Used primarily for testing to restore original behavior.
+        """
+        pass
+
+    def is_available(self) -> bool:
+        """
+        Check if the target library is available.
+
+        Override this to check if the library this instrumentor
+        targets is installed. By default returns True.
+        """
+        return True

gatewaysdk/instrumentation/litellm.py

@@ -0,0 +1,83 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""LiteLLM instrumentations.
+
+Patches LiteLLM's OpenTelemetry integration to add GatewaySDK-specific attributes
+to spans, including session_id, tags, environment, and call_type from context.
+
+The call_type is read from a ContextVar that is set by algorithm-specific
+decorators (e.g., @gepa.judge, @gepa.agent).
+
+[Related documentation](https://docs.litellm.ai/docs/observability/agentops_integration).
+"""
+
+import json
+from typing import Any, Optional
+
+from litellm.integrations.opentelemetry import OpenTelemetry
+
+from gatewaysdk.tracing.context import get_current_call_type
+
+__all__ = [
+    "instrument_litellm",
+    "uninstrument_litellm",
+]
+
+original_set_attributes = OpenTelemetry.set_attributes  # type: ignore
+
+
+def patched_set_attributes(self: Any, span: Any, kwargs: Any, response_obj: Optional[Any]):
+    """Enhanced set_attributes that adds GatewaySDK tracing metadata to spans.
+
+    Extracts session_id, tags, and environment from kwargs["metadata"] and sets
+    them as OTEL span attributes for visibility in Gateway/Langfuse.
+
+    Also reads the current call_type from context (set by decorators like
+    @gepa.judge) and adds it as a span attribute.
+    """
+    original_set_attributes(self, span, kwargs, response_obj)
+
+    # Add token IDs if available
+    if response_obj is not None and response_obj.get("prompt_token_ids"):
+        span.set_attribute("prompt_token_ids", list(response_obj.get("prompt_token_ids")))
+    if response_obj is not None and response_obj.get("response_token_ids"):
+        span.set_attribute("response_token_ids", list(response_obj.get("response_token_ids")[0]))
+
+    # Read call_type from context (set by @gepa.judge, @gepa.agent, etc.)
+    call_type = get_current_call_type()
+    if call_type:
+        span.set_attribute("gateway.call_type", call_type)
+
+    # Extract GatewaySDK tracing metadata from kwargs
+    metadata = kwargs.get("metadata", {}) if kwargs else {}
+    if not metadata:
+        return
+
+    # Set session_id as span attribute (standard Langfuse/Gateway attribute)
+    session_id = metadata.get("session_id")
+    if session_id:
+        span.set_attribute("session.id", session_id)
+
+    # Set tags as span attribute (JSON array for Langfuse compatibility)
+    tags = metadata.get("tags")
+    if tags and isinstance(tags, list):
+        # Langfuse expects tags as a JSON array string or individual attributes
+        span.set_attribute("tags", json.dumps(tags))
+        # Also set individual tag attributes for filtering
+        for i, tag in enumerate(tags):
+            span.set_attribute(f"tag.{i}", str(tag))
+
+    # Set environment as span attribute
+    environment = metadata.get("environment")
+    if environment:
+        span.set_attribute("environment", environment)
+
+
+def instrument_litellm():
+    """Instrument litellm to capture token IDs."""
+    OpenTelemetry.set_attributes = patched_set_attributes
+
+
+def uninstrument_litellm():
+    """Uninstrument litellm to stop capturing token IDs."""
+    OpenTelemetry.set_attributes = original_set_attributes

gatewaysdk/instrumentation/registry.py

@@ -0,0 +1,273 @@
+"""
+Instrumentor Registry for auto-instrumentation.
+
+The registry provides a central location to register and activate
+instrumentors. It supports both explicit registration and automatic
+discovery via Python entry points.
+"""
+
+import logging
+import sys
+from typing import Dict, List, Optional, Type
+
+from .base import Instrumentor, MetricSchema
+
+logger = logging.getLogger(__name__)
+
+
+class InstrumentorRegistry:
+    """
+    Global registry for metric instrumentors.
+
+    Instrumentors can be registered in two ways:
+    1. Explicit registration:
+        InstrumentorRegistry.register("openai", OpenAIInstrumentor)
+
+    2. Entry point discovery (in pyproject.toml):
+        [project.entry-points."gatewaysdk.instrumentors"]
+        openai = "gatewaysdk.instrumentation.openai:OpenAIInstrumentor"
+
+    Usage:
+        # Enable a specific instrumentor
+        gatewaysdk.instrument("openai")
+
+        # Enable multiple instrumentors
+        gatewaysdk.instrument("openai", "anthropic", "langchain")
+
+        # Get list of available instrumentors
+        InstrumentorRegistry.list_available()
+    """
+
+    _instrumentors: Dict[str, Type[Instrumentor]] = {}
+    _active: Dict[str, Instrumentor] = {}
+    _discovered: bool = False
+
+    @classmethod
+    def register(cls, name: str, instrumentor_cls: Type[Instrumentor]) -> None:
+        """
+        Register an instrumentor class.
+
+        This is typically called at import time by instrumentor modules,
+        or by entry point discovery.
+
+        Args:
+            name: Unique identifier for the instrumentor (e.g., "openai")
+            instrumentor_cls: The instrumentor class to register
+        """
+        cls._instrumentors[name] = instrumentor_cls
+        logger.debug(f"Registered instrumentor: {name}")
+
+    @classmethod
+    def instrument(cls, name: str, **kwargs) -> Optional[Instrumentor]:
+        """
+        Activate an instrumentor by name.
+
+        Args:
+            name: The instrumentor name to activate
+            **kwargs: Additional arguments passed to the instrumentor constructor
+
+        Returns:
+            The activated instrumentor instance, or None if not available
+
+        Raises:
+            ValueError: If the instrumentor name is not registered
+        """
+        # Auto-discover instrumentors if not already done
+        if not cls._discovered:
+            cls.discover()
+
+        if name not in cls._instrumentors:
+            available = list(cls._instrumentors.keys())
+            raise ValueError(f"Unknown instrumentor: {name}. Available: {available}")
+
+        # Return existing instance if already active
+        if name in cls._active:
+            return cls._active[name]
+
+        # Create and activate new instance
+        instrumentor_cls = cls._instrumentors[name]
+        instance = instrumentor_cls(**kwargs)
+
+        # Check if target library is available
+        if not instance.is_available():
+            logger.warning(f"Instrumentor '{name}' target library is not available")
+            return None
+
+        try:
+            instance.instrument()
+            cls._active[name] = instance
+            logger.info(f"Activated instrumentor: {name}")
+            return instance
+        except Exception as e:
+            logger.error(f"Failed to activate instrumentor '{name}': {e}")
+            return None
+
+    @classmethod
+    def uninstrument(cls, name: str) -> bool:
+        """
+        Deactivate an instrumentor by name.
+
+        Args:
+            name: The instrumentor name to deactivate
+
+        Returns:
+            True if successfully deactivated, False otherwise
+        """
+        if name not in cls._active:
+            logger.debug(f"Instrumentor '{name}' is not active")
+            return False
+
+        instance = cls._active[name]
+        try:
+            instance.uninstrument()
+            del cls._active[name]
+            logger.info(f"Deactivated instrumentor: {name}")
+            return True
+        except Exception as e:
+            logger.error(f"Failed to deactivate instrumentor '{name}': {e}")
+            return False
+
+    @classmethod
+    def uninstrument_all(cls) -> None:
+        """Deactivate all active instrumentors."""
+        for name in list(cls._active.keys()):
+            cls.uninstrument(name)
+
+    @classmethod
+    def discover(cls) -> None:
+        """
+        Auto-discover instrumentors via Python entry points.
+
+        Entry points should be defined in pyproject.toml:
+            [project.entry-points."gatewaysdk.instrumentors"]
+            openai = "gatewaysdk.instrumentation.openai:OpenAIInstrumentor"
+        """
+        if cls._discovered:
+            return
+
+        try:
+            # Python 3.10+ has importlib.metadata.entry_points with group param
+            if sys.version_info >= (3, 10):
+                from importlib.metadata import entry_points
+
+                eps = entry_points(group="gatewaysdk.instrumentors")
+            else:
+                # Python 3.9 compatibility
+                from importlib.metadata import entry_points
+
+                eps = entry_points().get("gatewaysdk.instrumentors", [])
+
+            for ep in eps:
+                try:
+                    instrumentor_cls = ep.load()
+                    cls.register(ep.name, instrumentor_cls)
+                except Exception as e:
+                    logger.warning(f"Failed to load instrumentor '{ep.name}': {e}")
+
+        except ImportError:
+            logger.debug("importlib.metadata not available, skipping entry point discovery")
+
+        cls._discovered = True
+
+    @classmethod
+    def list_available(cls) -> List[str]:
+        """
+        List all registered instrumentor names.
+
+        Returns:
+            List of instrumentor names
+        """
+        if not cls._discovered:
+            cls.discover()
+        return list(cls._instrumentors.keys())
+
+    @classmethod
+    def list_active(cls) -> List[str]:
+        """
+        List all currently active instrumentor names.
+
+        Returns:
+            List of active instrumentor names
+        """
+        return list(cls._active.keys())
+
+    @classmethod
+    def get_metric_schemas(cls) -> Dict[str, Dict[str, MetricSchema]]:
+        """
+        Get metric schemas from all registered instrumentors.
+
+        Returns:
+            Dict mapping instrumentor names to their metric schemas
+        """
+        if not cls._discovered:
+            cls.discover()
+
+        schemas = {}
+        for name, instrumentor_cls in cls._instrumentors.items():
+            instance = instrumentor_cls()
+            schema = instance.get_metric_schema()
+            if schema:
+                schemas[name] = schema
+        return schemas
+
+    @classmethod
+    def is_active(cls, name: str) -> bool:
+        """
+        Check if an instrumentor is currently active.
+
+        Args:
+            name: The instrumentor name to check
+
+        Returns:
+            True if the instrumentor is active
+        """
+        return name in cls._active
+
+    @classmethod
+    def get(cls, name: str) -> Optional[Instrumentor]:
+        """
+        Get an active instrumentor instance by name.
+
+        Args:
+            name: The instrumentor name
+
+        Returns:
+            The instrumentor instance if active, None otherwise
+        """
+        return cls._active.get(name)
+
+
+# Convenience function for module-level access
+def instrument(*names: str, **kwargs) -> List[Optional[Instrumentor]]:
+    """
+    Activate one or more instrumentors by name.
+
+    Args:
+        *names: Instrumentor names to activate
+        **kwargs: Additional arguments passed to instrumentor constructors
+
+    Returns:
+        List of activated instrumentor instances (None for failed activations)
+
+    Examples:
+        # Enable OpenAI instrumentation
+        gatewaysdk.instrument("openai")
+
+        # Enable multiple instrumentors
+        gatewaysdk.instrument("openai", "anthropic", "langchain")
+    """
+    return [InstrumentorRegistry.instrument(name, **kwargs) for name in names]
+
+
+def uninstrument(*names: str) -> None:
+    """
+    Deactivate one or more instrumentors by name.
+
+    Args:
+        *names: Instrumentor names to deactivate
+
+    Examples:
+        gatewaysdk.uninstrument("openai")
+    """
+    for name in names:
+        InstrumentorRegistry.uninstrument(name)

gatewaysdk/instrumentation/vllm.py

@@ -0,0 +1,81 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import warnings
+from typing import Any, List
+
+import vllm.entrypoints.openai.protocol
+from vllm.entrypoints.openai.protocol import ChatCompletionResponse
+from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
+
+__all__ = [
+    "instrument_vllm",
+    "uninstrument_vllm",
+]
+
+
+class ChatCompletionResponsePatched(ChatCompletionResponse):
+    prompt_token_ids: List[int] | None = None
+    response_token_ids: List[int] | None = None
+
+
+original_chat_completion_full_generator = OpenAIServingChat.chat_completion_full_generator
+
+
+async def chat_completion_full_generator(
+    self: Any,
+    request: Any,
+    result_generator: Any,
+    request_id: str,
+    model_name: str,
+    conversation: Any,
+    tokenizer: Any,
+    request_metadata: Any,
+) -> Any:
+    prompt_token_ids: List[int] | None = None
+    response_token_ids: List[List[int]] | None = None
+
+    async def _generate_inceptor():
+        nonlocal prompt_token_ids, response_token_ids
+        async for res in result_generator:
+            yield res
+            prompt_token_ids = res.prompt_token_ids
+            response_token_ids = [output.token_ids for output in res.outputs]
+
+    response = await original_chat_completion_full_generator(
+        self,
+        request,
+        _generate_inceptor(),
+        request_id,
+        model_name,
+        conversation,
+        tokenizer,
+        request_metadata,
+    )
+    response = response.model_copy(
+        update={
+            "prompt_token_ids": prompt_token_ids,
+            "response_token_ids": response_token_ids,
+        }
+    )
+
+    return response
+
+
+def instrument_vllm():
+    """Instrument vLLM to capture token IDs generated by engine.
+
+    This instrumentation has been merged to upstream vLLM since v0.10.2.
+    """
+    if vllm.entrypoints.openai.protocol.ChatCompletionResponse is ChatCompletionResponsePatched:
+        warnings.warn("vllm is already instrumented. Skip the instrumentation.")
+        return
+
+    vllm.entrypoints.openai.protocol.ChatCompletionResponse = ChatCompletionResponsePatched
+    OpenAIServingChat.chat_completion_full_generator = chat_completion_full_generator
+
+
+def uninstrument_vllm():
+    """Uninstrument vLLM to stop capturing token IDs generated by engine."""
+    OpenAIServingChat.chat_completion_full_generator = original_chat_completion_full_generator