opik 1.8.39__py3-none-any.whl → 1.9.71__py3-none-any.whl

opik/integrations/langchain/langgraph_async_context_bridge.py

@@ -0,0 +1,131 @@
+import logging
+from typing import Any, Dict, Optional
+
+import opik.api_objects.span as span
+import opik.integrations.langchain.opik_tracer as opik_tracer_module
+
+LOGGER = logging.getLogger(__name__)
+
+
+def extract_current_langgraph_span_data(
+    runnable_config: Dict[str, Any],
+) -> Optional[span.SpanData]:
+    """
+    Extract current span data for async LangGraph nodes.
+
+    This helper function is specifically designed for async LangGraph execution using `ainvoke()`.
+    Due to LangChain framework limitations in async scenarios, the execution context is not
+    automatically shared between callbacks (like OpikTracer) and node code. This function
+    extracts the current span data from the LangGraph config, allowing you to propagate
+    trace context to @track-decorated functions via distributed headers.
+
+    For synchronous execution using `invoke()`, this function is not needed as the context
+    is automatically shared.
+
+    Args:
+        runnable_config: The config dictionary automatically passed to LangGraph node functions,
+            containing an AsyncCallbackManager in the "callbacks" key.
+
+    Returns:
+        The current span data if found, None otherwise. Returns None if OpikTracer is not
+        configured in the callbacks or if the span data cannot be extracted.
+
+    Example:
+        ```python
+        from opik import track
+        from opik.integrations.langchain import OpikTracer, extract_current_langgraph_span_data
+        from langgraph.graph import StateGraph, START, END
+
+        @track
+        def process_data(value: int) -> int:
+            return value * 2
+
+        async def my_async_node(state, config):
+            # Extract current span data from LangGraph config
+            span_data = extract_current_langgraph_span_data(config)
+
+            if span_data is not None:
+                # Propagate trace context to tracked function
+                result = process_data(
+                    state["value"],
+                    opik_distributed_trace_headers=span_data.get_distributed_trace_headers()
+                )
+                return {"value": result}
+
+            return {"value": None}
+
+        # Build and execute graph
+        graph = StateGraph(dict)
+        graph.add_node("processor", my_async_node)
+        graph.add_edge(START, "processor")
+        graph.add_edge("processor", END)
+
+        app = graph.compile()
+        opik_tracer = OpikTracer()
+
+        # Asynchronous execution requires explicit trace context propagation
+        result = await app.ainvoke({"value": 21}, config={"callbacks": [opik_tracer]})
+        ```
+    """
+    try:
+        callback_manager = runnable_config.get("callbacks")
+        if callback_manager is None:
+            LOGGER.warning(
+                "Cannot extract span data: no callback manager found in LangGraph config. "
+                "Ensure OpikTracer is passed in the config: config={'callbacks': [opik_tracer]}"
+            )
+            return None
+
+        parent_run_id = getattr(callback_manager, "parent_run_id", None)
+        handlers = getattr(callback_manager, "handlers", [])
+
+        if parent_run_id is None or not handlers:
+            LOGGER.warning(
+                "Cannot extract span data: LangGraph callback manager is not properly initialized. "
+                "This may occur if the node is called outside of a LangGraph execution context."
+            )
+            return None
+
+        opik_tracer_instance = _find_opik_tracer(handlers)
+        if opik_tracer_instance is None:
+            LOGGER.warning(
+                "Cannot extract span data: OpikTracer not found in LangGraph callbacks. "
+                "Ensure OpikTracer is passed in the config: config={'callbacks': [opik_tracer]}"
+            )
+            return None
+
+        span_data = opik_tracer_instance.get_current_span_data_for_run(parent_run_id)
+        if span_data is None:
+            LOGGER.warning(
+                "Cannot extract span data: no span found for the current LangGraph node. "
+                "This may occur if OpikTracer was not properly initialized or if the node "
+                "is executing in an unexpected context."
+            )
+            return None
+
+        return span_data
+
+    except Exception as exception:
+        LOGGER.warning(
+            "Failed to extract span data from LangGraph config due to unexpected error",
+            exc_info=exception,
+        )
+        return None
+
+
+def _find_opik_tracer(
+    handlers: list,
+) -> Optional[opik_tracer_module.OpikTracer]:
+    """
+    Find OpikTracer instance in the list of handlers.
+
+    Args:
+        handlers: List of callback handlers.
+
+    Returns:
+        OpikTracer instance if found, None otherwise.
+    """
+    for handler in handlers:
+        if isinstance(handler, opik_tracer_module.OpikTracer):
+            return handler
+    return None

