Add usage stats to tool calls and handle multiple tool calls correctly (#1899)

87329c80ef9d9ad180406b30c8e3f457de06700e · 2026-04-15 19:40:00 +0100 · rltakashige

## Motivation

Tool calls are usually not end tokens, so they didn't have usage stats.

Files touched

M src/exo/worker/runner/llm_inference/model_output_parsers.py
M src/exo/worker/tests/unittests/test_runner/test_parse_tool_calls.py

Diff

commit 87329c80ef9d9ad180406b30c8e3f457de06700e
Author: rltakashige <rl.takashige@gmail.com>
Date:   Wed Apr 15 19:40:00 2026 +0100

    Add usage stats to tool calls and handle multiple tool calls correctly (#1899)
    
    ## Motivation
    
    Tool calls are usually not end tokens, so they didn't have usage stats.
---
 .../runner/llm_inference/model_output_parsers.py   | 31 +++++++++++++++++++---
 .../unittests/test_runner/test_parse_tool_calls.py | 11 +++-----
 2 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/src/exo/worker/runner/llm_inference/model_output_parsers.py b/src/exo/worker/runner/llm_inference/model_output_parsers.py
index a5cf19e6..27a24db0 100644
--- a/src/exo/worker/runner/llm_inference/model_output_parsers.py
+++ b/src/exo/worker/runner/llm_inference/model_output_parsers.py
@@ -379,6 +379,8 @@ def parse_tool_calls(
 ) -> Generator[GenerationResponse | ToolCallResponse | None]:
     in_tool_call = False
     tool_call_text_parts: list[str] = []
+    accumulated_tool_calls: list[ToolCallItem] = []
+
     for response in responses:
         if response is None:
             yield None
@@ -387,6 +389,19 @@ def parse_tool_calls(
         if not in_tool_call and response.text.startswith(tool_parser.start_parsing):
             in_tool_call = True
 
+        if (
+            not in_tool_call
+            and accumulated_tool_calls
+            and (response.stats is not None or response.finish_reason is not None)
+        ):
+            yield ToolCallResponse(
+                tool_calls=accumulated_tool_calls,
+                usage=response.usage,
+                stats=response.stats,
+            )
+            accumulated_tool_calls.clear()
+            continue
+
         if not in_tool_call:
             yield response
             continue
@@ -407,9 +422,16 @@ def parse_tool_calls(
                 )
                 break
 
-            yield ToolCallResponse(
-                tool_calls=parsed, usage=response.usage, stats=response.stats
-            )
+            accumulated_tool_calls.extend(parsed)
+            if accumulated_tool_calls and (
+                response.finish_reason is not None or response.stats is not None
+            ):
+                yield ToolCallResponse(
+                    tool_calls=accumulated_tool_calls,
+                    usage=response.usage,
+                    stats=response.stats,
+                )
+                accumulated_tool_calls.clear()
             continue
 
         if response.finish_reason is not None:
@@ -424,3 +446,6 @@ def parse_tool_calls(
                 }
             )
             yield response
+
+    if not accumulated_tool_calls:
+        logger.warning("Tool calls should have all been emitted but were not")
diff --git a/src/exo/worker/tests/unittests/test_runner/test_parse_tool_calls.py b/src/exo/worker/tests/unittests/test_runner/test_parse_tool_calls.py
index 49655f03..d69fbe09 100644
--- a/src/exo/worker/tests/unittests/test_runner/test_parse_tool_calls.py
+++ b/src/exo/worker/tests/unittests/test_runner/test_parse_tool_calls.py
@@ -9,17 +9,14 @@ from exo.worker.runner.llm_inference.model_output_parsers import parse_tool_call
 from exo.worker.runner.llm_inference.tool_parsers import make_mlx_parser
 
 
-def _make_responses(
-    texts: list[str],
-    finish_on_last: bool = True,
-) -> Generator[GenerationResponse]:
+def _make_responses(texts: list[str]) -> Generator[GenerationResponse]:
     """Create a sequence of GenerationResponses from text strings."""
     for i, text in enumerate(texts):
         is_last = i == len(texts) - 1
         yield GenerationResponse(
             text=text,
             token=i,
-            finish_reason="stop" if (is_last and finish_on_last) else None,
+            finish_reason="stop" if is_last else None,
             usage=None,
         )
 
@@ -39,7 +36,7 @@ class TestParseToolCalls:
         texts = ["<tool_call>", "test_fn", "</tool_call>"]
         results = list(
             parse_tool_calls(
-                _make_responses(texts, finish_on_last=False),
+                _make_responses(texts),
                 _dummy_parser,
                 tools=None,
             )
@@ -78,7 +75,7 @@ class TestParseToolCalls:
         texts = ["<tool_call>", "bad content", "</tool_call>"]
         results = list(
             parse_tool_calls(
-                _make_responses(texts, finish_on_last=False),
+                _make_responses(texts),
                 make_mlx_parser("<tool_call>", "</tool_call>", _failing_parser),
                 tools=None,
             )

← 8cdc8338 Drain tokens silently skipped in thinking parsing (#1898) · back to Exo · Better environment variables in MacOS app (#1901) 3eead802 →