diff --git a/py/src/braintrust/integrations/dspy/test_dspy.py b/py/src/braintrust/integrations/dspy/test_dspy.py index 57d74e84b..348e1afad 100644 --- a/py/src/braintrust/integrations/dspy/test_dspy.py +++ b/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() diff --git a/py/src/braintrust/integrations/dspy/tracing.py b/py/src/braintrust/integrations/dspy/tracing.py index 2348493d4..3a1c23818 100644 --- a/py/src/braintrust/integrations/dspy/tracing.py +++ b/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,