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
This commit is contained in:
parent
ecc1f31996
commit
29a542dc23
21 changed files with 888 additions and 65 deletions
|
|
@ -36,7 +36,7 @@ from pydantic import BaseModel, ConfigDict, Field, field_validator
|
|||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api.engine.db_utils import acquire_with_retry
|
||||
from hindsight_api.engine.memory_engine import Budget, fq_table
|
||||
from hindsight_api.engine.response_models import VALID_RECALL_FACT_TYPES
|
||||
from hindsight_api.engine.response_models import VALID_RECALL_FACT_TYPES, TokenUsage
|
||||
from hindsight_api.extensions import HttpExtension, OperationValidationError, load_extension
|
||||
from hindsight_api.metrics import create_metrics_collector, get_metrics_collector, initialize_metrics
|
||||
from hindsight_api.models import RequestContext
|
||||
|
|
@ -364,7 +364,15 @@ class RetainResponse(BaseModel):
|
|||
|
||||
model_config = ConfigDict(
|
||||
populate_by_name=True,
|
||||
json_schema_extra={"example": {"success": True, "bank_id": "user123", "items_count": 2, "async": False}},
|
||||
json_schema_extra={
|
||||
"example": {
|
||||
"success": True,
|
||||
"bank_id": "user123",
|
||||
"items_count": 2,
|
||||
"async": False,
|
||||
"usage": {"input_tokens": 500, "output_tokens": 100, "total_tokens": 600},
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
success: bool
|
||||
|
|
@ -373,6 +381,10 @@ class RetainResponse(BaseModel):
|
|||
is_async: bool = Field(
|
||||
alias="async", serialization_alias="async", description="Whether the operation was processed asynchronously"
|
||||
)
|
||||
usage: TokenUsage | None = Field(
|
||||
default=None,
|
||||
description="Token usage metrics for LLM calls during fact extraction (only present for synchronous operations)",
|
||||
)
|
||||
|
||||
|
||||
class FactsIncludeOptions(BaseModel):
|
||||
|
|
@ -472,6 +484,7 @@ class ReflectResponse(BaseModel):
|
|||
"summary": "AI is transformative",
|
||||
"key_points": ["Used in healthcare", "Discussed recently"],
|
||||
},
|
||||
"usage": {"input_tokens": 1500, "output_tokens": 500, "total_tokens": 2000},
|
||||
}
|
||||
}
|
||||
)
|
||||
|
|
@ -482,6 +495,10 @@ class ReflectResponse(BaseModel):
|
|||
default=None,
|
||||
description="Structured output parsed according to the request's response_schema. Only present when response_schema was provided in the request.",
|
||||
)
|
||||
usage: TokenUsage | None = Field(
|
||||
default=None,
|
||||
description="Token usage metrics for LLM calls during reflection.",
|
||||
)
|
||||
|
||||
|
||||
class BanksResponse(BaseModel):
|
||||
|
|
@ -1290,6 +1307,7 @@ def _register_routes(app: FastAPI):
|
|||
text=core_result.text,
|
||||
based_on=based_on_facts,
|
||||
structured_output=core_result.structured_output,
|
||||
usage=core_result.usage,
|
||||
)
|
||||
|
||||
except OperationValidationError as e:
|
||||
|
|
@ -2016,12 +2034,12 @@ def _register_routes(app: FastAPI):
|
|||
else:
|
||||
# Synchronous processing: wait for completion (record metrics)
|
||||
with metrics.record_operation("retain", bank_id=bank_id):
|
||||
result = await app.state.memory.retain_batch_async(
|
||||
bank_id=bank_id, contents=contents, request_context=request_context
|
||||
result, usage = await app.state.memory.retain_batch_async(
|
||||
bank_id=bank_id, contents=contents, request_context=request_context, return_usage=True
|
||||
)
|
||||
|
||||
return RetainResponse.model_validate(
|
||||
{"success": True, "bank_id": bank_id, "items_count": len(contents), "async": False}
|
||||
{"success": True, "bank_id": bank_id, "items_count": len(contents), "async": False, "usage": usage}
|
||||
)
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
|
|
|
|||
|
|
@ -23,6 +23,8 @@ from ..config import (
|
|||
ENV_LLM_MAX_CONCURRENT,
|
||||
ENV_LLM_TIMEOUT,
|
||||
)
|
||||
from ..metrics import get_metrics_collector
|
||||
from .response_models import TokenUsage
|
||||
|
||||
# Seed applied to every Groq request for deterministic behavior.
|
||||
DEFAULT_LLM_SEED = 4242
|
||||
|
|
@ -174,6 +176,7 @@ class LLMProvider:
|
|||
max_backoff: float = 60.0,
|
||||
skip_validation: bool = False,
|
||||
strict_schema: bool = False,
|
||||
return_usage: bool = False,
|
||||
) -> Any:
|
||||
"""
|
||||
Make an LLM API call with retry logic.
|
||||
|
|
@ -189,9 +192,11 @@ class LLMProvider:
|
|||
max_backoff: Maximum backoff time in seconds.
|
||||
skip_validation: Return raw JSON without Pydantic validation.
|
||||
strict_schema: Use strict JSON schema enforcement (OpenAI only). Guarantees all required fields.
|
||||
return_usage: If True, return tuple (result, TokenUsage) instead of just result.
|
||||
|
||||
Returns:
|
||||
Parsed response if response_format is provided, otherwise text content.
|
||||
If return_usage=False: Parsed response if response_format is provided, otherwise text content.
|
||||
If return_usage=True: Tuple of (result, TokenUsage) with token counts from the LLM call.
|
||||
|
||||
Raises:
|
||||
OutputTooLongError: If output exceeds token limits.
|
||||
|
|
@ -203,7 +208,14 @@ class LLMProvider:
|
|||
# Handle Gemini provider separately
|
||||
if self.provider == "gemini":
|
||||
return await self._call_gemini(
|
||||
messages, response_format, max_retries, initial_backoff, max_backoff, skip_validation, start_time
|
||||
messages,
|
||||
response_format,
|
||||
max_retries,
|
||||
initial_backoff,
|
||||
max_backoff,
|
||||
skip_validation,
|
||||
start_time,
|
||||
return_usage,
|
||||
)
|
||||
|
||||
# Handle Anthropic provider separately
|
||||
|
|
@ -217,6 +229,7 @@ class LLMProvider:
|
|||
max_backoff,
|
||||
skip_validation,
|
||||
start_time,
|
||||
return_usage,
|
||||
)
|
||||
|
||||
# Handle Ollama with native API for structured output (better schema enforcement)
|
||||
|
|
@ -231,6 +244,7 @@ class LLMProvider:
|
|||
max_backoff,
|
||||
skip_validation,
|
||||
start_time,
|
||||
return_usage,
|
||||
)
|
||||
|
||||
call_params = {
|
||||
|
|
@ -379,21 +393,41 @@ class LLMProvider:
|
|||
response = await self._client.chat.completions.create(**call_params)
|
||||
result = response.choices[0].message.content
|
||||
|
||||
# Log slow calls
|
||||
# Record token usage metrics
|
||||
duration = time.time() - start_time
|
||||
usage = response.usage
|
||||
if duration > 10.0:
|
||||
ratio = max(1, usage.completion_tokens) / usage.prompt_tokens
|
||||
input_tokens = usage.prompt_tokens or 0 if usage else 0
|
||||
output_tokens = usage.completion_tokens or 0 if usage else 0
|
||||
total_tokens = usage.total_tokens or 0 if usage else 0
|
||||
|
||||
if usage:
|
||||
get_metrics_collector().record_tokens(
|
||||
operation=scope,
|
||||
bank_id="llm",
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
)
|
||||
|
||||
# Log slow calls
|
||||
if duration > 10.0 and usage:
|
||||
ratio = max(1, output_tokens) / max(1, input_tokens)
|
||||
cached_tokens = 0
|
||||
if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details:
|
||||
cached_tokens = getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0
|
||||
cache_info = f", cached_tokens={cached_tokens}" if cached_tokens > 0 else ""
|
||||
logger.info(
|
||||
f"slow llm call: model={self.provider}/{self.model}, "
|
||||
f"input_tokens={usage.prompt_tokens}, output_tokens={usage.completion_tokens}, "
|
||||
f"total_tokens={usage.total_tokens}{cache_info}, time={duration:.3f}s, ratio out/in={ratio:.2f}"
|
||||
f"input_tokens={input_tokens}, output_tokens={output_tokens}, "
|
||||
f"total_tokens={total_tokens}{cache_info}, time={duration:.3f}s, ratio out/in={ratio:.2f}"
|
||||
)
|
||||
|
||||
if return_usage:
|
||||
token_usage = TokenUsage(
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
total_tokens=total_tokens,
|
||||
)
|
||||
return result, token_usage
|
||||
return result
|
||||
|
||||
except LengthFinishReasonError as e:
|
||||
|
|
@ -452,6 +486,7 @@ class LLMProvider:
|
|||
max_backoff: float,
|
||||
skip_validation: bool,
|
||||
start_time: float,
|
||||
return_usage: bool = False,
|
||||
) -> Any:
|
||||
"""Handle Anthropic-specific API calls."""
|
||||
from anthropic import APIConnectionError, APIStatusError, RateLimitError
|
||||
|
|
@ -524,17 +559,35 @@ class LLMProvider:
|
|||
else:
|
||||
result = content
|
||||
|
||||
# Log slow calls
|
||||
# Record token usage metrics
|
||||
duration = time.time() - start_time
|
||||
if duration > 10.0:
|
||||
input_tokens = response.usage.input_tokens
|
||||
output_tokens = response.usage.output_tokens
|
||||
input_tokens = response.usage.input_tokens or 0 if response.usage else 0
|
||||
output_tokens = response.usage.output_tokens or 0 if response.usage else 0
|
||||
total_tokens = input_tokens + output_tokens
|
||||
|
||||
if response.usage:
|
||||
get_metrics_collector().record_tokens(
|
||||
operation="memory",
|
||||
bank_id="llm",
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
)
|
||||
|
||||
# Log slow calls
|
||||
if duration > 10.0 and response.usage:
|
||||
logger.info(
|
||||
f"slow llm call: model={self.provider}/{self.model}, "
|
||||
f"input_tokens={input_tokens}, output_tokens={output_tokens}, "
|
||||
f"time={duration:.3f}s"
|
||||
)
|
||||
|
||||
if return_usage:
|
||||
token_usage = TokenUsage(
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
total_tokens=total_tokens,
|
||||
)
|
||||
return result, token_usage
|
||||
return result
|
||||
|
||||
except json.JSONDecodeError as e:
|
||||
|
|
@ -589,6 +642,7 @@ class LLMProvider:
|
|||
max_backoff: float,
|
||||
skip_validation: bool,
|
||||
start_time: float,
|
||||
return_usage: bool = False,
|
||||
) -> Any:
|
||||
"""
|
||||
Call Ollama using native API with JSON schema enforcement.
|
||||
|
|
@ -663,11 +717,35 @@ class LLMProvider:
|
|||
else:
|
||||
raise
|
||||
|
||||
# Extract token usage from Ollama response
|
||||
# Ollama returns prompt_eval_count (input) and eval_count (output)
|
||||
input_tokens = result.get("prompt_eval_count", 0) or 0
|
||||
output_tokens = result.get("eval_count", 0) or 0
|
||||
total_tokens = input_tokens + output_tokens
|
||||
|
||||
# Record to metrics
|
||||
if input_tokens > 0 or output_tokens > 0:
|
||||
get_metrics_collector().record_tokens(
|
||||
operation="memory",
|
||||
bank_id="llm",
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
)
|
||||
|
||||
# Validate against Pydantic model or return raw JSON
|
||||
if skip_validation:
|
||||
return json_data
|
||||
validated_result = json_data
|
||||
else:
|
||||
return response_format.model_validate(json_data)
|
||||
validated_result = response_format.model_validate(json_data)
|
||||
|
||||
if return_usage:
|
||||
token_usage = TokenUsage(
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
total_tokens=total_tokens,
|
||||
)
|
||||
return validated_result, token_usage
|
||||
return validated_result
|
||||
|
||||
except httpx.HTTPStatusError as e:
|
||||
last_exception = e
|
||||
|
|
@ -710,6 +788,7 @@ class LLMProvider:
|
|||
max_backoff: float,
|
||||
skip_validation: bool,
|
||||
start_time: float,
|
||||
return_usage: bool = False,
|
||||
) -> Any:
|
||||
"""Handle Gemini-specific API calls."""
|
||||
# Convert OpenAI-style messages to Gemini format
|
||||
|
|
@ -786,16 +865,36 @@ class LLMProvider:
|
|||
else:
|
||||
result = content
|
||||
|
||||
# Log slow calls
|
||||
# Record token usage metrics
|
||||
duration = time.time() - start_time
|
||||
if duration > 10.0 and hasattr(response, "usage_metadata") and response.usage_metadata:
|
||||
input_tokens = 0
|
||||
output_tokens = 0
|
||||
if hasattr(response, "usage_metadata") and response.usage_metadata:
|
||||
usage = response.usage_metadata
|
||||
logger.info(
|
||||
f"slow llm call: model={self.provider}/{self.model}, "
|
||||
f"input_tokens={usage.prompt_token_count}, output_tokens={usage.candidates_token_count}, "
|
||||
f"time={duration:.3f}s"
|
||||
input_tokens = usage.prompt_token_count or 0
|
||||
output_tokens = usage.candidates_token_count or 0
|
||||
get_metrics_collector().record_tokens(
|
||||
operation="memory",
|
||||
bank_id="llm",
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
)
|
||||
|
||||
# Log slow calls
|
||||
if duration > 10.0:
|
||||
logger.info(
|
||||
f"slow llm call: model={self.provider}/{self.model}, "
|
||||
f"input_tokens={input_tokens}, output_tokens={output_tokens}, "
|
||||
f"time={duration:.3f}s"
|
||||
)
|
||||
|
||||
if return_usage:
|
||||
token_usage = TokenUsage(
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
total_tokens=input_tokens + output_tokens,
|
||||
)
|
||||
return result, token_usage
|
||||
return result
|
||||
|
||||
except json.JSONDecodeError as e:
|
||||
|
|
|
|||
|
|
@ -136,7 +136,14 @@ from ..pg0 import EmbeddedPostgres, parse_pg0_url
|
|||
from .entity_resolver import EntityResolver
|
||||
from .llm_wrapper import LLMConfig
|
||||
from .query_analyzer import QueryAnalyzer
|
||||
from .response_models import VALID_RECALL_FACT_TYPES, EntityObservation, EntityState, MemoryFact, ReflectResult
|
||||
from .response_models import (
|
||||
VALID_RECALL_FACT_TYPES,
|
||||
EntityObservation,
|
||||
EntityState,
|
||||
MemoryFact,
|
||||
ReflectResult,
|
||||
TokenUsage,
|
||||
)
|
||||
from .response_models import RecallResult as RecallResultModel
|
||||
from .retain import bank_utils, embedding_utils
|
||||
from .retain.types import RetainContentDict
|
||||
|
|
@ -954,7 +961,8 @@ class MemoryEngine(MemoryEngineInterface):
|
|||
document_id: str | None = None,
|
||||
fact_type_override: str | None = None,
|
||||
confidence_score: float | None = None,
|
||||
) -> list[list[str]]:
|
||||
return_usage: bool = False,
|
||||
):
|
||||
"""
|
||||
Store multiple content items as memory units in ONE batch operation.
|
||||
|
||||
|
|
@ -975,9 +983,11 @@ class MemoryEngine(MemoryEngineInterface):
|
|||
Applies the same document_id to ALL content items that don't specify their own.
|
||||
fact_type_override: Override fact type for all facts ('world', 'experience', 'opinion')
|
||||
confidence_score: Confidence score for opinions (0.0 to 1.0)
|
||||
return_usage: If True, returns tuple of (unit_ids, TokenUsage). Default False for backward compatibility.
|
||||
|
||||
Returns:
|
||||
List of lists of unit IDs (one list per content item)
|
||||
If return_usage=False: List of lists of unit IDs (one list per content item)
|
||||
If return_usage=True: Tuple of (unit_ids, TokenUsage)
|
||||
|
||||
Example (new style - per-content document_id):
|
||||
unit_ids = await memory.retain_batch_async(
|
||||
|
|
@ -1004,6 +1014,8 @@ class MemoryEngine(MemoryEngineInterface):
|
|||
start_time = time.time()
|
||||
|
||||
if not contents:
|
||||
if return_usage:
|
||||
return [], TokenUsage()
|
||||
return []
|
||||
|
||||
# Authenticate tenant and set schema in context (for fq_table())
|
||||
|
|
@ -1033,6 +1045,7 @@ class MemoryEngine(MemoryEngineInterface):
|
|||
# Auto-chunk large batches by character count to avoid timeouts and memory issues
|
||||
# Calculate total character count
|
||||
total_chars = sum(len(item.get("content", "")) for item in contents)
|
||||
total_usage = TokenUsage()
|
||||
|
||||
CHARS_PER_BATCH = 600_000
|
||||
|
||||
|
|
@ -1073,7 +1086,7 @@ class MemoryEngine(MemoryEngineInterface):
|
|||
f"Processing sub-batch {i}/{len(sub_batches)}: {len(sub_batch)} items, {sub_batch_chars:,} chars"
|
||||
)
|
||||
|
||||
sub_results = await self._retain_batch_async_internal(
|
||||
sub_results, sub_usage = await self._retain_batch_async_internal(
|
||||
bank_id=bank_id,
|
||||
contents=sub_batch,
|
||||
document_id=document_id,
|
||||
|
|
@ -1082,6 +1095,7 @@ class MemoryEngine(MemoryEngineInterface):
|
|||
confidence_score=confidence_score,
|
||||
)
|
||||
all_results.extend(sub_results)
|
||||
total_usage = total_usage + sub_usage
|
||||
|
||||
total_time = time.time() - start_time
|
||||
logger.info(
|
||||
|
|
@ -1090,7 +1104,7 @@ class MemoryEngine(MemoryEngineInterface):
|
|||
result = all_results
|
||||
else:
|
||||
# Small batch - use internal method directly
|
||||
result = await self._retain_batch_async_internal(
|
||||
result, total_usage = await self._retain_batch_async_internal(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
document_id=document_id,
|
||||
|
|
@ -1119,6 +1133,8 @@ class MemoryEngine(MemoryEngineInterface):
|
|||
except Exception as e:
|
||||
logger.warning(f"Post-retain hook error (non-fatal): {e}")
|
||||
|
||||
if return_usage:
|
||||
return result, total_usage
|
||||
return result
|
||||
|
||||
async def _retain_batch_async_internal(
|
||||
|
|
@ -1129,7 +1145,7 @@ class MemoryEngine(MemoryEngineInterface):
|
|||
is_first_batch: bool = True,
|
||||
fact_type_override: str | None = None,
|
||||
confidence_score: float | None = None,
|
||||
) -> list[list[str]]:
|
||||
) -> tuple[list[list[str]], "TokenUsage"]:
|
||||
"""
|
||||
Internal method for batch processing without chunking logic.
|
||||
|
||||
|
|
@ -1145,6 +1161,9 @@ class MemoryEngine(MemoryEngineInterface):
|
|||
is_first_batch: Whether this is the first batch (for chunked operations, only delete on first batch)
|
||||
fact_type_override: Override fact type for all facts
|
||||
confidence_score: Confidence score for opinions
|
||||
|
||||
Returns:
|
||||
Tuple of (unit ID lists, token usage for fact extraction)
|
||||
"""
|
||||
# Backpressure: limit concurrent retains to prevent database contention
|
||||
async with self._put_semaphore:
|
||||
|
|
@ -3192,7 +3211,7 @@ Guidelines:
|
|||
response_format = JsonSchemaWrapper(response_schema)
|
||||
|
||||
llm_start = time.time()
|
||||
result = await self._llm_config.call(
|
||||
llm_result, usage = await self._llm_config.call(
|
||||
messages=messages,
|
||||
scope="memory_reflect",
|
||||
max_completion_tokens=max_tokens,
|
||||
|
|
@ -3201,17 +3220,18 @@ Guidelines:
|
|||
# Don't enforce strict_schema - not all providers support it and may retry forever
|
||||
# Soft enforcement (schema in prompt + json_object mode) is sufficient
|
||||
strict_schema=False,
|
||||
return_usage=True,
|
||||
)
|
||||
llm_time = time.time() - llm_start
|
||||
|
||||
# Handle response based on whether structured output was requested
|
||||
if response_schema is not None:
|
||||
structured_output = result
|
||||
structured_output = llm_result
|
||||
answer_text = "" # Empty for backward compatibility
|
||||
log_buffer.append(f"[REFLECT {reflect_id}] Structured output generated")
|
||||
else:
|
||||
structured_output = None
|
||||
answer_text = result.strip()
|
||||
answer_text = llm_result.strip()
|
||||
|
||||
# Submit form_opinion task for background processing
|
||||
# Pass tenant_id from request context for internal authentication in background task
|
||||
|
|
@ -3237,6 +3257,7 @@ Guidelines:
|
|||
based_on={"world": world_results, "experience": agent_results, "opinion": opinion_results},
|
||||
new_opinions=[], # Opinions are being extracted asynchronously
|
||||
structured_output=structured_output,
|
||||
usage=usage,
|
||||
)
|
||||
|
||||
# Call post-operation hook if validator is configured
|
||||
|
|
|
|||
|
|
@ -14,6 +14,37 @@ from pydantic import BaseModel, ConfigDict, Field
|
|||
VALID_RECALL_FACT_TYPES = frozenset(["world", "experience", "opinion"])
|
||||
|
||||
|
||||
class TokenUsage(BaseModel):
|
||||
"""
|
||||
Token usage metrics for LLM calls.
|
||||
|
||||
Tracks input/output tokens for a single request to enable
|
||||
per-request cost tracking and monitoring.
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(
|
||||
json_schema_extra={
|
||||
"example": {
|
||||
"input_tokens": 1500,
|
||||
"output_tokens": 500,
|
||||
"total_tokens": 2000,
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
input_tokens: int = Field(default=0, description="Number of input/prompt tokens consumed")
|
||||
output_tokens: int = Field(default=0, description="Number of output/completion tokens generated")
|
||||
total_tokens: int = Field(default=0, description="Total tokens (input + output)")
|
||||
|
||||
def __add__(self, other: "TokenUsage") -> "TokenUsage":
|
||||
"""Allow aggregating token usage from multiple calls."""
|
||||
return TokenUsage(
|
||||
input_tokens=self.input_tokens + other.input_tokens,
|
||||
output_tokens=self.output_tokens + other.output_tokens,
|
||||
total_tokens=self.total_tokens + other.total_tokens,
|
||||
)
|
||||
|
||||
|
||||
class DispositionTraits(BaseModel):
|
||||
"""
|
||||
Disposition traits for a memory bank.
|
||||
|
|
@ -147,6 +178,7 @@ class ReflectResult(BaseModel):
|
|||
},
|
||||
"new_opinions": ["Machine learning has great potential in healthcare"],
|
||||
"structured_output": {"summary": "ML in healthcare", "confidence": 0.9},
|
||||
"usage": {"input_tokens": 1500, "output_tokens": 500, "total_tokens": 2000},
|
||||
}
|
||||
}
|
||||
)
|
||||
|
|
@ -160,6 +192,10 @@ class ReflectResult(BaseModel):
|
|||
default=None,
|
||||
description="Structured output parsed according to the provided response schema. Only present when response_schema was provided.",
|
||||
)
|
||||
usage: TokenUsage | None = Field(
|
||||
default=None,
|
||||
description="Token usage metrics for the LLM calls made during this reflect operation.",
|
||||
)
|
||||
|
||||
|
||||
class Opinion(BaseModel):
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ from pydantic import BaseModel, ConfigDict, Field, field_validator
|
|||
|
||||
from ...config import get_config
|
||||
from ..llm_wrapper import LLMConfig, OutputTooLongError
|
||||
from ..response_models import TokenUsage
|
||||
|
||||
|
||||
def _infer_temporal_date(fact_text: str, event_date: datetime) -> str | None:
|
||||
|
|
@ -393,7 +394,7 @@ async def _extract_facts_from_chunk(
|
|||
llm_config: "LLMConfig",
|
||||
agent_name: str = None,
|
||||
extract_opinions: bool = False,
|
||||
) -> list[dict[str, str]]:
|
||||
) -> tuple[list[dict[str, str]], TokenUsage]:
|
||||
"""
|
||||
Extract facts from a single chunk (internal helper for parallel processing).
|
||||
|
||||
|
|
@ -685,16 +686,19 @@ Context: {sanitized_context}
|
|||
Text:
|
||||
{sanitized_chunk}"""
|
||||
|
||||
usage = TokenUsage() # Track cumulative usage across retries
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
extraction_response_json = await llm_config.call(
|
||||
extraction_response_json, call_usage = await llm_config.call(
|
||||
messages=[{"role": "system", "content": prompt}, {"role": "user", "content": user_message}],
|
||||
response_format=FactExtractionResponse,
|
||||
scope="memory_extract_facts",
|
||||
temperature=0.1,
|
||||
max_completion_tokens=config.retain_max_completion_tokens,
|
||||
skip_validation=True, # Get raw JSON, we'll validate leniently
|
||||
return_usage=True,
|
||||
)
|
||||
usage = usage + call_usage # Aggregate usage across retries
|
||||
|
||||
# Lenient parsing of facts from raw JSON
|
||||
chunk_facts = []
|
||||
|
|
@ -712,7 +716,7 @@ Text:
|
|||
f"LLM returned non-dict JSON after {max_retries} attempts: {type(extraction_response_json).__name__}. "
|
||||
f"Raw: {str(extraction_response_json)[:500]}"
|
||||
)
|
||||
return []
|
||||
return [], usage
|
||||
|
||||
raw_facts = extraction_response_json.get("facts", [])
|
||||
# Get top-level causal relationships (new schema)
|
||||
|
|
@ -927,7 +931,7 @@ Text:
|
|||
)
|
||||
continue
|
||||
|
||||
return chunk_facts
|
||||
return chunk_facts, usage
|
||||
|
||||
except BadRequestError as e:
|
||||
last_error = e
|
||||
|
|
@ -954,7 +958,7 @@ async def _extract_facts_with_auto_split(
|
|||
llm_config: LLMConfig,
|
||||
agent_name: str = None,
|
||||
extract_opinions: bool = False,
|
||||
) -> list[dict[str, str]]:
|
||||
) -> tuple[list[dict[str, str]], TokenUsage]:
|
||||
"""
|
||||
Extract facts from a chunk with automatic splitting if output exceeds token limits.
|
||||
|
||||
|
|
@ -972,7 +976,7 @@ async def _extract_facts_with_auto_split(
|
|||
extract_opinions: If True, extract ONLY opinions. If False, extract world and agent facts (no opinions)
|
||||
|
||||
Returns:
|
||||
List of fact dictionaries extracted from the chunk (possibly from sub-chunks)
|
||||
Tuple of (facts list, token usage) extracted from the chunk (possibly from sub-chunks)
|
||||
"""
|
||||
import logging
|
||||
|
||||
|
|
@ -1051,12 +1055,14 @@ async def _extract_facts_with_auto_split(
|
|||
|
||||
# Combine results from both halves
|
||||
all_facts = []
|
||||
for sub_result in sub_results:
|
||||
all_facts.extend(sub_result)
|
||||
total_usage = TokenUsage()
|
||||
for sub_facts, sub_usage in sub_results:
|
||||
all_facts.extend(sub_facts)
|
||||
total_usage = total_usage + sub_usage
|
||||
|
||||
logger.info(f"Successfully extracted {len(all_facts)} facts from split chunk {chunk_index + 1}")
|
||||
|
||||
return all_facts
|
||||
return all_facts, total_usage
|
||||
|
||||
|
||||
async def extract_facts_from_text(
|
||||
|
|
@ -1066,7 +1072,7 @@ async def extract_facts_from_text(
|
|||
agent_name: str,
|
||||
context: str = "",
|
||||
extract_opinions: bool = False,
|
||||
) -> tuple[list[Fact], list[tuple[str, int]]]:
|
||||
) -> tuple[list[Fact], list[tuple[str, int]], TokenUsage]:
|
||||
"""
|
||||
Extract semantic facts from conversational or narrative text using LLM.
|
||||
|
||||
|
|
@ -1085,9 +1091,10 @@ async def extract_facts_from_text(
|
|||
extract_opinions: If True, extract ONLY opinions. If False, extract world and bank facts (no opinions)
|
||||
|
||||
Returns:
|
||||
Tuple of (facts, chunks) where:
|
||||
Tuple of (facts, chunks, usage) where:
|
||||
- facts: List of Fact model instances
|
||||
- chunks: List of tuples (chunk_text, fact_count) for each chunk
|
||||
- usage: Aggregated token usage across all LLM calls
|
||||
"""
|
||||
chunks = chunk_text(text, max_chars=3000)
|
||||
tasks = [
|
||||
|
|
@ -1106,10 +1113,12 @@ async def extract_facts_from_text(
|
|||
chunk_results = await asyncio.gather(*tasks)
|
||||
all_facts = []
|
||||
chunk_metadata = [] # [(chunk_text, fact_count), ...]
|
||||
for chunk, chunk_facts in zip(chunks, chunk_results):
|
||||
total_usage = TokenUsage()
|
||||
for chunk, (chunk_facts, chunk_usage) in zip(chunks, chunk_results):
|
||||
all_facts.extend(chunk_facts)
|
||||
chunk_metadata.append((chunk, len(chunk_facts)))
|
||||
return all_facts, chunk_metadata
|
||||
total_usage = total_usage + chunk_usage
|
||||
return all_facts, chunk_metadata, total_usage
|
||||
|
||||
|
||||
# ============================================================================
|
||||
|
|
@ -1130,7 +1139,7 @@ SECONDS_PER_FACT = 10
|
|||
|
||||
async def extract_facts_from_contents(
|
||||
contents: list[RetainContent], llm_config, agent_name: str, extract_opinions: bool = False
|
||||
) -> tuple[list[ExtractedFactType], list[ChunkMetadata]]:
|
||||
) -> tuple[list[ExtractedFactType], list[ChunkMetadata], TokenUsage]:
|
||||
"""
|
||||
Extract facts from multiple content items in parallel.
|
||||
|
||||
|
|
@ -1147,10 +1156,10 @@ async def extract_facts_from_contents(
|
|||
extract_opinions: If True, extract only opinions; otherwise world/bank facts
|
||||
|
||||
Returns:
|
||||
Tuple of (extracted_facts, chunks_metadata)
|
||||
Tuple of (extracted_facts, chunks_metadata, usage)
|
||||
"""
|
||||
if not contents:
|
||||
return [], []
|
||||
return [], [], TokenUsage()
|
||||
|
||||
# Step 1: Create parallel fact extraction tasks
|
||||
fact_extraction_tasks = []
|
||||
|
|
@ -1173,11 +1182,15 @@ async def extract_facts_from_contents(
|
|||
# Step 3: Flatten and convert to typed objects
|
||||
extracted_facts: list[ExtractedFactType] = []
|
||||
chunks_metadata: list[ChunkMetadata] = []
|
||||
total_usage = TokenUsage()
|
||||
|
||||
global_chunk_idx = 0
|
||||
global_fact_idx = 0
|
||||
|
||||
for content_index, (content, (facts_from_llm, chunks_from_llm)) in enumerate(zip(contents, all_fact_results)):
|
||||
for content_index, (content, (facts_from_llm, chunks_from_llm, content_usage)) in enumerate(
|
||||
zip(contents, all_fact_results)
|
||||
):
|
||||
total_usage = total_usage + content_usage
|
||||
chunk_start_idx = global_chunk_idx
|
||||
|
||||
# Convert chunk tuples to ChunkMetadata objects
|
||||
|
|
@ -1231,7 +1244,7 @@ async def extract_facts_from_contents(
|
|||
# Step 4: Add time offsets to preserve ordering within each content
|
||||
_add_temporal_offsets(extracted_facts, contents)
|
||||
|
||||
return extracted_facts, chunks_metadata
|
||||
return extracted_facts, chunks_metadata, total_usage
|
||||
|
||||
|
||||
def _parse_datetime(date_str: str):
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ def utcnow():
|
|||
return datetime.now(UTC)
|
||||
|
||||
|
||||
from ..response_models import TokenUsage
|
||||
from . import (
|
||||
chunk_storage,
|
||||
deduplication,
|
||||
|
|
@ -47,7 +48,7 @@ async def retain_batch(
|
|||
is_first_batch: bool = True,
|
||||
fact_type_override: str | None = None,
|
||||
confidence_score: float | None = None,
|
||||
) -> list[list[str]]:
|
||||
) -> tuple[list[list[str]], TokenUsage]:
|
||||
"""
|
||||
Process a batch of content through the retain pipeline.
|
||||
|
||||
|
|
@ -67,7 +68,7 @@ async def retain_batch(
|
|||
confidence_score: Confidence score for opinions
|
||||
|
||||
Returns:
|
||||
List of unit ID lists (one list per content item)
|
||||
Tuple of (unit ID lists, token usage for fact extraction)
|
||||
"""
|
||||
start_time = time.time()
|
||||
total_chars = sum(len(item.get("content", "")) for item in contents_dicts)
|
||||
|
|
@ -99,7 +100,7 @@ async def retain_batch(
|
|||
step_start = time.time()
|
||||
extract_opinions = fact_type_override == "opinion"
|
||||
|
||||
extracted_facts, chunks = await fact_extraction.extract_facts_from_contents(
|
||||
extracted_facts, chunks, usage = await fact_extraction.extract_facts_from_contents(
|
||||
contents, llm_config, agent_name, extract_opinions
|
||||
)
|
||||
log_buffer.append(
|
||||
|
|
@ -164,7 +165,7 @@ async def retain_batch(
|
|||
logger.info(
|
||||
f"RETAIN_BATCH COMPLETE: 0 facts extracted from {len(contents)} contents in {total_time:.3f}s (document tracked, no facts)"
|
||||
)
|
||||
return [[] for _ in contents]
|
||||
return [[] for _ in contents], usage
|
||||
|
||||
# Apply fact_type_override if provided
|
||||
if fact_type_override:
|
||||
|
|
@ -344,7 +345,7 @@ async def retain_batch(
|
|||
non_duplicate_facts = deduplication.filter_duplicates(processed_facts, is_duplicate_flags)
|
||||
|
||||
if not non_duplicate_facts:
|
||||
return [[] for _ in contents]
|
||||
return [[] for _ in contents], usage
|
||||
|
||||
# Insert facts (document_id is now stored per-fact)
|
||||
step_start = time.time()
|
||||
|
|
@ -415,7 +416,7 @@ async def retain_batch(
|
|||
|
||||
logger.info("\n" + "\n".join(log_buffer) + "\n")
|
||||
|
||||
return result_unit_ids
|
||||
return result_unit_ids, usage
|
||||
|
||||
|
||||
def _map_results_to_contents(
|
||||
|
|
|
|||
|
|
@ -49,7 +49,7 @@ async def extract_facts(
|
|||
if not text or not text.strip():
|
||||
return [], []
|
||||
|
||||
facts, chunks = await extract_facts_from_text(
|
||||
facts, chunks, _ = await extract_facts_from_text(
|
||||
text,
|
||||
event_date,
|
||||
context=context,
|
||||
|
|
|
|||
|
|
@ -832,3 +832,134 @@ async def test_reflect_with_max_tokens(api_client):
|
|||
# Verify response has text
|
||||
assert "text" in result
|
||||
assert len(result["text"]) > 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reflect_returns_token_usage(api_client):
|
||||
"""Test that reflect endpoint returns token usage metrics.
|
||||
|
||||
The usage field should contain input_tokens, output_tokens, and total_tokens
|
||||
from the LLM call made during reflection.
|
||||
"""
|
||||
test_bank_id = f"reflect_usage_test_{datetime.now().timestamp()}"
|
||||
|
||||
# Store a memory to reflect on
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories",
|
||||
json={
|
||||
"items": [
|
||||
{
|
||||
"content": "The capital of France is Paris.",
|
||||
"context": "geography"
|
||||
}
|
||||
]
|
||||
}
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
# Call reflect
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/reflect",
|
||||
json={
|
||||
"query": "What is the capital of France?"
|
||||
}
|
||||
)
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
|
||||
# Verify response has text
|
||||
assert "text" in result
|
||||
assert len(result["text"]) > 0
|
||||
|
||||
# Verify usage field exists and has expected structure
|
||||
assert "usage" in result, "Response should include 'usage' field"
|
||||
usage = result["usage"]
|
||||
assert usage is not None, "Usage should not be None for reflect"
|
||||
assert "input_tokens" in usage, "Usage should have 'input_tokens'"
|
||||
assert "output_tokens" in usage, "Usage should have 'output_tokens'"
|
||||
assert "total_tokens" in usage, "Usage should have 'total_tokens'"
|
||||
|
||||
# Verify token counts are valid
|
||||
assert usage["input_tokens"] > 0, f"Expected input_tokens > 0, got {usage['input_tokens']}"
|
||||
assert usage["output_tokens"] >= 0, f"Expected output_tokens >= 0, got {usage['output_tokens']}"
|
||||
assert usage["total_tokens"] == usage["input_tokens"] + usage["output_tokens"]
|
||||
|
||||
print(f"Reflect token usage: input={usage['input_tokens']}, output={usage['output_tokens']}, total={usage['total_tokens']}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_returns_token_usage(api_client):
|
||||
"""Test that retain endpoint returns token usage metrics for synchronous operations.
|
||||
|
||||
The usage field should contain input_tokens, output_tokens, and total_tokens
|
||||
from the LLM calls made during fact extraction.
|
||||
"""
|
||||
test_bank_id = f"retain_usage_test_{datetime.now().timestamp()}"
|
||||
|
||||
# Store memory synchronously (async=false is default)
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories",
|
||||
json={
|
||||
"items": [
|
||||
{
|
||||
"content": "Alice is a software engineer at TechCorp. She specializes in machine learning.",
|
||||
"context": "team introduction"
|
||||
}
|
||||
]
|
||||
}
|
||||
)
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
|
||||
# Verify basic response
|
||||
assert result["success"] is True
|
||||
assert result["items_count"] == 1
|
||||
assert result["async"] is False
|
||||
|
||||
# Verify usage field exists and has expected structure
|
||||
assert "usage" in result, "Response should include 'usage' field"
|
||||
usage = result["usage"]
|
||||
assert usage is not None, "Usage should not be None for synchronous retain"
|
||||
assert "input_tokens" in usage, "Usage should have 'input_tokens'"
|
||||
assert "output_tokens" in usage, "Usage should have 'output_tokens'"
|
||||
assert "total_tokens" in usage, "Usage should have 'total_tokens'"
|
||||
|
||||
# Verify token counts are valid
|
||||
assert usage["input_tokens"] > 0, f"Expected input_tokens > 0, got {usage['input_tokens']}"
|
||||
assert usage["output_tokens"] >= 0, f"Expected output_tokens >= 0, got {usage['output_tokens']}"
|
||||
assert usage["total_tokens"] == usage["input_tokens"] + usage["output_tokens"]
|
||||
|
||||
print(f"Retain token usage: input={usage['input_tokens']}, output={usage['output_tokens']}, total={usage['total_tokens']}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_async_no_usage(api_client):
|
||||
"""Test that async retain does not return usage (as it's processed in background).
|
||||
|
||||
When async=true, the usage field should be None since the actual
|
||||
fact extraction happens asynchronously.
|
||||
"""
|
||||
test_bank_id = f"retain_async_no_usage_test_{datetime.now().timestamp()}"
|
||||
|
||||
# Store memory asynchronously
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories",
|
||||
json={
|
||||
"async": True,
|
||||
"items": [
|
||||
{
|
||||
"content": "Bob is a data scientist.",
|
||||
"context": "team introduction"
|
||||
}
|
||||
]
|
||||
}
|
||||
)
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
|
||||
# Verify async response
|
||||
assert result["success"] is True
|
||||
assert result["async"] is True
|
||||
|
||||
# Usage should be None for async operations
|
||||
assert result.get("usage") is None, "Async retain should not include usage"
|
||||
|
|
|
|||
242
hindsight-api/tests/test_llm_token_metrics.py
Normal file
242
hindsight-api/tests/test_llm_token_metrics.py
Normal file
|
|
@ -0,0 +1,242 @@
|
|||
"""
|
||||
Test that LLM calls record token metrics via the metrics collector.
|
||||
"""
|
||||
import os
|
||||
from unittest.mock import MagicMock, patch
|
||||
import pytest
|
||||
from hindsight_api.engine.llm_wrapper import LLMProvider
|
||||
from hindsight_api.metrics import (
|
||||
MetricsCollector,
|
||||
NoOpMetricsCollector,
|
||||
get_metrics_collector,
|
||||
initialize_metrics,
|
||||
create_metrics_collector,
|
||||
)
|
||||
|
||||
|
||||
def get_groq_api_key() -> str | None:
|
||||
"""Get Groq API key from environment."""
|
||||
return os.getenv("GROQ_API_KEY")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_token_metrics_recorded_for_groq():
|
||||
"""
|
||||
Test that token metrics are recorded when making LLM calls via Groq.
|
||||
Uses openai/gpt-oss-20b as recommended by Hindsight.
|
||||
"""
|
||||
api_key = get_groq_api_key()
|
||||
if not api_key:
|
||||
pytest.skip("Skipping: GROQ_API_KEY not set")
|
||||
|
||||
# Create a mock metrics collector to track record_tokens calls
|
||||
mock_collector = MagicMock(spec=MetricsCollector)
|
||||
|
||||
with patch("hindsight_api.engine.llm_wrapper.get_metrics_collector", return_value=mock_collector):
|
||||
llm = LLMProvider(
|
||||
provider="groq",
|
||||
api_key=api_key,
|
||||
base_url="",
|
||||
model="openai/gpt-oss-20b",
|
||||
)
|
||||
|
||||
# Make an LLM call with clear instruction
|
||||
response = await llm.call(
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant. Always respond."},
|
||||
{"role": "user", "content": "What is 2+2? Reply with just the number."}
|
||||
],
|
||||
max_completion_tokens=50,
|
||||
scope="test_metrics",
|
||||
)
|
||||
|
||||
# Verify record_tokens was called - this is the main test
|
||||
assert mock_collector.record_tokens.called, "record_tokens should have been called"
|
||||
|
||||
# Get the call arguments
|
||||
call_kwargs = mock_collector.record_tokens.call_args.kwargs
|
||||
|
||||
# Verify the call had correct structure
|
||||
assert call_kwargs["operation"] == "test_metrics", f"Expected operation='test_metrics', got {call_kwargs}"
|
||||
assert call_kwargs["bank_id"] == "llm", f"Expected bank_id='llm', got {call_kwargs}"
|
||||
assert call_kwargs["input_tokens"] > 0, f"Expected input_tokens > 0, got {call_kwargs['input_tokens']}"
|
||||
# Output tokens may be 0 for some edge cases, but input should always be > 0
|
||||
assert call_kwargs["output_tokens"] >= 0, f"Expected output_tokens >= 0, got {call_kwargs['output_tokens']}"
|
||||
|
||||
print(f"\nToken metrics recorded:")
|
||||
print(f" operation: {call_kwargs['operation']}")
|
||||
print(f" input_tokens: {call_kwargs['input_tokens']}")
|
||||
print(f" output_tokens: {call_kwargs['output_tokens']}")
|
||||
print(f" response: {response}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_token_metrics_recorded_for_structured_output():
|
||||
"""
|
||||
Test that token metrics are recorded for structured output (JSON) calls.
|
||||
"""
|
||||
api_key = get_groq_api_key()
|
||||
if not api_key:
|
||||
pytest.skip("Skipping: GROQ_API_KEY not set")
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
class SimpleResponse(BaseModel):
|
||||
greeting: str
|
||||
language: str
|
||||
|
||||
mock_collector = MagicMock(spec=MetricsCollector)
|
||||
|
||||
with patch("hindsight_api.engine.llm_wrapper.get_metrics_collector", return_value=mock_collector):
|
||||
llm = LLMProvider(
|
||||
provider="groq",
|
||||
api_key=api_key,
|
||||
base_url="",
|
||||
model="openai/gpt-oss-20b",
|
||||
)
|
||||
|
||||
# Make a structured output call
|
||||
response = await llm.call(
|
||||
messages=[{"role": "user", "content": "Say hello in French. Return greeting and language."}],
|
||||
response_format=SimpleResponse,
|
||||
max_completion_tokens=100,
|
||||
scope="structured_output_test",
|
||||
)
|
||||
|
||||
# Verify structured response
|
||||
assert isinstance(response, SimpleResponse)
|
||||
assert response.greeting is not None
|
||||
assert response.language is not None
|
||||
|
||||
# Verify record_tokens was called
|
||||
assert mock_collector.record_tokens.called, "record_tokens should have been called"
|
||||
|
||||
call_kwargs = mock_collector.record_tokens.call_args.kwargs
|
||||
assert call_kwargs["input_tokens"] > 0
|
||||
assert call_kwargs["output_tokens"] > 0
|
||||
|
||||
print(f"\nStructured output token metrics:")
|
||||
print(f" greeting: {response.greeting}")
|
||||
print(f" language: {response.language}")
|
||||
print(f" input_tokens: {call_kwargs['input_tokens']}")
|
||||
print(f" output_tokens: {call_kwargs['output_tokens']}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_noop_collector_when_metrics_disabled():
|
||||
"""
|
||||
Test that NoOpMetricsCollector is returned when metrics are not initialized.
|
||||
This verifies the fallback behavior doesn't break LLM calls.
|
||||
"""
|
||||
api_key = get_groq_api_key()
|
||||
if not api_key:
|
||||
pytest.skip("Skipping: GROQ_API_KEY not set")
|
||||
|
||||
# Without initializing metrics, get_metrics_collector returns NoOpMetricsCollector
|
||||
collector = get_metrics_collector()
|
||||
assert isinstance(collector, NoOpMetricsCollector), "Should return NoOpMetricsCollector when not initialized"
|
||||
|
||||
# Make an LLM call - should work fine with NoOp collector
|
||||
llm = LLMProvider(
|
||||
provider="groq",
|
||||
api_key=api_key,
|
||||
base_url="",
|
||||
model="openai/gpt-oss-20b",
|
||||
)
|
||||
|
||||
response = await llm.call(
|
||||
messages=[{"role": "user", "content": "Say 'test' in one word."}],
|
||||
max_completion_tokens=50,
|
||||
)
|
||||
|
||||
assert response is not None
|
||||
print(f"\nLLM call succeeded with NoOpMetricsCollector: {response}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_return_usage_returns_tuple():
|
||||
"""
|
||||
Test that return_usage=True returns (result, TokenUsage) tuple.
|
||||
"""
|
||||
from hindsight_api.engine.response_models import TokenUsage
|
||||
|
||||
api_key = get_groq_api_key()
|
||||
if not api_key:
|
||||
pytest.skip("Skipping: GROQ_API_KEY not set")
|
||||
|
||||
llm = LLMProvider(
|
||||
provider="groq",
|
||||
api_key=api_key,
|
||||
base_url="",
|
||||
model="openai/gpt-oss-20b",
|
||||
)
|
||||
|
||||
# Call with return_usage=True
|
||||
result, usage = await llm.call(
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "What is 2+2? Reply with just the number."}
|
||||
],
|
||||
max_completion_tokens=50,
|
||||
return_usage=True,
|
||||
)
|
||||
|
||||
# Verify result is the response text
|
||||
assert result is not None
|
||||
assert isinstance(result, str)
|
||||
|
||||
# Verify usage is TokenUsage model with valid counts
|
||||
assert isinstance(usage, TokenUsage)
|
||||
assert usage.input_tokens > 0, f"Expected input_tokens > 0, got {usage.input_tokens}"
|
||||
assert usage.output_tokens >= 0, f"Expected output_tokens >= 0, got {usage.output_tokens}"
|
||||
assert usage.total_tokens == usage.input_tokens + usage.output_tokens
|
||||
|
||||
print(f"\nreturn_usage=True test:")
|
||||
print(f" result: {result}")
|
||||
print(f" usage: {usage}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_return_usage_with_structured_output():
|
||||
"""
|
||||
Test that return_usage=True works with structured output (JSON).
|
||||
"""
|
||||
from pydantic import BaseModel
|
||||
from hindsight_api.engine.response_models import TokenUsage
|
||||
|
||||
api_key = get_groq_api_key()
|
||||
if not api_key:
|
||||
pytest.skip("Skipping: GROQ_API_KEY not set")
|
||||
|
||||
class MathAnswer(BaseModel):
|
||||
answer: int
|
||||
explanation: str
|
||||
|
||||
llm = LLMProvider(
|
||||
provider="groq",
|
||||
api_key=api_key,
|
||||
base_url="",
|
||||
model="openai/gpt-oss-20b",
|
||||
)
|
||||
|
||||
# Call with return_usage=True and structured output
|
||||
result, usage = await llm.call(
|
||||
messages=[{"role": "user", "content": "What is 5+3? Return the answer and a brief explanation."}],
|
||||
response_format=MathAnswer,
|
||||
max_completion_tokens=100,
|
||||
return_usage=True,
|
||||
)
|
||||
|
||||
# Verify result is the parsed response
|
||||
assert isinstance(result, MathAnswer)
|
||||
assert result.answer == 8
|
||||
assert result.explanation is not None
|
||||
|
||||
# Verify usage is TokenUsage model
|
||||
assert isinstance(usage, TokenUsage)
|
||||
assert usage.input_tokens > 0
|
||||
assert usage.output_tokens > 0
|
||||
|
||||
print(f"\nStructured output with return_usage=True:")
|
||||
print(f" result: {result}")
|
||||
print(f" usage: {usage}")
|
||||
|
|
@ -51,6 +51,7 @@ hindsight_client_api/models/reflect_request.py
|
|||
hindsight_client_api/models/reflect_response.py
|
||||
hindsight_client_api/models/retain_request.py
|
||||
hindsight_client_api/models/retain_response.py
|
||||
hindsight_client_api/models/token_usage.py
|
||||
hindsight_client_api/models/update_disposition_request.py
|
||||
hindsight_client_api/models/validation_error.py
|
||||
hindsight_client_api/models/validation_error_loc_inner.py
|
||||
|
|
|
|||
|
|
@ -76,6 +76,7 @@ from hindsight_client_api.models.reflect_request import ReflectRequest
|
|||
from hindsight_client_api.models.reflect_response import ReflectResponse
|
||||
from hindsight_client_api.models.retain_request import RetainRequest
|
||||
from hindsight_client_api.models.retain_response import RetainResponse
|
||||
from hindsight_client_api.models.token_usage import TokenUsage
|
||||
from hindsight_client_api.models.update_disposition_request import UpdateDispositionRequest
|
||||
from hindsight_client_api.models.validation_error import ValidationError
|
||||
from hindsight_client_api.models.validation_error_loc_inner import ValidationErrorLocInner
|
||||
|
|
|
|||
|
|
@ -54,6 +54,7 @@ from hindsight_client_api.models.reflect_request import ReflectRequest
|
|||
from hindsight_client_api.models.reflect_response import ReflectResponse
|
||||
from hindsight_client_api.models.retain_request import RetainRequest
|
||||
from hindsight_client_api.models.retain_response import RetainResponse
|
||||
from hindsight_client_api.models.token_usage import TokenUsage
|
||||
from hindsight_client_api.models.update_disposition_request import UpdateDispositionRequest
|
||||
from hindsight_client_api.models.validation_error import ValidationError
|
||||
from hindsight_client_api.models.validation_error_loc_inner import ValidationErrorLocInner
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ import json
|
|||
from pydantic import BaseModel, ConfigDict, StrictStr
|
||||
from typing import Any, ClassVar, Dict, List, Optional
|
||||
from hindsight_client_api.models.reflect_fact import ReflectFact
|
||||
from hindsight_client_api.models.token_usage import TokenUsage
|
||||
from typing import Optional, Set
|
||||
from typing_extensions import Self
|
||||
|
||||
|
|
@ -30,7 +31,8 @@ class ReflectResponse(BaseModel):
|
|||
text: StrictStr
|
||||
based_on: Optional[List[ReflectFact]] = None
|
||||
structured_output: Optional[Dict[str, Any]] = None
|
||||
__properties: ClassVar[List[str]] = ["text", "based_on", "structured_output"]
|
||||
usage: Optional[TokenUsage] = None
|
||||
__properties: ClassVar[List[str]] = ["text", "based_on", "structured_output", "usage"]
|
||||
|
||||
model_config = ConfigDict(
|
||||
populate_by_name=True,
|
||||
|
|
@ -78,11 +80,19 @@ class ReflectResponse(BaseModel):
|
|||
if _item_based_on:
|
||||
_items.append(_item_based_on.to_dict())
|
||||
_dict['based_on'] = _items
|
||||
# override the default output from pydantic by calling `to_dict()` of usage
|
||||
if self.usage:
|
||||
_dict['usage'] = self.usage.to_dict()
|
||||
# set to None if structured_output (nullable) is None
|
||||
# and model_fields_set contains the field
|
||||
if self.structured_output is None and "structured_output" in self.model_fields_set:
|
||||
_dict['structured_output'] = None
|
||||
|
||||
# set to None if usage (nullable) is None
|
||||
# and model_fields_set contains the field
|
||||
if self.usage is None and "usage" in self.model_fields_set:
|
||||
_dict['usage'] = None
|
||||
|
||||
return _dict
|
||||
|
||||
@classmethod
|
||||
|
|
@ -97,7 +107,8 @@ class ReflectResponse(BaseModel):
|
|||
_obj = cls.model_validate({
|
||||
"text": obj.get("text"),
|
||||
"based_on": [ReflectFact.from_dict(_item) for _item in obj["based_on"]] if obj.get("based_on") is not None else None,
|
||||
"structured_output": obj.get("structured_output")
|
||||
"structured_output": obj.get("structured_output"),
|
||||
"usage": TokenUsage.from_dict(obj["usage"]) if obj.get("usage") is not None else None
|
||||
})
|
||||
return _obj
|
||||
|
||||
|
|
|
|||
|
|
@ -18,7 +18,8 @@ import re # noqa: F401
|
|||
import json
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, StrictBool, StrictInt, StrictStr
|
||||
from typing import Any, ClassVar, Dict, List
|
||||
from typing import Any, ClassVar, Dict, List, Optional
|
||||
from hindsight_client_api.models.token_usage import TokenUsage
|
||||
from typing import Optional, Set
|
||||
from typing_extensions import Self
|
||||
|
||||
|
|
@ -30,7 +31,8 @@ class RetainResponse(BaseModel):
|
|||
bank_id: StrictStr
|
||||
items_count: StrictInt
|
||||
var_async: StrictBool = Field(description="Whether the operation was processed asynchronously", alias="async")
|
||||
__properties: ClassVar[List[str]] = ["success", "bank_id", "items_count", "async"]
|
||||
usage: Optional[TokenUsage] = None
|
||||
__properties: ClassVar[List[str]] = ["success", "bank_id", "items_count", "async", "usage"]
|
||||
|
||||
model_config = ConfigDict(
|
||||
populate_by_name=True,
|
||||
|
|
@ -71,6 +73,14 @@ class RetainResponse(BaseModel):
|
|||
exclude=excluded_fields,
|
||||
exclude_none=True,
|
||||
)
|
||||
# override the default output from pydantic by calling `to_dict()` of usage
|
||||
if self.usage:
|
||||
_dict['usage'] = self.usage.to_dict()
|
||||
# set to None if usage (nullable) is None
|
||||
# and model_fields_set contains the field
|
||||
if self.usage is None and "usage" in self.model_fields_set:
|
||||
_dict['usage'] = None
|
||||
|
||||
return _dict
|
||||
|
||||
@classmethod
|
||||
|
|
@ -86,7 +96,8 @@ class RetainResponse(BaseModel):
|
|||
"success": obj.get("success"),
|
||||
"bank_id": obj.get("bank_id"),
|
||||
"items_count": obj.get("items_count"),
|
||||
"async": obj.get("async")
|
||||
"async": obj.get("async"),
|
||||
"usage": TokenUsage.from_dict(obj["usage"]) if obj.get("usage") is not None else None
|
||||
})
|
||||
return _obj
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,91 @@
|
|||
# coding: utf-8
|
||||
|
||||
"""
|
||||
Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.1.0
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
""" # noqa: E501
|
||||
|
||||
|
||||
from __future__ import annotations
|
||||
import pprint
|
||||
import re # noqa: F401
|
||||
import json
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, StrictInt
|
||||
from typing import Any, ClassVar, Dict, List, Optional
|
||||
from typing import Optional, Set
|
||||
from typing_extensions import Self
|
||||
|
||||
class TokenUsage(BaseModel):
|
||||
"""
|
||||
Token usage metrics for LLM calls. Tracks input/output tokens for a single request to enable per-request cost tracking and monitoring.
|
||||
""" # noqa: E501
|
||||
input_tokens: Optional[StrictInt] = Field(default=0, description="Number of input/prompt tokens consumed")
|
||||
output_tokens: Optional[StrictInt] = Field(default=0, description="Number of output/completion tokens generated")
|
||||
total_tokens: Optional[StrictInt] = Field(default=0, description="Total tokens (input + output)")
|
||||
__properties: ClassVar[List[str]] = ["input_tokens", "output_tokens", "total_tokens"]
|
||||
|
||||
model_config = ConfigDict(
|
||||
populate_by_name=True,
|
||||
validate_assignment=True,
|
||||
protected_namespaces=(),
|
||||
)
|
||||
|
||||
|
||||
def to_str(self) -> str:
|
||||
"""Returns the string representation of the model using alias"""
|
||||
return pprint.pformat(self.model_dump(by_alias=True))
|
||||
|
||||
def to_json(self) -> str:
|
||||
"""Returns the JSON representation of the model using alias"""
|
||||
# TODO: pydantic v2: use .model_dump_json(by_alias=True, exclude_unset=True) instead
|
||||
return json.dumps(self.to_dict())
|
||||
|
||||
@classmethod
|
||||
def from_json(cls, json_str: str) -> Optional[Self]:
|
||||
"""Create an instance of TokenUsage from a JSON string"""
|
||||
return cls.from_dict(json.loads(json_str))
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Return the dictionary representation of the model using alias.
|
||||
|
||||
This has the following differences from calling pydantic's
|
||||
`self.model_dump(by_alias=True)`:
|
||||
|
||||
* `None` is only added to the output dict for nullable fields that
|
||||
were set at model initialization. Other fields with value `None`
|
||||
are ignored.
|
||||
"""
|
||||
excluded_fields: Set[str] = set([
|
||||
])
|
||||
|
||||
_dict = self.model_dump(
|
||||
by_alias=True,
|
||||
exclude=excluded_fields,
|
||||
exclude_none=True,
|
||||
)
|
||||
return _dict
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, obj: Optional[Dict[str, Any]]) -> Optional[Self]:
|
||||
"""Create an instance of TokenUsage from a dict"""
|
||||
if obj is None:
|
||||
return None
|
||||
|
||||
if not isinstance(obj, dict):
|
||||
return cls.model_validate(obj)
|
||||
|
||||
_obj = cls.model_validate({
|
||||
"input_tokens": obj.get("input_tokens") if obj.get("input_tokens") is not None else 0,
|
||||
"output_tokens": obj.get("output_tokens") if obj.get("output_tokens") is not None else 0,
|
||||
"total_tokens": obj.get("total_tokens") if obj.get("total_tokens") is not None else 0
|
||||
})
|
||||
return _obj
|
||||
|
||||
|
||||
8
hindsight-clients/typescript/.prettierrc
Normal file
8
hindsight-clients/typescript/.prettierrc
Normal file
|
|
@ -0,0 +1,8 @@
|
|||
{
|
||||
"tabWidth": 2,
|
||||
"useTabs": false,
|
||||
"semi": true,
|
||||
"singleQuote": false,
|
||||
"trailingComma": "all",
|
||||
"printWidth": 80
|
||||
}
|
||||
|
|
@ -966,6 +966,10 @@ export type ReflectResponse = {
|
|||
structured_output?: {
|
||||
[key: string]: unknown;
|
||||
} | null;
|
||||
/**
|
||||
* Token usage metrics for LLM calls during reflection.
|
||||
*/
|
||||
usage?: TokenUsage | null;
|
||||
};
|
||||
|
||||
/**
|
||||
|
|
@ -1010,6 +1014,39 @@ export type RetainResponse = {
|
|||
* Whether the operation was processed asynchronously
|
||||
*/
|
||||
async: boolean;
|
||||
/**
|
||||
* Token usage metrics for LLM calls during fact extraction (only present for synchronous operations)
|
||||
*/
|
||||
usage?: TokenUsage | null;
|
||||
};
|
||||
|
||||
/**
|
||||
* TokenUsage
|
||||
*
|
||||
* Token usage metrics for LLM calls.
|
||||
*
|
||||
* Tracks input/output tokens for a single request to enable
|
||||
* per-request cost tracking and monitoring.
|
||||
*/
|
||||
export type TokenUsage = {
|
||||
/**
|
||||
* Input Tokens
|
||||
*
|
||||
* Number of input/prompt tokens consumed
|
||||
*/
|
||||
input_tokens?: number;
|
||||
/**
|
||||
* Output Tokens
|
||||
*
|
||||
* Number of output/completion tokens generated
|
||||
*/
|
||||
output_tokens?: number;
|
||||
/**
|
||||
* Total Tokens
|
||||
*
|
||||
* Total tokens (input + output)
|
||||
*/
|
||||
total_tokens?: number;
|
||||
};
|
||||
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -8,6 +8,12 @@ This changelog highlights user-facing changes only. Internal maintenance, CI/CD,
|
|||
|
||||
For full release details, see [GitHub Releases](https://github.com/vectorize-io/hindsight/releases).
|
||||
|
||||
## [Unreleased]
|
||||
|
||||
**Features**
|
||||
|
||||
- Add per-request token usage tracking to retain and reflect endpoints for cost monitoring and billing integration.
|
||||
|
||||
## [0.2.0](https://github.com/vectorize-io/hindsight/releases/tag/v0.2.0)
|
||||
|
||||
**Features**
|
||||
|
|
|
|||
|
|
@ -55,6 +55,20 @@ Make sure you've completed the [Quick Start](./quickstart) to install the client
|
|||
| `max_tokens` | int | 4096 | Maximum tokens for the response |
|
||||
| `response_schema` | object | None | JSON Schema for [structured output](#structured-output) |
|
||||
|
||||
### Response Fields
|
||||
|
||||
| Field | Type | Description |
|
||||
|-------|------|-------------|
|
||||
| `text` | string | The generated answer text |
|
||||
| `based_on` | array | Facts used to generate the response |
|
||||
| `structured_output` | object | Parsed structured output (when `response_schema` provided) |
|
||||
| `usage` | TokenUsage | Token usage metrics for the LLM call |
|
||||
|
||||
The `usage` field contains:
|
||||
- `input_tokens`: Number of input/prompt tokens consumed
|
||||
- `output_tokens`: Number of output/completion tokens generated
|
||||
- `total_tokens`: Sum of input and output tokens
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="python" label="Python">
|
||||
<CodeSnippet code={reflectPy} section="reflect-with-params" language="python" />
|
||||
|
|
|
|||
|
|
@ -66,6 +66,25 @@ Always provide context and event dates for optimal memory extraction:
|
|||
|
||||
The `timestamp` defaults to the current time if not specified. Providing explicit timestamps enables temporal queries like "What happened last spring?"
|
||||
|
||||
### Response Fields
|
||||
|
||||
The retain response includes:
|
||||
|
||||
| Field | Type | Description |
|
||||
|-------|------|-------------|
|
||||
| `success` | bool | Whether the operation succeeded |
|
||||
| `bank_id` | string | The memory bank ID |
|
||||
| `items_count` | int | Number of items processed |
|
||||
| `async` | bool | Whether processed asynchronously |
|
||||
| `usage` | TokenUsage | Token usage metrics for LLM calls (synchronous only) |
|
||||
|
||||
The `usage` field contains token metrics for cost tracking:
|
||||
- `input_tokens`: Tokens consumed by prompts
|
||||
- `output_tokens`: Tokens generated by the LLM
|
||||
- `total_tokens`: Sum of input and output tokens
|
||||
|
||||
Note: `usage` is only present for synchronous operations. Async operations (`async: true`) do not return usage metrics.
|
||||
|
||||
## Batch Ingestion
|
||||
|
||||
Store multiple items in a single request. **Batch ingestion is the recommended approach** as it significantly improves performance by reducing network overhead and allowing Hindsight to optimize the memory extraction process across related content.
|
||||
|
|
|
|||
|
|
@ -3404,6 +3404,17 @@
|
|||
],
|
||||
"title": "Structured Output",
|
||||
"description": "Structured output parsed according to the request's response_schema. Only present when response_schema was provided in the request."
|
||||
},
|
||||
"usage": {
|
||||
"anyOf": [
|
||||
{
|
||||
"$ref": "#/components/schemas/TokenUsage"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"description": "Token usage metrics for LLM calls during reflection."
|
||||
}
|
||||
},
|
||||
"type": "object",
|
||||
|
|
@ -3432,7 +3443,12 @@
|
|||
],
|
||||
"summary": "AI is transformative"
|
||||
},
|
||||
"text": "Based on my understanding, AI is a transformative technology..."
|
||||
"text": "Based on my understanding, AI is a transformative technology...",
|
||||
"usage": {
|
||||
"input_tokens": 1500,
|
||||
"output_tokens": 500,
|
||||
"total_tokens": 2000
|
||||
}
|
||||
}
|
||||
},
|
||||
"RetainRequest": {
|
||||
|
|
@ -3491,6 +3507,17 @@
|
|||
"type": "boolean",
|
||||
"title": "Async",
|
||||
"description": "Whether the operation was processed asynchronously"
|
||||
},
|
||||
"usage": {
|
||||
"anyOf": [
|
||||
{
|
||||
"$ref": "#/components/schemas/TokenUsage"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"description": "Token usage metrics for LLM calls during fact extraction (only present for synchronous operations)"
|
||||
}
|
||||
},
|
||||
"type": "object",
|
||||
|
|
@ -3506,7 +3533,42 @@
|
|||
"async": false,
|
||||
"bank_id": "user123",
|
||||
"items_count": 2,
|
||||
"success": true
|
||||
"success": true,
|
||||
"usage": {
|
||||
"input_tokens": 500,
|
||||
"output_tokens": 100,
|
||||
"total_tokens": 600
|
||||
}
|
||||
}
|
||||
},
|
||||
"TokenUsage": {
|
||||
"properties": {
|
||||
"input_tokens": {
|
||||
"type": "integer",
|
||||
"title": "Input Tokens",
|
||||
"description": "Number of input/prompt tokens consumed",
|
||||
"default": 0
|
||||
},
|
||||
"output_tokens": {
|
||||
"type": "integer",
|
||||
"title": "Output Tokens",
|
||||
"description": "Number of output/completion tokens generated",
|
||||
"default": 0
|
||||
},
|
||||
"total_tokens": {
|
||||
"type": "integer",
|
||||
"title": "Total Tokens",
|
||||
"description": "Total tokens (input + output)",
|
||||
"default": 0
|
||||
}
|
||||
},
|
||||
"type": "object",
|
||||
"title": "TokenUsage",
|
||||
"description": "Token usage metrics for LLM calls.\n\nTracks input/output tokens for a single request to enable\nper-request cost tracking and monitoring.",
|
||||
"example": {
|
||||
"input_tokens": 1500,
|
||||
"output_tokens": 500,
|
||||
"total_tokens": 2000
|
||||
}
|
||||
},
|
||||
"UpdateDispositionRequest": {
|
||||
|
|
|
|||
Loading…
Reference in a new issue