genai-otel-instrument 0.1.24__py3-none-any.whl

genai_otel/instrumentors/ollama_server_metrics_poller.py

@@ -0,0 +1,336 @@
+"""Ollama server metrics poller for automatic KV cache and memory tracking.
+
+This module provides automatic collection of Ollama server metrics by polling
+the /api/ps endpoint to get running models and their VRAM usage.
+
+GPU VRAM is auto-detected using:
+1. pynvml (nvidia-ml-py) - preferred method
+2. nvidia-smi subprocess - fallback
+3. Environment variable - manual override
+"""
+
+import logging
+import subprocess  # nosec B404 - Only used for nvidia-smi with hardcoded args
+import threading
+import time
+from typing import Optional
+
+import requests
+
+from ..server_metrics import get_server_metrics
+
+logger = logging.getLogger(__name__)
+
+# Try to import nvidia-ml-py for GPU VRAM detection
+# Package: nvidia-ml-py, Import: pynvml
+try:
+    import pynvml
+
+    NVIDIA_ML_AVAILABLE = True
+except ImportError:
+    NVIDIA_ML_AVAILABLE = False
+
+
+class OllamaServerMetricsPoller:
+    """Polls Ollama server for metrics and updates ServerMetricsCollector.
+
+    This poller queries the /api/ps endpoint to get:
+    - Running models and their VRAM usage
+    - Total VRAM usage across all models
+    - Number of models currently loaded
+
+    It automatically updates the global ServerMetricsCollector with:
+    - KV cache usage per model (approximated from VRAM usage)
+    - Running models count
+    """
+
+    def __init__(
+        self,
+        base_url: str = "http://localhost:11434",
+        interval: float = 5.0,
+        max_vram_gb: Optional[float] = None,
+    ):
+        """Initialize the Ollama server metrics poller.
+
+        Args:
+            base_url: Base URL of the Ollama server (default: http://localhost:11434)
+            interval: Polling interval in seconds (default: 5.0)
+            max_vram_gb: Maximum VRAM in GB for percentage calculation.
+                        If None, will attempt auto-detection via nvidia-ml-py or nvidia-smi.
+        """
+        self.base_url = base_url.rstrip("/")
+        self.interval = interval
+
+        # Auto-detect GPU VRAM if not provided
+        if max_vram_gb is not None:
+            self.max_vram_bytes = max_vram_gb * 1024**3
+            logger.info(f"Using configured GPU VRAM: {max_vram_gb}GB")
+        else:
+            detected_vram_gb = self._detect_gpu_vram()
+            if detected_vram_gb:
+                self.max_vram_bytes = detected_vram_gb * 1024**3
+                logger.info(f"Auto-detected GPU VRAM: {detected_vram_gb}GB")
+            else:
+                self.max_vram_bytes = None
+                logger.info("GPU VRAM not detected, using heuristic-based percentages")
+
+        self._stop_event = threading.Event()
+        self._thread: Optional[threading.Thread] = None
+        self._running = False
+
+    def _detect_gpu_vram(self) -> Optional[float]:
+        """Auto-detect GPU VRAM in GB using nvidia-ml-py or nvidia-smi.
+
+        Returns:
+            GPU VRAM in GB, or None if detection failed
+        """
+        # Try nvidia-ml-py first (preferred method)
+        vram_gb = self._detect_vram_nvidia_ml()
+        if vram_gb:
+            return vram_gb
+
+        # Fallback to nvidia-smi
+        vram_gb = self._detect_vram_nvidia_smi()
+        if vram_gb:
+            return vram_gb
+
+        return None
+
+    def _detect_vram_nvidia_ml(self) -> Optional[float]:
+        """Detect GPU VRAM using nvidia-ml-py library.
+
+        Returns:
+            GPU VRAM in GB for the first GPU, or None if failed
+        """
+        if not NVIDIA_ML_AVAILABLE:
+            logger.debug("nvidia-ml-py not available for VRAM detection")
+            return None
+
+        try:
+            pynvml.nvmlInit()
+            device_count = pynvml.nvmlDeviceGetCount()
+
+            if device_count == 0:
+                logger.debug("No NVIDIA GPUs detected via nvidia-ml-py")
+                pynvml.nvmlShutdown()
+                return None
+
+            # Get VRAM from first GPU (Ollama typically uses GPU 0)
+            handle = pynvml.nvmlDeviceGetHandleByIndex(0)
+            mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
+            vram_bytes = mem_info.total
+            vram_gb = vram_bytes / (1024**3)
+
+            # Get GPU name for logging
+            gpu_name = pynvml.nvmlDeviceGetName(handle)
+            logger.debug(f"Detected via nvidia-ml-py: {gpu_name} with {vram_gb:.2f}GB VRAM")
+
+            pynvml.nvmlShutdown()
+            return vram_gb
+
+        except Exception as e:
+            logger.debug(f"Failed to detect VRAM via nvidia-ml-py: {e}")
+            try:
+                pynvml.nvmlShutdown()
+            except Exception:
+                pass
+            return None
+
+    def _detect_vram_nvidia_smi(self) -> Optional[float]:
+        """Detect GPU VRAM using nvidia-smi command.
+
+        Returns:
+            GPU VRAM in GB for the first GPU, or None if failed
+        """
+        try:
+            # Query nvidia-smi for total memory of GPU 0
+            result = subprocess.run(  # nosec B603 B607
+                [
+                    "nvidia-smi",
+                    "--query-gpu=memory.total",
+                    "--format=csv,noheader,nounits",
+                    "--id=0",
+                ],
+                capture_output=True,
+                text=True,
+                timeout=5,
+                check=True,
+            )
+
+            # Parse output (in MB)
+            vram_mb = float(result.stdout.strip())
+            vram_gb = vram_mb / 1024
+
+            logger.debug(f"Detected via nvidia-smi: {vram_gb:.2f}GB VRAM")
+            return vram_gb
+
+        except FileNotFoundError:
+            logger.debug("nvidia-smi command not found")
+            return None
+        except subprocess.TimeoutExpired:
+            logger.debug("nvidia-smi command timed out")
+            return None
+        except Exception as e:
+            logger.debug(f"Failed to detect VRAM via nvidia-smi: {e}")
+            return None
+
+    def start(self):
+        """Start the background polling thread."""
+        if self._running:
+            logger.warning("Ollama server metrics poller already running")
+            return
+
+        self._stop_event.clear()
+        self._thread = threading.Thread(target=self._poll_loop, daemon=True)
+        self._thread.start()
+        self._running = True
+        logger.info(
+            f"Ollama server metrics poller started (interval={self.interval}s, url={self.base_url})"
+        )
+
+    def stop(self):
+        """Stop the background polling thread."""
+        if not self._running:
+            return
+
+        self._stop_event.set()
+        if self._thread:
+            self._thread.join(timeout=self.interval + 1.0)
+        self._running = False
+        logger.info("Ollama server metrics poller stopped")
+
+    def _poll_loop(self):
+        """Background polling loop that runs in a separate thread."""
+        while not self._stop_event.is_set():
+            try:
+                self._collect_metrics()
+            except Exception as e:
+                logger.warning(f"Failed to collect Ollama server metrics: {e}")
+
+            # Wait for interval or until stop is requested
+            self._stop_event.wait(self.interval)
+
+    def _collect_metrics(self):
+        """Query Ollama /api/ps and update server metrics."""
+        try:
+            # Query /api/ps endpoint
+            response = requests.get(f"{self.base_url}/api/ps", timeout=2.0)
+            response.raise_for_status()
+            data = response.json()
+
+            # Get server metrics collector
+            server_metrics = get_server_metrics()
+            if not server_metrics:
+                logger.debug("Server metrics collector not initialized, skipping update")
+                return
+
+            # Extract running models
+            models = data.get("models", [])
+            num_models = len(models)
+
+            # Update the "max" capacity to reflect number of model slots
+            # This gives visibility into how many models can be loaded simultaneously
+            # Note: This is approximate - actual capacity depends on VRAM
+            if num_models > 0:
+                server_metrics.set_requests_max(num_models)
+
+            logger.debug(f"Ollama has {num_models} models loaded in memory")
+
+            # Process each model's VRAM usage and details
+            total_vram_bytes = 0
+            total_size_bytes = 0
+
+            for model_info in models:
+                model_name = model_info.get("name", "unknown")
+                size_vram = model_info.get("size_vram", 0)
+                size_total = model_info.get("size", 0)
+                total_vram_bytes += size_vram
+                total_size_bytes += size_total
+
+                # Calculate VRAM usage percentage for this model
+                if self.max_vram_bytes and self.max_vram_bytes > 0:
+                    # If we know max VRAM, calculate actual percentage
+                    vram_usage_pct = (size_vram / self.max_vram_bytes) * 100
+                    vram_usage_pct = min(100.0, vram_usage_pct)  # Cap at 100%
+                else:
+                    # If we don't know max VRAM, use a simple heuristic:
+                    # Models in memory are "using" the cache, so report a non-zero value
+                    # This is an approximation - actual KV cache varies with context
+                    vram_usage_pct = min(100.0, (size_vram / (8 * 1024**3)) * 100)
+
+                # Update KV cache usage for this model
+                server_metrics.set_kv_cache_usage(model_name, vram_usage_pct)
+
+                # Extract model details for logging
+                details = model_info.get("details", {})
+                param_size = details.get("parameter_size", "unknown")
+                quant_level = details.get("quantization_level", "unknown")
+                model_format = details.get("format", "unknown")
+
+                logger.debug(
+                    f"Model {model_name}: "
+                    f"VRAM={size_vram / 1024**3:.2f}GB ({vram_usage_pct:.1f}%), "
+                    f"Size={size_total / 1024**3:.2f}GB, "
+                    f"Params={param_size}, Quant={quant_level}, Format={model_format}"
+                )
+
+            # Log aggregate metrics
+            if total_vram_bytes > 0:
+                logger.debug(
+                    f"Total Ollama usage: VRAM={total_vram_bytes / 1024**3:.2f}GB, "
+                    f"Total={total_size_bytes / 1024**3:.2f}GB, "
+                    f"Models={num_models}"
+                )
+
+            # Note: We don't update requests.running here because that's tracked
+            # automatically by BaseInstrumentor when requests are made
+
+        except requests.exceptions.ConnectionError:
+            logger.debug(f"Cannot connect to Ollama server at {self.base_url}")
+        except requests.exceptions.Timeout:
+            logger.warning("Ollama server request timed out")
+        except requests.exceptions.RequestException as e:
+            logger.warning(f"Error querying Ollama server: {e}")
+        except Exception as e:
+            logger.error(f"Unexpected error collecting Ollama metrics: {e}", exc_info=True)
+
+
+# Global instance
+_global_poller: Optional[OllamaServerMetricsPoller] = None
+
+
+def start_ollama_metrics_poller(
+    base_url: str = "http://localhost:11434",
+    interval: float = 5.0,
+    max_vram_gb: Optional[float] = None,
+) -> OllamaServerMetricsPoller:
+    """Start the global Ollama server metrics poller.
+
+    Args:
+        base_url: Base URL of the Ollama server (default: http://localhost:11434)
+        interval: Polling interval in seconds (default: 5.0)
+        max_vram_gb: Maximum VRAM in GB for percentage calculation (optional)
+
+    Returns:
+        The global poller instance
+    """
+    global _global_poller
+
+    if _global_poller is not None and _global_poller._running:
+        logger.warning("Ollama metrics poller already running")
+        return _global_poller
+
+    _global_poller = OllamaServerMetricsPoller(
+        base_url=base_url, interval=interval, max_vram_gb=max_vram_gb
+    )
+    _global_poller.start()
+    return _global_poller
+
+
+def stop_ollama_metrics_poller():
+    """Stop the global Ollama server metrics poller."""
+    global _global_poller
+
+    if _global_poller is not None:
+        _global_poller.stop()
+        _global_poller = None

