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

genai_otel/instrumentors/huggingface_instrumentor.py

@@ -0,0 +1,399 @@
+"""OpenTelemetry instrumentor for HuggingFace Transformers and Inference API.
+
+This instrumentor automatically traces:
+1. HuggingFace Transformers pipelines (local model execution)
+2. HuggingFace Inference API calls via InferenceClient (used by smolagents)
+3. Direct model usage via AutoModelForCausalLM.generate() and forward()
+
+Note: Transformers runs models locally (no API costs), but InferenceClient makes
+API calls to HuggingFace endpoints which may have costs based on usage.
+Local model costs are estimated based on parameter count and token usage.
+"""
+
+import logging
+from typing import Dict, Optional
+
+from ..config import OTelConfig
+from .base import BaseInstrumentor
+
+logger = logging.getLogger(__name__)
+
+
+class HuggingFaceInstrumentor(BaseInstrumentor):
+    """Instrumentor for HuggingFace Transformers and Inference API.
+
+    Instruments:
+    - transformers.pipeline (local execution, estimated costs)
+    - transformers.AutoModelForCausalLM.generate() (local execution, estimated costs)
+    - transformers.AutoModelForCausalLM.forward() (local execution, estimated costs)
+    - huggingface_hub.InferenceClient (API calls, may have costs)
+    """
+
+    def __init__(self):
+        """Initialize the instrumentor."""
+        super().__init__()
+        self._transformers_available = False
+        self._inference_client_available = False
+        self._model_classes_instrumented = False
+        self._check_availability()
+
+    def _check_availability(self):
+        """Check if Transformers and InferenceClient libraries are available."""
+        try:
+            import transformers
+
+            self._transformers_available = True
+            logger.debug("Transformers library detected and available for instrumentation")
+        except ImportError:
+            logger.debug("Transformers library not installed, instrumentation will be skipped")
+            self._transformers_available = False
+
+        try:
+            from huggingface_hub import InferenceClient
+
+            self._inference_client_available = True
+            logger.debug("HuggingFace InferenceClient detected and available for instrumentation")
+        except ImportError:
+            logger.debug(
+                "huggingface_hub not installed, InferenceClient instrumentation will be skipped"
+            )
+            self._inference_client_available = False
+
+    def instrument(self, config: OTelConfig):
+        """Instrument HuggingFace Transformers pipelines, model classes, and InferenceClient."""
+        self._setup_config(config)
+
+        instrumented_count = 0
+
+        # Instrument transformers components if available
+        if self._transformers_available:
+            # Instrument pipeline
+            try:
+                self._instrument_transformers()
+                instrumented_count += 1
+            except Exception as e:
+                logger.error("Failed to instrument HuggingFace Transformers: %s", e, exc_info=True)
+                if config.fail_on_error:
+                    raise
+
+            # Instrument model classes (AutoModelForCausalLM, etc.)
+            try:
+                self._instrument_model_classes()
+                instrumented_count += 1
+            except Exception as e:
+                logger.error("Failed to instrument HuggingFace model classes: %s", e, exc_info=True)
+                if config.fail_on_error:
+                    raise
+
+        # Instrument InferenceClient if available
+        if self._inference_client_available:
+            try:
+                self._instrument_inference_client()
+                instrumented_count += 1
+            except Exception as e:
+                logger.error(
+                    "Failed to instrument HuggingFace InferenceClient: %s", e, exc_info=True
+                )
+                if config.fail_on_error:
+                    raise
+
+        if instrumented_count > 0:
+            self._instrumented = True
+            logger.info(f"HuggingFace instrumentation enabled ({instrumented_count} components)")
+
+    def _instrument_transformers(self):
+        """Instrument transformers.pipeline for local model execution."""
+        try:
+            import importlib
+
+            transformers_module = importlib.import_module("transformers")
+            original_pipeline = transformers_module.pipeline
+
+            # Capture self reference for use in nested classes
+            instrumentor = self
+
+            def wrapped_pipeline(*args, **kwargs):
+                pipe = original_pipeline(*args, **kwargs)
+
+                class WrappedPipeline:
+                    def __init__(self, original_pipe):
+                        self._original_pipe = original_pipe
+
+                    def __call__(self, *call_args, **call_kwargs):
+                        # Use instrumentor.tracer instead of config.tracer
+                        with instrumentor.tracer.start_span("huggingface.pipeline") as span:
+                            task = getattr(self._original_pipe, "task", "unknown")
+                            model = getattr(
+                                getattr(self._original_pipe, "model", None),
+                                "name_or_path",
+                                "unknown",
+                            )
+
+                            span.set_attribute("gen_ai.system", "huggingface")
+                            span.set_attribute("gen_ai.request.model", model)
+                            span.set_attribute("gen_ai.operation.name", task)
+                            span.set_attribute("huggingface.task", task)
+
+                            if instrumentor.request_counter:
+                                instrumentor.request_counter.add(
+                                    1, {"model": model, "provider": "huggingface"}
+                                )
+
+                            result = self._original_pipe(*call_args, **call_kwargs)
+
+                            # End span manually
+                            span.end()
+                            return result
+
+                    def __getattr__(self, name):
+                        # Delegate all other attribute access to the original pipe
+                        return getattr(self._original_pipe, name)
+
+                return WrappedPipeline(pipe)
+
+            transformers_module.pipeline = wrapped_pipeline
+            logger.debug("HuggingFace Transformers pipeline instrumented")
+
+        except Exception as e:
+            raise  # Re-raise to be caught by instrument() method
+
+    def _instrument_inference_client(self):
+        """Instrument HuggingFace InferenceClient for API calls."""
+        from huggingface_hub import InferenceClient
+
+        # Store original methods
+        original_chat_completion = InferenceClient.chat_completion
+        original_text_generation = InferenceClient.text_generation
+
+        # Wrap chat_completion method
+        wrapped_chat_completion = self.create_span_wrapper(
+            span_name="huggingface.inference.chat_completion",
+            extract_attributes=self._extract_inference_client_attributes,
+        )(original_chat_completion)
+
+        # Wrap text_generation method
+        wrapped_text_generation = self.create_span_wrapper(
+            span_name="huggingface.inference.text_generation",
+            extract_attributes=self._extract_inference_client_attributes,
+        )(original_text_generation)
+
+        InferenceClient.chat_completion = wrapped_chat_completion
+        InferenceClient.text_generation = wrapped_text_generation
+        logger.debug("HuggingFace InferenceClient instrumented")
+
+    def _instrument_model_classes(self):
+        """Instrument HuggingFace model classes for direct model usage."""
+        try:
+            import wrapt
+
+            # Import GenerationMixin - the base class that provides generate() method
+            # All generative models (AutoModelForCausalLM, AutoModelForSeq2SeqLM, etc.) inherit from it
+            try:
+                from transformers.generation.utils import GenerationMixin
+            except ImportError:
+                # Fallback for older transformers versions
+                from transformers.generation import GenerationMixin
+
+            # Store reference to instrumentor for use in wrapper
+            instrumentor = self
+
+            # Wrap the generate() method at GenerationMixin level (all models inherit from this)
+            original_generate = GenerationMixin.generate
+
+            @wrapt.decorator
+            def generate_wrapper(wrapped, instance, args, kwargs):
+                """Wrapper for model.generate() method."""
+                # Extract model info
+                model_name = getattr(instance, "name_or_path", "unknown")
+                if hasattr(instance.config, "_name_or_path"):
+                    model_name = instance.config._name_or_path
+
+                # Get input token count
+                input_ids = kwargs.get("input_ids") or (args[0] if args else None)
+                prompt_tokens = 0
+                if input_ids is not None:
+                    if hasattr(input_ids, "shape"):
+                        prompt_tokens = int(input_ids.shape[-1])
+                    elif isinstance(input_ids, (list, tuple)):
+                        prompt_tokens = len(input_ids[0]) if input_ids else 0
+
+                # Create span
+                with instrumentor.tracer.start_as_current_span(
+                    "huggingface.model.generate"
+                ) as span:
+                    # Set attributes
+                    span.set_attribute("gen_ai.system", "huggingface")
+                    span.set_attribute("gen_ai.request.model", model_name)
+                    span.set_attribute("gen_ai.operation.name", "text_generation")
+                    span.set_attribute("gen_ai.request.type", "chat")
+
+                    # Extract generation parameters
+                    if "max_length" in kwargs:
+                        span.set_attribute("gen_ai.request.max_tokens", kwargs["max_length"])
+                    if "max_new_tokens" in kwargs:
+                        span.set_attribute("gen_ai.request.max_tokens", kwargs["max_new_tokens"])
+                    if "temperature" in kwargs:
+                        span.set_attribute("gen_ai.request.temperature", kwargs["temperature"])
+                    if "top_p" in kwargs:
+                        span.set_attribute("gen_ai.request.top_p", kwargs["top_p"])
+
+                    # Call original generate
+                    import time
+
+                    start_time = time.time()
+                    result = wrapped(*args, **kwargs)
+                    duration = time.time() - start_time
+
+                    # Extract output token count
+                    completion_tokens = 0
+                    if hasattr(result, "shape"):
+                        # result is a tensor
+                        total_length = int(result.shape[-1])
+                        completion_tokens = max(0, total_length - prompt_tokens)
+                    elif isinstance(result, (list, tuple)):
+                        # result is a list of sequences
+                        if result and hasattr(result[0], "shape"):
+                            total_length = int(result[0].shape[-1])
+                            completion_tokens = max(0, total_length - prompt_tokens)
+
+                    total_tokens = prompt_tokens + completion_tokens
+
+                    # Set token usage attributes
+                    if prompt_tokens > 0:
+                        span.set_attribute("gen_ai.usage.prompt_tokens", prompt_tokens)
+                    if completion_tokens > 0:
+                        span.set_attribute("gen_ai.usage.completion_tokens", completion_tokens)
+                    if total_tokens > 0:
+                        span.set_attribute("gen_ai.usage.total_tokens", total_tokens)
+
+                    # Record metrics
+                    if instrumentor.request_counter:
+                        instrumentor.request_counter.add(
+                            1, {"model": model_name, "provider": "huggingface"}
+                        )
+
+                    if instrumentor.token_counter and total_tokens > 0:
+                        if prompt_tokens > 0:
+                            instrumentor.token_counter.add(
+                                prompt_tokens, {"token_type": "prompt", "operation": span.name}
+                            )
+                        if completion_tokens > 0:
+                            instrumentor.token_counter.add(
+                                completion_tokens,
+                                {"token_type": "completion", "operation": span.name},
+                            )
+
+                    if instrumentor.latency_histogram:
+                        instrumentor.latency_histogram.record(duration, {"operation": span.name})
+
+                    # Calculate and record cost if enabled
+                    if (
+                        instrumentor.config
+                        and instrumentor.config.enable_cost_tracking
+                        and total_tokens > 0
+                    ):
+                        try:
+                            usage = {
+                                "prompt_tokens": prompt_tokens,
+                                "completion_tokens": completion_tokens,
+                                "total_tokens": total_tokens,
+                            }
+
+                            costs = instrumentor.cost_calculator.calculate_granular_cost(
+                                model=model_name, usage=usage, call_type="chat"
+                            )
+
+                            if costs["total"] > 0:
+                                if instrumentor.cost_counter:
+                                    instrumentor.cost_counter.add(
+                                        costs["total"], {"model": model_name}
+                                    )
+                                span.set_attribute("gen_ai.usage.cost.total", costs["total"])
+                                if costs["prompt"] > 0:
+                                    span.set_attribute("gen_ai.usage.cost.prompt", costs["prompt"])
+                                if costs["completion"] > 0:
+                                    span.set_attribute(
+                                        "gen_ai.usage.cost.completion", costs["completion"]
+                                    )
+
+                                logger.debug(
+                                    f"HuggingFace model {model_name}: {total_tokens} tokens, "
+                                    f"cost: ${costs['total']:.6f}"
+                                )
+                        except Exception as e:
+                            logger.warning(f"Failed to calculate cost: {e}")
+
+                    return result
+
+            # Apply wrapper to GenerationMixin.generate (all models inherit this)
+            GenerationMixin.generate = generate_wrapper(original_generate)
+
+            self._model_classes_instrumented = True
+            logger.debug(
+                "HuggingFace GenerationMixin.generate() instrumented "
+                "(covers all models: AutoModelForCausalLM, AutoModelForSeq2SeqLM, etc.)"
+            )
+
+        except ImportError as e:
+            logger.debug(f"Could not import model classes for instrumentation: {e}")
+        except Exception as e:
+            raise  # Re-raise to be caught by instrument() method
+
+    def _extract_inference_client_attributes(self, instance, args, kwargs) -> Dict[str, str]:
+        """Extract attributes from Inference API call."""
+        attrs = {}
+        model = kwargs.get("model") or (args[0] if args else "unknown")
+
+        attrs["gen_ai.system"] = "huggingface"
+        attrs["gen_ai.request.model"] = str(model)
+        attrs["gen_ai.operation.name"] = "chat"  # Default to chat
+
+        # Extract parameters if available
+        if "max_tokens" in kwargs:
+            attrs["gen_ai.request.max_tokens"] = kwargs["max_tokens"]
+        if "temperature" in kwargs:
+            attrs["gen_ai.request.temperature"] = kwargs["temperature"]
+        if "top_p" in kwargs:
+            attrs["gen_ai.request.top_p"] = kwargs["top_p"]
+
+        return attrs
+
+    def _extract_usage(self, result) -> Optional[Dict[str, int]]:
+        """Extract token usage from HuggingFace response.
+
+        Handles both:
+        1. Transformers pipeline (local execution) - returns None
+        2. InferenceClient API calls - extracts token usage from response
+
+        Args:
+            result: The pipeline output or InferenceClient response.
+
+        Returns:
+            Dict with token counts for InferenceClient calls, None for local execution.
+        """
+        # Check if this is an InferenceClient API response
+        if result is not None and hasattr(result, "usage"):
+            usage = result.usage
+
+            # Extract token counts from usage object
+            prompt_tokens = getattr(usage, "prompt_tokens", None)
+            completion_tokens = getattr(usage, "completion_tokens", None)
+            total_tokens = getattr(usage, "total_tokens", None)
+
+            # If usage is a dict instead of object
+            if isinstance(usage, dict):
+                prompt_tokens = usage.get("prompt_tokens")
+                completion_tokens = usage.get("completion_tokens")
+                total_tokens = usage.get("total_tokens")
+
+            # Return token counts if available
+            if prompt_tokens is not None or completion_tokens is not None:
+                return {
+                    "prompt_tokens": prompt_tokens or 0,
+                    "completion_tokens": completion_tokens or 0,
+                    "total_tokens": total_tokens or (prompt_tokens or 0) + (completion_tokens or 0),
+                }
+
+        # HuggingFace Transformers is free (local execution)
+        # No token-based costs to track
+        return None

