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

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,245 @@ class MistralAIInstrumentor(BaseInstrumentor):
             import wrapt
             from mistralai import Mistral
 
-            # Wrap the Mistral client __init__ to instrument each instance
-            original_init = Mistral.__init__
-
-            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)")
+            # 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
+            # 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."""

genai_otel/instrumentors/ollama_instrumentor.py

@@ -2,11 +2,11 @@
 
 This instrumentor automatically traces calls to Ollama models for both
 generation and chat functionalities, capturing relevant attributes such as
-the model name.
+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
@@ -22,8 +22,8 @@ class OllamaInstrumentor(BaseInstrumentor):
         super().__init__()
         self._ollama_available = False
         self._ollama_module = None
-        self._original_generate = None  # Add this
-        self._original_chat = None  # Add this
+        self._original_generate = None
+        self._original_chat = None
         self._check_availability()
 
     def _check_availability(self):
@@ -46,38 +46,107 @@ class OllamaInstrumentor(BaseInstrumentor):
         if not self._ollama_available or self._ollama_module is None:
             return
 
-        # Store original methods
-        self._original_generate = self._ollama_module.generate
-        self._original_chat = self._ollama_module.chat
-
-        def wrapped_generate(*args, **kwargs):
-            with self.tracer.start_as_current_span("ollama.generate") as span:
-                model = kwargs.get("model", "unknown")
-
-                span.set_attribute("gen_ai.system", "ollama")
-                span.set_attribute("gen_ai.request.model", model)
-
-                if self.request_counter:
-                    self.request_counter.add(1, {"model": model, "provider": "ollama"})
-
-                result = self._original_generate(*args, **kwargs)
-                return result
-
-        def wrapped_chat(*args, **kwargs):
-            with self.tracer.start_as_current_span("ollama.chat") as span:
-                model = kwargs.get("model", "unknown")
-
-                span.set_attribute("gen_ai.system", "ollama")
-                span.set_attribute("gen_ai.request.model", model)
+        try:
+            # Store original methods and wrap them
+            self._original_generate = self._ollama_module.generate
+            self._original_chat = self._ollama_module.chat
+
+            # Wrap generate method
+            wrapped_generate = self.create_span_wrapper(
+                span_name="ollama.generate",
+                extract_attributes=self._extract_generate_attributes,
+            )(self._original_generate)
+            self._ollama_module.generate = wrapped_generate
+
+            # Wrap chat method
+            wrapped_chat = self.create_span_wrapper(
+                span_name="ollama.chat",
+                extract_attributes=self._extract_chat_attributes,
+            )(self._original_chat)
+            self._ollama_module.chat = wrapped_chat
+
+            self._instrumented = True
+            logger.info("Ollama instrumentation enabled")
+
+        except Exception as e:
+            logger.error("Failed to instrument Ollama: %s", e, exc_info=True)
+            if config.fail_on_error:
+                raise
+
+    def _extract_generate_attributes(self, instance: Any, args: Any, kwargs: Any) -> Dict[str, Any]:
+        """Extract attributes from Ollama generate call.
+
+        Args:
+            instance: The client instance (None for module-level functions).
+            args: Positional arguments.
+            kwargs: Keyword arguments.
+
+        Returns:
+            Dict[str, Any]: Dictionary of attributes to set on the span.
+        """
+        attrs = {}
+        model = kwargs.get("model", "unknown")
+
+        attrs["gen_ai.system"] = "ollama"
+        attrs["gen_ai.request.model"] = model
+        attrs["gen_ai.operation.name"] = "generate"
+
+        return attrs
+
+    def _extract_chat_attributes(self, instance: Any, args: Any, kwargs: Any) -> Dict[str, Any]:
+        """Extract attributes from Ollama chat call.
+
+        Args:
+            instance: The client instance (None for module-level functions).
+            args: Positional arguments.
+            kwargs: Keyword arguments.
+
+        Returns:
+            Dict[str, Any]: Dictionary of attributes to set on the span.
+        """
+        attrs = {}
+        model = kwargs.get("model", "unknown")
+        messages = kwargs.get("messages", [])
+
+        attrs["gen_ai.system"] = "ollama"
+        attrs["gen_ai.request.model"] = model
+        attrs["gen_ai.operation.name"] = "chat"
+        attrs["gen_ai.request.message_count"] = len(messages)
+
+        return attrs
 
