Skip to content

Commit 29a542d

Browse files
authored
feat: Add per-request LLM token usage metrics (#117)
* feat: Record LLM token metrics via Prometheus Wire up the existing token metrics infrastructure to actually record token usage from LLM calls. The MetricsCollector already had record_tokens() method and Prometheus counters (hindsight.tokens.input, hindsight.tokens.output), but they were never being populated. Changes: - Import get_metrics_collector in llm_wrapper.py - Call record_tokens() after successful LLM calls for: - OpenAI/Groq (using response.usage.prompt_tokens, completion_tokens) - Anthropic (using response.usage.input_tokens, output_tokens) - Gemini (using response.usage_metadata.prompt_token_count, candidates_token_count) - Add test file to verify token metrics are recorded Note: Ollama's native API doesn't return token usage, so metrics are not recorded for that provider. The token metrics will now be available via /metrics endpoint: - hindsight_tokens_input_total - hindsight_tokens_output_total * feat: add per-request token usage tracking to retain and reflect endpoints - Add TokenUsage model with input_tokens, output_tokens, total_tokens - Return usage metrics in retain response (sync operations only) - Return usage metrics in reflect response - Update Python, TypeScript, and Rust clients - Add API documentation for usage fields - Add changelog entry
1 parent ecc1f31 commit 29a542d

21 files changed

Lines changed: 888 additions & 65 deletions

File tree

hindsight-api/hindsight_api/api/http.py

Lines changed: 23 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ def _parse_metadata(metadata: Any) -> dict[str, Any]:
3636
from hindsight_api import MemoryEngine
3737
from hindsight_api.engine.db_utils import acquire_with_retry
3838
from hindsight_api.engine.memory_engine import Budget, fq_table
39-
from hindsight_api.engine.response_models import VALID_RECALL_FACT_TYPES
39+
from hindsight_api.engine.response_models import VALID_RECALL_FACT_TYPES, TokenUsage
4040
from hindsight_api.extensions import HttpExtension, OperationValidationError, load_extension
4141
from hindsight_api.metrics import create_metrics_collector, get_metrics_collector, initialize_metrics
4242
from hindsight_api.models import RequestContext
@@ -364,7 +364,15 @@ class RetainResponse(BaseModel):
364364

365365
model_config = ConfigDict(
366366
populate_by_name=True,
367-
json_schema_extra={"example": {"success": True, "bank_id": "user123", "items_count": 2, "async": False}},
367+
json_schema_extra={
368+
"example": {
369+
"success": True,
370+
"bank_id": "user123",
371+
"items_count": 2,
372+
"async": False,
373+
"usage": {"input_tokens": 500, "output_tokens": 100, "total_tokens": 600},
374+
}
375+
},
368376
)
369377

370378
success: bool
@@ -373,6 +381,10 @@ class RetainResponse(BaseModel):
373381
is_async: bool = Field(
374382
alias="async", serialization_alias="async", description="Whether the operation was processed asynchronously"
375383
)
384+
usage: TokenUsage | None = Field(
385+
default=None,
386+
description="Token usage metrics for LLM calls during fact extraction (only present for synchronous operations)",
387+
)
376388

377389

378390
class FactsIncludeOptions(BaseModel):
@@ -472,6 +484,7 @@ class ReflectResponse(BaseModel):
472484
"summary": "AI is transformative",
473485
"key_points": ["Used in healthcare", "Discussed recently"],
474486
},
487+
"usage": {"input_tokens": 1500, "output_tokens": 500, "total_tokens": 2000},
475488
}
476489
}
477490
)
@@ -482,6 +495,10 @@ class ReflectResponse(BaseModel):
482495
default=None,
483496
description="Structured output parsed according to the request's response_schema. Only present when response_schema was provided in the request.",
484497
)
498+
usage: TokenUsage | None = Field(
499+
default=None,
500+
description="Token usage metrics for LLM calls during reflection.",
501+
)
485502

486503

487504
class BanksResponse(BaseModel):
@@ -1290,6 +1307,7 @@ async def api_reflect(
12901307
text=core_result.text,
12911308
based_on=based_on_facts,
12921309
structured_output=core_result.structured_output,
1310+
usage=core_result.usage,
12931311
)
12941312

