From adb963acb1986dafd31cbb648b1eb735b8e4924b Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 19 Jan 2026 03:02:21 +0000
Subject: [PATCH 1/9] Initial plan


From 7f381e84be4f0e45a9b7b3ee4986c38fb8cdcfad Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 19 Jan 2026 03:09:11 +0000
Subject: [PATCH 2/9] Fix dspy.Parallel usage tracking inside dspy.Module

Modified _set_lm_usage() to handle lists of Predictions, not just single Prediction or tuple

Co-authored-by: TomeHirata <33407409+TomeHirata@users.noreply.github.com>
---
 dspy/primitives/module.py | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/dspy/primitives/module.py b/dspy/primitives/module.py
index 94d48cae63..413fdd6728 100644
--- a/dspy/primitives/module.py
+++ b/dspy/primitives/module.py
@@ -186,13 +186,17 @@ def _set_lm_usage(self, tokens: dict[str, Any], output: Any):
         # module.forward to return a tuple: (prediction, trace).
         # When usage tracking is enabled, ensure we attach usage to the
         # prediction object if present.
-        prediction_in_output = None
+        predictions_in_output = []
         if isinstance(output, Prediction):
-            prediction_in_output = output
+            predictions_in_output = [output]
         elif isinstance(output, tuple) and len(output) > 0 and isinstance(output[0], Prediction):
-            prediction_in_output = output[0]
-        if prediction_in_output:
-            prediction_in_output.set_lm_usage(tokens)
+            predictions_in_output = [output[0]]
+        elif isinstance(output, list) and all(isinstance(item, Prediction) for item in output):
+            predictions_in_output = output
+        
+        if predictions_in_output:
+            for prediction in predictions_in_output:
+                prediction.set_lm_usage(tokens)
         else:
             logger.warning("Failed to set LM usage. Please return `dspy.Prediction` object from dspy.Module to enable usage tracking.")
 

From 0fd3fcc03c64c12dbe2cc38b865a488c7910bd43 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 19 Jan 2026 03:10:50 +0000
Subject: [PATCH 3/9] Add test for parallel usage tracking fix

Co-authored-by: TomeHirata <33407409+TomeHirata@users.noreply.github.com>
---
 tests/predict/test_parallel.py | 34 ++++++++++++++++++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/tests/predict/test_parallel.py b/tests/predict/test_parallel.py
index f8524d4d25..458e610fc7 100644
--- a/tests/predict/test_parallel.py
+++ b/tests/predict/test_parallel.py
@@ -194,3 +194,37 @@ def forward(self, value: int) -> str:
     assert len(exceptions) == 1
     assert isinstance(exceptions[0], ValueError)
     assert str(exceptions[0]) == "test error"
+
+
+def test_parallel_with_track_usage():
+    """Test that usage tracking works correctly when using Parallel inside a Module."""
+    lm = DummyLM(
+        [
+            {"output": "test output 1"},
+            {"output": "test output 2"},
+        ]
+    )
+
+    class MyModule(dspy.Module):
+        def forward(self):
+            predictor = dspy.Predict("input -> output")
+            parallel = dspy.Parallel(num_threads=2)
+            return parallel(
+                [
+                    (predictor, {"input": "test input 1"}),
+                    (predictor, {"input": "test input 2"}),
+                ]
+            )
+
+    # Test with track_usage enabled
+    with dspy.settings.context(lm=lm, track_usage=True):
+        results = MyModule()()
+
+    # Verify that results are Prediction objects with usage data set
+    assert len(results) == 2
+    assert all(isinstance(r, dspy.Prediction) for r in results)
+    assert all(r.get_lm_usage() is not None for r in results)
+
+    # Verify predictions have the expected outputs
+    expected_outputs = {"test output 1", "test output 2"}
+    assert {r.output for r in results} == expected_outputs

From db542516252f3665505adc8d8c88bd4104c7bbc5 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 19 Jan 2026 03:12:20 +0000
Subject: [PATCH 4/9] Handle tuple with list of Predictions in usage tracking

Extended _set_lm_usage to handle tuple containing list of Predictions, such as ([Prediction, ...], trace)

