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

opik/integrations/dspy/callback.py

@@ -1,35 +1,53 @@
-from contextvars import ContextVar, Token
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict, Optional, Tuple, Union
+import logging
 
 import dspy
-from dspy.utils.callback import BaseCallback
+from dspy.utils import callback as dspy_callback
 
-from opik import opik_context, types
-from opik.api_objects import helpers, span, trace
-from opik.api_objects.opik_client import get_client_cached
+from opik import context_storage, opik_context, tracing_runtime_config
+from opik import llm_usage
+from opik.api_objects import helpers, span, trace, opik_client
 from opik.decorator import error_info_collector
 
-ContextType = Union[span.SpanData, trace.TraceData]
+from .graph import build_mermaid_graph_from_module
+from .parsers import LMHistoryInfo, extract_lm_info_from_history, get_span_type
 
+LOGGER = logging.getLogger(__name__)
+
+SpanOrTraceData = Union[span.SpanData, trace.TraceData]
+
+
+class OpikCallback(dspy_callback.BaseCallback):
+    """
+    Callback for DSPy Opik logging.
+
+    Args:
+        project_name: The name of the Opik project to log data.
+        log_graph: If True, will log a mermaid diagram for each
+            module
+    """
 
-class OpikCallback(BaseCallback):
     def __init__(
         self,
         project_name: Optional[str] = None,
+        log_graph: bool = False,
     ):
         self._map_call_id_to_span_data: Dict[str, span.SpanData] = {}
         self._map_call_id_to_trace_data: Dict[str, trace.TraceData] = {}
-        self._map_span_id_or_trace_id_to_token: Dict[str, Token] = {}
+        # Store (lm_instance, expected_messages) for extracting usage and verifying correct history entry
+        self._map_call_id_to_lm_info: Dict[str, Tuple[Any, Optional[Any]]] = {}
 
-        self._origins_metadata = {"created_from": "dspy"}
+        self._origins_metadata: Dict[str, Any] = {"created_from": "dspy"}
 
-        self._current_callback_context: ContextVar[Optional[ContextType]] = ContextVar(
-            "opik_context", default=None
-        )
+        self._context_storage = context_storage.OpikContextStorage()
 
         self._project_name = project_name
+        self.log_graph = log_graph
+
+        self._opik_client = opik_client.get_client_cached()
 
-        self._opik_client = get_client_cached()
+    def _skip_tracking(self) -> bool:
+        return not tracing_runtime_config.is_tracing_active()
 
     def on_module_start(
         self,
@@ -37,50 +55,46 @@ class OpikCallback(BaseCallback):
         instance: Any,
         inputs: Dict[str, Any],
     ) -> None:
-        if current_callback_context_data := self._current_callback_context.get():
-            if isinstance(current_callback_context_data, span.SpanData):
-                self._attach_span_to_existing_span(
-                    call_id=call_id,
-                    current_span_data=current_callback_context_data,
-                    instance=instance,
-                    inputs=inputs,
-                )
-            else:
-                self._attach_span_to_existing_trace(
-                    call_id=call_id,
-                    current_trace_data=current_callback_context_data,
-                    instance=instance,
-                    inputs=inputs,
-                )
+        if self._skip_tracking():
             return
 
-        if current_span_data := opik_context.get_current_span_data():
+        # First we check the callback's context
+        if (current_span_data := self._context_storage.top_span_data()) is not None:
             self._attach_span_to_existing_span(
                 call_id=call_id,
                 current_span_data=current_span_data,
                 instance=instance,
                 inputs=inputs,
             )
-            new_span_data = self._map_call_id_to_span_data[call_id]
-            self._callback_context_set(new_span_data)
-            return
-
-        if current_trace_data := opik_context.get_current_trace_data():
+        elif (current_trace_data := self._context_storage.get_trace_data()) is not None:
             self._attach_span_to_existing_trace(
                 call_id=call_id,
                 current_trace_data=current_trace_data,
                 instance=instance,
                 inputs=inputs,
             )
