genai-otel-instrument 0.1.4.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.4.dev0'
-__version_tuple__ = version_tuple = (0, 1, 4, 'dev0')
+__version__ = version = '0.1.7.dev0'
+__version_tuple__ = version_tuple = (0, 1, 7, 'dev0')
 
 __commit_id__ = commit_id = None

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__)
 
@@ -104,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
@@ -120,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

@@ -13,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."""
@@ -60,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,
@@ -150,7 +216,7 @@ class CostCalculator:
                     model,
                     param_count,
                     pricing["promptPrice"],
-                    pricing["completionPrice"]
+                    pricing["completionPrice"],
                 )
             else:
                 logger.debug("Pricing not found for chat model: %s", model)
@@ -319,14 +385,14 @@ class CostCalculator:
 
         # 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|:|$|-)'
+        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':
+            if unit == "m":
                 return value / 1000  # Convert millions to billions
-            elif unit == 'b':
+            elif unit == "b":
                 return value
 
         # Fallback to common model size indicators for HuggingFace models

genai_otel/cost_enrichment_processor.py

@@ -1,177 +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
+"""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}")