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

genai_otel/instrumentors/cohere_instrumentor.py

@@ -1,11 +1,11 @@
 """OpenTelemetry instrumentor for the Cohere SDK.
 
 This instrumentor automatically traces calls to Cohere models, capturing
-relevant attributes such as the model name.
+relevant attributes such as the model name and token usage.
 """
 
 import logging
-from typing import Dict, Optional
+from typing import Any, Dict, Optional
 
 from ..config import OTelConfig
 from .base import BaseInstrumentor
@@ -34,7 +34,7 @@ class CohereInstrumentor(BaseInstrumentor):
             self._cohere_available = False
 
     def instrument(self, config: OTelConfig):
-        """Instrument  cohere available if available."""
+        """Instrument cohere if available."""
         if not self._cohere_available:
             logger.debug("Skipping instrumentation - library not available")
             return
@@ -50,27 +50,91 @@ class CohereInstrumentor(BaseInstrumentor):
                 self._instrument_client(instance)
 
             cohere.Client.__init__ = wrapped_init
+            self._instrumented = True
+            logger.info("Cohere instrumentation enabled")
 
-        except ImportError:
-            pass
+        except Exception as e:
+            logger.error("Failed to instrument Cohere: %s", e, exc_info=True)
+            if config.fail_on_error:
+                raise
 
     def _instrument_client(self, client):
+        """Instrument Cohere client methods."""
         original_generate = client.generate
 
-        def wrapped_generate(*args, **kwargs):
-            with self.tracer.start_as_current_span("cohere.generate") as span:
-                model = kwargs.get("model", "command")
+        # Wrap using create_span_wrapper
+        wrapped_generate = self.create_span_wrapper(
+            span_name="cohere.generate",
+            extract_attributes=self._extract_generate_attributes,
+        )(original_generate)
 
-                span.set_attribute("gen_ai.system", "cohere")
-                span.set_attribute("gen_ai.request.model", model)
+        client.generate = wrapped_generate
 
-                if self.request_counter:
-                    self.request_counter.add(1, {"model": model, "provider": "cohere"})
+    def _extract_generate_attributes(self, instance: Any, args: Any, kwargs: Any) -> Dict[str, Any]:
+        """Extract attributes from Cohere generate call.
 
-                result = original_generate(*args, **kwargs)
-                return result
+        Args:
+            instance: The client instance.
+            args: Positional arguments.
+            kwargs: Keyword arguments.
 
-        client.generate = wrapped_generate
+        Returns:
+            Dict[str, Any]: Dictionary of attributes to set on the span.
+        """
+        attrs = {}
+        model = kwargs.get("model", "command")
+        prompt = kwargs.get("prompt", "")
+
+        attrs["gen_ai.system"] = "cohere"
+        attrs["gen_ai.request.model"] = model
+        attrs["gen_ai.operation.name"] = "generate"
+        attrs["gen_ai.request.message_count"] = 1 if prompt else 0
+
+        return attrs
 
     def _extract_usage(self, result) -> Optional[Dict[str, int]]:
-        return None
+        """Extract token usage from Cohere response.
+
+        Cohere responses include meta.tokens with:
+        - input_tokens: Input tokens
+        - output_tokens: Output tokens
+
+        Args:
+            result: The API response object.
+
+        Returns:
+            Optional[Dict[str, int]]: Dictionary with token counts or None.
+        """
+        try:
+            # Handle object response
+            if hasattr(result, "meta") and result.meta:
+                meta = result.meta
+                # Check for tokens object
+                if hasattr(meta, "tokens") and meta.tokens:
+                    tokens = meta.tokens
+                    input_tokens = getattr(tokens, "input_tokens", 0)
+                    output_tokens = getattr(tokens, "output_tokens", 0)
+
+                    if input_tokens or output_tokens:
+                        return {
+                            "prompt_tokens": int(input_tokens) if input_tokens else 0,
+                            "completion_tokens": int(output_tokens) if output_tokens else 0,
+                            "total_tokens": int(input_tokens or 0) + int(output_tokens or 0),
+                        }
+                # Fallback to billed_units
+                elif hasattr(meta, "billed_units") and meta.billed_units:
+                    billed = meta.billed_units
+                    input_tokens = getattr(billed, "input_tokens", 0)
+                    output_tokens = getattr(billed, "output_tokens", 0)
+
+                    if input_tokens or output_tokens:
+                        return {
+                            "prompt_tokens": int(input_tokens) if input_tokens else 0,
+                            "completion_tokens": int(output_tokens) if output_tokens else 0,
+                            "total_tokens": int(input_tokens or 0) + int(output_tokens or 0),
+                        }
+
+            return None
+        except Exception as e:
+            logger.debug("Failed to extract usage from Cohere response: %s", e)
+            return None