genai_otel/instrumentors/openai_agents_instrumentor.py

@@ -0,0 +1,291 @@
+"""OpenTelemetry instrumentor for the OpenAI Agents SDK.
+
+This instrumentor automatically traces agent execution, handoffs, sessions,
+and guardrails using the OpenAI Agents SDK (openai-agents package).
+"""
+
+import json
+import logging
+from typing import Any, Dict, Optional
+
+from ..config import OTelConfig
+from .base import BaseInstrumentor
+
+logger = logging.getLogger(__name__)
+
+
+class OpenAIAgentsInstrumentor(BaseInstrumentor):
+    """Instrumentor for OpenAI Agents SDK"""
+
+    def __init__(self):
+        """Initialize the instrumentor."""
+        super().__init__()
+        self._agents_available = False
+        self._check_availability()
+
+    def _check_availability(self):
+        """Check if OpenAI Agents SDK is available."""
+        try:
+            import agents
+
+            self._agents_available = True
+            logger.debug("OpenAI Agents SDK detected and available for instrumentation")
+        except ImportError:
+            logger.debug("OpenAI Agents SDK not installed, instrumentation will be skipped")
+            self._agents_available = False
+
+    def instrument(self, config: OTelConfig):
+        """Instrument OpenAI Agents SDK if available.
+
+        Args:
+            config (OTelConfig): The OpenTelemetry configuration object.
+        """
+        if not self._agents_available:
+            logger.debug("Skipping OpenAI Agents instrumentation - library not available")
+            return
+
+        self.config = config
+
+        try:
+            import agents
+            import wrapt
+
+            # Instrument Runner.run() (async) and Runner.run_sync() (sync)
+            if hasattr(agents, "Runner"):
+                # Instrument async run method
+                if hasattr(agents.Runner, "run"):
+                    original_run = agents.Runner.run
+                    agents.Runner.run = wrapt.FunctionWrapper(original_run, self._wrap_runner_run)
+
+                # Instrument sync run method
+                if hasattr(agents.Runner, "run_sync"):
+                    original_run_sync = agents.Runner.run_sync
+                    agents.Runner.run_sync = wrapt.FunctionWrapper(
+                        original_run_sync, self._wrap_runner_run_sync
+                    )
+
+                self._instrumented = True
+                logger.info("OpenAI Agents SDK instrumentation enabled")
+
+        except Exception as e:
+            logger.error("Failed to instrument OpenAI Agents SDK: %s", e, exc_info=True)
+            if config.fail_on_error:
+                raise
+
+    def _wrap_runner_run(self, wrapped, instance, args, kwargs):
+        """Wrap Runner.run() async method with span.
+
+        Args:
+            wrapped: The original method.
+            instance: The Runner instance.
+            args: Positional arguments.
+            kwargs: Keyword arguments.
+        """
+        return self.create_span_wrapper(
+            span_name="openai_agents.runner.run",
+            extract_attributes=self._extract_runner_attributes,
+        )(wrapped)(instance, *args, **kwargs)
+
+    def _wrap_runner_run_sync(self, wrapped, instance, args, kwargs):
+        """Wrap Runner.run_sync() sync method with span.
+
+        Args:
+            wrapped: The original method.
+            instance: The Runner instance.
+            args: Positional arguments.
+            kwargs: Keyword arguments.
+        """
+        return self.create_span_wrapper(
+            span_name="openai_agents.runner.run_sync",
+            extract_attributes=self._extract_runner_attributes,
+        )(wrapped)(instance, *args, **kwargs)
+
+    def _extract_runner_attributes(self, instance: Any, args: Any, kwargs: Any) -> Dict[str, Any]:
+        """Extract attributes from Runner.run() or Runner.run_sync() call.
+
+        Args:
+            instance: The Runner instance.
+            args: Positional arguments (agent, input_data, session, etc.).
+            kwargs: Keyword arguments.
+
+        Returns:
+            Dict[str, Any]: Dictionary of attributes to set on the span.
+        """
+        attrs = {}
+
+        # Core attributes
+        attrs["gen_ai.system"] = "openai_agents"
+        attrs["gen_ai.operation.name"] = "agent.run"
+
+        # Extract agent from args (first positional argument)
+        agent = None
+        if len(args) > 0:
+            agent = args[0]
+        else:
+            agent = kwargs.get("agent")
+
+        if agent:
+            # Agent attributes
+            if hasattr(agent, "name"):
+                attrs["openai.agent.name"] = agent.name
+                attrs["gen_ai.request.model"] = agent.name  # Use agent name as "model"
+
+            if hasattr(agent, "model"):
+                attrs["openai.agent.model"] = agent.model
+
+            if hasattr(agent, "instructions"):
+                # Truncate instructions to avoid span size issues
+                instructions = str(agent.instructions)[:500]
+                attrs["openai.agent.instructions"] = instructions
+
+            # Extract tools
+            if hasattr(agent, "tools") and agent.tools:
+                try:
+                    tool_names = [getattr(tool, "name", str(tool)[:50]) for tool in agent.tools]
+                    attrs["openai.agent.tools"] = tool_names
+                    attrs["openai.agent.tool_count"] = len(agent.tools)
+                except Exception as e:
+                    logger.debug("Failed to extract agent tools: %s", e)
+
+            # Extract handoffs
+            if hasattr(agent, "handoffs") and agent.handoffs:
+                try:
+                    handoff_names = [getattr(h, "name", str(h)[:50]) for h in agent.handoffs]
+                    attrs["openai.agent.handoffs"] = handoff_names
+                    attrs["openai.agent.handoff_count"] = len(agent.handoffs)
+                except Exception as e:
+                    logger.debug("Failed to extract agent handoffs: %s", e)
+
+            # Extract guardrails
+            if hasattr(agent, "guardrails") and agent.guardrails:
+                try:
+                    attrs["openai.agent.guardrails_enabled"] = True
+                    attrs["openai.agent.guardrail_count"] = len(agent.guardrails)
+                except Exception as e:
+                    logger.debug("Failed to extract agent guardrails: %s", e)
+
+        # Extract input data (second positional argument)
+        input_data = None
+        if len(args) > 1:
+            input_data = args[1]
+        else:
+            input_data = kwargs.get("input_data")
+
+        if input_data:
+            # Truncate input to avoid span size issues
+            input_str = str(input_data)[:500]
+            attrs["openai.agent.input"] = input_str
+            attrs["openai.agent.input_length"] = len(str(input_data))
+
+        # Extract session (third positional argument or kwarg)
+        session = None
+        if len(args) > 2:
+            session = args[2]
+        else:
+            session = kwargs.get("session")
+
+        if session:
+            # Session attributes
+            if hasattr(session, "session_id"):
+                attrs["openai.session.id"] = session.session_id
+                attrs["session.id"] = session.session_id  # Generic session ID
+
+            # Detect session type
+            session_type = type(session).__name__
+            attrs["openai.session.type"] = session_type
+
+        return attrs
+
+    def _extract_usage(self, result) -> Optional[Dict[str, int]]:
+        """Extract token usage from agent run result.
+
+        Note: The OpenAI Agents SDK may not directly expose token usage in the result.
+        Token usage is captured by the underlying OpenAI SDK instrumentor.
+
+        Args:
+            result: The agent run result object.
+
+        Returns:
+            Optional[Dict[str, int]]: Dictionary with token counts or None.
+        """
+        # OpenAI Agents SDK doesn't directly expose usage in the result
+        # Token usage is captured by the OpenAI SDK instrumentor for underlying LLM calls
+        # We can try to extract if the result has usage information
+        if hasattr(result, "usage") and result.usage:
+            usage = result.usage
+            return {
+                "prompt_tokens": getattr(usage, "prompt_tokens", 0),
+                "completion_tokens": getattr(usage, "completion_tokens", 0),
+                "total_tokens": getattr(usage, "total_tokens", 0),
+            }
+        return None
+
+    def _extract_response_attributes(self, result) -> Dict[str, Any]:
+        """Extract response attributes from agent run result.
+
+        Args:
+            result: The agent run result object.
+
+        Returns:
+            Dict[str, Any]: Dictionary of response attributes.
+        """
+        attrs = {}
+
+        # Extract final output
+        if hasattr(result, "final_output"):
+            output = str(result.final_output)[:500]  # Truncate to avoid span size issues
+            attrs["openai.agent.output"] = output
+            attrs["openai.agent.output_length"] = len(str(result.final_output))
+
+        # Extract handoff information if the agent handed off to another agent
+        if hasattr(result, "handoff") and result.handoff:
+            attrs["openai.handoff.occurred"] = True
+            if hasattr(result.handoff, "target_agent"):
+                attrs["openai.handoff.to_agent"] = result.handoff.target_agent
+
+        # Extract guardrail validation results
+        if hasattr(result, "guardrail_results"):
+            try:
+                guardrail_results = result.guardrail_results
+                if guardrail_results:
+                    # Count violations
+                    violation_count = sum(
+                        1 for r in guardrail_results if not getattr(r, "passed", True)
+                    )
+                    attrs["openai.guardrail.violations"] = violation_count
+                    attrs["openai.guardrail.validated"] = True
+            except Exception as e:
+                logger.debug("Failed to extract guardrail results: %s", e)
+
+        # Extract metadata if available
+        if hasattr(result, "metadata"):
+            try:
+                metadata = result.metadata
+                if isinstance(metadata, dict):
+                    # Add selected metadata fields
+                    for key in ["run_id", "agent_id", "session_id"]:
+                        if key in metadata:
+                            attrs[f"openai.agent.metadata.{key}"] = metadata[key]
+            except Exception as e:
+                logger.debug("Failed to extract result metadata: %s", e)
+
+        return attrs
+
+    def _extract_finish_reason(self, result) -> Optional[str]:
+        """Extract finish reason from agent run result.
+
+        Args:
+            result: The agent run result object.
+
+        Returns:
+            Optional[str]: The finish reason string or None if not available.
+        """
+        # Try to extract finish reason if available
+        if hasattr(result, "finish_reason"):
+            return result.finish_reason
+
+        # Check if result has a status
+        if hasattr(result, "status"):
+            return result.status
+
+        return None