Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
dc30ba63f5 | ||
|
|
78ae3bbbff |
@@ -3383,6 +3383,8 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
llm_input_tokens=total_usage.input_tokens,
|
||||
llm_output_tokens=total_usage.output_tokens,
|
||||
llm_total_tokens=total_usage.total_tokens,
|
||||
llm_cached_input_tokens=getattr(total_usage, "cached_tokens", 0) or 0,
|
||||
llm_thoughts_tokens=getattr(total_usage, "thoughts_tokens", 0) or 0,
|
||||
processed_content_tokens=total_processed_content_tokens,
|
||||
)
|
||||
try:
|
||||
|
||||
@@ -12,6 +12,7 @@ import io
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
import traceback
|
||||
from contextvars import ContextVar
|
||||
from typing import Any
|
||||
|
||||
@@ -27,6 +28,79 @@ from hindsight_api.worker.stage import set_stage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Dedicated structured logger for per-call attribution. Operators can route
|
||||
# this to BigQuery / Loki for offline reconciliation against provider billing
|
||||
# without polluting the general application log. Each line is JSON.
|
||||
_call_audit_logger = logging.getLogger("hindsight.llm.gemini.calls")
|
||||
|
||||
|
||||
def _truncated_caller_stack(max_frames: int = 8) -> list[str]:
|
||||
"""Return up to ``max_frames`` of caller info, project files only.
|
||||
|
||||
Used by the per-call audit log to attribute LLM spend back to the source
|
||||
code path. Stdlib + 3p frames are skipped so the output stays small and
|
||||
actionable. Only file:lineno:func is captured — no locals, no source
|
||||
lines, no message content — so this is safe to ship to long-term log
|
||||
storage.
|
||||
"""
|
||||
frames: list[str] = []
|
||||
# ``[:-2]`` skips the helper itself and the immediate gemini-provider frame.
|
||||
for frame in traceback.extract_stack(limit=64)[:-2]:
|
||||
if "/site-packages/" in frame.filename:
|
||||
continue
|
||||
if "/lib/python" in frame.filename:
|
||||
continue
|
||||
# Trim absolute path to the last three path segments so the line
|
||||
# stays short and the project structure is still recognisable.
|
||||
segments = frame.filename.rsplit("/", 3)
|
||||
path = "/".join(segments[-3:]) if len(segments) > 3 else frame.filename
|
||||
frames.append(f"{path}:{frame.lineno}:{frame.name}")
|
||||
if len(frames) >= max_frames:
|
||||
break
|
||||
return frames
|
||||
|
||||
|
||||
def _emit_call_audit(
|
||||
*,
|
||||
provider: str,
|
||||
model: str,
|
||||
scope: str | None,
|
||||
input_tokens: int,
|
||||
cached_input_tokens: int,
|
||||
output_tokens: int,
|
||||
thoughts_tokens: int,
|
||||
duration_ms: int,
|
||||
finish_reason: str | None,
|
||||
) -> None:
|
||||
"""Emit one JSON line per Gemini call for offline attribution analysis.
|
||||
|
||||
Goes to a dedicated logger so operators can route it independently of
|
||||
the application log. The cost is ~400 bytes per call. At 200k
|
||||
calls/day that's ~80 MB/day — well within any normal log retention
|
||||
budget.
|
||||
"""
|
||||
try:
|
||||
_call_audit_logger.info(
|
||||
json.dumps(
|
||||
{
|
||||
"provider": provider,
|
||||
"model": model,
|
||||
"scope": scope,
|
||||
"input_tokens": input_tokens,
|
||||
"cached_input_tokens": cached_input_tokens,
|
||||
"output_tokens": output_tokens,
|
||||
"thoughts_tokens": thoughts_tokens,
|
||||
"duration_ms": duration_ms,
|
||||
"finish_reason": finish_reason,
|
||||
"caller_stack": _truncated_caller_stack(),
|
||||
}
|
||||
)
|
||||
)
|
||||
except Exception:
|
||||
# Audit must never fail the request path. Swallow and move on; the
|
||||
# Prometheus metrics path is the durable signal.
|
||||
pass
|
||||
|
||||
# Per-request Gemini safety settings override.
|
||||
# Set exclusively by ConfiguredLLMProvider.call() / call_with_tools() via token-based
|
||||
# set/reset, so it is properly scoped to each individual LLM call and never leaks.
|
||||
@@ -406,12 +480,25 @@ class GeminiLLM(LLMInterface):
|
||||
f"time={duration:.3f}s"
|
||||
)
|
||||
|
||||
_emit_call_audit(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
input_tokens=input_tokens,
|
||||
cached_input_tokens=cached_input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
thoughts_tokens=thoughts_tokens,
|
||||
duration_ms=int(duration * 1000),
|
||||
finish_reason=None,
|
||||
)
|
||||
|
||||
if return_usage:
|
||||
token_usage = TokenUsage(
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
total_tokens=input_tokens + output_tokens,
|
||||
cached_tokens=cached_tokens,
|
||||
thoughts_tokens=thoughts_tokens,
|
||||
)
|
||||
return result, token_usage
|
||||
return result
|
||||
@@ -743,12 +830,26 @@ class GeminiLLM(LLMInterface):
|
||||
cached_tokens=cached_input_tokens,
|
||||
)
|
||||
|
||||
_emit_call_audit(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
input_tokens=input_tokens,
|
||||
cached_input_tokens=cached_input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
thoughts_tokens=thoughts_tokens,
|
||||
duration_ms=int(duration * 1000),
|
||||
finish_reason=finish_reason,
|
||||
)
|
||||
|
||||
return LLMToolCallResult(
|
||||
content=content,
|
||||
tool_calls=tool_calls,
|
||||
finish_reason=finish_reason,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
cached_tokens=cached_input_tokens,
|
||||
thoughts_tokens=thoughts_tokens,
|
||||
)
|
||||
|
||||
except genai_errors.APIError as e:
|
||||
|
||||
@@ -141,7 +141,7 @@ async def _generate_structured_output(
|
||||
response_schema: dict,
|
||||
llm_config: "LLMProvider",
|
||||
reflect_id: str,
|
||||
) -> tuple[dict[str, Any] | None, int, int]:
|
||||
) -> tuple[dict[str, Any] | None, int, int, int, int]:
|
||||
"""Generate structured output from an answer using the provided JSON schema.
|
||||
|
||||
Args:
|
||||
@@ -151,7 +151,7 @@ async def _generate_structured_output(
|
||||
reflect_id: Reflect ID for logging
|
||||
|
||||
Returns:
|
||||
Tuple of (structured_output, input_tokens, output_tokens).
|
||||
Tuple of (structured_output, input_tokens, output_tokens, cached_tokens, thoughts_tokens).
|
||||
structured_output is None if generation fails.
|
||||
"""
|
||||
try:
|
||||
@@ -186,7 +186,7 @@ async def _generate_structured_output(
|
||||
|
||||
if not fields:
|
||||
logger.warning(f"[REFLECT {reflect_id}] No fields found in response_schema, skipping structured output")
|
||||
return None, 0, 0
|
||||
return None, 0, 0, 0, 0
|
||||
|
||||
DynamicModel = create_model("StructuredResponse", **fields)
|
||||
|
||||
@@ -259,11 +259,17 @@ OUTPUT:"""
|
||||
logger.warning(f"[REFLECT {reflect_id}] Required field '{field_name}' is empty in structured output")
|
||||
|
||||
logger.info(f"[REFLECT {reflect_id}] Generated structured output with {len(structured_output)} fields")
|
||||
return structured_output, usage.input_tokens, usage.output_tokens
|
||||
return (
|
||||
structured_output,
|
||||
usage.input_tokens,
|
||||
usage.output_tokens,
|
||||
getattr(usage, "cached_tokens", 0) or 0,
|
||||
getattr(usage, "thoughts_tokens", 0) or 0,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"[REFLECT {reflect_id}] Failed to generate structured output: {e}")
|
||||
return None, 0, 0
|
||||
return None, 0, 0, 0, 0
|
||||
|
||||
|
||||
def _count_messages_tokens(messages: list[dict[str, Any]]) -> int:
|
||||
@@ -435,9 +441,15 @@ async def run_reflect_agent(
|
||||
llm_trace: list[dict[str, Any]] = []
|
||||
context_history: list[dict[str, Any]] = [] # For final prompt fallback
|
||||
|
||||
# Token usage tracking - accumulate across all LLM calls
|
||||
# Token usage tracking - accumulate across all LLM calls.
|
||||
# cached_tokens and thoughts_tokens are surfaced so downstream metering
|
||||
# can split COGS by cached/fresh/visible-output/reasoning. Both are
|
||||
# subsets of (or parallel to) input/output and are not double-counted
|
||||
# in total_tokens.
|
||||
total_input_tokens = 0
|
||||
total_output_tokens = 0
|
||||
total_cached_tokens = 0
|
||||
total_thoughts_tokens = 0
|
||||
|
||||
# Track available IDs for validation (prevents hallucinated citations)
|
||||
available_memory_ids: set[str] = set()
|
||||
@@ -460,6 +472,8 @@ async def run_reflect_agent(
|
||||
input_tokens=total_input_tokens,
|
||||
output_tokens=total_output_tokens,
|
||||
total_tokens=total_input_tokens + total_output_tokens,
|
||||
cached_tokens=total_cached_tokens,
|
||||
thoughts_tokens=total_thoughts_tokens,
|
||||
)
|
||||
|
||||
def _log_completion(answer: str, iterations: int, forced: bool = False):
|
||||
@@ -526,6 +540,8 @@ async def run_reflect_agent(
|
||||
llm_duration = int((time.time() - llm_start) * 1000)
|
||||
total_input_tokens += usage.input_tokens
|
||||
total_output_tokens += usage.output_tokens
|
||||
total_cached_tokens += getattr(usage, "cached_tokens", 0) or 0
|
||||
total_thoughts_tokens += getattr(usage, "thoughts_tokens", 0) or 0
|
||||
llm_trace.append(
|
||||
{
|
||||
"scope": "final",
|
||||
@@ -539,11 +555,13 @@ async def run_reflect_agent(
|
||||
# Generate structured output if schema provided
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
structured_output, struct_in, struct_out = await _generate_structured_output(
|
||||
structured_output, struct_in, struct_out, struct_cached, struct_thoughts = await _generate_structured_output(
|
||||
answer, response_schema, llm_config, reflect_id
|
||||
)
|
||||
total_input_tokens += struct_in
|
||||
total_output_tokens += struct_out
|
||||
total_cached_tokens += struct_cached
|
||||
total_thoughts_tokens += struct_thoughts
|
||||
|
||||
_log_completion(answer, iteration + 1, forced=True)
|
||||
return ReflectAgentResult(
|
||||
@@ -588,6 +606,8 @@ async def run_reflect_agent(
|
||||
llm_duration = int((time.time() - llm_start) * 1000)
|
||||
total_input_tokens += usage.input_tokens
|
||||
total_output_tokens += usage.output_tokens
|
||||
total_cached_tokens += getattr(usage, "cached_tokens", 0) or 0
|
||||
total_thoughts_tokens += getattr(usage, "thoughts_tokens", 0) or 0
|
||||
llm_trace.append(
|
||||
{
|
||||
"scope": "final",
|
||||
@@ -600,11 +620,13 @@ async def run_reflect_agent(
|
||||
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
structured_output, struct_in, struct_out = await _generate_structured_output(
|
||||
structured_output, struct_in, struct_out, struct_cached, struct_thoughts = await _generate_structured_output(
|
||||
answer, response_schema, llm_config, reflect_id
|
||||
)
|
||||
total_input_tokens += struct_in
|
||||
total_output_tokens += struct_out
|
||||
total_cached_tokens += struct_cached
|
||||
total_thoughts_tokens += struct_thoughts
|
||||
|
||||
_log_completion(answer, iteration + 1, forced=True)
|
||||
return ReflectAgentResult(
|
||||
@@ -661,6 +683,8 @@ async def run_reflect_agent(
|
||||
consecutive_errors = 0
|
||||
total_input_tokens += result.input_tokens
|
||||
total_output_tokens += result.output_tokens
|
||||
total_cached_tokens += getattr(result, "cached_tokens", 0) or 0
|
||||
total_thoughts_tokens += getattr(result, "thoughts_tokens", 0) or 0
|
||||
llm_trace.append(
|
||||
{
|
||||
"scope": f"agent_{iteration + 1}",
|
||||
@@ -709,6 +733,8 @@ async def run_reflect_agent(
|
||||
llm_duration = int((time.time() - llm_start) * 1000)
|
||||
total_input_tokens += usage.input_tokens
|
||||
total_output_tokens += usage.output_tokens
|
||||
total_cached_tokens += getattr(usage, "cached_tokens", 0) or 0
|
||||
total_thoughts_tokens += getattr(usage, "thoughts_tokens", 0) or 0
|
||||
llm_trace.append(
|
||||
{
|
||||
"scope": "final",
|
||||
@@ -722,11 +748,13 @@ async def run_reflect_agent(
|
||||
# Generate structured output if schema provided
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
structured_output, struct_in, struct_out = await _generate_structured_output(
|
||||
structured_output, struct_in, struct_out, struct_cached, struct_thoughts = await _generate_structured_output(
|
||||
answer, response_schema, llm_config, reflect_id
|
||||
)
|
||||
total_input_tokens += struct_in
|
||||
total_output_tokens += struct_out
|
||||
total_cached_tokens += struct_cached
|
||||
total_thoughts_tokens += struct_thoughts
|
||||
|
||||
_log_completion(answer, iteration + 1, forced=True)
|
||||
return ReflectAgentResult(
|
||||
@@ -796,11 +824,13 @@ async def run_reflect_agent(
|
||||
# Generate structured output if schema provided
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
structured_output, struct_in, struct_out = await _generate_structured_output(
|
||||
structured_output, struct_in, struct_out, struct_cached, struct_thoughts = await _generate_structured_output(
|
||||
answer, response_schema, llm_config, reflect_id
|
||||
)
|
||||
total_input_tokens += struct_in
|
||||
total_output_tokens += struct_out
|
||||
total_cached_tokens += struct_cached
|
||||
total_thoughts_tokens += struct_thoughts
|
||||
|
||||
_log_completion(answer, iteration + 1)
|
||||
return ReflectAgentResult(
|
||||
@@ -835,6 +865,8 @@ async def run_reflect_agent(
|
||||
llm_duration = int((time.time() - llm_start) * 1000)
|
||||
total_input_tokens += usage.input_tokens
|
||||
total_output_tokens += usage.output_tokens
|
||||
total_cached_tokens += getattr(usage, "cached_tokens", 0) or 0
|
||||
total_thoughts_tokens += getattr(usage, "thoughts_tokens", 0) or 0
|
||||
llm_trace.append(
|
||||
{
|
||||
"scope": "final",
|
||||
@@ -848,11 +880,13 @@ async def run_reflect_agent(
|
||||
# Generate structured output if schema provided
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
structured_output, struct_in, struct_out = await _generate_structured_output(
|
||||
structured_output, struct_in, struct_out, struct_cached, struct_thoughts = await _generate_structured_output(
|
||||
answer, response_schema, llm_config, reflect_id
|
||||
)
|
||||
total_input_tokens += struct_in
|
||||
total_output_tokens += struct_out
|
||||
total_cached_tokens += struct_cached
|
||||
total_thoughts_tokens += struct_thoughts
|
||||
|
||||
_log_completion(answer, iteration + 1, forced=True)
|
||||
return ReflectAgentResult(
|
||||
@@ -1147,7 +1181,7 @@ async def _process_done_tool(
|
||||
structured_output = None
|
||||
final_usage = usage
|
||||
if response_schema and llm_config and answer:
|
||||
structured_output, struct_in, struct_out = await _generate_structured_output(
|
||||
structured_output, struct_in, struct_out, struct_cached, struct_thoughts = await _generate_structured_output(
|
||||
answer, response_schema, llm_config, reflect_id
|
||||
)
|
||||
# Add structured output tokens to usage
|
||||
@@ -1155,6 +1189,8 @@ async def _process_done_tool(
|
||||
input_tokens=usage.input_tokens + struct_in,
|
||||
output_tokens=usage.output_tokens + struct_out,
|
||||
total_tokens=usage.total_tokens + struct_in + struct_out,
|
||||
cached_tokens=usage.cached_tokens + struct_cached,
|
||||
thoughts_tokens=usage.thoughts_tokens + struct_thoughts,
|
||||
)
|
||||
|
||||
log_completion(answer, iterations)
|
||||
|
||||
@@ -78,9 +78,19 @@ class DirectiveInfo(BaseModel):
|
||||
class TokenUsageSummary(BaseModel):
|
||||
"""Total token usage across all LLM calls."""
|
||||
|
||||
input_tokens: int = Field(default=0, description="Total input tokens used")
|
||||
output_tokens: int = Field(default=0, description="Total output tokens used")
|
||||
total_tokens: int = Field(default=0, description="Total tokens (input + output)")
|
||||
input_tokens: int = Field(default=0, description="Total input tokens used (includes any cached prefix tokens)")
|
||||
output_tokens: int = Field(
|
||||
default=0, description="Total visible output tokens used (excludes reasoning/thoughts)"
|
||||
)
|
||||
total_tokens: int = Field(default=0, description="Total tokens (input + output, excludes thoughts)")
|
||||
cached_tokens: int = Field(
|
||||
default=0,
|
||||
description="Cached/cache-read prompt tokens summed across calls. Subset of input_tokens.",
|
||||
)
|
||||
thoughts_tokens: int = Field(
|
||||
default=0,
|
||||
description="Reasoning/thinking tokens summed across calls. Billed at the output rate by some providers but not part of visible output.",
|
||||
)
|
||||
|
||||
|
||||
class ReflectAgentResult(BaseModel):
|
||||
|
||||
@@ -31,8 +31,10 @@ class LLMToolCallResult(BaseModel):
|
||||
content: str | None = Field(default=None, description="Text content if any")
|
||||
tool_calls: list[LLMToolCall] = Field(default_factory=list, description="Tool calls requested by the LLM")
|
||||
finish_reason: str | None = Field(default=None, description="Reason the LLM stopped: 'stop', 'tool_calls', etc.")
|
||||
input_tokens: int = Field(default=0, description="Input tokens used in this call")
|
||||
output_tokens: int = Field(default=0, description="Output tokens used in this call")
|
||||
input_tokens: int = Field(default=0, description="Input tokens used in this call (includes any cached prefix tokens reported by the provider)")
|
||||
output_tokens: int = Field(default=0, description="Visible output tokens used in this call (excludes reasoning/thoughts)")
|
||||
cached_tokens: int = Field(default=0, description="Cached prefix tokens, when reported by the provider. Subset of input_tokens.")
|
||||
thoughts_tokens: int = Field(default=0, description="Reasoning/thinking tokens. Billed at the output rate by some providers but not part of visible output.")
|
||||
|
||||
|
||||
class ToolCallTrace(BaseModel):
|
||||
@@ -91,9 +93,10 @@ class TokenUsage(BaseModel):
|
||||
)
|
||||
|
||||
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)")
|
||||
output_tokens: int = Field(default=0, description="Number of visible output/completion tokens generated (excludes reasoning/thoughts)")
|
||||
total_tokens: int = Field(default=0, description="Total tokens (input + output, excludes thoughts)")
|
||||
cached_tokens: int = Field(default=0, description="Cached/cache-read prompt tokens, when reported by the provider")
|
||||
thoughts_tokens: int = Field(default=0, description="Reasoning/thinking tokens generated by the model. Billed at the output rate by some providers (e.g. Gemini 2.5+ family) but not surfaced in the visible response.")
|
||||
|
||||
def __add__(self, other: "TokenUsage") -> "TokenUsage":
|
||||
"""Allow aggregating token usage from multiple calls."""
|
||||
@@ -102,6 +105,7 @@ class TokenUsage(BaseModel):
|
||||
output_tokens=self.output_tokens + other.output_tokens,
|
||||
total_tokens=self.total_tokens + other.total_tokens,
|
||||
cached_tokens=self.cached_tokens + other.cached_tokens,
|
||||
thoughts_tokens=self.thoughts_tokens + other.thoughts_tokens,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -203,6 +203,13 @@ class RetainResult:
|
||||
llm_input_tokens: int | None = None
|
||||
llm_output_tokens: int | None = None
|
||||
llm_total_tokens: int | None = None
|
||||
# Diagnostic token splits surfaced for COGS reconciliation. ``llm_cached_input_tokens``
|
||||
# is the subset of llm_input_tokens that the provider served from a cache (e.g. Gemini's
|
||||
# cached_content_token_count). ``llm_thoughts_tokens`` is reasoning/thinking tokens that
|
||||
# are billed at the output rate by some providers (Gemini 2.5+) but not part of the visible
|
||||
# response. Both default to None when the engine/provider didn't report them.
|
||||
llm_cached_input_tokens: int | None = None
|
||||
llm_thoughts_tokens: int | None = None
|
||||
# Content tokens the retain pipeline actually processed, after
|
||||
# chunk-level content-hash deduplication. Semantics:
|
||||
# None — no dedup signal available (e.g. a first-time retain or a
|
||||
|
||||
@@ -0,0 +1,142 @@
|
||||
"""Tests for the Gemini per-call audit log and thoughts_tokens plumbing.
|
||||
|
||||
The audit log is the diagnostic instrument we use to attribute LLM spend
|
||||
back to a specific code path when the application-layer metering shows a
|
||||
gap vs provider billing. These tests pin the audit-log shape so a refactor
|
||||
doesn't silently drop fields that an external reconciliation pipeline is
|
||||
relying on.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
|
||||
from hindsight_api.engine.providers.gemini_llm import (
|
||||
_emit_call_audit,
|
||||
_truncated_caller_stack,
|
||||
)
|
||||
from hindsight_api.engine.response_models import LLMToolCallResult, TokenUsage
|
||||
|
||||
|
||||
def test_token_usage_aggregates_thoughts_tokens():
|
||||
"""Aggregating two TokenUsage entries sums thoughts_tokens alongside the others.
|
||||
|
||||
The reflect agent and retain orchestrator both accumulate per-call
|
||||
usage via ``+=``. If thoughts_tokens is not summed, the per-op total
|
||||
will undercount reasoning spend by a factor of N (where N is the
|
||||
number of LLM sub-calls per op).
|
||||
"""
|
||||
a = TokenUsage(input_tokens=10, output_tokens=5, total_tokens=15, cached_tokens=2, thoughts_tokens=7)
|
||||
b = TokenUsage(input_tokens=20, output_tokens=8, total_tokens=28, cached_tokens=3, thoughts_tokens=11)
|
||||
c = a + b
|
||||
assert c.input_tokens == 30
|
||||
assert c.output_tokens == 13
|
||||
assert c.total_tokens == 43
|
||||
assert c.cached_tokens == 5
|
||||
assert c.thoughts_tokens == 18
|
||||
|
||||
|
||||
def test_token_usage_thoughts_defaults_zero():
|
||||
"""Existing callers that construct TokenUsage without thoughts still work."""
|
||||
u = TokenUsage(input_tokens=100, output_tokens=50, total_tokens=150)
|
||||
assert u.thoughts_tokens == 0
|
||||
|
||||
|
||||
def test_llm_tool_call_result_carries_thoughts_and_cached():
|
||||
"""call_with_tools returns LLMToolCallResult; downstream reflect-agent
|
||||
aggregation needs both new fields visible at this layer."""
|
||||
r = LLMToolCallResult(
|
||||
content="ok",
|
||||
input_tokens=1234,
|
||||
output_tokens=56,
|
||||
cached_tokens=200,
|
||||
thoughts_tokens=78,
|
||||
)
|
||||
assert r.cached_tokens == 200
|
||||
assert r.thoughts_tokens == 78
|
||||
|
||||
|
||||
def test_truncated_caller_stack_filters_stdlib_and_packages():
|
||||
"""The stack must only include project frames, capped at max_frames."""
|
||||
frames = _truncated_caller_stack(max_frames=4)
|
||||
# All returned entries follow ``path:lineno:func`` shape.
|
||||
for f in frames:
|
||||
parts = f.rsplit(":", 2)
|
||||
assert len(parts) == 3
|
||||
assert parts[1].isdigit()
|
||||
# No stdlib or 3p frames leaked through.
|
||||
for f in frames:
|
||||
assert "/site-packages/" not in f
|
||||
assert "/lib/python" not in f
|
||||
# Cap honored.
|
||||
assert len(frames) <= 4
|
||||
|
||||
|
||||
def test_emit_call_audit_writes_valid_json_with_all_fields(caplog):
|
||||
"""The audit log line must be JSON-parseable and contain every field
|
||||
the reconciliation script expects. If a future refactor renames or
|
||||
drops a field, this test fails before the rename ships."""
|
||||
caplog.set_level(logging.INFO, logger="hindsight.llm.gemini.calls")
|
||||
_emit_call_audit(
|
||||
provider="gemini",
|
||||
model="gemini-3.1-flash-lite",
|
||||
scope="test_scope",
|
||||
input_tokens=1500,
|
||||
cached_input_tokens=300,
|
||||
output_tokens=100,
|
||||
thoughts_tokens=50,
|
||||
duration_ms=420,
|
||||
finish_reason="stop",
|
||||
)
|
||||
|
||||
audit_records = [r for r in caplog.records if r.name == "hindsight.llm.gemini.calls"]
|
||||
assert len(audit_records) == 1
|
||||
payload = json.loads(audit_records[0].message)
|
||||
|
||||
expected_keys = {
|
||||
"provider",
|
||||
"model",
|
||||
"scope",
|
||||
"input_tokens",
|
||||
"cached_input_tokens",
|
||||
"output_tokens",
|
||||
"thoughts_tokens",
|
||||
"duration_ms",
|
||||
"finish_reason",
|
||||
"caller_stack",
|
||||
}
|
||||
assert set(payload.keys()) == expected_keys
|
||||
assert payload["input_tokens"] == 1500
|
||||
assert payload["cached_input_tokens"] == 300
|
||||
assert payload["output_tokens"] == 100
|
||||
assert payload["thoughts_tokens"] == 50
|
||||
assert payload["duration_ms"] == 420
|
||||
assert isinstance(payload["caller_stack"], list)
|
||||
|
||||
|
||||
def test_emit_call_audit_never_raises_on_internal_error(monkeypatch, caplog):
|
||||
"""Audit emission must NEVER fail the request path. If the logger
|
||||
itself raises (e.g. log handler corruption), we swallow silently."""
|
||||
|
||||
class _BoomLogger:
|
||||
def info(self, *_args, **_kwargs):
|
||||
raise RuntimeError("log backend down")
|
||||
|
||||
monkeypatch.setattr(
|
||||
"hindsight_api.engine.providers.gemini_llm._call_audit_logger",
|
||||
_BoomLogger(),
|
||||
)
|
||||
|
||||
# Must not raise.
|
||||
_emit_call_audit(
|
||||
provider="gemini",
|
||||
model="gemini-3.1-flash-lite",
|
||||
scope=None,
|
||||
input_tokens=1,
|
||||
cached_input_tokens=0,
|
||||
output_tokens=1,
|
||||
thoughts_tokens=0,
|
||||
duration_ms=1,
|
||||
finish_reason="stop",
|
||||
)
|
||||
Reference in New Issue
Block a user