genai_otel/instrumentors/hyperbolic_instrumentor.py

@@ -0,0 +1,236 @@
+"""OpenTelemetry instrumentor for Hyperbolic API calls.
+
+This instrumentor automatically traces HTTP requests to Hyperbolic's API,
+capturing relevant LLM attributes such as model name and token usage from
+the raw HTTP response.
+"""
+
+import json
+import logging
+from typing import Any, Dict, Optional
+
+import wrapt
+
+from ..config import OTelConfig
+from .base import BaseInstrumentor
+
+logger = logging.getLogger(__name__)
+
+
+class HyperbolicInstrumentor(BaseInstrumentor):
+    """Instrumentor for Hyperbolic API (raw HTTP requests)"""
+
+    HYPERBOLIC_API_BASE = "https://api.hyperbolic.xyz"
+
+    def __init__(self):
+        """Initialize the instrumentor."""
+        super().__init__()
+        self._requests_available = False
+        self._check_availability()
+
+    def _check_availability(self):
+        """Check if requests library is available."""
+        try:
+            import requests
+
+            self._requests_available = True
+            logger.debug("Requests library detected, Hyperbolic instrumentation available")
+        except ImportError:
+            logger.debug("Requests library not installed, Hyperbolic instrumentation skipped")
+            self._requests_available = False
+
+    def instrument(self, config: OTelConfig):
+        """Instrument requests library for Hyperbolic API calls.
+
+        Args:
+            config (OTelConfig): The OpenTelemetry configuration object.
+        """
+        if not self._requests_available:
+            logger.debug("Skipping Hyperbolic instrumentation - requests library not available")
+            return
+
+        self.config = config
+
+        try:
+            import requests
+
+            # Wrap requests.post to intercept Hyperbolic API calls
+            original_post = requests.post
+
+            @wrapt.decorator
+            def hyperbolic_post_wrapper(wrapped, instance, args, kwargs):
+                # Check if this is a Hyperbolic API call
+                url = args[0] if args else kwargs.get("url", "")
+                if not url.startswith(self.HYPERBOLIC_API_BASE):
+                    # Not a Hyperbolic call, pass through
+                    return wrapped(*args, **kwargs)
+
+                # Extract attributes before the call
+                request_data = kwargs.get("json", {})
+                attrs = self._extract_request_attributes(request_data)
+
+                # Create span wrapper
+                with self.tracer.start_as_current_span("hyperbolic.chat.completion") as span:
+                    # Set request attributes
+                    for key, value in attrs.items():
+                        span.set_attribute(key, value)
+
+                    # Record request metric
+                    model = attrs.get("gen_ai.request.model", "unknown")
+                    if self.request_counter:
+                        self.request_counter.add(1, {"model": model, "provider": "hyperbolic"})
+
+                    try:
+                        # Make the actual API call
+                        response = wrapped(*args, **kwargs)
+
+                        # Extract response attributes
+                        if response.status_code == 200:
+                            response_data = response.json()
+                            self._extract_and_record_response(span, response_data)
+                        else:
+                            span.set_attribute("error", True)
+                            span.set_attribute("http.status_code", response.status_code)
+
+                        return response
+
+                    except Exception as e:
+                        span.set_attribute("error", True)
+                        span.record_exception(e)
+                        if self.error_counter:
+                            self.error_counter.add(
+                                1,
+                                {
+                                    "operation": "chat.completion",
+                                    "error.type": type(e).__name__,
+                                    "provider": "hyperbolic",
+                                },
+                            )
+                        raise
+
+            # Apply the wrapper
+            requests.post = hyperbolic_post_wrapper(original_post)
+            self._instrumented = True
+            logger.info("Hyperbolic instrumentation enabled")
+
+        except Exception as e:
+            logger.error("Failed to instrument Hyperbolic: %s", e, exc_info=True)
+            if config.fail_on_error:
+                raise
+
+    def _extract_request_attributes(self, request_data: Dict[str, Any]) -> Dict[str, Any]:
+        """Extract attributes from Hyperbolic API request.
+
+        Args:
+            request_data: The JSON request payload.
+
+        Returns:
+            Dict[str, Any]: Dictionary of attributes to set on the span.
+        """
+        attrs = {}
+
+        # Core attributes
+        attrs["gen_ai.system"] = "hyperbolic"
+        attrs["gen_ai.request.model"] = request_data.get("model", "unknown")
+        attrs["gen_ai.operation.name"] = "chat"
+
+        messages = request_data.get("messages", [])
+        attrs["gen_ai.request.message_count"] = len(messages)
+
+        # Request parameters
+        if "temperature" in request_data:
+            attrs["gen_ai.request.temperature"] = request_data["temperature"]
+        if "top_p" in request_data:
+            attrs["gen_ai.request.top_p"] = request_data["top_p"]
+        if "max_tokens" in request_data:
+            attrs["gen_ai.request.max_tokens"] = request_data["max_tokens"]
+
+        # First message preview
+        if messages:
+            first_message = str(messages[0])[:200]
+            attrs["gen_ai.request.first_message"] = first_message
+
+        return attrs
+
+    def _extract_and_record_response(self, span, response_data: Dict[str, Any]):
+        """Extract response attributes and record metrics.
+
+        Args:
+            span: The OpenTelemetry span.
+            response_data: The JSON response from Hyperbolic API.
+        """
+        # Response ID
+        if "id" in response_data:
+            span.set_attribute("gen_ai.response.id", response_data["id"])
+
+        # Response model
+        if "model" in response_data:
+            span.set_attribute("gen_ai.response.model", response_data["model"])
+
+        # Finish reasons
+        choices = response_data.get("choices", [])
+        if choices:
+            finish_reasons = [
+                choice.get("finish_reason") for choice in choices if "finish_reason" in choice
+            ]
+            if finish_reasons:
+                span.set_attribute("gen_ai.response.finish_reasons", finish_reasons)
+
+        # Extract token usage
+        usage_data = response_data.get("usage", {})
+        if usage_data:
+            usage_dict = {
+                "prompt_tokens": usage_data.get("prompt_tokens", 0),
+                "completion_tokens": usage_data.get("completion_tokens", 0),
+                "total_tokens": usage_data.get("total_tokens", 0),
+            }
+
+            # Record token usage as span attributes
+            span.set_attribute("gen_ai.usage.prompt_tokens", usage_dict["prompt_tokens"])
+            span.set_attribute("gen_ai.usage.completion_tokens", usage_dict["completion_tokens"])
+            span.set_attribute("gen_ai.usage.total_tokens", usage_dict["total_tokens"])
+
+            # Record token metrics
+            if self.token_counter:
+                model = span.attributes.get("gen_ai.request.model", "unknown")
+                self.token_counter.add(
+                    usage_dict["prompt_tokens"],
+                    {"token_type": "prompt", "model": model, "provider": "hyperbolic"},
+                )
+                self.token_counter.add(
+                    usage_dict["completion_tokens"],
+                    {"token_type": "completion", "model": model, "provider": "hyperbolic"},
+                )
+
+            # Calculate and record cost
+            if self.config.enable_cost_tracking:
+                from ..cost_calculator import CostCalculator
+
+                cost_calc = CostCalculator(custom_pricing_json=self.config.custom_pricing_json)
+                model = span.attributes.get("gen_ai.request.model", "unknown")
+                cost = cost_calc.calculate_cost(
+                    model_name=model,
+                    prompt_tokens=usage_dict["prompt_tokens"],
+                    completion_tokens=usage_dict["completion_tokens"],
+                    call_type="chat",
+                )
+
+                if cost > 0 and self.cost_counter:
+                    span.set_attribute("gen_ai.cost.amount", cost)
+                    self.cost_counter.add(
+                        cost, {"model": model, "provider": "hyperbolic", "call_type": "chat"}
+                    )
+
+    def _extract_usage(self, result) -> Optional[Dict[str, int]]:
+        """Extract token usage from response.
+
+        Note: This method is required by BaseInstrumentor but not used for HTTP-based
+        instrumentation. Token extraction is handled in _extract_and_record_response.
+
+        Args:
+            result: The API response (unused for HTTP instrumentation).
+
+        Returns:
+            None
+        """
+        return None