From 83ca6690119f1409725a2ba99eab6c73e00c94ce Mon Sep 17 00:00:00 2001 From: Chris Bartholomew Date: Wed, 11 Feb 2026 04:41:41 -0500 Subject: [PATCH] Add actual LLM token usage fields to RetainResult (#342) * Add actual LLM token usage fields to RetainResult RetainResult now carries llm_input_tokens, llm_output_tokens, and llm_total_tokens populated from the engine's TokenUsage, so downstream operation validator extensions can access actual LLM token counts. * Test that RetainResult includes actual LLM token usage --- hindsight-api/hindsight_api/engine/memory_engine.py | 3 +++ .../hindsight_api/extensions/operation_validator.py | 4 ++++ hindsight-api/tests/test_extensions.py | 8 ++++++++ 3 files changed, 15 insertions(+) diff --git a/hindsight-api/hindsight_api/engine/memory_engine.py b/hindsight-api/hindsight_api/engine/memory_engine.py index b2786e32..0d49139d 100644 --- a/hindsight-api/hindsight_api/engine/memory_engine.py +++ b/hindsight-api/hindsight_api/engine/memory_engine.py @@ -1487,6 +1487,9 @@ class MemoryEngine(MemoryEngineInterface): unit_ids=result, success=True, error=None, + llm_input_tokens=total_usage.input_tokens, + llm_output_tokens=total_usage.output_tokens, + llm_total_tokens=total_usage.total_tokens, ) try: await self._operation_validator.on_retain_complete(result_ctx) diff --git a/hindsight-api/hindsight_api/extensions/operation_validator.py b/hindsight-api/hindsight_api/extensions/operation_validator.py index e5b34557..15eb1b22 100644 --- a/hindsight-api/hindsight_api/extensions/operation_validator.py +++ b/hindsight-api/hindsight_api/extensions/operation_validator.py @@ -132,6 +132,10 @@ class RetainResult: unit_ids: list[list[str]] # List of unit IDs per content item success: bool = True error: str | None = None + # Actual LLM token usage (populated by engine when available) + llm_input_tokens: int | None = None + llm_output_tokens: int | None = None + llm_total_tokens: int | None = None @dataclass diff --git a/hindsight-api/tests/test_extensions.py b/hindsight-api/tests/test_extensions.py index 9956a509..39290447 100644 --- a/hindsight-api/tests/test_extensions.py +++ b/hindsight-api/tests/test_extensions.py @@ -353,6 +353,14 @@ class TestOperationHooksParameters: assert post_result.error is None assert post_result.unit_ids == result # Should match the return value + # Verify actual LLM token usage is populated + assert post_result.llm_input_tokens is not None + assert post_result.llm_input_tokens > 0 + assert post_result.llm_output_tokens is not None + assert post_result.llm_output_tokens > 0 + assert post_result.llm_total_tokens is not None + assert post_result.llm_total_tokens == post_result.llm_input_tokens + post_result.llm_output_tokens + @pytest.mark.asyncio async def test_recall_pre_hook_receives_all_parameters(self, memory_with_tracking_validator): """Pre-recall hook receives all user-provided parameters."""