genai-otel-instrument 0.1.2.dev0__py3-none-any.whl → 0.1.7.dev0__py3-none-any.whl

genai_otel/__version__.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.1.2.dev0'
-__version_tuple__ = version_tuple = (0, 1, 2, 'dev0')
+__version__ = version = '0.1.7.dev0'
+__version_tuple__ = version_tuple = (0, 1, 7, 'dev0')
 
 __commit_id__ = commit_id = None

genai_otel/auto_instrument.py

@@ -17,6 +17,8 @@ from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter
 
 from .config import OTelConfig
+from .cost_calculator import CostCalculator
+from .cost_enrichment_processor import CostEnrichmentSpanProcessor
 from .gpu_metrics import GPUMetricsCollector
 from .mcp_instrumentors import MCPInstrumentorManager
 from .metrics import (
@@ -117,12 +119,16 @@ INSTRUMENTORS = {
 }
 
 # Add OpenInference instrumentors if available (requires Python >= 3.10)
+# IMPORTANT: Order matters! Load in this specific sequence:
+# 1. smolagents - instruments the agent framework
+# 2. litellm - instruments LLM calls made by agents
+# 3. mcp - instruments Model Context Protocol tools
 if OPENINFERENCE_AVAILABLE:
     INSTRUMENTORS.update(
         {
             "smolagents": SmolagentsInstrumentor,
-            "mcp": MCPInstrumentor,
             "litellm": LiteLLMInstrumentor,
+            "mcp": MCPInstrumentor,
         }
     )
 
@@ -163,6 +169,17 @@ def setup_auto_instrumentation(config: OTelConfig):
 
     set_global_textmap(TraceContextTextMapPropagator())
 
+    # Add cost enrichment processor for OpenInference instrumentors
+    # This enriches spans from smolagents, litellm, mcp with cost attributes
+    if config.enable_cost_tracking:
+        try:
+            cost_calculator = CostCalculator()
+            cost_processor = CostEnrichmentSpanProcessor(cost_calculator)
+            tracer_provider.add_span_processor(cost_processor)
+            logger.info("Cost enrichment processor added for OpenInference instrumentors")
+        except Exception as e:
+            logger.warning(f"Failed to add cost enrichment processor: {e}", exc_info=True)
+
     logger.debug(f"OTelConfig endpoint: {config.endpoint}")
     if config.endpoint:
         # Convert timeout to float safely

genai_otel/config.py

@@ -11,7 +11,7 @@ import logging
 import os
 import sys
 from dataclasses import dataclass, field
-from typing import Dict, List, Optional
+from typing import Any, Callable, Dict, List, Optional, Tuple
 
 logger = logging.getLogger(__name__)
 
@@ -44,6 +44,9 @@ DEFAULT_INSTRUMENTORS = [
 ]
 
 # Add OpenInference instrumentors only for Python >= 3.10
+# IMPORTANT: Order matters! Load in this specific sequence:
+# 1. smolagents - instruments the agent framework
+# 2. litellm - instruments the LLM calls made by agents
 if sys.version_info >= (3, 10):
     DEFAULT_INSTRUMENTORS.extend(["smolagents", "litellm"])
 
@@ -101,6 +104,10 @@ class OTelConfig:
         default_factory=lambda: float(os.getenv("GENAI_CARBON_INTENSITY", "475.0"))
     )  # gCO2e/kWh
 
+    power_cost_per_kwh: float = field(
+        default_factory=lambda: float(os.getenv("GENAI_POWER_COST_PER_KWH", "0.12"))
+    )  # USD per kWh - electricity cost for power consumption tracking
+
     gpu_collection_interval: int = field(
         default_factory=lambda: int(os.getenv("GENAI_GPU_COLLECTION_INTERVAL", "5"))
     )  # seconds - how often to collect GPU metrics and CO2 emissions
@@ -117,6 +124,20 @@ class OTelConfig:
         default_factory=lambda: os.getenv("GENAI_ENABLE_CONTENT_CAPTURE", "false").lower() == "true"
     )
 
