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

genai_otel/instrumentors/cohere_instrumentor.py

@@ -1,76 +1,140 @@
-"""OpenTelemetry instrumentor for the Cohere SDK.
-
-This instrumentor automatically traces calls to Cohere models, capturing
-relevant attributes such as the model name.
-"""
-
-import logging
-from typing import Dict, Optional
-
-from ..config import OTelConfig
-from .base import BaseInstrumentor
-
-logger = logging.getLogger(__name__)
-
-
-class CohereInstrumentor(BaseInstrumentor):
-    """Instrumentor for Cohere"""
-
-    def __init__(self):
-        """Initialize the instrumentor."""
-        super().__init__()
-        self._cohere_available = False
-        self._check_availability()
-
-    def _check_availability(self):
-        """Check if cohere library is available."""
-        try:
-            import cohere
-
-            self._cohere_available = True
-            logger.debug("cohere library detected and available for instrumentation")
-        except ImportError:
-            logger.debug("cohere library not installed, instrumentation will be skipped")
-            self._cohere_available = False
-
-    def instrument(self, config: OTelConfig):
-        """Instrument  cohere available if available."""
-        if not self._cohere_available:
-            logger.debug("Skipping instrumentation - library not available")
-            return
-
-        self.config = config
-        try:
-            import cohere
-
-            original_init = cohere.Client.__init__
-
-            def wrapped_init(instance, *args, **kwargs):
-                original_init(instance, *args, **kwargs)
-                self._instrument_client(instance)
-
-            cohere.Client.__init__ = wrapped_init
-
-        except ImportError:
-            pass
-
-    def _instrument_client(self, client):
-        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")
-
-                span.set_attribute("gen_ai.system", "cohere")
-                span.set_attribute("gen_ai.request.model", model)
-
-                if self.request_counter:
-                    self.request_counter.add(1, {"model": model, "provider": "cohere"})
-
-                result = original_generate(*args, **kwargs)
-                return result
-
-        client.generate = wrapped_generate
-
-    def _extract_usage(self, result) -> Optional[Dict[str, int]]:
-        return None
+"""OpenTelemetry instrumentor for the Cohere SDK.
+
+This instrumentor automatically traces calls to Cohere models, capturing
+relevant attributes such as the model name and token usage.
+"""
+
+import logging
+from typing import Any, Dict, Optional
+
+from ..config import OTelConfig
+from .base import BaseInstrumentor
+
+logger = logging.getLogger(__name__)
+
+
+class CohereInstrumentor(BaseInstrumentor):
+    """Instrumentor for Cohere"""
+
+    def __init__(self):
+        """Initialize the instrumentor."""
+        super().__init__()
+        self._cohere_available = False
+        self._check_availability()
+
+    def _check_availability(self):
+        """Check if cohere library is available."""
+        try:
+            import cohere
+
+            self._cohere_available = True
+            logger.debug("cohere library detected and available for instrumentation")
+        except ImportError:
+            logger.debug("cohere library not installed, instrumentation will be skipped")
+            self._cohere_available = False
+
+    def instrument(self, config: OTelConfig):
+        """Instrument cohere if available."""
+        if not self._cohere_available:
+            logger.debug("Skipping instrumentation - library not available")
+            return
+
+        self.config = config
+        try:
+            import cohere
+
+            original_init = cohere.Client.__init__
+
+            def wrapped_init(instance, *args, **kwargs):
+                original_init(instance, *args, **kwargs)
+                self._instrument_client(instance)
+
+            cohere.Client.__init__ = wrapped_init
+            self._instrumented = True
+            logger.info("Cohere instrumentation enabled")
+
+        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
+
+        # Wrap using create_span_wrapper
+        wrapped_generate = self.create_span_wrapper(
+            span_name="cohere.generate",
+            extract_attributes=self._extract_generate_attributes,
+        )(original_generate)
+
+        client.generate = wrapped_generate
+
+    def _extract_generate_attributes(self, instance: Any, args: Any, kwargs: Any) -> Dict[str, Any]:
+        """Extract attributes from Cohere generate call.
+
+        Args:
+            instance: The client instance.
+            args: Positional arguments.
+            kwargs: Keyword arguments.
+
+        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]]:
+        """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,51 @@ 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 +116,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 +137,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/langchain_instrumentor.py

