diff --git a/hindsight-api/hindsight_api/api/http.py b/hindsight-api/hindsight_api/api/http.py index 20958ae9..71cf122d 100644 --- a/hindsight-api/hindsight_api/api/http.py +++ b/hindsight-api/hindsight_api/api/http.py @@ -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) diff --git a/hindsight-api/hindsight_api/engine/llm_wrapper.py b/hindsight-api/hindsight_api/engine/llm_wrapper.py index fd8f436a..65ba39ed 100644 --- a/hindsight-api/hindsight_api/engine/llm_wrapper.py +++ b/hindsight-api/hindsight_api/engine/llm_wrapper.py @@ -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: diff --git a/hindsight-api/hindsight_api/engine/memory_engine.py b/hindsight-api/hindsight_api/engine/memory_engine.py index a94f5dcd..a0c43fd7 100644 --- a/hindsight-api/hindsight_api/engine/memory_engine.py +++ b/hindsight-api/hindsight_api/engine/memory_engine.py @@ -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 diff --git a/hindsight-api/hindsight_api/engine/response_models.py b/hindsight-api/hindsight_api/engine/response_models.py index 9607848b..c9f7568d 100644 --- a/hindsight-api/hindsight_api/engine/response_models.py +++ b/hindsight-api/hindsight_api/engine/response_models.py @@ -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): diff --git a/hindsight-api/hindsight_api/engine/retain/fact_extraction.py b/hindsight-api/hindsight_api/engine/retain/fact_extraction.py index ab3494b6..b51c52a5 100644 --- a/hindsight-api/hindsight_api/engine/retain/fact_extraction.py +++ b/hindsight-api/hindsight_api/engine/retain/fact_extraction.py @@ -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): diff --git a/hindsight-api/hindsight_api/engine/retain/orchestrator.py b/hindsight-api/hindsight_api/engine/retain/orchestrator.py index 7cebf666..05b2b0be 100644 --- a/hindsight-api/hindsight_api/engine/retain/orchestrator.py +++ b/hindsight-api/hindsight_api/engine/retain/orchestrator.py @@ -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( diff --git a/hindsight-api/hindsight_api/engine/utils.py b/hindsight-api/hindsight_api/engine/utils.py index 1d1a132b..87f0824b 100644 --- a/hindsight-api/hindsight_api/engine/utils.py +++ b/hindsight-api/hindsight_api/engine/utils.py @@ -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, diff --git a/hindsight-api/tests/test_http_api_integration.py b/hindsight-api/tests/test_http_api_integration.py index 49b22364..5c3b7b15 100644 --- a/hindsight-api/tests/test_http_api_integration.py +++ b/hindsight-api/tests/test_http_api_integration.py @@ -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" diff --git a/hindsight-api/tests/test_llm_token_metrics.py b/hindsight-api/tests/test_llm_token_metrics.py new file mode 100644 index 00000000..d4a81e67 --- /dev/null +++ b/hindsight-api/tests/test_llm_token_metrics.py @@ -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}") diff --git a/hindsight-clients/python/.openapi-generator/FILES b/hindsight-clients/python/.openapi-generator/FILES index af6a82fc..6ddbc545 100644 --- a/hindsight-clients/python/.openapi-generator/FILES +++ b/hindsight-clients/python/.openapi-generator/FILES @@ -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 diff --git a/hindsight-clients/python/hindsight_client_api/__init__.py b/hindsight-clients/python/hindsight_client_api/__init__.py index 14ac3d49..76992a52 100644 --- a/hindsight-clients/python/hindsight_client_api/__init__.py +++ b/hindsight-clients/python/hindsight_client_api/__init__.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 diff --git a/hindsight-clients/python/hindsight_client_api/models/__init__.py b/hindsight-clients/python/hindsight_client_api/models/__init__.py index 5cb0bab9..3c7ecc33 100644 --- a/hindsight-clients/python/hindsight_client_api/models/__init__.py +++ b/hindsight-clients/python/hindsight_client_api/models/__init__.py @@ -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 diff --git a/hindsight-clients/python/hindsight_client_api/models/reflect_response.py b/hindsight-clients/python/hindsight_client_api/models/reflect_response.py index 74773142..177ae36a 100644 --- a/hindsight-clients/python/hindsight_client_api/models/reflect_response.py +++ b/hindsight-clients/python/hindsight_client_api/models/reflect_response.py @@ -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 diff --git a/hindsight-clients/python/hindsight_client_api/models/retain_response.py b/hindsight-clients/python/hindsight_client_api/models/retain_response.py index 3ca9de2f..10fe8a67 100644 --- a/hindsight-clients/python/hindsight_client_api/models/retain_response.py +++ b/hindsight-clients/python/hindsight_client_api/models/retain_response.py @@ -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 diff --git a/hindsight-clients/python/hindsight_client_api/models/token_usage.py b/hindsight-clients/python/hindsight_client_api/models/token_usage.py new file mode 100644 index 00000000..2b19d301 --- /dev/null +++ b/hindsight-clients/python/hindsight_client_api/models/token_usage.py @@ -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 + + diff --git a/hindsight-clients/typescript/.prettierrc b/hindsight-clients/typescript/.prettierrc new file mode 100644 index 00000000..a6d2a48e --- /dev/null +++ b/hindsight-clients/typescript/.prettierrc @@ -0,0 +1,8 @@ +{ + "tabWidth": 2, + "useTabs": false, + "semi": true, + "singleQuote": false, + "trailingComma": "all", + "printWidth": 80 +} diff --git a/hindsight-clients/typescript/generated/types.gen.ts b/hindsight-clients/typescript/generated/types.gen.ts index 6eb760e8..30f160cc 100644 --- a/hindsight-clients/typescript/generated/types.gen.ts +++ b/hindsight-clients/typescript/generated/types.gen.ts @@ -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; }; /** diff --git a/hindsight-docs/docs/changelog/index.md b/hindsight-docs/docs/changelog/index.md index a00a5c1a..5c048043 100644 --- a/hindsight-docs/docs/changelog/index.md +++ b/hindsight-docs/docs/changelog/index.md @@ -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** diff --git a/hindsight-docs/docs/developer/api/reflect.mdx b/hindsight-docs/docs/developer/api/reflect.mdx index 751ea08f..9944fe99 100644 --- a/hindsight-docs/docs/developer/api/reflect.mdx +++ b/hindsight-docs/docs/developer/api/reflect.mdx @@ -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 + diff --git a/hindsight-docs/docs/developer/api/retain.mdx b/hindsight-docs/docs/developer/api/retain.mdx index f3cf9bad..4116dbfe 100644 --- a/hindsight-docs/docs/developer/api/retain.mdx +++ b/hindsight-docs/docs/developer/api/retain.mdx @@ -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. diff --git a/hindsight-docs/static/openapi.json b/hindsight-docs/static/openapi.json index 820f628b..ef44a9a1 100644 --- a/hindsight-docs/static/openapi.json +++ b/hindsight-docs/static/openapi.json @@ -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": {