genai-otel-instrument 0.1.4.dev0__py3-none-any.whl → 0.1.9.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.9.dev0'
+__version_tuple__ = version_tuple = (0, 1, 9, 'dev0')
 
 __commit_id__ = commit_id = None

genai_otel/auto_instrument.py

@@ -19,6 +19,7 @@ from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExport
 from .config import OTelConfig
 from .cost_calculator import CostCalculator
 from .cost_enrichment_processor import CostEnrichmentSpanProcessor
+from .cost_enriching_exporter import CostEnrichingSpanExporter
 from .gpu_metrics import GPUMetricsCollector
 from .mcp_instrumentors import MCPInstrumentorManager
 from .metrics import (
@@ -169,14 +170,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
+    # Add cost enrichment processor for custom instrumentors (OpenAI, Ollama, etc.)
+    # These instrumentors set cost attributes directly, so processor is mainly for logging
+    # Also attempts to enrich OpenInference spans (smolagents, litellm, mcp), though
+    # the processor can't modify ReadableSpan - the exporter below handles that
+    cost_calculator = None
     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")
+            logger.info("Cost enrichment processor added")
         except Exception as e:
             logger.warning(f"Failed to add cost enrichment processor: {e}", exc_info=True)
 

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_enriching_exporter.py

@@ -0,0 +1,207 @@
+"""Custom SpanExporter that enriches spans with cost attributes before export.
+
+This exporter wraps another exporter (like OTLPSpanExporter) and adds cost
+attributes to spans before passing them to the wrapped exporter.
+"""
+
+import logging
+from typing import Optional, Sequence
+
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
+
+from .cost_calculator import CostCalculator
+
+logger = logging.getLogger(__name__)
+
+
+class CostEnrichingSpanExporter(SpanExporter):
+    """Wraps a SpanExporter and enriches spans with cost attributes before export.
+
+    This exporter:
+    1. Receives ReadableSpan objects from the SDK
+    2. Extracts model name and token usage from span attributes
+    3. Calculates cost using CostCalculator
+    4. Creates enriched span data with cost attributes
+    5. Exports to the wrapped exporter (e.g., OTLP)
+    """
+
+    def __init__(
+        self, wrapped_exporter: SpanExporter, cost_calculator: Optional[CostCalculator] = None
+    ):
+        """Initialize the cost enriching exporter.
+
+        Args:
+            wrapped_exporter: The underlying exporter to send enriched spans to.
+            cost_calculator: CostCalculator instance to use for cost calculations.
+                           If None, creates a new instance.
+        """
+        self.wrapped_exporter = wrapped_exporter
+        self.cost_calculator = cost_calculator or CostCalculator()
+        logger.info(
+            f"CostEnrichingSpanExporter initialized, wrapping {type(wrapped_exporter).__name__}"
+        )
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        """Export spans after enriching them with cost attributes.
+
+        Args:
+            spans: Sequence of ReadableSpan objects to export.
+
+        Returns:
+            SpanExportResult from the wrapped exporter.
+        """
+        try:
+            # Enrich spans with cost attributes
+            enriched_spans = []
+            for span in spans:
+                enriched_span = self._enrich_span(span)
+                enriched_spans.append(enriched_span)
+
+            # Export to wrapped exporter
+            return self.wrapped_exporter.export(enriched_spans)
+
+        except Exception as e:
+            logger.error(f"Failed to export spans: {e}", exc_info=True)
+            return SpanExportResult.FAILURE
+
+    def _enrich_span(self, span: ReadableSpan) -> ReadableSpan:
+        """Enrich a span with cost attributes if applicable.
+
+        Args:
+            span: The original ReadableSpan.
+
+        Returns:
+            A new ReadableSpan with cost attributes added (or the original if not applicable).
+        """
+        try:
+            # Check if span has LLM-related attributes
+            if not span.attributes:
+                return span
+
+            attributes = dict(span.attributes)  # Make a mutable copy
+
+            # 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 span
+
+            # Skip if cost attributes are already present
+            if "gen_ai.usage.cost.total" in attributes:
+                logger.debug(f"Span '{span.name}' already has cost attributes, skipping enrichment")
+                return span
+
+            # 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 span
+
+            # Get call type - support both GenAI and OpenInference conventions
+            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
+            call_type_mapping = {
+                "chat": "chat",
+                "completion": "chat",
+                "embedding": "embedding",
+                "embeddings": "embedding",
+                "text_generation": "chat",
+                "image_generation": "image",
+                "audio": "audio",
+                "llm": "chat",
+                "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),
+            }
+
+            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 mutable copy
+                attributes["gen_ai.usage.cost.total"] = cost_info["total"]
+
+                if cost_info.get("prompt", 0.0) > 0:
+                    attributes["gen_ai.usage.cost.prompt"] = cost_info["prompt"]
+                if cost_info.get("completion", 0.0) > 0:
+                    attributes["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)"
+                )
+
+                # Create a new ReadableSpan with enriched attributes
+                # ReadableSpan is a NamedTuple, so we need to replace it
+                from opentelemetry.sdk.trace import ReadableSpan as RS
+
+                enriched_span = RS(
+                    name=span.name,
+                    context=span.context,
+                    kind=span.kind,
+                    parent=span.parent,
+                    start_time=span.start_time,
+                    end_time=span.end_time,
+                    status=span.status,
+                    attributes=attributes,  # Use enriched attributes
+                    events=span.events,
+                    links=span.links,
+                    resource=span.resource,
+                    instrumentation_scope=span.instrumentation_scope,
+                )
+                return enriched_span
+
+        except Exception as e:
+            logger.warning(
+                f"Failed to enrich span '{getattr(span, 'name', 'unknown')}' with cost: {e}",
+                exc_info=True,
+            )
+
+        return span
+
+    def shutdown(self) -> None:
+        """Shutdown the wrapped exporter."""
+        logger.info("CostEnrichingSpanExporter shutting down")
+        self.wrapped_exporter.shutdown()
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        """Force flush the wrapped exporter.
+
+        Args:
+            timeout_millis: Timeout in milliseconds.
+
+        Returns:
+            True if flush succeeded.
+        """
+        return self.wrapped_exporter.force_flush(timeout_millis)