fix(dspy): record LM history token metrics

py/src/braintrust/integrations/dspy/test_dspy.py

@@ -3,6 +3,8 @@
 """
 
 import inspect
+import os
+from types import SimpleNamespace
 
 import dspy
 import pytest
@@ -33,7 +35,7 @@ def test_dspy_callback(memory_logger):
     assert not memory_logger.pop()
 
     # Configure DSPy with Braintrust callback
-    lm = dspy.LM(MODEL, **LM_KWARGS)
+    lm = dspy.LM(MODEL, cache=False, **LM_KWARGS)
     dspy.configure(lm=lm, callbacks=[BraintrustDSpyCallback()])
 
     # Use ChainOfThought for a more interesting test
@@ -61,6 +63,12 @@ def test_dspy_callback(memory_logger):
     assert lm_span["metadata"]["provider"] == "openai"
     assert "input" in lm_span
     assert "output" in lm_span
+    # The latest cassette contains a real OpenAI usage payload. DSPy 2.6's
+    # recorded response predates usage being included.
+    if os.environ.get("BRAINTRUST_TEST_PACKAGE_VERSION") == "latest":
+        assert lm_span["metrics"]["prompt_tokens"] == 170
+        assert lm_span["metrics"]["completion_tokens"] == 53
+        assert lm_span["metrics"]["tokens"] == 223
 
     format_span = spans_by_name["dspy.adapter.format"]
     assert format_span["span_attributes"]["type"] == "task"
@@ -87,6 +95,54 @@ def test_dspy_callback(memory_logger):
     assert parse_span["span_parents"][0] in span_ids
 
 
+def test_dspy_callback_ignores_usage_from_previous_lm_call(memory_logger):
+    """A failed call must not inherit usage from an older history entry."""
+    instance = SimpleNamespace(
+        model=MODEL,
+        history=[
+            {
+                "prompt": "previous request",
+                "outputs": ["previous response"],
+                "usage": {"prompt_tokens": 20, "completion_tokens": 4, "total_tokens": 24},
+            }
+        ],
+    )
+    callback = BraintrustDSpyCallback()
+    callback.on_lm_start("failed-call", instance, {"prompt": "new request"})
+    callback.on_lm_end("failed-call", None, RuntimeError("provider failed"))
+
+    span = next(span for span in memory_logger.pop() if span["span_attributes"]["name"] == "dspy.lm")
+    assert not {"prompt_tokens", "completion_tokens", "tokens"} & span.get("metrics", {}).keys()
+
+
+def test_dspy_callback_correlates_concurrent_lm_usage(memory_logger):
+    """Concurrent calls on one LM must read their own new history entries."""
+    instance = SimpleNamespace(model=MODEL, history=[])
+    callback = BraintrustDSpyCallback()
+    callback.on_lm_start("call-a", instance, {"prompt": "same request"})
+    callback.on_lm_start("call-b", instance, {"prompt": "same request"})
+
+    outputs_a = ["same response"]
+    outputs_b = ["same response"]
+    entry_a = {
+        "prompt": "same request",
+        "outputs": outputs_a,
+        "usage": {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12},
+    }
+    entry_b = {
+        "prompt": "same request",
+        "outputs": outputs_b,
+        "usage": {"prompt_tokens": 30, "completion_tokens": 5, "total_tokens": 35},
+    }
+    # Both provider calls can append history before either end callback runs.
+    instance.history.extend([entry_a, entry_b])
+    callback.on_lm_end("call-a", outputs_a)
+    callback.on_lm_end("call-b", outputs_b)
+
+    spans = [span for span in memory_logger.pop() if span["span_attributes"]["name"] == "dspy.lm"]
+    assert [span["metrics"]["tokens"] for span in spans] == [12, 35]
+
+
 def test_dspy_adapter_callbacks(memory_logger):
     """Adapter format/parse callbacks should log spans without an LM call."""
     assert not memory_logger.pop()

py/src/braintrust/integrations/dspy/tracing.py

@@ -68,8 +68,8 @@ class BraintrustDSpyCallback(BaseCallback):
         ```
 
     Advanced Example with LiteLLM Patching:
-        For additional detailed token metrics from LiteLLM's wrapper, patch before importing DSPy
-        and disable DSPy's disk cache:
+        To capture additional provider-level spans from LiteLLM, patch before importing DSPy and
+        disable DSPy's disk cache:
 
         ```python
         from braintrust.integrations.litellm import patch_litellm
@@ -95,8 +95,8 @@ class BraintrustDSpyCallback(BaseCallback):
     - Tool calls
     - Evaluation runs
 
-    For detailed token usage and cost metrics, use LiteLLM patching (see Advanced Example above).
-    The patched LiteLLM wrapper will create additional "Completion" spans with comprehensive metrics.
+    Token usage metrics are read from the DSPy LM's history. LiteLLM patching (see Advanced Example
+    above) can additionally create "Completion" spans with provider-level details.
 
     Spans are automatically nested based on the execution hierarchy.
     """
