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

genai_otel/instrumentors/langchain_instrumentor.py

@@ -0,0 +1,340 @@
+"""OpenTelemetry instrumentor for the LangChain framework.
+
+This instrumentor automatically traces various components within LangChain,
+including chains, agents, and chat models, capturing relevant attributes for observability.
+"""
+
+import asyncio
+import functools
+import logging
+from typing import Any, 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._langchain_core_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
+
+        # Check for langchain_core (required for chat model instrumentation)
+        try:
+            import langchain_core
+
+            self._langchain_core_available = True
+            logger.debug("langchain_core library detected and available for instrumentation")
+        except ImportError:
+            logger.debug(
+                "langchain_core library not installed, chat model instrumentation will be skipped"
+            )
+            self._langchain_core_available = False
+
+    def instrument(self, config: OTelConfig):
+        """Instrument langchain components if available."""
+        if not self._langchain_available:
+            logger.debug("Skipping instrumentation - library not available")
+            return
+
+        self.config = config
+
+        # Instrument chains and agents
+        self._instrument_chains_and_agents()
+
+        # Instrument chat models if langchain_core is available
+        if self._langchain_core_available:
+            self._instrument_chat_models()
+
+    def _instrument_chains_and_agents(self):
+        """Instrument LangChain chains and agents."""
+        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
+            logger.debug("Chains and agents instrumentation completed")
+
+        except ImportError:
+            logger.debug("Could not import chains or agents, skipping instrumentation")
+
+    def _instrument_chat_models(self):
+        """Instrument LangChain chat models."""
+        try:
+            from langchain_core.language_models.chat_models import BaseChatModel
+
+            # Instrument invoke method
+            original_invoke = BaseChatModel.invoke
+
+            @functools.wraps(original_invoke)
+            def wrapped_invoke(instance, *args, **kwargs):
+                import time
+
+                model_name = self._get_model_name(instance)
+                with self.tracer.start_as_current_span("langchain.chat_model.invoke") as span:
+                    start_time = time.time()
+                    self._set_chat_attributes(span, instance, args, kwargs, model_name)
+
+                    result = original_invoke(instance, *args, **kwargs)
+
+                    # Use standard metrics recording from BaseInstrumentor
+                    # This will extract usage, calculate costs, and record metrics
+                    self._record_result_metrics(span, result, start_time, kwargs)
+
+                    return result
+
+            BaseChatModel.invoke = wrapped_invoke
+
+            # Instrument ainvoke (async invoke) method
+            original_ainvoke = BaseChatModel.ainvoke
+
+            @functools.wraps(original_ainvoke)
+            async def wrapped_ainvoke(instance, *args, **kwargs):
+                import time
+
+                model_name = self._get_model_name(instance)
+                with self.tracer.start_as_current_span("langchain.chat_model.ainvoke") as span:
+                    start_time = time.time()
+                    self._set_chat_attributes(span, instance, args, kwargs, model_name)
+
+                    result = await original_ainvoke(instance, *args, **kwargs)
+
+                    # Use standard metrics recording from BaseInstrumentor
+                    # This will extract usage, calculate costs, and record metrics
+                    self._record_result_metrics(span, result, start_time, kwargs)
+
+                    return result
+
+            BaseChatModel.ainvoke = wrapped_ainvoke
+
+            # Instrument batch method
+            original_batch = BaseChatModel.batch
+
+            @functools.wraps(original_batch)
+            def wrapped_batch(instance, *args, **kwargs):
+                import time
+
+                model_name = self._get_model_name(instance)
+                with self.tracer.start_as_current_span("langchain.chat_model.batch") as span:
+                    start_time = time.time()
+
+                    # Set standard GenAI attributes
+                    provider = self._extract_provider(instance)
+                    if provider:
+                        span.set_attribute("gen_ai.system", provider)
+                    else:
+                        span.set_attribute("gen_ai.system", "langchain")
+
+                    span.set_attribute("gen_ai.request.model", model_name)
+                    span.set_attribute("gen_ai.operation.name", "batch")
+
+                    # Also set LangChain-specific attributes
+                    span.set_attribute("langchain.chat_model.name", model_name)
+                    span.set_attribute("langchain.chat_model.operation", "batch")
+
+                    # Get batch size
+                    if args and len(args) > 0:
+                        batch_size = len(args[0]) if hasattr(args[0], "__len__") else 1
+                        span.set_attribute("langchain.chat_model.batch_size", batch_size)
+
+                    result = original_batch(instance, *args, **kwargs)
+
+                    # Record metrics (though batch results may not have usage info)
+                    self._record_result_metrics(span, result, start_time, kwargs)
+
+                    return result
+
+            BaseChatModel.batch = wrapped_batch
+
+            # Instrument abatch (async batch) method
+            original_abatch = BaseChatModel.abatch
+
+            @functools.wraps(original_abatch)
+            async def wrapped_abatch(instance, *args, **kwargs):
+                import time
+
+                model_name = self._get_model_name(instance)
+                with self.tracer.start_as_current_span("langchain.chat_model.abatch") as span:
+                    start_time = time.time()
+
+                    # Set standard GenAI attributes
+                    provider = self._extract_provider(instance)
+                    if provider:
+                        span.set_attribute("gen_ai.system", provider)
+                    else:
+                        span.set_attribute("gen_ai.system", "langchain")
+
+                    span.set_attribute("gen_ai.request.model", model_name)
+                    span.set_attribute("gen_ai.operation.name", "batch")
+
+                    # Also set LangChain-specific attributes
+                    span.set_attribute("langchain.chat_model.name", model_name)
+                    span.set_attribute("langchain.chat_model.operation", "abatch")
+
+                    # Get batch size
+                    if args and len(args) > 0:
+                        batch_size = len(args[0]) if hasattr(args[0], "__len__") else 1
+                        span.set_attribute("langchain.chat_model.batch_size", batch_size)
+
+                    result = await original_abatch(instance, *args, **kwargs)
+
+                    # Record metrics (though batch results may not have usage info)
+                    self._record_result_metrics(span, result, start_time, kwargs)
+
+                    return result
+
+            BaseChatModel.abatch = wrapped_abatch
+
+            logger.info("LangChain chat models instrumentation completed")
+
+        except ImportError as e:
+            logger.debug(f"Could not import langchain_core chat models: {e}")
+        except Exception as e:
+            logger.error(f"Error instrumenting chat models: {e}", exc_info=True)
+
+    def _get_model_name(self, instance: Any) -> str:
+        """Extract model name from chat model instance."""
+        # Try common attribute names for model name
+        for attr in ["model_name", "model", "model_id"]:
+            if hasattr(instance, attr):
+                value = getattr(instance, attr)
+                if value:
+                    return str(value)
+
+        # Fallback to class name
+        return instance.__class__.__name__
+
+    def _set_chat_attributes(self, span, instance: Any, args: tuple, kwargs: dict, model_name: str):
+        """Set span attributes for chat model invocations."""
+        # Set standard GenAI semantic convention attributes
+        provider = self._extract_provider(instance)
+        if provider:
+            span.set_attribute("gen_ai.system", provider)
+        else:
+            span.set_attribute("gen_ai.system", "langchain")
+
+        span.set_attribute("gen_ai.request.model", model_name)
+        span.set_attribute("gen_ai.operation.name", "chat")
+
+        # Also set LangChain-specific attributes for backward compatibility
+        span.set_attribute("langchain.chat_model.name", model_name)
+        span.set_attribute("langchain.chat_model.operation", "invoke")
+        if provider:
+            span.set_attribute("langchain.chat_model.provider", provider)
+
+        # Count messages if available
+        if args and len(args) > 0:
+            messages = args[0]
+            if hasattr(messages, "__len__"):
+                message_count = len(messages)
+                span.set_attribute("gen_ai.request.message_count", message_count)
+                span.set_attribute("langchain.chat_model.message_count", message_count)
+
+    def _extract_provider(self, instance: Any) -> Optional[str]:
+        """Extract provider name from chat model instance."""
+        class_name = instance.__class__.__name__.lower()
+        module_name = instance.__class__.__module__.lower()
+
+        # Map class names to providers
+        provider_mapping = {
+            "openai": "openai",
+            "anthropic": "anthropic",
+            "google": "google",
+            "ollama": "ollama",
+            "bedrock": "bedrock",
+            "cohere": "cohere",
+            "groq": "groq",
+            "mistral": "mistral",
+        }
+
+        # Check class name
+        for key, value in provider_mapping.items():
+            if key in class_name:
+                return value
+
+        # Check module name
+        for key, value in provider_mapping.items():
+            if key in module_name:
+                return value
+
+        return None
+
+    def _extract_usage(self, result) -> Optional[Dict[str, int]]:
+        """Extract usage information for BaseInstrumentor compatibility."""
+        try:
+            usage_data = None
+
+            if hasattr(result, "usage_metadata") and result.usage_metadata:
+                usage_data = result.usage_metadata
+            elif hasattr(result, "response_metadata") and result.response_metadata:
+                metadata = result.response_metadata
+                if "token_usage" in metadata:
+                    usage_data = metadata["token_usage"]
+                elif "usage" in metadata:
+                    usage_data = metadata["usage"]
+
+            if usage_data:
+                if isinstance(usage_data, dict):
+                    prompt_tokens = usage_data.get("input_tokens") or usage_data.get(
+                        "prompt_tokens"
+                    )
+                    completion_tokens = usage_data.get("output_tokens") or usage_data.get(
+                        "completion_tokens"
+                    )
+                else:
+                    prompt_tokens = getattr(usage_data, "input_tokens", None) or getattr(
+                        usage_data, "prompt_tokens", None
+                    )
+                    completion_tokens = getattr(usage_data, "output_tokens", None) or getattr(
+                        usage_data, "completion_tokens", None
+                    )
+
+                if prompt_tokens or completion_tokens:
+                    return {
+                        "prompt_tokens": int(prompt_tokens) if prompt_tokens else 0,
+                        "completion_tokens": int(completion_tokens) if completion_tokens else 0,
+                        "total_tokens": int(prompt_tokens or 0) + int(completion_tokens or 0),
+                    }
+        except Exception:
+            pass
+
+        return None