@@ -1,75 +1,75 @@
-"""OpenTelemetry instrumentor for the LangChain framework.
-
-This instrumentor automatically traces various components within LangChain,
-including chains and agents, capturing relevant attributes for observability.
-"""
-
-import logging
-from typing import Dict, Optional
-
-from ..config import OTelConfig
-from .base import BaseInstrumentor
-
-logger = logging.getLogger(__name__)
-
-
-class LangChainInstrumentor(BaseInstrumentor):
-    """Instrumentor for LangChain"""
-
-    def __init__(self):
-        """Initialize the instrumentor."""
-        super().__init__()
-        self._langchain_available = False
-        self._check_availability()
-
-    def _check_availability(self):
-        """Check if langchain library is available."""
-        try:
-            import langchain
-
-            self._langchain_available = True
-            logger.debug("langchain library detected and available for instrumentation")
-        except ImportError:
-            logger.debug("langchain library not installed, instrumentation will be skipped")
-            self._langchain_available = False
-
-    def instrument(self, config: OTelConfig):
-        """Instrument  langchain available if available."""
-        if not self._langchain_available:
-            logger.debug("Skipping instrumentation - library not available")
-            return
-
-        self.config = config
-        try:
-            from langchain.agents.agent import AgentExecutor
-            from langchain.chains.base import Chain
-
-            # Instrument Chains
-            original_call = Chain.__call__
-
-            def wrapped_call(instance, *args, **kwargs):
-                chain_type = instance.__class__.__name__
-                with self.tracer.start_as_current_span(f"langchain.chain.{chain_type}") as span:
-                    span.set_attribute("langchain.chain.type", chain_type)
-                    result = original_call(instance, *args, **kwargs)
-                    return result
-
-            Chain.__call__ = wrapped_call
-
-            # Instrument Agents
-            original_agent_call = AgentExecutor.__call__
-
-            def wrapped_agent_call(instance, *args, **kwargs):
-                with self.tracer.start_as_current_span("langchain.agent.execute") as span:
-                    agent_name = getattr(instance, "agent", {}).get("name", "unknown")
-                    span.set_attribute("langchain.agent.name", agent_name)
-                    result = original_agent_call(instance, *args, **kwargs)
-                    return result
-
-            AgentExecutor.__call__ = wrapped_agent_call
-
-        except ImportError:
-            pass
-
-    def _extract_usage(self, result) -> Optional[Dict[str, int]]:
-        return None
+"""OpenTelemetry instrumentor for the LangChain framework.
+
+This instrumentor automatically traces various components within LangChain,
+including chains and agents, capturing relevant attributes for observability.
+"""
+
+import logging
+from typing import Dict, Optional
+
+from ..config import OTelConfig
+from .base import BaseInstrumentor
+
+logger = logging.getLogger(__name__)
+
+
+class LangChainInstrumentor(BaseInstrumentor):
+    """Instrumentor for LangChain"""
+
+    def __init__(self):
+        """Initialize the instrumentor."""
+        super().__init__()
+        self._langchain_available = False
+        self._check_availability()
+
+    def _check_availability(self):
+        """Check if langchain library is available."""
+        try:
+            import langchain
+
+            self._langchain_available = True
+            logger.debug("langchain library detected and available for instrumentation")
+        except ImportError:
+            logger.debug("langchain library not installed, instrumentation will be skipped")
+            self._langchain_available = False
+
+    def instrument(self, config: OTelConfig):
+        """Instrument  langchain available if available."""
+        if not self._langchain_available:
+            logger.debug("Skipping instrumentation - library not available")
+            return
+
+        self.config = config
+        try:
+            from langchain.agents.agent import AgentExecutor
+            from langchain.chains.base import Chain
+
+            # Instrument Chains
+            original_call = Chain.__call__
+
+            def wrapped_call(instance, *args, **kwargs):
+                chain_type = instance.__class__.__name__
+                with self.tracer.start_as_current_span(f"langchain.chain.{chain_type}") as span:
+                    span.set_attribute("langchain.chain.type", chain_type)
+                    result = original_call(instance, *args, **kwargs)
+                    return result
+
+            Chain.__call__ = wrapped_call
+
+            # Instrument Agents
+            original_agent_call = AgentExecutor.__call__
+
+            def wrapped_agent_call(instance, *args, **kwargs):
+                with self.tracer.start_as_current_span("langchain.agent.execute") as span:
+                    agent_name = getattr(instance, "agent", {}).get("name", "unknown")
+                    span.set_attribute("langchain.agent.name", agent_name)
+                    result = original_agent_call(instance, *args, **kwargs)
+                    return result
+
+            AgentExecutor.__call__ = wrapped_agent_call
+
+        except ImportError:
+            pass
+
+    def _extract_usage(self, result) -> Optional[Dict[str, int]]:
+        return None