-            new_span_data = self._map_call_id_to_span_data[call_id]
-            self._callback_context_set(new_span_data)
-            return
-
-        self._start_trace(
-            call_id=call_id,
-            instance=instance,
-            inputs=inputs,
-        )
+        # Callback's context is empty, we check opik's context
+        elif (current_span_data := opik_context.get_current_span_data()) is not None:
+            self._attach_span_to_existing_span(
+                call_id=call_id,
+                current_span_data=current_span_data,
+                instance=instance,
+                inputs=inputs,
+            )
+        elif (current_trace_data := opik_context.get_current_trace_data()) is not None:
+            self._attach_span_to_existing_trace(
+                call_id=call_id,
+                current_trace_data=current_trace_data,
+                instance=instance,
+                inputs=inputs,
+            )
+        else:
+            # Both callback's and opik's context are empty
+            self._start_trace(
+                call_id=call_id,
+                instance=instance,
+                inputs=inputs,
+            )
 
     def _attach_span_to_existing_span(
         self,
@@ -93,7 +107,7 @@ class OpikCallback(BaseCallback):
             parent_project_name=current_span_data.project_name,
             child_project_name=self._project_name,
         )
-        span_type = self._get_span_type(instance)
+        span_type = get_span_type(instance)
 
         span_data = span.SpanData(
             trace_id=current_span_data.trace_id,
@@ -102,9 +116,9 @@ class OpikCallback(BaseCallback):
             input=inputs,
             type=span_type,
             project_name=project_name,
-            metadata=self._origins_metadata,
+            metadata=self._get_opik_metadata(instance),
         )
-        self._map_call_id_to_span_data[call_id] = span_data
+        self._start_span(call_id=call_id, span_data=span_data)
 
     def _attach_span_to_existing_trace(
         self,
@@ -117,7 +131,7 @@ class OpikCallback(BaseCallback):
             current_trace_data.project_name,
             self._project_name,
         )
-        span_type = self._get_span_type(instance)
+        span_type = get_span_type(instance)
 
         span_data = span.SpanData(
             trace_id=current_trace_data.id,
@@ -126,9 +140,19 @@ class OpikCallback(BaseCallback):
             input=inputs,
             type=span_type,
             project_name=project_name,
-            metadata=self._origins_metadata,
+            metadata=self._get_opik_metadata(instance),
         )
+        self._start_span(call_id=call_id, span_data=span_data)
+
+    def _start_span(self, call_id: str, span_data: span.SpanData) -> None:
         self._map_call_id_to_span_data[call_id] = span_data
