genai-otel-instrument 0.1.1.dev0__py3-none-any.whl

genai_otel/instrumentors/groq_instrumentor.py

@@ -0,0 +1,106 @@
+"""OpenTelemetry instrumentor for the Groq SDK.
+
+This instrumentor automatically traces chat completion calls to Groq models,
+capturing relevant attributes such as the model name and token usage.
+"""
+
+import logging
+from typing import Dict, Optional
+
+from ..config import OTelConfig
+from .base import BaseInstrumentor
+
+logger = logging.getLogger(__name__)
+
+
+class GroqInstrumentor(BaseInstrumentor):
+    """Instrumentor for Groq"""
+
+    def __init__(self):
+        """Initialize the instrumentor."""
+        super().__init__()
+        self._groq_available = False
+        self._check_availability()
+
+    def _check_availability(self):
+        """Check if Groq library is available."""
+        try:
+            import groq
+
+            self._groq_available = True
+            logger.debug("Groq library detected and available for instrumentation")
+        except ImportError:
+            logger.debug("Groq library not installed, instrumentation will be skipped")
+            self._groq_available = False
+
+    def instrument(self, config: OTelConfig):
+        """Instrument Groq SDK if available.
+
+        Args:
+            config (OTelConfig): The OpenTelemetry configuration object.
+        """
+        if not self._groq_available:
+            logger.debug("Skipping Groq instrumentation - library not available")
+            return
+
+        self.config = config
+
+        try:
+            import groq
+
+            original_init = groq.Groq.__init__
+
+            def wrapped_init(instance, *args, **kwargs):
+                original_init(instance, *args, **kwargs)
+                self._instrument_client(instance)
+                return instance
+
+            groq.Groq.__init__ = wrapped_init
+            self._instrumented = True
+            logger.info("Groq instrumentation enabled")
+
+        except Exception as e:
+            logger.error("Failed to instrument Groq: %s", e, exc_info=True)
+            if config.fail_on_error:
+                raise
+
+    def _instrument_client(self, client):
+        """Instrument Groq client methods.
+
+        Args:
+            client: The Groq client instance to instrument.
+        """
+        original_create = client.chat.completions.create
+
+        def wrapped_create(*args, **kwargs):
+            with self.tracer.start_as_current_span("groq.chat.completions") as span:
+                model = kwargs.get("model", "unknown")
+
+                span.set_attribute("gen_ai.system", "groq")
+                span.set_attribute("gen_ai.request.model", model)
+
+                if self.request_counter:
+                    self.request_counter.add(1, {"model": model, "provider": "groq"})
+
+                result = original_create(*args, **kwargs)
+                self._record_result_metrics(span, result, 0)
+                return result
+
+        client.chat.completions.create = wrapped_create
+
+    def _extract_usage(self, result) -> Optional[Dict[str, int]]:
+        """Extract token usage from Groq response.
+
+        Args:
+            result: The API response object.
+
+        Returns:
+            Optional[Dict[str, int]]: Dictionary with token counts or None.
+        """
+        if hasattr(result, "usage"):
+            return {
+                "prompt_tokens": result.usage.prompt_tokens,
+                "completion_tokens": result.usage.completion_tokens,
+                "total_tokens": result.usage.total_tokens,
+            }
+        return None

genai_otel/instrumentors/huggingface_instrumentor.py