+    # Custom pricing configuration for models not in llm_pricing.json
+    # Format: JSON string with same structure as llm_pricing.json
+    # Example: {"chat": {"custom-model": {"promptPrice": 0.001, "completionPrice": 0.002}}}
+    custom_pricing_json: Optional[str] = field(
+        default_factory=lambda: os.getenv("GENAI_CUSTOM_PRICING_JSON")
+    )
+
+    # Session and user tracking (Phase 4.1)
+    # Optional callable functions to extract session_id and user_id from requests
+    # Signature: (instance, args, kwargs) -> Optional[str]
+    # Example: lambda instance, args, kwargs: kwargs.get("metadata", {}).get("session_id")
+    session_id_extractor: Optional[Callable[[Any, Tuple, Dict], Optional[str]]] = None
+    user_id_extractor: Optional[Callable[[Any, Tuple, Dict], Optional[str]]] = None
+
 
 import os
 

genai_otel/cost_calculator.py

@@ -2,6 +2,7 @@
 
 import json
 import logging
+import re
 from typing import Any, Dict, Optional
 
 logger = logging.getLogger(__name__)
@@ -12,10 +13,18 @@ class CostCalculator:
 
     DEFAULT_PRICING_FILE = "llm_pricing.json"
 
-    def __init__(self):
-        """Initializes the CostCalculator by loading pricing data from a JSON file."""
+    def __init__(self, custom_pricing_json: Optional[str] = None):
+        """Initializes the CostCalculator by loading pricing data from a JSON file.
+
+        Args:
+            custom_pricing_json: Optional JSON string with custom model pricing.
+                Format: {"chat": {"model-name": {"promptPrice": 0.001, "completionPrice": 0.002}}}
+                Custom prices will be merged with default pricing, with custom taking precedence.
+        """
         self.pricing_data: Dict[str, Any] = {}
         self._load_pricing()
+        if custom_pricing_json:
+            self._merge_custom_pricing(custom_pricing_json)
 
     def _load_pricing(self):
         """Load pricing data from the JSON configuration file."""
@@ -59,6 +68,64 @@ class CostCalculator:
         except Exception as e:
             logger.error("An unexpected error occurred while loading pricing: %s", e, exc_info=True)
 