+        self._set_current_context_data(span_data)
+
+        if (
+            self._opik_client.config.log_start_trace_span
+            and tracing_runtime_config.is_tracing_active()
+        ):
+            self._opik_client.span(**span_data.as_start_parameters)
 
     def _start_trace(
         self,
@@ -139,11 +163,17 @@ class OpikCallback(BaseCallback):
         trace_data = trace.TraceData(
             name=instance.__class__.__name__,
             input=inputs,
-            metadata={"created_from": "dspy"},
+            metadata=self._get_opik_metadata(instance),
             project_name=self._project_name,
         )
         self._map_call_id_to_trace_data[call_id] = trace_data
-        self._callback_context_set(trace_data)
+        self._set_current_context_data(trace_data)
+
+        if (
+            self._opik_client.config.log_start_trace_span
+            and tracing_runtime_config.is_tracing_active()
+        ):
+            self._opik_client.trace(**trace_data.as_start_parameters)
 
     def on_module_end(
         self,
@@ -160,38 +190,63 @@ class OpikCallback(BaseCallback):
 
     def _end_trace(self, call_id: str) -> None:
         if trace_data := self._map_call_id_to_trace_data.pop(call_id, None):
-            trace_data.init_end_time()
-            self._opik_client.trace(**trace_data.__dict__)
+            if tracing_runtime_config.is_tracing_active():
+                trace_data.init_end_time()
+                self._opik_client.trace(**trace_data.as_parameters)
 
-            # remove trace data from context
-            if token := self._map_span_id_or_trace_id_to_token.pop(trace_data.id, None):
-                self._current_callback_context.reset(token)
+            if self._context_storage.get_trace_data() == trace_data:
+                self._context_storage.set_trace_data(None)
 
     def _end_span(
         self,
         call_id: str,
         outputs: Optional[Any],
         exception: Optional[Exception] = None,
+        usage: Optional[llm_usage.OpikUsage] = None,
+        extra_metadata: Optional[Dict[str, Any]] = None,
+        actual_provider: Optional[str] = None,
+        total_cost: Optional[float] = None,
     ) -> None:
         if span_data := self._map_call_id_to_span_data.pop(call_id, None):
             if exception:
                 error_info = error_info_collector.collect(exception)
                 span_data.update(error_info=error_info)
 
-            span_data.update(output={"output": outputs}).init_end_time()
-            self._opik_client.span(**span_data.__dict__)
+            # Prepare the update dict
+            update_kwargs: Dict[str, Any] = {
+                "output": {"output": outputs},
+                "usage": usage,
+                "total_cost": total_cost,
+            }
+
+            # Handle LLM routers like OpenRouter that return the actual serving provider
+            if (
+                actual_provider is not None
+                and span_data.provider is not None
+                and span_data.provider.lower() != actual_provider.lower()
+            ):
+                # Store the original provider (e.g., "openrouter") in metadata
+                if extra_metadata is None:
+                    extra_metadata = {}
+                extra_metadata["llm_router"] = span_data.provider
+                # Update to the actual provider for accurate cost tracking
+                update_kwargs["provider"] = actual_provider
+
+            update_kwargs["metadata"] = extra_metadata
+
+            span_data.update(**update_kwargs).init_end_time()
+            if tracing_runtime_config.is_tracing_active():
+                self._opik_client.span(**span_data.as_parameters)
 
             # remove span data from context
-            if token := self._map_span_id_or_trace_id_to_token.pop(span_data.id, None):
-                self._current_callback_context.reset(token)
-
-    def on_lm_start(
-        self,
-        call_id: str,
-        instance: Any,
-        inputs: Dict[str, Any],
-    ) -> None:
-        current_callback_context_data = self._current_callback_context.get()
+            current_span = self._context_storage.top_span_data()
+            if current_span and current_span.id == span_data.id:
+                self._context_storage.pop_span_data()
+
+    def _collect_common_span_data(
+        self, instance: Any, inputs: Dict[str, Any]
+    ) -> span.SpanData:
+        current_callback_context_data = self._get_current_context_data()
         assert current_callback_context_data is not None
 
         project_name = helpers.resolve_child_span_project_name(
@@ -206,46 +261,153 @@ class OpikCallback(BaseCallback):
             trace_id = current_callback_context_data.id
             parent_span_id = None
 
-        provider, model = instance.model.split(r"/", 1)
-        span_type = self._get_span_type(instance)
+        span_type = get_span_type(instance)
 
-        span_data = span.SpanData(
+        return span.SpanData(
             trace_id=trace_id,
-            name=instance.__class__.__name__,
             parent_span_id=parent_span_id,
-            type=span_type,
+            name=(
+                instance.name
+                if hasattr(instance, "name")
+                else instance.__class__.__name__
+            ),
             input=inputs,
+            type=span_type,
             project_name=project_name,
+            metadata=self._get_opik_metadata(instance),
+        )
+
+    def on_lm_start(
+        self,
+        call_id: str,
+        instance: Any,
+        inputs: Dict[str, Any],
+    ) -> None:
+        span_data = self._collect_common_span_data(instance, inputs)
+
+        provider, model = instance.model.split(r"/", 1)
+
+        span_data.update(
             provider=provider,
             model=model,
-            metadata=self._origins_metadata,
+            name=f"{span_data.name}: {provider} - {model}",
         )
         self._map_call_id_to_span_data[call_id] = span_data
 
+        # Store LM instance and expected messages for extracting usage
+        self._map_call_id_to_lm_info[call_id] = (
+            instance,
+            inputs.get("messages"),
+        )
+
+        self._set_current_context_data(span_data)
+
     def on_lm_end(
         self,
         call_id: str,
         outputs: Optional[Dict[str, Any]],
         exception: Optional[Exception] = None,
+    ) -> None:
+        lm_info = self._extract_lm_info_from_history(call_id)
+
+        # Add cache_hit to span metadata only when we have a definitive value
+        extra_metadata = (
+            {"cache_hit": lm_info.cache_hit} if lm_info.cache_hit is not None else None
+        )
+
+        self._end_span(
+            call_id=call_id,
+            exception=exception,
+            outputs=outputs,
+            usage=lm_info.usage,
+            extra_metadata=extra_metadata,
+            actual_provider=lm_info.actual_provider,
+            total_cost=lm_info.total_cost,
+        )
+
+    def on_tool_start(
+        self,
+        call_id: str,
+        instance: Any,
+        inputs: Dict[str, Any],
+    ) -> None:
+        span_data = self._collect_common_span_data(instance, inputs)
+        self._map_call_id_to_span_data[call_id] = span_data
+        self._set_current_context_data(span_data)
+
+    def on_tool_end(
+        self,
+        call_id: str,
+        outputs: Optional[Dict[str, Any]],
+        exception: Optional[Exception] = None,
     ) -> None:
         self._end_span(
             call_id=call_id,
             exception=exception,
             outputs=outputs,
         )
-        self._end_trace(call_id=call_id)
 
     def flush(self) -> None:
         """Sends pending Opik data to the backend"""
         self._opik_client.flush()
 
-    def _callback_context_set(self, value: ContextType) -> None:
-        token = self._current_callback_context.set(value)
-        self._map_span_id_or_trace_id_to_token[value.id] = token
+    def _set_current_context_data(self, value: SpanOrTraceData) -> None:
+        if isinstance(value, span.SpanData):
+            self._context_storage.add_span_data(value)
+        elif isinstance(value, trace.TraceData):
+            self._context_storage.set_trace_data(value)
+        else:
+            raise ValueError(f"Invalid context type: {type(value)}")
+
+    def _get_current_context_data(self) -> Optional[SpanOrTraceData]:
+        if span_data := self._context_storage.top_span_data():
+            return span_data
+        return self._context_storage.get_trace_data()
+
+    def _extract_lm_info_from_history(self, call_id: str) -> LMHistoryInfo:
+        """
+        Extract token usage, cache status, actual provider, and cost from the LM's history.
+
+        DSPy stores usage information in the LM's history after each call.
+        We verify the history entry matches our expected messages to handle
+        potential race conditions with concurrent LM calls.
+
+        For routers like OpenRouter, the response contains the actual provider
+        that served the request (e.g., "Novita", "Together"), which differs from
+        the router name used in the model string (e.g., "openrouter").
+
+        The cost field is provided by providers like OpenRouter and includes
+        accurate pricing for all token types (reasoning, cache, multimodal).
+
+        Returns:
+            LMHistoryInfo containing usage, cache_hit, actual_provider, and total_cost.
+        """
+        lm_info = self._map_call_id_to_lm_info.pop(call_id, None)
+        if lm_info is None:
+            return LMHistoryInfo(
+                usage=None, cache_hit=None, actual_provider=None, total_cost=None
+            )
 
-    def _get_span_type(self, instance: Any) -> types.SpanType:
-        if isinstance(instance, dspy.Predict):
-            return "llm"
-        elif isinstance(instance, dspy.LM):
-            return "llm"
-        return "general"
+        lm_instance, expected_messages = lm_info
+        return extract_lm_info_from_history(lm_instance, expected_messages)
+
+    def _get_opik_metadata(self, instance: Any) -> Dict[str, Any]:
+        graph = None
+        if self.log_graph and isinstance(instance, dspy.Module):
+            try:
+                graph = build_mermaid_graph_from_module(instance)
+            except Exception:
+                LOGGER.warning("Unable to generate graph from DSPy module")
+
+        if graph:
+            return {
+                **self._origins_metadata,
+                **{
+                    "_opik_graph_definition": {
+                        "format": "mermaid",
+                        "data": graph,
+                    }
+                },
+            }
+        else:
+            return self._origins_metadata

opik/integrations/dspy/graph.py

@@ -0,0 +1,88 @@
+from typing import Any, Dict, List
+import dspy
+import string
+
+STYLES = """
+classDef ReAct fill:#90EE90
+classDef Predict fill:#F08080
+classDef ChainOfThought fill:#ADD8E6
+classDef Tools fill:##D3D3D3
+"""
+
+
+def build_mermaid_graph_from_module(module: dspy.Module) -> str:
+    modules: Dict[str, Any] = {}
+    _get_dspy_module_heirarchy(module, modules)
+    graph = ""
+    queue = [modules]
+    states: Dict[str, Any] = {}
+    while queue:
+        current = queue.pop()
+        graph += _get_mermaid_arrows(queue, current, states)
+
+    if graph:
+        for key in states:
+            graph += "class %s %s\n" % (states[key]["state_name"], states[key]["class"])
+        graph += "class Tools Tools\n"
+
+        return "graph TD\n" + graph + STYLES
+    return ""
+
+
+def _get_mermaid_arrows(
+    queue: List[dspy.Module], current: Dict[str, Any], states: Dict[str, Any]
+) -> str:
+    start = _get_mermaid_state(current, states)
+    arrows = []
+    for module in current["sub_data"]:
+        end = _get_mermaid_state(module, states)
+        arrows.append("%s --> %s\n" % (start["text"], end["text"]))
+        queue.append(module)
+    if "tools" in current["details"]:
+        arrows.append("%s --> Tools\n" % (start["text"],))
+        arrows.append("subgraph Tools[<b>Tools</b>]\n")
+        for tool in current["details"]["tools"]:
+            arrows.append("    %s(<b>%s</b>)\n" % (tool, tool))
+        arrows.append("end\n")
+    return "".join(arrows)
+
+
+def _get_mermaid_state(
+    current: Dict[str, Any], states: Dict[str, Any]
+) -> Dict[str, Any]:
+    if current["id"] not in states:
+        state_name = string.ascii_uppercase[len(states)]
+        state_text = "<b>%s</b>" % current["name"]
+        if "instructions" in current["details"]:
+            state_text += "<br><i>%s</i>" % current["details"]["instructions"][
+                :100
+            ].replace("`", "'")
+        states[current["id"]] = {
+            "state_name": state_name,
+            "text": "%s(%s)" % (state_name, state_text),
+            "class": current["name"],
+        }
+    return states[current["id"]]
+
+
+def _get_dspy_module_heirarchy(module: dspy.Module, data: Dict[str, Any]) -> None:
+    data["name"] = module.__class__.__name__
+    data["id"] = id(module)
+    data["details"] = {}
+    if hasattr(module, "tools"):
+        data["details"]["tools"] = [name for name in module.tools]
+    if hasattr(module, "signature"):
+        data["details"]["instructions"] = module.signature.instructions
+    data["sub_data"] = []
+    for name, attribute in module.__dict__.items():
+        if isinstance(attribute, dspy.Module):
+            sub_data: Dict[str, Any] = {}
+            _get_dspy_module_heirarchy(attribute, sub_data)
+            data["sub_data"].append(sub_data)
+        elif name == "lm":
+            lm_data: Dict[str, Any] = {}
+            lm_data["name"] = "LM"
+            lm_data["id"] = id(attribute)
+            lm_data["details"] = {}
+            lm_data["sub_data"] = []
+            data["sub_data"].append(lm_data)