Co-authored-by: TomeHirata <33407409+TomeHirata@users.noreply.github.com>
---
 dspy/primitives/module.py      |  9 +++++--
 tests/predict/test_parallel.py | 44 ++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/dspy/primitives/module.py b/dspy/primitives/module.py
index 413fdd6728..2bbecc2e49 100644
--- a/dspy/primitives/module.py
+++ b/dspy/primitives/module.py
@@ -189,8 +189,13 @@ def _set_lm_usage(self, tokens: dict[str, Any], output: Any):
         predictions_in_output = []
         if isinstance(output, Prediction):
             predictions_in_output = [output]
-        elif isinstance(output, tuple) and len(output) > 0 and isinstance(output[0], Prediction):
-            predictions_in_output = [output[0]]
+        elif isinstance(output, tuple) and len(output) > 0:
+            # Handle tuple cases: (Prediction, trace) or ([Prediction, ...], trace)
+            first_element = output[0]
+            if isinstance(first_element, Prediction):
+                predictions_in_output = [first_element]
+            elif isinstance(first_element, list) and all(isinstance(item, Prediction) for item in first_element):
+                predictions_in_output = first_element
         elif isinstance(output, list) and all(isinstance(item, Prediction) for item in output):
             predictions_in_output = output
         
diff --git a/tests/predict/test_parallel.py b/tests/predict/test_parallel.py
index 458e610fc7..01c211187a 100644
--- a/tests/predict/test_parallel.py
+++ b/tests/predict/test_parallel.py
@@ -228,3 +228,47 @@ def forward(self):
     # Verify predictions have the expected outputs
     expected_outputs = {"test output 1", "test output 2"}
     assert {r.output for r in results} == expected_outputs
+
+
+def test_parallel_with_track_usage_and_tuple_return():
+    """Test that usage tracking works when Module returns a tuple with list of Predictions."""
+    lm = DummyLM(
+        [
+            {"output": "test output 1"},
+            {"output": "test output 2"},
+        ]
+    )
+
+    class MyModule(dspy.Module):
+        def forward(self):
+            predictor = dspy.Predict("input -> output")
+            parallel = dspy.Parallel(num_threads=2)
+            results = parallel(
+                [
+                    (predictor, {"input": "test input 1"}),
+                    (predictor, {"input": "test input 2"}),
+                ]
+            )
+            # Return tuple mimicking optimizer behavior (e.g., GEPA bootstrap tracing)
+            return (results, {"trace": "some_trace_data"})
+
+    # Test with track_usage enabled
+    with dspy.settings.context(lm=lm, track_usage=True):
+        output = MyModule()()
+
+    # Verify that output is a tuple
+    assert isinstance(output, tuple)
+    assert len(output) == 2
+    results, trace = output
+
+    # Verify that results are Prediction objects with usage data set
+    assert len(results) == 2
+    assert all(isinstance(r, dspy.Prediction) for r in results)
+    assert all(r.get_lm_usage() is not None for r in results)
+
+    # Verify predictions have the expected outputs
+    expected_outputs = {"test output 1", "test output 2"}
+    assert {r.output for r in results} == expected_outputs
+    
+    # Verify trace is preserved
+    assert trace == {"trace": "some_trace_data"}

From df5efbf99efbb1d583638e529b2ed6402469ba89 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 19 Jan 2026 03:14:29 +0000
Subject: [PATCH 5/9] Add non-empty list checks in _set_lm_usage

Added checks to ensure lists are non-empty before iterating with all(), improving performance

Co-authored-by: TomeHirata <33407409+TomeHirata@users.noreply.github.com>
---
 dspy/primitives/module.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/dspy/primitives/module.py b/dspy/primitives/module.py
index 2bbecc2e49..f9fd3962a3 100644
--- a/dspy/primitives/module.py
+++ b/dspy/primitives/module.py
@@ -194,9 +194,9 @@ def _set_lm_usage(self, tokens: dict[str, Any], output: Any):
             first_element = output[0]
             if isinstance(first_element, Prediction):
                 predictions_in_output = [first_element]
-            elif isinstance(first_element, list) and all(isinstance(item, Prediction) for item in first_element):
+            elif isinstance(first_element, list) and first_element and all(isinstance(item, Prediction) for item in first_element):
                 predictions_in_output = first_element