@@ -0,0 +1,97 @@
+"""OpenTelemetry instrumentor for HuggingFace Transformers library.
+
+This instrumentor automatically traces calls made through HuggingFace pipelines,
+capturing relevant attributes such as the model name and task type.
+"""
+
+import logging
+import types
+from typing import Dict, Optional
+
+from ..config import OTelConfig
+from .base import BaseInstrumentor
+
+logger = logging.getLogger(__name__)
+
+
+class HuggingFaceInstrumentor(BaseInstrumentor):
+    """Instrumentor for HuggingFace Transformers"""
+
+    def __init__(self):
+        """Initialize the instrumentor."""
+        super().__init__()
+        self._transformers_available = False
+        self._check_availability()
+
+    def _check_availability(self):
+        """Check if Transformers library is 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
+
+    def instrument(self, config: OTelConfig):
+        self.config = config
+
+        if not self._transformers_available:
+            return
+
+        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("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.info("HuggingFace instrumentation enabled")
+
+        except ImportError:
+            pass
+
+    def _extract_usage(self, result) -> Optional[Dict[str, int]]:
+        return None

genai_otel/instrumentors/langchain_instrumentor.py

@@ -0,0 +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

genai_otel/instrumentors/llamaindex_instrumentor.py

@@ -0,0 +1,36 @@
+"""OpenTelemetry instrumentor for the LlamaIndex framework.
+
+This instrumentor automatically traces query engine operations within LlamaIndex,
+capturing relevant attributes such as the query text.
+"""
+
+from typing import Dict, Optional
+
+from ..config import OTelConfig
+from .base import BaseInstrumentor
+
+
+class LlamaIndexInstrumentor(BaseInstrumentor):
+    """Instrumentor for LlamaIndex"""
+
+    def instrument(self, config: OTelConfig):
+        self.config = config
+        try:
+            from llama_index.core.query_engine import BaseQueryEngine
+
+            original_query = BaseQueryEngine.query
+
+            def wrapped_query(instance, *args, **kwargs):
+                with self.tracer.start_as_current_span("llamaindex.query_engine") as span:
+                    query_text = args[0] if args else kwargs.get("query_str", "")
+                    span.set_attribute("llamaindex.query", str(query_text)[:200])
+                    result = original_query(instance, *args, **kwargs)
+                    return result
+
+            BaseQueryEngine.query = wrapped_query
+
+        except ImportError:
+            pass
+
+    def _extract_usage(self, result) -> Optional[Dict[str, int]]:
+        return None

genai_otel/instrumentors/mistralai_instrumentor.py

@@ -0,0 +1,119 @@
+"""OpenTelemetry instrumentor for the Mistral AI SDK (v1.0+).
+
+This instrumentor automatically traces chat calls to Mistral AI models,
+capturing relevant attributes such as the model name and token usage.
+
+Supports Mistral SDK v1.0+ with the new API structure:
+- Mistral.chat.complete()
+- Mistral.chat.stream()
+- Mistral.embeddings.create()
+"""
+
+import logging
+from typing import Any, Dict, Optional
+
+from ..config import OTelConfig
+from .base import BaseInstrumentor
+
+logger = logging.getLogger(__name__)
+
+
+class MistralAIInstrumentor(BaseInstrumentor):
+    """Instrumentor for Mistral AI SDK v1.0+"""
+
+    def instrument(self, config: OTelConfig):
+        self.config = config
+        try:
+            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)")
+
+        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)
+
+    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 _extract_chat_attributes(self, instance: Any, args: Any, kwargs: Any) -> Dict[str, Any]:
+        """Extract attributes from chat.complete() or chat.stream() call."""
+        model = kwargs.get("model", "unknown")
+        attributes = {
+            "gen_ai.system": "mistralai",
+            "gen_ai.request.model": model,
+            "gen_ai.request.type": "chat",
+        }
+
+        # Add optional parameters
+        if "temperature" in kwargs and kwargs["temperature"] is not None:
+            attributes["gen_ai.request.temperature"] = kwargs["temperature"]
+        if "top_p" in kwargs and kwargs["top_p"] is not None:
+            attributes["gen_ai.request.top_p"] = kwargs["top_p"]
+        if "max_tokens" in kwargs and kwargs["max_tokens"] is not None:
+            attributes["gen_ai.request.max_tokens"] = kwargs["max_tokens"]
+
+        return attributes
+
+    def _extract_embeddings_attributes(
+        self, instance: Any, args: Any, kwargs: Any
+    ) -> Dict[str, Any]:
+        """Extract attributes from embeddings.create() call."""
+        model = kwargs.get("model", "mistral-embed")
+        attributes = {
+            "gen_ai.system": "mistralai",
+            "gen_ai.request.model": model,
+            "gen_ai.request.type": "embedding",
+        }
+        return attributes
+
+    def _extract_usage(self, result) -> Optional[Dict[str, int]]:
+        """Extract usage information from Mistral AI response"""
+        try:
+            if hasattr(result, "usage"):
+                usage = result.usage
+                return {
+                    "prompt_tokens": getattr(usage, "prompt_tokens", 0),
+                    "completion_tokens": getattr(usage, "completion_tokens", 0),
+                    "total_tokens": getattr(usage, "total_tokens", 0),
+                }
+        except Exception as e:
+            logger.debug(f"Could not extract usage from MistralAI response: {e}")
+
+        return None

genai_otel/instrumentors/ollama_instrumentor.py

@@ -0,0 +1,83 @@
+"""OpenTelemetry instrumentor for the Ollama library.
+
+This instrumentor automatically traces calls to Ollama models for both
+generation and chat functionalities, 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 OllamaInstrumentor(BaseInstrumentor):
+    """Instrumentor for Ollama"""
+
+    def __init__(self):
+        """Initialize the instrumentor."""
+        super().__init__()
+        self._ollama_available = False
+        self._ollama_module = None
+        self._original_generate = None  # Add this
+        self._original_chat = None  # Add this
+        self._check_availability()
+
+    def _check_availability(self):
+        """Check if Ollama library is available."""
+        try:
+            import ollama
+
+            self._ollama_available = True
+            self._ollama_module = ollama
+            logger.debug("Ollama library detected and available for instrumentation")
+        except ImportError:
+            logger.debug("Ollama library not installed, instrumentation will be skipped")
+            self._ollama_available = False
+            self._ollama_module = None
+
+    def instrument(self, config: OTelConfig):
+        """Instrument the Ollama library."""
+        self.config = config
+
+        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)
+
+                if self.request_counter:
+                    self.request_counter.add(1, {"model": model, "provider": "ollama"})
+
+                result = self._original_chat(*args, **kwargs)
+                return result
+
+        self._ollama_module.generate = wrapped_generate
+        self._ollama_module.chat = wrapped_chat
+
+    def _extract_usage(self, result) -> Optional[Dict[str, int]]:
+        return None