opik 1.9.41__py3-none-any.whl → 1.9.86__py3-none-any.whl

opik/integrations/langchain/opik_tracer.py

@@ -1,5 +1,6 @@
 import logging
 import datetime
+import re
 from typing import (
     Any,
     Dict,
@@ -21,6 +22,7 @@ from langchain_core.tracers.schemas import Run
 
 from opik import context_storage, dict_utils, llm_usage, tracing_runtime_config
 from opik.api_objects import span, trace
+from opik.decorator import arguments_helpers, span_creation_handler
 from opik.types import DistributedTraceHeadersDict, ErrorInfoDict
 from opik.validation import parameters_validator
 from . import (
@@ -53,6 +55,11 @@ SkipErrorCallback = Callable[[str], bool]
 # due to a handled/ignored error during execution.
 ERROR_SKIPPED_OUTPUTS = {"warning": "Error output skipped by skip_error_callback."}
 
+# Constants for LangGraph interrupt/resume functionality
+LANGGRAPH_INTERRUPT_OUTPUT_KEY = "__interrupt__"
+LANGGRAPH_RESUME_INPUT_KEY = "__resume__"
+LANGGRAPH_INTERRUPT_METADATA_KEY = "_langgraph_interrupt"
+
 
 class TrackRootRunResult(NamedTuple):
     new_trace_data: Optional[trace.TraceData]
@@ -77,6 +84,117 @@ def _get_run_metadata(run_dict: Dict[str, Any]) -> Dict[str, Any]:
     return run_dict["extra"].get("metadata", {})
 
 
+def _parse_graph_interrupt_value(error_traceback: str) -> Optional[str]:
+    """
+    Parse GraphInterrupt error traceback to extract the interrupt value as a string.
+
+    The function extracts the value from the Interrupt object representation in the traceback.
+    It handles both string values (with quotes) and non-string values, including nested structures.
+    For string values, escape sequences are decoded (e.g., \\n becomes a newline character).
+
+    Args:
+        error_traceback: The error traceback string containing GraphInterrupt information.
+
+    Returns:
+        The interrupt value as a string if found, None otherwise.
+    """
+    # Search for GraphInterrupt( anywhere in the traceback
+    match = re.search(
+        r"GraphInterrupt\(.*?Interrupt\(value=",
+        error_traceback,
+        re.DOTALL,
+    )
+    if not match:
+        return None
+
+    # Start parsing from after "value="
+    start_pos = match.end()
+    value_str = error_traceback[start_pos:]
+
+    # Extract the value, handling nested parentheses and brackets
+    paren_depth = 0
+    bracket_depth = 0
+    brace_depth = 0
+    in_string = False
+    string_char = None
+    i = 0
+
+    for i, char in enumerate(value_str):
+        # Handle string boundaries
+        if char in ('"', "'") and (i == 0 or value_str[i - 1] != "\\"):
+            if not in_string:
+                in_string = True
+                string_char = char
+            elif char == string_char:
+                in_string = False
+                string_char = None
+
+        # Skip counting brackets/parens inside strings
+        if in_string:
+            continue
+
+        # Track nesting depth
+        if char == "(":
+            paren_depth += 1
+        elif char == ")":
+            if paren_depth > 0:
+                paren_depth -= 1
+            else:
+                # Found the closing paren of Interrupt(...), stop here
+                break
+        elif char == "[":
+            bracket_depth += 1
+        elif char == "]":
+            bracket_depth -= 1
+        elif char == "{":
+            brace_depth += 1
+        elif char == "}":
+            brace_depth -= 1
+        elif (
+            char == "," and paren_depth == 0 and bracket_depth == 0 and brace_depth == 0
+        ):
+            # Found a comma at the top level, stop here
+            break
+
+    # Extract and clean the value
+    value = value_str[:i].strip()
+
+    # Check if the value was originally a quoted string
+    was_quoted_string = False
+    if len(value) >= 2 and value[0] in ('"', "'") and value[-1] == value[0]:
+        was_quoted_string = True
+        value = value[1:-1]
+
+    # Decode escape sequences for string values
+    if was_quoted_string:
+        try:
+            value = value.encode("utf-8").decode("unicode_escape")
+        except (UnicodeDecodeError, AttributeError):
+            # If decoding fails, return the original value
+            pass
+
+    return value
+
+
+def _extract_resume_value_from_command(obj: Any) -> Optional[str]:
+    """
+    Extract the resume value from a LangGraph Command object or serialized Command dict.
+
+    Args:
+        obj: A Command object or dict representing a serialized Command object (from run.dict()).
+
+    Returns:
+        The resume value as a string if found, None otherwise.
+    """
+    # Check if it's a Command object (has a resume attribute)
+    if hasattr(obj, "resume") and obj.resume is not None:
+        return str(obj.resume)
+    # Check if it's a serialized Command dict
+    if obj is not None and isinstance(obj, dict) and "resume" in obj:
+        return str(obj["resume"])
+    return None
+
+
 class OpikTracer(BaseTracer):
     """Langchain Opik Tracer."""
 
@@ -89,6 +207,7 @@ class OpikTracer(BaseTracer):
         distributed_headers: Optional[DistributedTraceHeadersDict] = None,
         thread_id: Optional[str] = None,
         skip_error_callback: Optional[SkipErrorCallback] = None,
+        opik_context_read_only_mode: bool = False,
         **kwargs: Any,
     ) -> None:
         """
@@ -108,6 +227,12 @@ class OpikTracer(BaseTracer):
                 Please note that in traces/spans where errors are intentionally skipped,
                 the output will be replaced with `ERROR_SKIPPED_OUTPUTS`. You can provide
                 the output manually using `opik_context.get_current_span_data().update(output=...)`.
+            opik_context_read_only_mode: Whether to adding/popping spans/traces to/from the context storage.
+                * If False (default), OpikTracer will add created spans/traces to the opik context, so if there is a @track-decorated
+                  function called inside the LangChain runnable, it will be attached to it's parent span from LangChain automatically.
+                * If True, OpikTracer will not modify the context storage and only create spans/traces from LangChain's Run objects.
+                  This might be useful when the environment doesn't support proper context isolation for concurrent operations and you
+                  want to avoid modifying the Opik context stack due to unsafety.
             **kwargs: Additional arguments passed to the parent class constructor.
         """
         validator = parameters_validator.create_validator(
@@ -125,10 +250,7 @@ class OpikTracer(BaseTracer):
         self._trace_default_metadata["created_from"] = "langchain"
 
         if graph:
-            self._trace_default_metadata["_opik_graph_definition"] = {
-                "format": "mermaid",
-                "data": graph.draw_mermaid(),
-            }
+            self.set_graph(graph)
 
         self._trace_default_tags = tags
 
@@ -164,6 +286,23 @@ class OpikTracer(BaseTracer):
 
         self._skip_error_callback = skip_error_callback
 
+        self._opik_context_read_only_mode = opik_context_read_only_mode
+
+    def set_graph(self, graph: "Graph") -> None:
+        """
+        Set the LangGraph graph structure for visualization in Opik traces.
+
+        This method extracts the graph structure and stores it in trace metadata,
+        allowing the graph to be visualized in the Opik UI.
+
+        Args:
+            graph: A LangGraph Graph object (typically obtained via graph.get_graph(xray=True)).
+        """
+        self._trace_default_metadata["_opik_graph_definition"] = {
+            "format": "mermaid",
+            "data": graph.draw_mermaid(),
+        }
+
     def _is_opik_span_created_by_this_tracer(self, span_id: str) -> bool:
         return any(span_.id == span_id for span_ in self._span_data_map.values())
 
@@ -179,11 +318,16 @@ class OpikTracer(BaseTracer):
         trace_additional_metadata: Dict[str, Any] = {}
 
         error_str = run_dict.get("error")
-        outputs = None
+        outputs: Optional[Dict[str, Any]] = None
         error_info = None
 
         if error_str is not None:
-            if not self._should_skip_error(error_str):
+            # GraphInterrupt is not an error - it's a normal control flow for LangGraph
+            if interrupt_value := _parse_graph_interrupt_value(error_str):
+                outputs = {LANGGRAPH_INTERRUPT_OUTPUT_KEY: interrupt_value}
+                trace_additional_metadata[LANGGRAPH_INTERRUPT_METADATA_KEY] = True
+                # Don't set error_info - this is not an error
+            elif not self._should_skip_error(error_str):
                 error_info = ErrorInfoDict(
                     exception_type="Exception",
                     traceback=error_str,
@@ -195,7 +339,8 @@ class OpikTracer(BaseTracer):
                 langchain_helpers.split_big_langgraph_outputs(outputs)
             )
 
-        self._ensure_no_hanging_opik_tracer_spans()
+        if not self._opik_context_read_only_mode:
+            self._ensure_no_hanging_opik_tracer_spans()
 
         span_data = self._span_data_map.get(run.id)
         if (
@@ -228,6 +373,25 @@ class OpikTracer(BaseTracer):
         # workaround for `.astream()` method usage
         if trace_data.input == {"input": ""}:
             trace_data.input = run_dict["inputs"]
+        elif isinstance(trace_data.input, dict) and "input" in trace_data.input:
+            input_value = trace_data.input.get("input")
+            if resume_value := _extract_resume_value_from_command(input_value):
+                trace_data.input = {LANGGRAPH_RESUME_INPUT_KEY: resume_value}
+
+        # Check if any child span has a GraphInterrupt output and use it for trace output
+        for _, span_data in self._span_data_map.items():
+            if (
+                span_data.trace_id == trace_data.id
+                and span_data.metadata is not None
+                and span_data.metadata.get(LANGGRAPH_INTERRUPT_METADATA_KEY) is True
+            ):
+                # Use the interrupt output from the child span
+                outputs = span_data.output
+                # Also propagate the interrupt metadata to trace
+                if trace_additional_metadata is None:
+                    trace_additional_metadata = {}
+                trace_additional_metadata[LANGGRAPH_INTERRUPT_METADATA_KEY] = True
+                break
 
         if trace_additional_metadata:
             trace_data.update(metadata=trace_additional_metadata)
@@ -237,7 +401,8 @@ class OpikTracer(BaseTracer):
 
         assert trace_ is not None
         self._created_traces.append(trace_)
-        self._opik_context_storage.pop_trace_data(ensure_id=trace_data.id)
+        if not self._opik_context_read_only_mode:
+            self._opik_context_storage.pop_trace_data(ensure_id=trace_data.id)
 
     def _ensure_no_hanging_opik_tracer_spans(self) -> None:
         root_run_external_parent_span_id = self._root_run_external_parent_span_id.get()
@@ -260,19 +425,7 @@ class OpikTracer(BaseTracer):
         root_metadata = dict_utils.deepmerge(self._trace_default_metadata, run_metadata)
         self._update_thread_id_from_metadata(run_dict)
 
-        # Skip creating a span for root runs only when creating a new trace
-        # Keep the span when invoked from a tracked function, existing trace or distributed headers
-
-        if self._distributed_headers:
-            new_span_data = self._attach_span_to_distributed_headers(
-                run_dict=run_dict,
-                metadata=root_metadata,
-            )
-            return TrackRootRunResult(
-                new_trace_data=None,
-                new_span_data=new_span_data,
-            )
-
+        # Track the parent span ID for LangGraph cleanup later
         current_span_data = self._opik_context_storage.top_span_data()
         parent_span_id_when_langgraph_started = (
             current_span_data.id if current_span_data is not None else None
@@ -280,146 +433,49 @@ class OpikTracer(BaseTracer):
         self._root_run_external_parent_span_id.set(
             parent_span_id_when_langgraph_started
         )
-        if current_span_data is not None:
-            # When invoked from a tracked function, keep the root span
-            # and attach it to the parent span (don't skip it)
-            new_span_data = self._attach_span_to_external_span(
-                run_dict=run_dict,
-                current_span_data=current_span_data,
-                root_metadata=root_metadata,
-            )
-            return TrackRootRunResult(
-                new_trace_data=None,
-                new_span_data=new_span_data,
-            )
-
-        current_trace_data = self._opik_context_storage.get_trace_data()
-        if current_trace_data is not None:
-            # When invoked under an existing trace, keep the root span
-            # and attach it to the parent trace (don't skip it)
-            new_span_data = self._attach_span_to_external_trace(
-                run_dict=run_dict,
-                current_trace_data=current_trace_data,
-                root_metadata=root_metadata,
-            )
-            return TrackRootRunResult(
-                new_trace_data=None,
-                new_span_data=new_span_data,
-            )
-
-        return self._initialize_span_and_trace_from_scratch(
-            run_dict=run_dict,
-            root_metadata=root_metadata,
-            allow_duplicating_root_span=allow_duplicating_root_span,
-        )
-
-    def _initialize_span_and_trace_from_scratch(
-        self,
-        run_dict: Dict[str, Any],
-        root_metadata: Dict[str, Any],
-        allow_duplicating_root_span: bool,
-    ) -> TrackRootRunResult:
-        trace_data = trace.TraceData(
-            name=run_dict["name"],
-            input=run_dict["inputs"],
-            metadata=root_metadata,
-            tags=self._trace_default_tags,
-            project_name=self._project_name,
-            thread_id=self._thread_id,
-        )
 
-        # Skip creating a span for LangGraph root runs - children will be attached directly to trace
-        if _is_root_run(run_dict) and not allow_duplicating_root_span:
-            return TrackRootRunResult(
-                new_trace_data=trace_data,
-                new_span_data=None,
-            )
-
-        span_data = span.SpanData(
-            trace_id=trace_data.id,
-            parent_span_id=None,
+        start_span_arguments = arguments_helpers.StartSpanParameters(
             name=run_dict["name"],
             input=run_dict["inputs"],
             type=_get_span_type(run_dict),
-            metadata=root_metadata,
             tags=self._trace_default_tags,
+            metadata=root_metadata,
             project_name=self._project_name,
-        )
-        return TrackRootRunResult(new_trace_data=trace_data, new_span_data=span_data)
-
-    def _attach_span_to_external_span(
-        self,
-        run_dict: Dict[str, Any],
-        current_span_data: span.SpanData,
-        root_metadata: Dict[str, Any],
-    ) -> span.SpanData:
-        project_name = helpers.resolve_child_span_project_name(
-            current_span_data.project_name,
-            self._project_name,
+            thread_id=self._thread_id,
         )
 
-        span_data = span.SpanData(
-            trace_id=current_span_data.trace_id,
-            parent_span_id=current_span_data.id,
-            name=run_dict["name"],
-            input=run_dict["inputs"],
-            metadata=root_metadata,
-            tags=self._trace_default_tags,
-            project_name=project_name,
-            type=_get_span_type(run_dict),
+        span_creation_result = span_creation_handler.create_span_respecting_context(
+            start_span_arguments=start_span_arguments,
+            distributed_trace_headers=self._distributed_headers,
+            opik_context_storage=self._opik_context_storage,
         )
-        if not self._is_opik_trace_created_by_this_tracer(span_data.trace_id):
-            self._externally_created_traces_ids.add(span_data.trace_id)
-
-        return span_data
 
-    def _attach_span_to_external_trace(
-        self,
-        run_dict: Dict[str, Any],
-        current_trace_data: trace.TraceData,
-        root_metadata: Dict[str, Any],
-    ) -> span.SpanData:
-        project_name = helpers.resolve_child_span_project_name(
-            current_trace_data.project_name,
-            self._project_name,
+        trace_created_externally = (
+            span_creation_result.trace_data is None
+            and not self._is_opik_trace_created_by_this_tracer(
+                span_creation_result.span_data.trace_id
+            )
         )
+        if trace_created_externally:
+            self._externally_created_traces_ids.add(
+                span_creation_result.span_data.trace_id
+            )
 
-        span_data = span.SpanData(
-            trace_id=current_trace_data.id,
-            parent_span_id=None,
-            name=run_dict["name"],
-            input=run_dict["inputs"],
-            metadata=root_metadata,
-            tags=self._trace_default_tags,
-            project_name=project_name,
-            type=_get_span_type(run_dict),
+        should_skip_root_span_creation = (
+            span_creation_result.trace_data is not None
+            and _is_root_run(run_dict)
+            and not allow_duplicating_root_span
         )
-        span_data.update(metadata={"created_from": "langchain"})
-
-        if not self._is_opik_trace_created_by_this_tracer(current_trace_data.id):
-            self._externally_created_traces_ids.add(current_trace_data.id)
-        return span_data
+        if should_skip_root_span_creation:
+            return TrackRootRunResult(
+                new_trace_data=span_creation_result.trace_data,
+                new_span_data=None,
+            )
 
-    def _attach_span_to_distributed_headers(
-        self,
-        run_dict: Dict[str, Any],
-        metadata: Dict[str, Any],
-    ) -> span.SpanData:
-        if self._distributed_headers is None:
-            raise ValueError("Distributed headers are not set")
-
-        span_data = span.SpanData(
-            trace_id=self._distributed_headers["opik_trace_id"],
-            parent_span_id=self._distributed_headers["opik_parent_span_id"],
-            name=run_dict["name"],
-            input=run_dict["inputs"],
-            metadata=metadata,
-            tags=self._trace_default_tags,
-            project_name=self._project_name,
-            type=_get_span_type(run_dict),
+        return TrackRootRunResult(
+            new_trace_data=span_creation_result.trace_data,
+            new_span_data=span_creation_result.span_data,
         )
-        self._externally_created_traces_ids.add(span_data.trace_id)
-        return span_data
 
     def _process_start_span(self, run: Run, allow_duplicating_root_span: bool) -> None:
         try:
@@ -468,7 +524,11 @@ class OpikTracer(BaseTracer):
         # This is the first run for the chain.
         root_run_result = self._track_root_run(run_dict, allow_duplicating_root_span)
         if root_run_result.new_trace_data is not None:
-            self._opik_context_storage.set_trace_data(root_run_result.new_trace_data)
+            if not self._opik_context_read_only_mode:
+                self._opik_context_storage.set_trace_data(
+                    root_run_result.new_trace_data
+                )
+
             if (
                 self._opik_client.config.log_start_trace_span
                 and tracing_runtime_config.is_tracing_active()
@@ -501,7 +561,9 @@ class OpikTracer(BaseTracer):
                 trace_data=root_run_result.new_trace_data,
             )
 
-            self._opik_context_storage.add_span_data(root_run_result.new_span_data)
+            if not self._opik_context_read_only_mode:
+                self._opik_context_storage.add_span_data(root_run_result.new_span_data)
+
             if (
                 self._opik_client.config.log_start_trace_span
                 and tracing_runtime_config.is_tracing_active()
@@ -549,7 +611,9 @@ class OpikTracer(BaseTracer):
                 parent_run_id
             ]
 
-        self._opik_context_storage.add_span_data(new_span_data)
+        if not self._opik_context_read_only_mode:
+            self._opik_context_storage.add_span_data(new_span_data)
+
         if (
             self._opik_client.config.log_start_trace_span
             and tracing_runtime_config.is_tracing_active()
@@ -586,19 +650,40 @@ class OpikTracer(BaseTracer):
 
         elif self._distributed_headers:
             # LangGraph with distributed headers - attach to distributed trace
-            new_span_data = self._attach_span_to_distributed_headers(
-                run_dict=run_dict,
+            new_span_data = span.SpanData(
+                trace_id=self._distributed_headers["opik_trace_id"],
+                parent_span_id=self._distributed_headers["opik_parent_span_id"],
+                name=run_dict["name"],
+                input=run_dict["inputs"],
                 metadata=_get_run_metadata(run_dict),
+                tags=self._trace_default_tags,
+                project_name=self._project_name,
+                type=_get_span_type(run_dict),
             )
+            self._externally_created_traces_ids.add(new_span_data.trace_id)
+
         elif (
             current_trace_data := self._opik_context_storage.get_trace_data()
         ) is not None:
             # LangGraph attached to existing trace - attach children directly to trace
-            new_span_data = self._attach_span_to_external_trace(
-                run_dict=run_dict,
-                current_trace_data=current_trace_data,
-                root_metadata=_get_run_metadata(run_dict),
+            project_name = helpers.resolve_child_span_project_name(
+                current_trace_data.project_name,
+                self._project_name,
             )
+
+            new_span_data = span.SpanData(
+                trace_id=current_trace_data.id,
+                parent_span_id=None,
+                name=run_dict["name"],
+                input=run_dict["inputs"],
+                metadata=_get_run_metadata(run_dict),
+                tags=self._trace_default_tags,
+                project_name=project_name,
+                type=_get_span_type(run_dict),
+            )
+
+            if not self._is_opik_trace_created_by_this_tracer(current_trace_data.id):
+                self._externally_created_traces_ids.add(current_trace_data.id)
         else:
             LOGGER.warning(
                 f"Cannot find trace data or distributed headers for LangGraph child run '{run_id}'"
@@ -612,7 +697,9 @@ class OpikTracer(BaseTracer):
             trace_data=None,
         )
 
-        self._opik_context_storage.add_span_data(new_span_data)
+        if not self._opik_context_read_only_mode:
+            self._opik_context_storage.add_span_data(new_span_data)
+
         if (
             self._opik_client.config.log_start_trace_span
             and tracing_runtime_config.is_tracing_active()
@@ -641,8 +728,12 @@ class OpikTracer(BaseTracer):
                 usage_info = llm_usage.LLMUsageInfo()
 
             # workaround for `.astream()` method usage
-            if span_data.input == {"input": ""}:
+            if span_data.input == {"input": ""} or span_data.input == {"input": {}}:
                 span_data.input = run_dict["inputs"]
+            elif isinstance(span_data.input, dict):
+                input_value = span_data.input.get("input")
+                if resume_value := _extract_resume_value_from_command(input_value):
+                    span_data.input = {LANGGRAPH_RESUME_INPUT_KEY: resume_value}
 
             filtered_output, additional_metadata = (
                 langchain_helpers.split_big_langgraph_outputs(run_dict["outputs"])
@@ -667,7 +758,7 @@ class OpikTracer(BaseTracer):
         except Exception as e:
             LOGGER.error(f"Failed during _process_end_span: {e}", exc_info=True)
         finally:
-            if span_data is not None:
+            if span_data is not None and not self._opik_context_read_only_mode:
                 self._opik_context_storage.trim_span_data_stack_to_certain_span(
                     span_id=span_data.id
                 )
@@ -696,7 +787,14 @@ class OpikTracer(BaseTracer):
             span_data = self._span_data_map[run.id]
             error_str = run_dict["error"]
 
-            if self._should_skip_error(error_str):
+            # GraphInterrupt is not an error - it's a normal control flow for LangGraph
+            if interrupt_value := _parse_graph_interrupt_value(error_str):
+                span_data.init_end_time().update(
+                    metadata={LANGGRAPH_INTERRUPT_METADATA_KEY: True},
+                    output={LANGGRAPH_INTERRUPT_OUTPUT_KEY: interrupt_value},
+                )
+            # Don't set error_info - this is not an error
+            elif self._should_skip_error(error_str):
                 span_data.init_end_time().update(output=ERROR_SKIPPED_OUTPUTS)
             else:
                 error_info = ErrorInfoDict(
@@ -713,7 +811,7 @@ class OpikTracer(BaseTracer):
         except Exception as e:
             LOGGER.debug(f"Failed during _process_end_span_with_error: {e}")
         finally:
-            if span_data is not None:
+            if span_data is not None and not self._opik_context_read_only_mode:
                 self._opik_context_storage.trim_span_data_stack_to_certain_span(
                     span_id=span_data.id
                 )

opik/integrations/langchain/provider_usage_extractors/langchain_run_helpers/helpers.py

@@ -99,6 +99,9 @@ def find_token_usage_dict(
     that includes one or more keys from the specified candidate keys and returns it.
     If no such dictionary is found, the function returns None.
 
+    Searches in reverse order to prioritize more recent data (e.g., in multi-turn conversations,
+    the most recent turn's usage data comes last in the structure).
+
     Args:
         all_keys_should_match: if True, all candidate keys must be present in the dictionary.
         data: A nested data structure containing dictionaries, lists, or tuples to search through.
@@ -117,15 +120,15 @@ def find_token_usage_dict(
         elif not all_keys_should_match and len(matched_keys) > 0:
             return data
 
-        # Recursively search through dictionary values
-        for value in data.values():
+        # Recursively search through dictionary values in reverse order
+        for value in list(data.values())[::-1]:
             result = find_token_usage_dict(value, candidate_keys, all_keys_should_match)
             if result is not None:
                 return result
 
-    # Handle list and tuple cases
+    # Handle list and tuple cases - search in reverse order
     elif isinstance(data, (list, tuple)):
-        for item in data:
+        for item in reversed(data):
             result = find_token_usage_dict(item, candidate_keys, all_keys_should_match)
             if result is not None:
                 return result