-        elif isinstance(output, list) and all(isinstance(item, Prediction) for item in output):
+        elif isinstance(output, list) and output and all(isinstance(item, Prediction) for item in output):
             predictions_in_output = output
         
         if predictions_in_output:

From adc30e3b2545cc42f3bed06e96573b18875cc892 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 19 Jan 2026 15:24:54 +0000
Subject: [PATCH 6/9] Refactor _set_lm_usage to use recursive collection

Changed from type-specific checks to a recursive function that finds all Prediction objects in any structure (lists, tuples, dicts, nested combinations). This makes the implementation more robust and future-proof for user-defined return types.

Co-authored-by: TomeHirata <33407409+TomeHirata@users.noreply.github.com>
---
 dspy/primitives/module.py      | 43 ++++++++++++++++++++++------------
 tests/predict/test_parallel.py | 36 ++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 15 deletions(-)

diff --git a/dspy/primitives/module.py b/dspy/primitives/module.py
index f9fd3962a3..50d141c0fd 100644
--- a/dspy/primitives/module.py
+++ b/dspy/primitives/module.py
@@ -182,22 +182,35 @@ def batch(
             return results
 
     def _set_lm_usage(self, tokens: dict[str, Any], output: Any):
-        # Some optimizers (e.g., GEPA bootstrap tracing) temporarily patch
-        # module.forward to return a tuple: (prediction, trace).
-        # When usage tracking is enabled, ensure we attach usage to the
-        # prediction object if present.
+        # When usage tracking is enabled, recursively find all Prediction objects
+        # in the output and set usage on them. This handles any structure:
+        # - Single Prediction
+        # - List/tuple of Predictions
+        # - Nested structures (e.g., [[Prediction, ...], ...])
+        # - Tuples from optimizers (e.g., (prediction, trace))
         predictions_in_output = []
-        if isinstance(output, Prediction):
-            predictions_in_output = [output]
-        elif isinstance(output, tuple) and len(output) > 0:
-            # Handle tuple cases: (Prediction, trace) or ([Prediction, ...], trace)
-            first_element = output[0]
-            if isinstance(first_element, Prediction):
-                predictions_in_output = [first_element]
-            elif isinstance(first_element, list) and first_element and all(isinstance(item, Prediction) for item in first_element):
-                predictions_in_output = first_element
-        elif isinstance(output, list) and output and all(isinstance(item, Prediction) for item in output):
-            predictions_in_output = output
+        
+        def collect_predictions(obj, seen=None):
+            """Recursively collect all Prediction objects from any structure."""
+            if seen is None:
+                seen = set()
+            
+            # Avoid infinite recursion on circular references
+            obj_id = id(obj)
+            if obj_id in seen:
+                return
+            seen.add(obj_id)
+            
+            if isinstance(obj, Prediction):
+                predictions_in_output.append(obj)
+            elif isinstance(obj, (list, tuple)):
+                for item in obj:
+                    collect_predictions(item, seen)
+            elif isinstance(obj, dict):
+                for value in obj.values():
+                    collect_predictions(value, seen)
+        
+        collect_predictions(output)
         
         if predictions_in_output:
             for prediction in predictions_in_output:
diff --git a/tests/predict/test_parallel.py b/tests/predict/test_parallel.py
index 01c211187a..1df0a55395 100644
--- a/tests/predict/test_parallel.py
+++ b/tests/predict/test_parallel.py
@@ -272,3 +272,39 @@ def forward(self):
     
     # Verify trace is preserved
     assert trace == {"trace": "some_trace_data"}
+
+
+def test_track_usage_with_nested_and_dict_structures():
+    """Test that usage tracking works with nested lists, dicts, and mixed structures."""
+    lm = DummyLM([{"output": "test1"}, {"output": "test2"}, {"output": "test3"}])
+
+    class NestedStructureModule(dspy.Module):
+        def forward(self):
+            predictor = dspy.Predict("input -> output")
+            p1 = predictor(input="input1")
+            p2 = predictor(input="input2")
+            p3 = predictor(input="input3")
+            # Return a complex nested structure
+            return {
+                "nested_list": [[p1, p2]],
+                "single": p3,
+            }
+
+    # Test with track_usage enabled
+    with dspy.settings.context(lm=lm, track_usage=True):
+        result = NestedStructureModule()()
+
+    # Verify structure
+    assert isinstance(result, dict)
+    assert "nested_list" in result
+    assert "single" in result
+
+    # Verify all predictions have usage data
+    all_predictions = result["nested_list"][0] + [result["single"]]
+    assert len(all_predictions) == 3
+    assert all(isinstance(p, dspy.Prediction) for p in all_predictions)
+    assert all(p.get_lm_usage() is not None for p in all_predictions)
+
+    # Verify predictions have the expected outputs
+    expected_outputs = {"test1", "test2", "test3"}
+    assert {p.output for p in all_predictions} == expected_outputs

From d4614e61678ce509d8cb72d60184fca2ab00c721 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 19 Jan 2026 15:26:29 +0000
Subject: [PATCH 7/9] Improve documentation for collect_predictions function

Added detailed docstring explaining parameters, behavior, and side effects. Clarified circular reference prevention logic.

Co-authored-by: TomeHirata <33407409+TomeHirata@users.noreply.github.com>
---
 dspy/primitives/module.py | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/dspy/primitives/module.py b/dspy/primitives/module.py
index 50d141c0fd..aa464a864e 100644
--- a/dspy/primitives/module.py
+++ b/dspy/primitives/module.py
@@ -188,10 +188,20 @@ def _set_lm_usage(self, tokens: dict[str, Any], output: Any):
         # - List/tuple of Predictions
         # - Nested structures (e.g., [[Prediction, ...], ...])
         # - Tuples from optimizers (e.g., (prediction, trace))
+        # - Dicts with Prediction values
+        # - Any combination of the above
         predictions_in_output = []
         
         def collect_predictions(obj, seen=None):
-            """Recursively collect all Prediction objects from any structure."""
+            """Recursively collect all Prediction objects from any structure.
+            
+            Args:
+                obj: The object to search for Predictions (can be any type)
+                seen: Set of object ids already visited (prevents infinite recursion on circular refs)
+            
+            Side effects:
+                Appends found Prediction objects to predictions_in_output list in outer scope
+            """
             if seen is None:
                 seen = set()
             

From e8c0ee00c461416f417bee3077ce338a0d7dd0c1 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 19 Jan 2026 15:34:41 +0000
Subject: [PATCH 8/9] Fix nested module usage tracking in Module.__call__

Modified Module.__call__ to always call _set_lm_usage when track_usage is enabled, even for nested module calls. Previously, only the root module call would set usage. Now nested modules also get their usage set by computing the delta between token counts before and after the forward call.

Added _compute_token_delta helper method to calculate the difference in token usage.
Added test for nested module usage tracking.

Co-authored-by: TomeHirata <33407409+TomeHirata@users.noreply.github.com>
---
 dspy/primitives/module.py      | 49 ++++++++++++++++++++++++++++++++--
 tests/predict/test_parallel.py | 39 +++++++++++++++++++++++++++
 2 files changed, 86 insertions(+), 2 deletions(-)

diff --git a/dspy/primitives/module.py b/dspy/primitives/module.py
index aa464a864e..bf59cf7c6e 100644
--- a/dspy/primitives/module.py
+++ b/dspy/primitives/module.py
@@ -72,14 +72,38 @@ def __call__(self, *args, **kwargs) -> Prediction:
 
         with settings.context(caller_modules=caller_modules):
             if settings.track_usage and thread_local_overrides.get().get("usage_tracker") is None:
+                # No existing tracker - create a new one for this module call
                 with track_usage() as usage_tracker:
                     output = self.forward(*args, **kwargs)
                 tokens = usage_tracker.get_total_tokens()
                 self._set_lm_usage(tokens, output)
 
                 return output
-
-            return self.forward(*args, **kwargs)
+            elif settings.track_usage:
+                # Tracker exists (nested call) - capture usage for this specific module call
+                usage_tracker = thread_local_overrides.get().get("usage_tracker")
+                if usage_tracker:
+                    # Get current token count before calling forward
+                    tokens_before = usage_tracker.get_total_tokens()
+                    output = self.forward(*args, **kwargs)
+                    # Get tokens used by this module call
+                    tokens_after = usage_tracker.get_total_tokens()
+                    
+                    # Calculate the delta (tokens used by this module call only)
+                    tokens_delta = {}
+                    for lm_name in tokens_after:
+                        if lm_name in tokens_before:
+                            # Compute difference for each token type
+                            tokens_delta[lm_name] = self._compute_token_delta(tokens_before[lm_name], tokens_after[lm_name])
+                        else:
+                            tokens_delta[lm_name] = tokens_after[lm_name]
+                    
+                    self._set_lm_usage(tokens_delta, output)
+                    return output
+                else:
+                    return self.forward(*args, **kwargs)
+            else:
+                return self.forward(*args, **kwargs)
 
     @with_callbacks
     async def acall(self, *args, **kwargs) -> Prediction:
@@ -181,6 +205,27 @@ def batch(
             results = parallel_executor.forward(exec_pairs)
             return results
 
+    def _compute_token_delta(self, tokens_before: dict[str, Any], tokens_after: dict[str, Any]) -> dict[str, Any]:
+        """Compute the difference in token usage between two snapshots.
+        
+        Args:
+            tokens_before: Token usage before the operation
+            tokens_after: Token usage after the operation
+            
+        Returns:
+            Dictionary with the delta (difference) in token usage
+        """
+        delta = {}
+        for key, value_after in tokens_after.items():
+            value_before = tokens_before.get(key, 0)
+            if isinstance(value_after, dict) and isinstance(value_before, dict):
+                # Recursively compute delta for nested dicts
+                delta[key] = self._compute_token_delta(value_before, value_after)
+            elif value_after is not None or value_before is not None:
+                # Compute numeric difference
+                delta[key] = (value_after or 0) - (value_before or 0)
+        return delta
+
     def _set_lm_usage(self, tokens: dict[str, Any], output: Any):
         # When usage tracking is enabled, recursively find all Prediction objects
         # in the output and set usage on them. This handles any structure:
diff --git a/tests/predict/test_parallel.py b/tests/predict/test_parallel.py
index 1df0a55395..ac90eabe52 100644
--- a/tests/predict/test_parallel.py
+++ b/tests/predict/test_parallel.py
@@ -308,3 +308,42 @@ def forward(self):
     # Verify predictions have the expected outputs
     expected_outputs = {"test1", "test2", "test3"}
     assert {p.output for p in all_predictions} == expected_outputs
+
+
+def test_track_usage_with_nested_modules():
+    """Test that usage tracking works correctly for nested module calls."""
+    lm = DummyLM([{"output": "inner1"}, {"output": "inner2"}])
+
+    class InnerModule(dspy.Module):
+        def forward(self, input):
+            predictor = dspy.Predict("input -> output")
+            result = predictor(input=input)
+            # Verify usage is set inside the nested module
+            assert result.get_lm_usage() is not None, "Nested module should have usage set"
+            return result
+
+    class OuterModule(dspy.Module):
+        def __init__(self):
+            super().__init__()
+            self.inner = InnerModule()
+
+        def forward(self):
+            r1 = self.inner(input="test1")
+            r2 = self.inner(input="test2")
+            # Verify usage is set on nested module results
+            assert r1.get_lm_usage() is not None, "Nested call result should have usage set"
+            assert r2.get_lm_usage() is not None, "Nested call result should have usage set"
+            return [r1, r2]
+
+    # Test with track_usage enabled
+    with dspy.settings.context(lm=lm, track_usage=True):
+        results = OuterModule()()
+
+    # Verify all results have usage data
+    assert len(results) == 2
+    assert all(isinstance(r, dspy.Prediction) for r in results)
+    assert all(r.get_lm_usage() is not None for r in results)
+
+    # Verify predictions have the expected outputs
+    expected_outputs = {"inner1", "inner2"}
+    assert {r.output for r in results} == expected_outputs

From 10966b260d1a3e00c6a9213005323f2df494e5bb Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 19 Jan 2026 15:37:14 +0000
Subject: [PATCH 9/9] Improve robustness of Module.__call__ and
 _compute_token_delta

- Removed redundant usage_tracker None check in nested call branch
- Handle all LM names from both before and after token snapshots in delta calculation
- Improved _compute_token_delta to handle non-numeric values and process all keys from both dicts
- Added type checking to prevent TypeError on non-numeric types

Co-authored-by: TomeHirata <33407409+TomeHirata@users.noreply.github.com>
---
 dspy/primitives/module.py | 63 +++++++++++++++++++++++----------------
 1 file changed, 37 insertions(+), 26 deletions(-)

diff --git a/dspy/primitives/module.py b/dspy/primitives/module.py
index bf59cf7c6e..3202ca4c6e 100644
--- a/dspy/primitives/module.py
+++ b/dspy/primitives/module.py
@@ -82,26 +82,27 @@ def __call__(self, *args, **kwargs) -> Prediction:
             elif settings.track_usage:
                 # Tracker exists (nested call) - capture usage for this specific module call
                 usage_tracker = thread_local_overrides.get().get("usage_tracker")
-                if usage_tracker:
-                    # Get current token count before calling forward
-                    tokens_before = usage_tracker.get_total_tokens()
-                    output = self.forward(*args, **kwargs)
-                    # Get tokens used by this module call
-                    tokens_after = usage_tracker.get_total_tokens()
-                    
-                    # Calculate the delta (tokens used by this module call only)
-                    tokens_delta = {}
-                    for lm_name in tokens_after:
-                        if lm_name in tokens_before:
-                            # Compute difference for each token type
-                            tokens_delta[lm_name] = self._compute_token_delta(tokens_before[lm_name], tokens_after[lm_name])
-                        else:
-                            tokens_delta[lm_name] = tokens_after[lm_name]
-                    
-                    self._set_lm_usage(tokens_delta, output)
-                    return output
-                else:
-                    return self.forward(*args, **kwargs)
+                # Get current token count before calling forward
+                tokens_before = usage_tracker.get_total_tokens()
+                output = self.forward(*args, **kwargs)
+                # Get tokens used by this module call
+                tokens_after = usage_tracker.get_total_tokens()
+                
+                # Calculate the delta (tokens used by this module call only)
+                tokens_delta = {}
+                # Process all LM names from both before and after
+                all_lm_names = set(tokens_before.keys()) | set(tokens_after.keys())
+                for lm_name in all_lm_names:
+                    if lm_name in tokens_before and lm_name in tokens_after:
+                        # Compute difference for each token type
+                        tokens_delta[lm_name] = self._compute_token_delta(tokens_before[lm_name], tokens_after[lm_name])
+                    elif lm_name in tokens_after:
+                        # New LM used in this call
+                        tokens_delta[lm_name] = tokens_after[lm_name]
+                    # If lm_name only in tokens_before, it means no usage in this call, so skip
+                
+                self._set_lm_usage(tokens_delta, output)
+                return output
             else:
                 return self.forward(*args, **kwargs)
 
@@ -216,14 +217,24 @@ def _compute_token_delta(self, tokens_before: dict[str, Any], tokens_after: dict
             Dictionary with the delta (difference) in token usage
         """
         delta = {}
-        for key, value_after in tokens_after.items():
-            value_before = tokens_before.get(key, 0)
-            if isinstance(value_after, dict) and isinstance(value_before, dict):
+        # Process all keys from both dicts
+        all_keys = set(tokens_before.keys()) | set(tokens_after.keys())
+        for key in all_keys:
+            value_before = tokens_before.get(key)
+            value_after = tokens_after.get(key)
+            
+            if isinstance(value_after, dict) or isinstance(value_before, dict):
                 # Recursively compute delta for nested dicts
-                delta[key] = self._compute_token_delta(value_before, value_after)
-            elif value_after is not None or value_before is not None:
-                # Compute numeric difference
+                delta[key] = self._compute_token_delta(
+                    value_before if isinstance(value_before, dict) else {},
+                    value_after if isinstance(value_after, dict) else {}
+                )
+            elif isinstance(value_after, (int, float)) or isinstance(value_before, (int, float)):
+                # Compute numeric difference only for numeric types
                 delta[key] = (value_after or 0) - (value_before or 0)
+            elif value_after is not None:
+                # For non-numeric types (strings, etc.), just use the after value
+                delta[key] = value_after
         return delta
 
     def _set_lm_usage(self, tokens: dict[str, Any], output: Any):