@@ -108,6 +108,7 @@ def __init__(self):
         super().__init__()
         # Map call_id to span objects for proper nesting
         self._spans: dict[str, Any] = {}
+        self._lm_calls: dict[str, tuple[Any, dict[str, Any], Any | None]] = {}
 
     def on_lm_start(
         self,
@@ -141,12 +142,19 @@ def on_lm_start(
         )
         span.set_current()
         self._spans[call_id] = span
+        try:
+            history = getattr(instance, "history", None)
+            last_history_entry = history[-1] if isinstance(history, list) and history else None
+        except Exception:
+            last_history_entry = None
+        self._lm_calls[call_id] = (instance, dict(inputs), last_history_entry)
 
     def _end_span(
         self,
         call_id: str,
         outputs: Any | None,
         exception: Exception | None = None,
+        metrics: dict[str, int] | None = None,
     ):
         """Pop span by call_id, log outputs/exception, and end it."""
         span = self._spans.pop(call_id, None)
@@ -159,6 +167,8 @@ def _end_span(
                 log_data["error"] = exception
             if outputs is not None:
                 log_data["output"] = outputs
+            if metrics:
+                log_data["metrics"] = metrics
 
             if log_data:
                 span.log(**log_data)
@@ -179,7 +189,48 @@ def on_lm_end(
             outputs: Output from the LM, or None if there was an exception
             exception: Exception raised during execution, if any
         """
-        self._end_span(call_id, outputs, exception)
+        call = self._lm_calls.pop(call_id, None)
+        metrics: dict[str, int] = {}
+        try:
+            if call is not None and outputs is not None:
+                instance, inputs, last_history_entry = call
+                history = getattr(instance, "history", None)
+                if isinstance(history, list):
+                    if last_history_entry is None:
+                        new_entries = history
+                    else:
+                        marker_index = next(
+                            (index for index, entry in enumerate(history) if entry is last_history_entry), None
+                        )
+                        new_entries = history[marker_index + 1 :] if marker_index is not None else []
+
+                    input_matches = [
+                        entry
+                        for entry in new_entries
+                        if isinstance(entry, dict)
+                        and entry.get("prompt") == inputs.get("prompt")
+                        and entry.get("messages") == inputs.get("messages")
+                    ]
+                    # DSPy stores the same outputs object it returns to on_lm_end. Identity
+                    # disambiguates concurrent calls with equal input and output values.
+                    matching_entries = [entry for entry in input_matches if entry.get("outputs") is outputs]
+                    if not matching_entries:
+                        matching_entries = [entry for entry in input_matches if entry.get("outputs") == outputs]
+                    usage = matching_entries[0].get("usage") if len(matching_entries) == 1 else None
+                    if isinstance(usage, dict):
+                        for source, target in (
+                            ("prompt_tokens", "prompt_tokens"),
+                            ("completion_tokens", "completion_tokens"),
+                            ("total_tokens", "tokens"),
+                        ):
+                            value = usage.get(source)
+                            if isinstance(value, int) and not isinstance(value, bool) and value >= 0:
+                                metrics[target] = value
+        except Exception:
+            # Usage extraction is best-effort and must not affect the DSPy call.
+            pass
+
+        self._end_span(call_id, outputs, exception, metrics or None)
 
     def on_module_start(
         self,