From 9cd139c3178c3cf7e344fd00ec184a1f82cdb7a4 Mon Sep 17 00:00:00 2001 From: Declan Brady Date: Wed, 19 Nov 2025 16:57:39 -0500 Subject: [PATCH 1/2] Add concat functionality for usage --- src/agentex/lib/utils/completions.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/src/agentex/lib/utils/completions.py b/src/agentex/lib/utils/completions.py index f33495c1..24ef2c1f 100644 --- a/src/agentex/lib/utils/completions.py +++ b/src/agentex/lib/utils/completions.py @@ -9,6 +9,7 @@ Choice, ToolCall, Completion, + Usage, ToolCallRequest, ) @@ -21,6 +22,8 @@ def _concat_chunks(_a: None, b: Any): @_concat_chunks.register def _(a: Completion, b: Completion) -> Completion: a.choices = [_concat_chunks(*c) for c in zip(a.choices, b.choices, strict=False)] + a.usage = _concat_chunks(a.usage, b.usage) + return a @@ -35,6 +38,17 @@ def _(a: Choice, b: Choice) -> Choice: a.finish_reason = a.finish_reason or b.finish_reason return a +@_concat_chunks.register +def _(a: Usage | None, b: Usage | None) -> Usage | None: + if a is not None and b is not None: + return Usage( + prompt_tokens=a.prompt_tokens + b.prompt_tokens, + completion_tokens=a.completion_tokens + b.completion_tokens, + total_tokens=a.total_tokens + b.total_tokens, + ) + else: + return a or b + @_concat_chunks.register def _(a: Delta, b: Delta) -> Delta: From 80e35471c61edbd0178d1dd88f1d406135f7f9c3 Mon Sep 17 00:00:00 2001 From: Declan Brady Date: Wed, 19 Nov 2025 17:01:23 -0500 Subject: [PATCH 2/2] Fix lint --- src/agentex/lib/utils/completions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/agentex/lib/utils/completions.py b/src/agentex/lib/utils/completions.py index 24ef2c1f..fe62c7d1 100644 --- a/src/agentex/lib/utils/completions.py +++ b/src/agentex/lib/utils/completions.py @@ -6,10 +6,10 @@ from agentex.lib.types.llm_messages import ( Delta, + Usage, Choice, ToolCall, Completion, - Usage, ToolCallRequest, )