opik/integrations/langchain/langgraph_tracer_injector.py

@@ -0,0 +1,88 @@
+import logging
+from typing import Any, Dict, List, TypeVar
+
+from langchain_core.runnables import base as runnables_base
+
+from . import opik_tracer as opik_tracer_module
+
+LOGGER = logging.getLogger(__name__)
+
+CompiledGraphType = TypeVar("CompiledGraphType", bound=runnables_base.Runnable)
+
+
+def track_langgraph(
+    graph: CompiledGraphType,
+    opik_tracer: opik_tracer_module.OpikTracer,
+) -> CompiledGraphType:
+    """
+    Adds Opik tracking to a compiled LangGraph graph by injecting OpikTracer into its default config.
+
+    After calling this function, all subsequent invocations of the graph will automatically
+    be tracked without needing to pass the OpikTracer in the config parameter.
+
+    The function will automatically extract the graph structure visualization from the compiled
+    graph if it wasn't already provided when creating the OpikTracer. This visualization will
+    be included in the trace metadata in the Opik UI.
+
+    Args:
+        graph: A compiled LangGraph graph (result of StateGraph.compile()).
+        opik_tracer: An OpikTracer instance to use for tracking the graph.
+
+    Returns:
+        The modified graph with Opik tracking enabled.
+
+    Example:
+        ```python
+        from langgraph.graph import StateGraph, START, END
+        from opik.integrations.langchain import OpikTracer, track_langgraph
+
+        # Build your graph
+        builder = StateGraph(State)
+        builder.add_node("my_node", my_node_function)
+        builder.add_edge(START, "my_node")
+        builder.add_edge("my_node", END)
+
+        # Compile the graph
+        graph = builder.compile()
+
+        # Create OpikTracer and track the graph once
+        # No need to manually extract the graph - it's done automatically!
+        opik_tracer = OpikTracer(
+            tags=["production"],
+            metadata={"version": "1.0"}
+        )
+        graph = track_langgraph(graph, opik_tracer)
+
+        # Now all invocations are tracked automatically
+        result = graph.invoke({"message": "Hello"})
+        # No need to pass config={"callbacks": [opik_tracer]} anymore!
+        ```
+
+    Note:
+        - The graph visualization is automatically extracted and added to trace metadata
+          if not already provided in the OpikTracer constructor.
+        - If you need to customize the OpikTracer for specific invocations, you can still
+          pass it explicitly in the config parameter, which will override the default.
+        - The graph object is modified in-place and also returned for convenience.
+        - For async invocations using `ainvoke()`, you may still need to use
+          `extract_current_langgraph_span_data()` to propagate context to @track-decorated
+          functions within async nodes.
+    """
+    graph_structure = graph.get_graph(xray=True)
+    opik_tracer.set_graph(graph_structure)
+
+    # Inject the callback into the graph's default config
+    config: Dict[str, Any] = getattr(graph, "config", None) or {}
+    graph.config = config  # type: ignore[attr-defined]
+    callbacks: List[Any] = config.setdefault("callbacks", [])
+
+    if any(isinstance(cb, opik_tracer_module.OpikTracer) for cb in callbacks):
+        LOGGER.warning(
+            "Graph already has an OpikTracer callback injected. "
+            "Skipping re-tracking to avoid duplicate callbacks."
+        )
+        return graph
+
+    callbacks.append(opik_tracer)
+
+    return graph

opik/integrations/langchain/opik_encoder_extension.py

@@ -1,5 +1,5 @@
 from typing import Any, Dict
-from langchain.load import serializable
+from langchain_core.load import serializable
 import opik.jsonable_encoder as jsonable_encoder