-                if self.request_counter:
-                    self.request_counter.add(1, {"model": model, "provider": "ollama"})
+    def _extract_usage(self, result) -> Optional[Dict[str, int]]:
+        """Extract token usage from Ollama response.
 
-                result = self._original_chat(*args, **kwargs)
-                return result
+        Ollama responses include:
+        - prompt_eval_count: Input tokens
+        - eval_count: Output tokens
 
-        self._ollama_module.generate = wrapped_generate
-        self._ollama_module.chat = wrapped_chat
+        Args:
+            result: The API response object or dictionary.
 
-    def _extract_usage(self, result) -> Optional[Dict[str, int]]:
-        return None
+        Returns:
+            Optional[Dict[str, int]]: Dictionary with token counts or None.
+        """
+        try:
+            # Handle both dict and object responses
+            if isinstance(result, dict):
+                prompt_tokens = result.get("prompt_eval_count", 0)
+                completion_tokens = result.get("eval_count", 0)
+            elif hasattr(result, "prompt_eval_count") and hasattr(result, "eval_count"):
+                prompt_tokens = getattr(result, "prompt_eval_count", 0)
+                completion_tokens = getattr(result, "eval_count", 0)
+            else:
+                return None
+
+            if prompt_tokens == 0 and completion_tokens == 0:
+                return None
+
+            return {
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "total_tokens": prompt_tokens + completion_tokens,
+            }
+        except Exception as e:
+            logger.debug("Failed to extract usage from Ollama response: %s", e)
+            return None

genai_otel/instrumentors/replicate_instrumentor.py

@@ -2,41 +2,86 @@
 
 This instrumentor automatically traces calls to Replicate models, capturing
 relevant attributes such as the model name.
+
+Note: Replicate uses hardware-based pricing (per second of GPU/CPU time),
+not token-based pricing. Cost tracking is not applicable as the pricing model
+is fundamentally different from token-based LLM APIs.
 """
 
-from typing import Dict, Optional
+import logging
+from typing import Any, Dict, Optional
 
 from ..config import OTelConfig
 from .base import BaseInstrumentor
 
+logger = logging.getLogger(__name__)
+
 
 class ReplicateInstrumentor(BaseInstrumentor):
-    """Instrumentor for Replicate"""
+    """Instrumentor for Replicate.
+
+    Note: Replicate uses hardware-based pricing ($/second), not token-based.
+    Cost tracking returns None as pricing is based on execution time and hardware type.
+    """
 
     def instrument(self, config: OTelConfig):
+        """Instrument Replicate SDK if available."""
         self.config = config
         try:
             import replicate
 
             original_run = replicate.run
 
-            def wrapped_run(*args, **kwargs):
-                with self.tracer.start_as_current_span("replicate.run") as span:
-                    model = args[0] if args else "unknown"
+            # Wrap using create_span_wrapper
+            wrapped_run = self.create_span_wrapper(
+                span_name="replicate.run",
+                extract_attributes=self._extract_run_attributes,
+            )(original_run)
 
-                    span.set_attribute("gen_ai.system", "replicate")
-                    span.set_attribute("gen_ai.request.model", model)
+            replicate.run = wrapped_run
+            self._instrumented = True
+            logger.info("Replicate instrumentation enabled")
 
-                    if self.request_counter:
-                        self.request_counter.add(1, {"model": model, "provider": "replicate"})
+        except ImportError:
+            logger.debug("Replicate library not installed, instrumentation will be skipped")
+        except Exception as e:
+            logger.error("Failed to instrument Replicate: %s", e, exc_info=True)
+            if config.fail_on_error:
+                raise
 
-                    result = original_run(*args, **kwargs)
-                    return result
+    def _extract_run_attributes(self, instance: Any, args: Any, kwargs: Any) -> Dict[str, Any]:
+        """Extract attributes from Replicate run call.
 
-            replicate.run = wrapped_run
+        Args:
+            instance: The instance (None for module-level functions).
+            args: Positional arguments (first arg is typically the model).
+            kwargs: Keyword arguments.
 
-        except ImportError:
-            pass
+        Returns:
+            Dict[str, Any]: Dictionary of attributes to set on the span.
+        """
+        attrs = {}
+        model = args[0] if args else kwargs.get("model", "unknown")
+
+        attrs["gen_ai.system"] = "replicate"
+        attrs["gen_ai.request.model"] = model
+        attrs["gen_ai.operation.name"] = "run"
+
+        return attrs
 
     def _extract_usage(self, result) -> Optional[Dict[str, int]]:
+        """Extract token usage from Replicate response.
+
+        Note: Replicate uses hardware-based pricing ($/second of GPU/CPU time),
+        not token-based pricing. Returns None as the pricing model is incompatible
+        with token-based cost calculation.
+
+        Args:
+            result: The API response.
+
+        Returns:
+            None: Replicate uses hardware-based pricing, not token-based.
+        """
+        # Replicate uses hardware-based pricing ($/second), not tokens
+        # Cannot track costs with token-based calculator
         return None