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
This commit is contained in:
parent
e798979733
commit
83ca669011
3 changed files with 15 additions and 0 deletions
|
|
@ -1487,6 +1487,9 @@ class MemoryEngine(MemoryEngineInterface):
|
||||||
unit_ids=result,
|
unit_ids=result,
|
||||||
success=True,
|
success=True,
|
||||||
error=None,
|
error=None,
|
||||||
|
llm_input_tokens=total_usage.input_tokens,
|
||||||
|
llm_output_tokens=total_usage.output_tokens,
|
||||||
|
llm_total_tokens=total_usage.total_tokens,
|
||||||
)
|
)
|
||||||
try:
|
try:
|
||||||
await self._operation_validator.on_retain_complete(result_ctx)
|
await self._operation_validator.on_retain_complete(result_ctx)
|
||||||
|
|
|
||||||
|
|
@ -132,6 +132,10 @@ class RetainResult:
|
||||||
unit_ids: list[list[str]] # List of unit IDs per content item
|
unit_ids: list[list[str]] # List of unit IDs per content item
|
||||||
success: bool = True
|
success: bool = True
|
||||||
error: str | None = None
|
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
|
@dataclass
|
||||||
|
|
|
||||||
|
|
@ -353,6 +353,14 @@ class TestOperationHooksParameters:
|
||||||
assert post_result.error is None
|
assert post_result.error is None
|
||||||
assert post_result.unit_ids == result # Should match the return value
|
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
|
@pytest.mark.asyncio
|
||||||
async def test_recall_pre_hook_receives_all_parameters(self, memory_with_tracking_validator):
|
async def test_recall_pre_hook_receives_all_parameters(self, memory_with_tracking_validator):
|
||||||
"""Pre-recall hook receives all user-provided parameters."""
|
"""Pre-recall hook receives all user-provided parameters."""
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue