feat(gemini): per-call audit log + thoughts/cached token plumbing
The Gemini 2.5+ family reports four token-count fields on every response:
prompt_token_count, candidates_token_count, cached_content_token_count,
and thoughts_token_count. The provider was capturing all four for
Prometheus and OpenTelemetry spans but only propagating input + output
(via TokenUsage and LLMToolCallResult) to callers. As a result,
downstream accumulators (reflect agent, retain orchestrator) had no way
to attribute reasoning-token spend per operation, and reconciling
recorded token counts against the provider bill required reading the
metrics layer instead of the application-layer data we already track.
This change:
- Adds thoughts_tokens to TokenUsage and propagates it through
TokenUsage.__add__ so per-call totals aggregate correctly across
multi-iteration agentic loops.
- Adds cached_tokens and thoughts_tokens to LLMToolCallResult so
callers of call_with_tools see the full breakdown.
- Threads both fields through the Gemini provider's two call paths
(call and call_with_tools).
- Emits a per-call structured JSON line on a dedicated logger
(hindsight.llm.gemini.calls) capturing provider, model, scope,
every token field, duration, finish reason, and a truncated
project-only caller stack. Operators can route this logger to
BigQuery / Loki and reconcile against provider billing offline
without polluting the general application log.
- Audit emission is fail-safe: a broken log backend swallows silently
so the request path is never blocked.
Backward compatible: TokenUsage and LLMToolCallResult default the new
fields to 0, and output_tokens continues to mean "visible output"
(excludes reasoning) so any caller using it for usage billing sees no
behavior change.
This commit is contained in:
Chris Bartholomew
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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