+    def _merge_custom_pricing(self, custom_pricing_json: str):
+        """Merge custom pricing from JSON string into existing pricing data.
+
+        Args:
+            custom_pricing_json: JSON string with custom model pricing.
+                Format: {"chat": {"model-name": {"promptPrice": 0.001, "completionPrice": 0.002}}}
+        """
+        try:
+            custom_pricing = json.loads(custom_pricing_json)
+
+            if not isinstance(custom_pricing, dict):
+                logger.error(
+                    "Custom pricing must be a JSON object/dict. Got: %s",
+                    type(custom_pricing).__name__,
+                )
+                return
+
+            # Merge custom pricing into each category (chat, embeddings, images, audio)
+            for category, models in custom_pricing.items():
+                if category not in ["chat", "embeddings", "images", "audio"]:
+                    logger.warning(
+                        "Unknown pricing category '%s' in custom pricing. Valid categories: "
+                        "chat, embeddings, images, audio",
+                        category,
+                    )
+                    continue
+
+                if not isinstance(models, dict):
+                    logger.error(
+                        "Custom pricing for category '%s' must be a dict. Got: %s",
+                        category,
+                        type(models).__name__,
+                    )
+                    continue
+
+                # Initialize category if it doesn't exist
+                if category not in self.pricing_data:
+                    self.pricing_data[category] = {}
+
+                # Merge models into the category
+                for model_name, pricing in models.items():
+                    self.pricing_data[category][model_name] = pricing
+                    logger.info(
+                        "Added custom pricing for %s model '%s': %s",
+                        category,
+                        model_name,
+                        pricing,
+                    )
+
+        except json.JSONDecodeError as e:
+            logger.error(
+                "Failed to decode custom pricing JSON: %s. Custom pricing will be ignored.", e
+            )
+        except Exception as e:
+            logger.error(
+                "An unexpected error occurred while merging custom pricing: %s", e, exc_info=True
+            )
+
     def calculate_cost(
         self,
         model: str,
@@ -137,18 +204,32 @@ class CostCalculator:
             Dict with keys: total, prompt, completion, reasoning, cache_read, cache_write
         """
         model_key = self._normalize_model_name(model, "chat")
-        if not model_key:
-            logger.debug("Pricing not found for chat model: %s", model)
-            return {
-                "total": 0.0,
-                "prompt": 0.0,
-                "completion": 0.0,
-                "reasoning": 0.0,
-                "cache_read": 0.0,
-                "cache_write": 0.0,
-            }
 
-        pricing = self.pricing_data["chat"][model_key]
+        # Fallback for unknown local models (Ollama, HuggingFace): estimate pricing based on parameter count
+        if not model_key:
+            param_count = self._extract_param_count_from_model_name(model)
+            if param_count is not None:
+                pricing = self._get_local_model_price_tier(param_count)
+                logger.info(
+                    "Using fallback pricing for unknown local model '%s' with %.2fB parameters: "
+                    "$%.4f prompt / $%.4f completion per 1k tokens",
+                    model,
+                    param_count,
+                    pricing["promptPrice"],
+                    pricing["completionPrice"],
+                )
+            else:
+                logger.debug("Pricing not found for chat model: %s", model)
+                return {
+                    "total": 0.0,
+                    "prompt": 0.0,
+                    "completion": 0.0,
+                    "reasoning": 0.0,
+                    "cache_read": 0.0,
+                    "cache_write": 0.0,
+                }
+        else:
+            pricing = self.pricing_data["chat"][model_key]
 
         # Standard prompt and completion tokens
         prompt_tokens = usage.get("prompt_tokens", 0)
@@ -274,3 +355,113 @@ class CostCalculator:
             if key.lower() in normalized_model:
                 return key
         return None
+
+    def _extract_param_count_from_model_name(self, model: str) -> Optional[float]:
+        """Extract parameter count from Ollama or HuggingFace model name.
+
+        Supports both explicit size indicators and common model size names.
+
+        Examples:
+            Ollama models:
+            "smollm2:360m" -> 0.36
+            "llama3:7b" -> 7.0
+            "llama3.1:70b" -> 70.0
+            "deepseek-r1:32b" -> 32.0
+
+            HuggingFace models:
+            "gpt2" -> 0.124 (base)
+            "gpt2-xl" -> 1.5
+            "bert-base-uncased" -> 0.11
+            "bert-large-uncased" -> 0.34
+            "t5-small" -> 0.06
+            "t5-xxl" -> 11.0
+            "llama-2-7b" -> 7.0
+            "mistral-7b-v0.1" -> 7.0
+
+        Returns:
+            Parameter count in billions, or None if not parseable.
+        """
+        model_lower = model.lower()
+
+        # First try explicit parameter count patterns (e.g., 135m, 7b, 70b)
+        # Matches: digits followed by optional decimal, then 'm' or 'b'
+        pattern = r"(\d+(?:\.\d+)?)(m|b)(?:\s|:|$|-)"
+        match = re.search(pattern, model_lower)
+        if match:
+            value = float(match.group(1))
+            unit = match.group(2)
+            if unit == "m":
+                return value / 1000  # Convert millions to billions
+            elif unit == "b":
+                return value
+
+        # Fallback to common model size indicators for HuggingFace models
+        # These are approximate values based on typical model sizes
+        size_map = {
+            # T5 family
+            "t5-small": 0.06,
+            "t5-base": 0.22,
+            "t5-large": 0.77,
+            "t5-xl": 3.0,
+            "t5-xxl": 11.0,
+            # GPT-2 family
+            "gpt2-small": 0.124,
+            "gpt2-medium": 0.355,
+            "gpt2-large": 0.774,
+            "gpt2-xl": 1.5,
+            "gpt2": 0.124,  # default GPT-2 is small
+            # BERT family
+            "bert-tiny": 0.004,
+            "bert-mini": 0.011,
+            "bert-small": 0.029,
+            "bert-medium": 0.041,
+            "bert-base": 0.11,
+            "bert-large": 0.34,
+            # Generic size indicators (fallback)
+            "tiny": 0.01,
+            "mini": 0.02,
+            "small": 0.06,
+            "base": 0.11,
+            "medium": 0.35,
+            "large": 0.77,
+            "xl": 1.5,
+            "xxl": 11.0,
+        }
+
+        # Check for size indicators in the model name
+        for size_key, param_count in size_map.items():
+            if size_key in model_lower:
+                return param_count
+
+        return None
+
+    def _get_local_model_price_tier(self, param_count_billions: float) -> Dict[str, float]:
+        """Get pricing tier based on parameter count for local models (Ollama, HuggingFace).
+
+        Local models (Ollama, HuggingFace Transformers) are free but consume GPU power
+        and electricity. We estimate costs based on parameter count and comparable
+        cloud API pricing.
+
+        Price Tiers (based on parameter count):
+        - Tiny (< 1B params): $0.0001 / $0.0002 (prompt/completion)
+        - Small (1-10B): $0.0003 / $0.0006
+        - Medium (10-20B): $0.0005 / $0.001
+        - Large (20-80B): $0.0008 / $0.0008
+        - XLarge (80B+): $0.0012 / $0.0012
+
+        Args:
+            param_count_billions: Model parameter count in billions
+
+        Returns:
+            Dict with promptPrice and completionPrice
+        """
+        if param_count_billions < 1.0:
+            return {"promptPrice": 0.0001, "completionPrice": 0.0002}
+        elif param_count_billions < 10.0:
+            return {"promptPrice": 0.0003, "completionPrice": 0.0006}
+        elif param_count_billions < 20.0:
+            return {"promptPrice": 0.0005, "completionPrice": 0.001}
+        elif param_count_billions < 80.0:
+            return {"promptPrice": 0.0008, "completionPrice": 0.0008}
+        else:
+            return {"promptPrice": 0.0012, "completionPrice": 0.0012}

genai_otel/cost_enrichment_processor.py

@@ -0,0 +1,175 @@
+"""Custom SpanProcessor to enrich OpenInference spans with cost tracking.
+
+This processor adds cost attributes to spans created by OpenInference instrumentors
+(smolagents, litellm, mcp) by extracting token usage and model information from
+existing span attributes and calculating costs using our CostCalculator.
+
+Supports both OpenTelemetry GenAI and OpenInference semantic conventions:
+- GenAI: gen_ai.request.model, gen_ai.usage.{prompt_tokens,completion_tokens}
+- OpenInference: llm.model_name, llm.token_count.{prompt,completion}
+"""
+
+import logging
+from typing import Optional
+
+from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor
+from opentelemetry.trace import SpanContext
+
+from .cost_calculator import CostCalculator
+
+logger = logging.getLogger(__name__)
+
+
+class CostEnrichmentSpanProcessor(SpanProcessor):
+    """Enriches spans with cost tracking attributes.
+
+    This processor:
+    1. Identifies spans from OpenInference instrumentors (smolagents, litellm, mcp)
+    2. Extracts model name and token usage from span attributes
+    3. Calculates cost using CostCalculator
+    4. Adds cost attributes (gen_ai.usage.cost.total, etc.) to the span
+    """
+
+    def __init__(self, cost_calculator: Optional[CostCalculator] = None):
+        """Initialize the cost enrichment processor.
+
+        Args:
+            cost_calculator: CostCalculator instance to use for cost calculations.
+                           If None, creates a new instance.
+        """
+        self.cost_calculator = cost_calculator or CostCalculator()
+        logger.info("CostEnrichmentSpanProcessor initialized")
+
+    def on_start(self, span: Span, parent_context: Optional[SpanContext] = None) -> None:
+        """Called when a span starts. No action needed."""
+        pass
+
+    def on_end(self, span: ReadableSpan) -> None:
+        """Called when a span ends. Enriches with cost attributes if applicable.
+
+        Args:
+            span: The span that just ended.
+        """
+        try:
+            # Only process spans that have LLM-related attributes
+            if not span.attributes:
+                return
+
+            attributes = span.attributes
+
+            # Check for model name - support both GenAI and OpenInference conventions
+            model = (
+                attributes.get("gen_ai.request.model")
+                or attributes.get("llm.model_name")
+                or attributes.get("embedding.model_name")
+            )
+            if not model:
+                return
+
+            # Skip if cost attributes are already present (added by instrumentor)
+            if "gen_ai.usage.cost.total" in attributes:
+                logger.debug(f"Span '{span.name}' already has cost attributes, skipping enrichment")
+                return
+
+            # Extract token usage - support GenAI, OpenInference, and legacy conventions
+            prompt_tokens = (
+                attributes.get("gen_ai.usage.prompt_tokens")
+                or attributes.get("gen_ai.usage.input_tokens")
+                or attributes.get("llm.token_count.prompt")  # OpenInference
+                or 0
+            )
+            completion_tokens = (
+                attributes.get("gen_ai.usage.completion_tokens")
+                or attributes.get("gen_ai.usage.output_tokens")
+                or attributes.get("llm.token_count.completion")  # OpenInference
+                or 0
+            )
+
+            # Skip if no tokens recorded
+            if prompt_tokens == 0 and completion_tokens == 0:
+                return
+
+            # Get call type - support both GenAI and OpenInference conventions
+            # OpenInference uses openinference.span.kind (values: LLM, EMBEDDING, etc.)
+            span_kind = attributes.get("openinference.span.kind", "").upper()
+            call_type = attributes.get("gen_ai.operation.name") or span_kind.lower() or "chat"
+
+            # Map operation names to call types for cost calculator
+            # Supports both GenAI and OpenInference conventions
+            call_type_mapping = {
+                # GenAI conventions
+                "chat": "chat",
+                "completion": "chat",
+                "embedding": "embedding",
+                "embeddings": "embedding",
+                "text_generation": "chat",
+                "image_generation": "image",
+                "audio": "audio",
+                # OpenInference conventions (span.kind values)
+                "llm": "chat",
+                "embedding": "embedding",
+                "chain": "chat",
+                "retriever": "embedding",
+                "reranker": "embedding",
+                "tool": "chat",
+                "agent": "chat",
+            }
+            normalized_call_type = call_type_mapping.get(str(call_type).lower(), "chat")
+
+            # Calculate cost
+            usage = {
+                "prompt_tokens": int(prompt_tokens),
+                "completion_tokens": int(completion_tokens),
+                "total_tokens": int(prompt_tokens) + int(completion_tokens),
+            }
+
+            # Use calculate_granular_cost to get detailed breakdown
+            cost_info = self.cost_calculator.calculate_granular_cost(
+                model=str(model),
+                usage=usage,
+                call_type=normalized_call_type,
+            )
+
+            if cost_info and cost_info.get("total", 0.0) > 0:
+                # Add cost attributes to the span
+                # Note: We can't modify ReadableSpan attributes directly,
+                # but we can if span is still a Span instance
+                if isinstance(span, Span):
+                    span.set_attribute("gen_ai.usage.cost.total", cost_info["total"])
+
+                    if cost_info.get("prompt", 0.0) > 0:
+                        span.set_attribute("gen_ai.usage.cost.prompt", cost_info["prompt"])
+                    if cost_info.get("completion", 0.0) > 0:
+                        span.set_attribute("gen_ai.usage.cost.completion", cost_info["completion"])
+
+                    logger.info(
+                        f"Enriched span '{span.name}' with cost: {cost_info['total']:.6f} USD "
+                        f"for model {model} ({usage['total_tokens']} tokens)"
+                    )
+                else:
+                    logger.warning(
+                        f"Span '{span.name}' is not mutable (type: {type(span).__name__}), "
+                        "cannot add cost attributes"
+                    )
+
+        except Exception as e:
+            # Don't fail span processing due to cost enrichment errors
+            logger.warning(
+                f"Failed to enrich span '{getattr(span, 'name', 'unknown')}' with cost: {e}",
+                exc_info=True,
+            )
+
+    def shutdown(self) -> None:
+        """Called when the processor is shutdown."""
+        logger.info("CostEnrichmentSpanProcessor shutdown")
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        """Force flush any pending spans.
+
+        Args:
+            timeout_millis: Timeout in milliseconds.
+
+        Returns:
+            True if flush succeeded.
+        """
+        return True

genai_otel/gpu_metrics.py

@@ -42,6 +42,7 @@ class GPUMetricsCollector:
         self._stop_event = threading.Event()
         self.gpu_utilization_counter: Optional[ObservableCounter] = None
         self.gpu_memory_used_gauge: Optional[ObservableGauge] = None
+        self.gpu_memory_total_gauge: Optional[ObservableGauge] = None
         self.gpu_temperature_gauge: Optional[ObservableGauge] = None
         self.gpu_power_gauge: Optional[ObservableGauge] = None
         self.config = config
@@ -67,6 +68,11 @@ class GPUMetricsCollector:
             description="Cumulative CO2 equivalent emissions in grams",
             unit="gCO2e",
         )
+        self.power_cost_counter = meter.create_counter(
+            "gen_ai.power.cost",  # New metric name
+            description="Cumulative electricity cost in USD based on GPU power consumption",
+            unit="USD",
+        )
         if not NVML_AVAILABLE:
             logger.warning(
                 "GPU metrics collection not available - nvidia-ml-py not installed. "
@@ -88,6 +94,12 @@ class GPUMetricsCollector:
                 description="GPU memory used in MiB",
                 unit="MiB",
             )
+            self.gpu_memory_total_gauge = self.meter.create_observable_gauge(
+                "gen_ai.gpu.memory.total",  # Fixed metric name
+                callbacks=[self._observe_gpu_memory_total],
+                description="Total GPU memory capacity in MiB",
+                unit="MiB",
+            )
             self.gpu_temperature_gauge = self.meter.create_observable_gauge(
                 "gen_ai.gpu.temperature",  # Fixed metric name
                 callbacks=[self._observe_gpu_temperature],
@@ -167,6 +179,33 @@ class GPUMetricsCollector:
         except Exception as e:
             logger.error("Error observing GPU memory: %s", e)
 
+    def _observe_gpu_memory_total(self, options):
+        """Observable callback for total GPU memory capacity."""
+        if not NVML_AVAILABLE or not self.gpu_available:
+            return
+
+        try:
+            pynvml.nvmlInit()
+            device_count = pynvml.nvmlDeviceGetCount()
+
+            for i in range(device_count):
+                handle = pynvml.nvmlDeviceGetHandleByIndex(i)
+                device_name = self._get_device_name(handle, i)
+
+                try:
+                    memory_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
+                    gpu_memory_total = memory_info.total / (1024**2)  # Convert to MiB
+                    yield Observation(
+                        value=gpu_memory_total,
+                        attributes={"gpu_id": str(i), "gpu_name": device_name},
+                    )
+                except Exception as e:
+                    logger.debug("Failed to get total GPU memory for GPU %d: %s", i, e)
+
+            pynvml.nvmlShutdown()
+        except Exception as e:
+            logger.error("Error observing total GPU memory: %s", e)
+
     def _observe_gpu_temperature(self, options):
         """Observable callback for GPU temperature."""
         if not NVML_AVAILABLE or not self.gpu_available:
@@ -249,11 +288,22 @@ class GPUMetricsCollector:
                         delta_time_hours * 3600.0
                     )  # Wh (power in kW * hours = kWh, but track in Wh for precision)
                     self.cumulative_energy_wh[i] += delta_energy_wh
+
+                    # Calculate and record CO2 emissions if enabled
                     if self.config.enable_co2_tracking:
                         delta_co2_g = (
                             delta_energy_wh / 1000.0
                         ) * self.config.carbon_intensity  # gCO2e
                         self.co2_counter.add(delta_co2_g, {"gpu_id": str(i)})
+
+                    # Calculate and record power cost
+                    # delta_energy_wh is in Wh, convert to kWh and multiply by cost per kWh
+                    delta_cost_usd = (delta_energy_wh / 1000.0) * self.config.power_cost_per_kwh
+                    device_name = self._get_device_name(handle, i)
+                    self.power_cost_counter.add(
+                        delta_cost_usd, {"gpu_id": str(i), "gpu_name": device_name}
+                    )
+
                     self.last_timestamp[i] = current_time
                 except Exception as e:
                     logger.error(f"Error collecting GPU {i} metrics: {e}")