genai_otel/instrumentors/langgraph_instrumentor.py

@@ -0,0 +1,328 @@
+"""OpenTelemetry instrumentor for the LangGraph framework.
+
+This instrumentor automatically traces graph execution, nodes, edges, state updates,
+and checkpoints using the LangGraph stateful workflow framework.
+"""
+
+import json
+import logging
+from typing import Any, Dict, Optional
+
+from ..config import OTelConfig
+from .base import BaseInstrumentor
+
+logger = logging.getLogger(__name__)
+
+
+class LangGraphInstrumentor(BaseInstrumentor):
+    """Instrumentor for LangGraph stateful workflow framework"""
+
+    def __init__(self):
+        """Initialize the instrumentor."""
+        super().__init__()
+        self._langgraph_available = False
+        self._check_availability()
+
+    def _check_availability(self):
+        """Check if LangGraph library is available."""
+        try:
+            import langgraph
+
+            self._langgraph_available = True
+            logger.debug("LangGraph library detected and available for instrumentation")
+        except ImportError:
+            logger.debug("LangGraph library not installed, instrumentation will be skipped")
+            self._langgraph_available = False
+
+    def instrument(self, config: OTelConfig):
+        """Instrument LangGraph framework if available.
+
+        Args:
+            config (OTelConfig): The OpenTelemetry configuration object.
+        """
+        if not self._langgraph_available:
+            logger.debug("Skipping LangGraph instrumentation - library not available")
+            return
+
+        self.config = config
+
+        try:
+            import wrapt
+            from langgraph.graph import StateGraph
+
+            # Instrument StateGraph.compile() to wrap the resulting CompiledGraph
+            if hasattr(StateGraph, "compile"):
+                original_compile = StateGraph.compile
+
+                def wrapped_compile(wrapped, instance, args, kwargs):
+                    # Get the compiled graph
+                    compiled_graph = wrapped(*args, **kwargs)
+
+                    # Instrument the compiled graph's execution methods
+                    self._instrument_compiled_graph(compiled_graph, instance)
+
+                    return compiled_graph
+
+                StateGraph.compile = wrapt.FunctionWrapper(original_compile, wrapped_compile)
+
+                self._instrumented = True
+                logger.info("LangGraph instrumentation enabled")
+
+        except Exception as e:
+            logger.error("Failed to instrument LangGraph: %s", e, exc_info=True)
+            if config.fail_on_error:
+                raise
+
+    def _instrument_compiled_graph(self, compiled_graph, state_graph):
+        """Instrument a compiled graph's execution methods.
+
+        Args:
+            compiled_graph: The compiled graph object.
+            state_graph: The original StateGraph instance.
+        """
+        import wrapt
+
+        # Instrument invoke method (synchronous execution)
+        if hasattr(compiled_graph, "invoke"):
+            original_invoke = compiled_graph.invoke
+            compiled_graph.invoke = wrapt.FunctionWrapper(
+                original_invoke,
+                lambda w, i, a, kw: self._wrap_graph_invoke(
+                    w, i, a, kw, state_graph, is_async=False
+                ),
+            )
+
+        # Instrument stream method (synchronous streaming)
+        if hasattr(compiled_graph, "stream"):
+            original_stream = compiled_graph.stream
+            compiled_graph.stream = wrapt.FunctionWrapper(
+                original_stream,
+                lambda w, i, a, kw: self._wrap_graph_stream(
+                    w, i, a, kw, state_graph, is_async=False
+                ),
+            )
+
+        # Instrument ainvoke method (asynchronous execution)
+        if hasattr(compiled_graph, "ainvoke"):
+            original_ainvoke = compiled_graph.ainvoke
+            compiled_graph.ainvoke = wrapt.FunctionWrapper(
+                original_ainvoke,
+                lambda w, i, a, kw: self._wrap_graph_invoke(
+                    w, i, a, kw, state_graph, is_async=True
+                ),
+            )
+
+        # Instrument astream method (asynchronous streaming)
+        if hasattr(compiled_graph, "astream"):
+            original_astream = compiled_graph.astream
+            compiled_graph.astream = wrapt.FunctionWrapper(
+                original_astream,
+                lambda w, i, a, kw: self._wrap_graph_stream(
+                    w, i, a, kw, state_graph, is_async=True
+                ),
+            )
+
+    def _wrap_graph_invoke(self, wrapped, instance, args, kwargs, state_graph, is_async):
+        """Wrap graph invoke/ainvoke method with span.
+
+        Args:
+            wrapped: The original method.
+            instance: The compiled graph instance.
+            args: Positional arguments.
+            kwargs: Keyword arguments.
+            state_graph: The original StateGraph instance.
+            is_async: Whether this is async invocation.
+        """
+        operation_name = "graph.ainvoke" if is_async else "graph.invoke"
+        return self.create_span_wrapper(
+            span_name=f"langgraph.{operation_name}",
+            extract_attributes=lambda i, a, kw: self._extract_graph_attributes(
+                i, a, kw, state_graph
+            ),
+        )(wrapped)(instance, *args, **kwargs)
+
+    def _wrap_graph_stream(self, wrapped, instance, args, kwargs, state_graph, is_async):
+        """Wrap graph stream/astream method with span.
+
+        Args:
+            wrapped: The original method.
+            instance: The compiled graph instance.
+            args: Positional arguments.
+            kwargs: Keyword arguments.
+            state_graph: The original StateGraph instance.
+            is_async: Whether this is async streaming.
+        """
+        operation_name = "graph.astream" if is_async else "graph.stream"
+        return self.create_span_wrapper(
+            span_name=f"langgraph.{operation_name}",
+            extract_attributes=lambda i, a, kw: self._extract_graph_attributes(
+                i, a, kw, state_graph
+            ),
+        )(wrapped)(instance, *args, **kwargs)
+
+    def _extract_graph_attributes(
+        self, instance: Any, args: Any, kwargs: Any, state_graph: Any
+    ) -> Dict[str, Any]:
+        """Extract attributes from graph execution.
+
+        Args:
+            instance: The compiled graph instance.
+            args: Positional arguments (input state).
+            kwargs: Keyword arguments.
+            state_graph: The original StateGraph instance.
+
+        Returns:
+            Dict[str, Any]: Dictionary of attributes to set on the span.
+        """
+        attrs = {}
+
+        # Core attributes
+        attrs["gen_ai.system"] = "langgraph"
+        attrs["gen_ai.operation.name"] = "graph.execution"
+
+        # Extract graph structure information from StateGraph
+        try:
+            # Get nodes from the graph
+            if hasattr(state_graph, "nodes"):
+                nodes = state_graph.nodes
+                if nodes:
+                    node_names = list(nodes.keys())
+                    attrs["langgraph.node_count"] = len(node_names)
+                    attrs["langgraph.nodes"] = node_names[:10]  # Limit to 10 nodes
+
+            # Get edges from the graph
+            if hasattr(state_graph, "edges"):
+                edges = state_graph.edges
+                if edges:
+                    attrs["langgraph.edge_count"] = len(edges)
+
+            # Get channels (state schema) if available
+            if hasattr(state_graph, "channels"):
+                channels = state_graph.channels
+                if channels:
+                    channel_names = list(channels.keys())
+                    attrs["langgraph.channels"] = channel_names[:10]
+                    attrs["langgraph.channel_count"] = len(channel_names)
+
+        except Exception as e:
+            logger.debug("Failed to extract graph structure: %s", e)
+
+        # Extract input state (first positional argument)
+        input_state = None
+        if len(args) > 0:
+            input_state = args[0]
+        elif "input" in kwargs:
+            input_state = kwargs["input"]
+
+        if input_state:
+            try:
+                if isinstance(input_state, dict):
+                    # Store input state keys
+                    state_keys = list(input_state.keys())
+                    attrs["langgraph.input.keys"] = state_keys[:10]
+
+                    # Store truncated values for important keys
+                    for key in ["messages", "query", "question", "input"][:3]:
+                        if key in input_state:
+                            value = str(input_state[key])[:200]
+                            attrs[f"langgraph.input.{key}"] = value
+                else:
+                    # Non-dict input
+                    attrs["langgraph.input"] = str(input_state)[:200]
+            except Exception as e:
+                logger.debug("Failed to extract input state: %s", e)
+
+        # Extract config if provided
+        config = kwargs.get("config")
+        if config:
+            try:
+                # Extract configurable values
+                if isinstance(config, dict):
+                    if "configurable" in config:
+                        configurable = config["configurable"]
+                        if "thread_id" in configurable:
+                            attrs["langgraph.thread_id"] = configurable["thread_id"]
+                        if "checkpoint_id" in configurable:
+                            attrs["langgraph.checkpoint_id"] = configurable["checkpoint_id"]
+
+                    # Extract recursion limit
+                    if "recursion_limit" in config:
+                        attrs["langgraph.recursion_limit"] = config["recursion_limit"]
+
+            except Exception as e:
+                logger.debug("Failed to extract config: %s", e)
+
+        return attrs
+
+    def _extract_usage(self, result) -> Optional[Dict[str, int]]:
+        """Extract token usage from graph execution result.
+
+        Note: LangGraph doesn't directly expose token usage in the result.
+        Token usage is captured by underlying LLM provider instrumentors.
+
+        Args:
+            result: The graph execution result.
+
+        Returns:
+            Optional[Dict[str, int]]: Dictionary with token counts or None.
+        """
+        # LangGraph doesn't directly expose usage
+        # Token usage is captured by LLM provider instrumentors (OpenAI, Anthropic, etc.)
+        return None
+
+    def _extract_response_attributes(self, result) -> Dict[str, Any]:
+        """Extract response attributes from graph execution result.
+
+        Args:
+            result: The graph execution result.
+
+        Returns:
+            Dict[str, Any]: Dictionary of response attributes.
+        """
+        attrs = {}
+
+        try:
+            # LangGraph result is typically a dict (the final state)
+            if isinstance(result, dict):
+                # Store output state keys
+                state_keys = list(result.keys())
+                attrs["langgraph.output.keys"] = state_keys[:10]
+
+                # Store truncated values for important keys
+                for key in ["messages", "answer", "output", "result"][:3]:
+                    if key in result:
+                        value_str = str(result[key])[:500]
+                        attrs[f"langgraph.output.{key}"] = value_str
+
+                # Count the number of state updates/steps
+                if "messages" in result and isinstance(result["messages"], list):
+                    attrs["langgraph.message_count"] = len(result["messages"])
+
+            # Try to extract metadata if available
+            if hasattr(result, "__metadata__"):
+                try:
+                    metadata = result.__metadata__
+                    if "step" in metadata:
+                        attrs["langgraph.steps"] = metadata["step"]
+                except Exception as e:
+                    logger.debug("Failed to extract metadata: %s", e)
+
+        except Exception as e:
+            logger.debug("Failed to extract response attributes: %s", e)
+
+        return attrs
+
+    def _extract_finish_reason(self, result) -> Optional[str]:
+        """Extract finish reason from graph execution result.
+
+        Args:
+            result: The graph execution result.
+
+        Returns:
+            Optional[str]: The finish reason string or None if not available.
+        """
+        # LangGraph doesn't typically provide a finish_reason
+        # We could infer completion status
+        if result:
+            return "completed"
+        return None