genai_otel/instrumentors/huggingface_instrumentor.py

@@ -1,11 +1,14 @@
-"""OpenTelemetry instrumentor for HuggingFace Transformers library.
+"""OpenTelemetry instrumentor for HuggingFace Transformers and Inference API.
 
-This instrumentor automatically traces calls made through HuggingFace pipelines,
-capturing relevant attributes such as the model name and task type.
+This instrumentor automatically traces:
+1. HuggingFace Transformers pipelines (local model execution)
+2. HuggingFace Inference API calls via InferenceClient (used by smolagents)
+
+Note: Transformers runs models locally (no API costs), but InferenceClient makes
+API calls to HuggingFace endpoints which may have costs based on usage.
 """
 
 import logging
-import types
 from typing import Dict, Optional
 
 from ..config import OTelConfig
@@ -15,16 +18,22 @@ logger = logging.getLogger(__name__)
 
 
 class HuggingFaceInstrumentor(BaseInstrumentor):
-    """Instrumentor for HuggingFace Transformers"""
+    """Instrumentor for HuggingFace Transformers and Inference API.
+
+    Instruments both:
+    - transformers.pipeline (local execution, no API 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._check_availability()
 
     def _check_availability(self):
-        """Check if Transformers library is available."""
+        """Check if Transformers and InferenceClient libraries are available."""
         try:
             import transformers
 
@@ -34,12 +43,47 @@ class HuggingFaceInstrumentor(BaseInstrumentor):
             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 and InferenceClient."""
         self.config = config
 
-        if not self._transformers_available:
-            return
-
+        instrumented_count = 0
+
+        # Instrument transformers.pipeline if available
+        if self._transformers_available:
+            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 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
 
@@ -68,6 +112,7 @@ class HuggingFaceInstrumentor(BaseInstrumentor):
 
                             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:
@@ -88,10 +133,90 @@ class HuggingFaceInstrumentor(BaseInstrumentor):
                 return WrappedPipeline(pipe)
 
             transformers_module.pipeline = wrapped_pipeline
-            logger.info("HuggingFace instrumentation enabled")
-
-        except ImportError:
-            pass
+            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 _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/mistralai_instrumentor.py

@@ -10,6 +10,7 @@ Supports Mistral SDK v1.0+ with the new API structure:
 """
 
 import logging
+import time
 from typing import Any, Dict, Optional
 
 from ..config import OTelConfig
@@ -27,50 +28,261 @@ class MistralAIInstrumentor(BaseInstrumentor):
             import wrapt
             from mistralai import Mistral
 
-            # Wrap the Mistral client __init__ to instrument each instance
-            original_init = Mistral.__init__
+            # Get access to the chat and embeddings modules
+            # In Mistral SDK v1.0+, structure is:
+            # - Mistral client has .chat and .embeddings properties
+            # - These are bound methods that call internal APIs
 
-            def wrapped_init(wrapped, instance, args, kwargs):
-                result = wrapped(*args, **kwargs)
-                self._instrument_client(instance)
-                return result
-
-            Mistral.__init__ = wrapt.FunctionWrapper(original_init, wrapped_init)
-            logger.info("MistralAI instrumentation enabled (v1.0+ SDK)")
+            # Store original methods at module level before any instances are created
+            if not hasattr(Mistral, '_genai_otel_instrumented'):
+                self._wrap_mistral_methods(Mistral, wrapt)
+                Mistral._genai_otel_instrumented = True
+                logger.info("MistralAI instrumentation enabled (v1.0+ SDK)")
 
         except ImportError:
             logger.warning("mistralai package not available, skipping instrumentation")
         except Exception as e:
             logger.error(f"Failed to instrument mistralai: {e}", exc_info=True)
+            if config.fail_on_error:
+                raise
+
+    def _wrap_mistral_methods(self, Mistral, wrapt):
+        """Wrap Mistral client methods at the class level."""
+        # Import the internal classes that handle chat and embeddings
+        try:
+            from mistralai.chat import Chat
+            from mistralai.embeddings import Embeddings
+
+            # Wrap Chat.complete method
+            if hasattr(Chat, 'complete'):
+                wrapt.wrap_function_wrapper(
+                    'mistralai.chat',
+                    'Chat.complete',
+                    self._wrap_chat_complete
+                )
+                logger.debug("Wrapped Mistral Chat.complete")
+
+            # Wrap Chat.stream method
+            if hasattr(Chat, 'stream'):
+                wrapt.wrap_function_wrapper(
+                    'mistralai.chat',
+                    'Chat.stream',
+                    self._wrap_chat_stream
+                )
+                logger.debug("Wrapped Mistral Chat.stream")
+
+            # Wrap Embeddings.create method
+            if hasattr(Embeddings, 'create'):
+                wrapt.wrap_function_wrapper(
+                    'mistralai.embeddings',
+                    'Embeddings.create',
+                    self._wrap_embeddings_create
+                )
+                logger.debug("Wrapped Mistral Embeddings.create")
+
+        except (ImportError, AttributeError) as e:
+            logger.warning(f"Could not access Mistral internal classes: {e}")
+
+    def _wrap_chat_complete(self, wrapped, instance, args, kwargs):
+        """Wrapper for chat.complete() method."""
+        model = kwargs.get("model", "mistral-small-latest")
+        span_name = f"mistralai.chat.complete {model}"
+
+        with self.tracer.start_span(span_name) as span:
+            # Set attributes
+            attributes = self._extract_chat_attributes(instance, args, kwargs)
+            for key, value in attributes.items():
+                span.set_attribute(key, value)
+
+            # Record request metric
+            if self.request_counter:
+                self.request_counter.add(1, {"model": model, "provider": "mistralai"})
+
+            # Execute the call
+            start_time = time.time()
+            try:
+                response = wrapped(*args, **kwargs)
+
+                # Record metrics from response
+                self._record_result_metrics(span, response, start_time, kwargs)
+
+                return response
+
+            except Exception as e:
+                if self.error_counter:
+                    self.error_counter.add(
+                        1, {"operation": span_name, "error.type": type(e).__name__}
+                    )
+                span.record_exception(e)
+                raise
+
+    def _wrap_chat_stream(self, wrapped, instance, args, kwargs):
+        """Wrapper for chat.stream() method - handles streaming responses."""
+        model = kwargs.get("model", "mistral-small-latest")
+        span_name = f"mistralai.chat.stream {model}"
+
+        # Start the span
+        span = self.tracer.start_span(span_name)
+
+        # Set attributes
+        attributes = self._extract_chat_attributes(instance, args, kwargs)
+        for key, value in attributes.items():
+            span.set_attribute(key, value)
+
+        # Record request metric
+        if self.request_counter:
+            self.request_counter.add(1, {"model": model, "provider": "mistralai"})
+
+        start_time = time.time()
+
+        # Execute and get the stream
+        try:
+            stream = wrapped(*args, **kwargs)
+
+            # Wrap the stream with our tracking wrapper
+            return self._StreamWrapper(
+                stream, span, self, model, start_time, span_name
+            )
+
+        except Exception as e:
+            if self.error_counter:
+                self.error_counter.add(
+                    1, {"operation": span_name, "error.type": type(e).__name__}
+                )
+            span.record_exception(e)
+            span.end()
+            raise
+
+    def _wrap_embeddings_create(self, wrapped, instance, args, kwargs):
+        """Wrapper for embeddings.create() method."""
+        model = kwargs.get("model", "mistral-embed")
+        span_name = f"mistralai.embeddings.create {model}"
+
+        with self.tracer.start_span(span_name) as span:
+            # Set attributes
+            attributes = self._extract_embeddings_attributes(instance, args, kwargs)
+            for key, value in attributes.items():
+                span.set_attribute(key, value)
+
+            # Record request metric
+            if self.request_counter:
+                self.request_counter.add(1, {"model": model, "provider": "mistralai"})
+
+            # Execute the call
+            start_time = time.time()
+            try:
+                response = wrapped(*args, **kwargs)
+
+                # Record metrics from response
+                self._record_result_metrics(span, response, start_time, kwargs)
+
+                return response
+
+            except Exception as e:
+                if self.error_counter:
+                    self.error_counter.add(
+                        1, {"operation": span_name, "error.type": type(e).__name__}
+                    )
+                span.record_exception(e)
+                raise
+
+    class _StreamWrapper:
+        """Wrapper for streaming responses that collects metrics."""
+
+        def __init__(self, stream, span, instrumentor, model, start_time, span_name):
+            self._stream = stream
+            self._span = span
+            self._instrumentor = instrumentor
+            self._model = model
+            self._start_time = start_time
+            self._span_name = span_name
+            self._usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
+            self._response_text = ""
+            self._first_chunk = True
+            self._ttft = None
+
+        def __iter__(self):
+            return self
+
+        def __next__(self):
+            try:
+                chunk = next(self._stream)
+
+                # Record time to first token
+                if self._first_chunk:
+                    self._ttft = time.time() - self._start_time
+                    self._first_chunk = False
+
+                # Process chunk to extract usage and content
+                self._process_chunk(chunk)
+
+                return chunk
+
+            except StopIteration:
+                # Stream completed - record final metrics
+                try:
+                    # Set TTFT if we got any chunks
+                    if self._ttft is not None:
+                        self._span.set_attribute("gen_ai.server.ttft", self._ttft)
+
+                    # Record usage metrics if available
+                    if self._usage["total_tokens"] > 0:
+                        # Create a mock response object with usage for _record_result_metrics
+                        class MockUsage:
+                            def __init__(self, usage_dict):
+                                self.prompt_tokens = usage_dict["prompt_tokens"]
+                                self.completion_tokens = usage_dict["completion_tokens"]
+                                self.total_tokens = usage_dict["total_tokens"]
+
+                        class MockResponse:
+                            def __init__(self, usage_dict):
+                                self.usage = MockUsage(usage_dict)
+
+                        mock_response = MockResponse(self._usage)
+                        self._instrumentor._record_result_metrics(
+                            self._span,
+                            mock_response,
+                            self._start_time,
+                            {"model": self._model}
+                        )
+
+                finally:
+                    self._span.end()
+
+                raise
+
+        def _process_chunk(self, chunk):
+            """Process a streaming chunk to extract usage."""
+            try:
+                # Mistral streaming chunks have: data.choices[0].delta.content
+                if hasattr(chunk, 'data'):
+                    data = chunk.data
+                    if hasattr(data, 'choices') and len(data.choices) > 0:
+                        delta = data.choices[0].delta
+                        if hasattr(delta, 'content') and delta.content:
+                            self._response_text += delta.content
+
+                    # Extract usage if available on final chunk
+                    if hasattr(data, 'usage') and data.usage:
+                        usage = data.usage
+                        if hasattr(usage, 'prompt_tokens'):
+                            self._usage["prompt_tokens"] = usage.prompt_tokens
+                        if hasattr(usage, 'completion_tokens'):
+                            self._usage["completion_tokens"] = usage.completion_tokens
+                        if hasattr(usage, 'total_tokens'):
+                            self._usage["total_tokens"] = usage.total_tokens
+
+            except Exception as e:
+                logger.debug(f"Error processing Mistral stream chunk: {e}")
+
+        def __enter__(self):
+            return self
 
-    def _instrument_client(self, client):
-        """Instrument Mistral client instance methods."""
-        # Instrument chat.complete()
-        if hasattr(client, "chat") and hasattr(client.chat, "complete"):
-            original_complete = client.chat.complete
-            instrumented_complete = self.create_span_wrapper(
-                span_name="mistralai.chat.complete",
-                extract_attributes=self._extract_chat_attributes,
-            )(original_complete)
-            client.chat.complete = instrumented_complete
-
-        # Instrument chat.stream()
-        if hasattr(client, "chat") and hasattr(client.chat, "stream"):
-            original_stream = client.chat.stream
-            instrumented_stream = self.create_span_wrapper(
-                span_name="mistralai.chat.stream",
-                extract_attributes=self._extract_chat_attributes,
-            )(original_stream)
-            client.chat.stream = instrumented_stream
-
-        # Instrument embeddings.create()
-        if hasattr(client, "embeddings") and hasattr(client.embeddings, "create"):
-            original_embeddings = client.embeddings.create
-            instrumented_embeddings = self.create_span_wrapper(
-                span_name="mistralai.embeddings.create",
-                extract_attributes=self._extract_embeddings_attributes,
-            )(original_embeddings)
-            client.embeddings.create = instrumented_embeddings
+        def __exit__(self, exc_type, exc_val, exc_tb):
+            if exc_type is not None:
+                self._span.record_exception(exc_val)
+            self._span.end()
+            return False
 
     def _extract_chat_attributes(self, instance: Any, args: Any, kwargs: Any) -> Dict[str, Any]:
         """Extract attributes from chat.complete() or chat.stream() call."""