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."""