12951313
except OperationValidationError as e:
@@ -2016,12 +2034,12 @@ async def api_retain(
20162034
else:
20172035
# Synchronous processing: wait for completion (record metrics)
20182036
with metrics.record_operation("retain", bank_id=bank_id):
2019-
result = await app.state.memory.retain_batch_async(
2020-
bank_id=bank_id, contents=contents, request_context=request_context
2037+
result, usage = await app.state.memory.retain_batch_async(
2038+
bank_id=bank_id, contents=contents, request_context=request_context, return_usage=True
20212039
)
20222040

20232041
return RetainResponse.model_validate(
2024-
{"success": True, "bank_id": bank_id, "items_count": len(contents), "async": False}
2042+
{"success": True, "bank_id": bank_id, "items_count": len(contents), "async": False, "usage": usage}
20252043
)
20262044
except OperationValidationError as e:
20272045
raise HTTPException(status_code=e.status_code, detail=e.reason)

hindsight-api/hindsight_api/engine/llm_wrapper.py

Lines changed: 118 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@
2323
ENV_LLM_MAX_CONCURRENT,
2424
ENV_LLM_TIMEOUT,
2525
)
26+
from ..metrics import get_metrics_collector
27+
from .response_models import TokenUsage
2628

2729
# Seed applied to every Groq request for deterministic behavior.
2830
DEFAULT_LLM_SEED = 4242
@@ -174,6 +176,7 @@ async def call(
174176
max_backoff: float = 60.0,
175177
skip_validation: bool = False,
176178
strict_schema: bool = False,
179+
return_usage: bool = False,
177180
) -> Any:
178181
"""
179182
Make an LLM API call with retry logic.
@@ -189,9 +192,11 @@ async def call(
189192
max_backoff: Maximum backoff time in seconds.
190193
skip_validation: Return raw JSON without Pydantic validation.
191194
strict_schema: Use strict JSON schema enforcement (OpenAI only). Guarantees all required fields.
195+
return_usage: If True, return tuple (result, TokenUsage) instead of just result.
192196
193197
Returns:
194-
Parsed response if response_format is provided, otherwise text content.
198+
If return_usage=False: Parsed response if response_format is provided, otherwise text content.
199+
If return_usage=True: Tuple of (result, TokenUsage) with token counts from the LLM call.
195200
196201
Raises:
197202
OutputTooLongError: If output exceeds token limits.
@@ -203,7 +208,14 @@ async def call(
203208
# Handle Gemini provider separately
204209
if self.provider == "gemini":
205210
return await self._call_gemini(
206-
messages, response_format, max_retries, initial_backoff, max_backoff, skip_validation, start_time
211+
messages,
212+
response_format,
213+
max_retries,
214+
initial_backoff,
215+
max_backoff,
216+
skip_validation,
217+
start_time,
218+
return_usage,
207219
)
208220

209221
# Handle Anthropic provider separately
@@ -217,6 +229,7 @@ async def call(
217229
max_backoff,
218230
skip_validation,
219231
start_time,
232+
return_usage,
220233
)
221234

222235
# Handle Ollama with native API for structured output (better schema enforcement)
@@ -231,6 +244,7 @@ async def call(
231244
max_backoff,
232245
skip_validation,
233246
start_time,
247+
return_usage,
234248
)
235249

236250
call_params = {
@@ -379,21 +393,41 @@ async def call(
379393
response = await self._client.chat.completions.create(**call_params)
380394
result = response.choices[0].message.content
381395

382-
# Log slow calls
396+
# Record token usage metrics
383397
duration = time.time() - start_time
384398
usage = response.usage
385-
if duration > 10.0:
386-
ratio = max(1, usage.completion_tokens) / usage.prompt_tokens
399+
input_tokens = usage.prompt_tokens or 0 if usage else 0
400+
output_tokens = usage.completion_tokens or 0 if usage else 0
401+
total_tokens = usage.total_tokens or 0 if usage else 0
402+
403+
if usage:
404+
get_metrics_collector().record_tokens(
405+
operation=scope,
406+
bank_id="llm",
407+
input_tokens=input_tokens,
408+
output_tokens=output_tokens,
409+
)
410+
411+
# Log slow calls
412+
if duration > 10.0 and usage:
413+
ratio = max(1, output_tokens) / max(1, input_tokens)
387414
cached_tokens = 0
388415
if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details:
389416
cached_tokens = getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0
390417
cache_info = f", cached_tokens={cached_tokens}" if cached_tokens > 0 else ""
391418
logger.info(
392419
f"slow llm call: model={self.provider}/{self.model}, "
393-
f"input_tokens={usage.prompt_tokens}, output_tokens={usage.completion_tokens}, "
394-
f"total_tokens={usage.total_tokens}{cache_info}, time={duration:.3f}s, ratio out/in={ratio:.2f}"
420+
f"input_tokens={input_tokens}, output_tokens={output_tokens}, "
421+
f"total_tokens={total_tokens}{cache_info}, time={duration:.3f}s, ratio out/in={ratio:.2f}"
395422
)
396423

424+
if return_usage:
425+
token_usage = TokenUsage(
426+
input_tokens=input_tokens,
427+
output_tokens=output_tokens,
428+
total_tokens=total_tokens,
429+
)
430+
return result, token_usage
397431
return result
398432

399433
except LengthFinishReasonError as e:
@@ -452,6 +486,7 @@ async def _call_anthropic(
452486
max_backoff: float,
453487
skip_validation: bool,
454488
start_time: float,
489+
return_usage: bool = False,
455490
) -> Any:
456491
"""Handle Anthropic-specific API calls."""
457492
from anthropic import APIConnectionError, APIStatusError, RateLimitError
@@ -524,17 +559,35 @@ async def _call_anthropic(
524559
else:
525560
result = content
526561

527-
# Log slow calls
562+
# Record token usage metrics
528563
duration = time.time() - start_time
529-
if duration > 10.0:
530-
input_tokens = response.usage.input_tokens
531-
output_tokens = response.usage.output_tokens
564+
input_tokens = response.usage.input_tokens or 0 if response.usage else 0
565+
output_tokens = response.usage.output_tokens or 0 if response.usage else 0
566+
total_tokens = input_tokens + output_tokens
567+
568+
if response.usage:
569+
get_metrics_collector().record_tokens(
570+
operation="memory",
571+
bank_id="llm",
572+
input_tokens=input_tokens,
573+
output_tokens=output_tokens,
574+
)
575+
576+
# Log slow calls
577+
if duration > 10.0 and response.usage:
532578
logger.info(
533579
f"slow llm call: model={self.provider}/{self.model}, "
534580
f"input_tokens={input_tokens}, output_tokens={output_tokens}, "
535581
f"time={duration:.3f}s"
536582
)
537583

584+
if return_usage:
585+
token_usage = TokenUsage(
586+
input_tokens=input_tokens,
587+
output_tokens=output_tokens,
588+
total_tokens=total_tokens,
589+
)
590+
return result, token_usage
538591
return result
539592

540593
except json.JSONDecodeError as e:
@@ -589,6 +642,7 @@ async def _call_ollama_native(
589642
max_backoff: float,
590643
skip_validation: bool,
591644
start_time: float,
645+
return_usage: bool = False,
592646
) -> Any:
593647
"""
594648
Call Ollama using native API with JSON schema enforcement.
@@ -663,11 +717,35 @@ async def _call_ollama_native(
663717
else:
664718
raise
665719

720+
# Extract token usage from Ollama response
721+
# Ollama returns prompt_eval_count (input) and eval_count (output)
722+
input_tokens = result.get("prompt_eval_count", 0) or 0
723+
output_tokens = result.get("eval_count", 0) or 0
724+
total_tokens = input_tokens + output_tokens
725+
726+
# Record to metrics
727+
if input_tokens > 0 or output_tokens > 0:
728+
get_metrics_collector().record_tokens(
729+
operation="memory",
730+
bank_id="llm",
731+
input_tokens=input_tokens,
732+
output_tokens=output_tokens,
733+
)
734+
666735
# Validate against Pydantic model or return raw JSON
667736
if skip_validation:
668-
return json_data
737+
validated_result = json_data
669738
else:
670-
return response_format.model_validate(json_data)
739+
validated_result = response_format.model_validate(json_data)
740+
741+
if return_usage:
742+
token_usage = TokenUsage(
743+
input_tokens=input_tokens,
744+
output_tokens=output_tokens,
745+
total_tokens=total_tokens,
746+
)
747+
return validated_result, token_usage
748+
return validated_result
671749

672750
except httpx.HTTPStatusError as e:
673751
last_exception = e
@@ -710,6 +788,7 @@ async def _call_gemini(
710788
max_backoff: float,
711789
skip_validation: bool,
712790
start_time: float,
791+
return_usage: bool = False,
713792
) -> Any:
714793
"""Handle Gemini-specific API calls."""
715794
# Convert OpenAI-style messages to Gemini format
@@ -786,16 +865,36 @@ async def _call_gemini(
786865
else:
787866
result = content
788867

789-
# Log slow calls
868+
# Record token usage metrics
790869
duration = time.time() - start_time
791-
if duration > 10.0 and hasattr(response, "usage_metadata") and response.usage_metadata:
870+
input_tokens = 0
871+
output_tokens = 0
872+
if hasattr(response, "usage_metadata") and response.usage_metadata:
792873
usage = response.usage_metadata
793-
logger.info(
794-
f"slow llm call: model={self.provider}/{self.model}, "
795-
f"input_tokens={usage.prompt_token_count}, output_tokens={usage.candidates_token_count}, "
796-
f"time={duration:.3f}s"
874+
input_tokens = usage.prompt_token_count or 0
875+
output_tokens = usage.candidates_token_count or 0
876+
get_metrics_collector().record_tokens(
877+
operation="memory",
878+
bank_id="llm",
879+
input_tokens=input_tokens,
880+
output_tokens=output_tokens,
797881
)
798882

883+
# Log slow calls
884+
if duration > 10.0:
885+
logger.info(
886+
f"slow llm call: model={self.provider}/{self.model}, "
887+
f"input_tokens={input_tokens}, output_tokens={output_tokens}, "
888+
f"time={duration:.3f}s"
889+
)
890+
891+
if return_usage:
892+
token_usage = TokenUsage(
893+
input_tokens=input_tokens,
894+
output_tokens=output_tokens,
895+
total_tokens=input_tokens + output_tokens,
896+
)
897+
return result, token_usage
799898
return result
800899

801900
except json.JSONDecodeError as e:

0 commit comments

Comments
 (0)