296 lines
15 KiB
Python
296 lines
15 KiB
Python
"""Per-response usage accounting for the conversation turn loop.
|
|
|
|
After every successful model API call, ``record_response_usage`` folds ``response.usage``
|
|
into: the context compressor (``update_from_response`` + the compression-budget rearm
|
|
latch), the usage anchor for display/compression math, per-session token/cost counters,
|
|
the state.db token-delta queue, and the observability log line. MoA sessions additionally
|
|
fold advisor fan-out usage into the reported counts and price the aggregator at its REAL
|
|
model/provider. Logger name stays ``agent.conversation_loop`` for caplog parity.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from contextlib import suppress
|
|
from dataclasses import dataclass
|
|
from typing import Any, Dict, List
|
|
|
|
from agent.image_token_cost import calibrate_from_usage
|
|
from agent.usage_anchor import capture_usage_anchor, set_usage_anchor
|
|
from agent.usage_pricing import estimate_usage_cost, normalize_usage
|
|
|
|
logger = logging.getLogger("agent.conversation_loop")
|
|
|
|
|
|
def _agent_session_source(agent: Any) -> str:
|
|
"""The surface the agent's own row create would stamp (``_ensure_db_session``), so an
|
|
accounting guard that wins the row-creation race never mints an anonymous session."""
|
|
from run_agent import _session_source_for_agent # late: run_agent imports this module
|
|
return _session_source_for_agent(getattr(agent, "platform", None))
|
|
|
|
|
|
@dataclass
|
|
class ResponseUsageOutcome:
|
|
"""``compression_attempts`` is the (possibly rearmed-to-zero) budget counter;
|
|
``rearmed`` tells the loop to also clear its preflight-block latch."""
|
|
|
|
compression_attempts: int
|
|
rearmed: bool = False
|
|
|
|
|
|
def _loop_mod():
|
|
"""Lazy ``agent.conversation_loop`` import (avoids an import cycle)."""
|
|
import agent.conversation_loop as _cl
|
|
|
|
return _cl
|
|
|
|
|
|
def _fold_moa_usage(agent, canonical_usage):
|
|
"""MoA: fold advisor fan-out usage into REPORTED token counts (only aggregator usage is
|
|
returned, so advisor spend would be invisible) and flush the full-turn trace when
|
|
``moa.save_traces`` is on. Returns ``(client, canonical_usage, advisor_cost)``."""
|
|
_moa_ref_cost = None
|
|
_moa_client = getattr(agent, "client", None)
|
|
if _moa_client is not None and hasattr(_moa_client, "consume_reference_usage"):
|
|
try:
|
|
_ref_usage, _moa_ref_cost = _moa_client.consume_reference_usage()
|
|
if _ref_usage is not None:
|
|
canonical_usage = canonical_usage + _ref_usage
|
|
except Exception as _moa_acct_exc: # pragma: no cover - defensive
|
|
logger.debug("MoA reference usage accounting failed: %s", _moa_acct_exc)
|
|
if _moa_client is not None and hasattr(_moa_client, "consume_and_save_trace"):
|
|
try:
|
|
# Streaming path: pass the streamed acting text so the trace is self-contained.
|
|
_agg_streamed_text = getattr(agent, "_current_streamed_assistant_text", "") or ""
|
|
_moa_client.consume_and_save_trace(
|
|
agent.session_id, aggregator_output_fallback=_agg_streamed_text or None
|
|
)
|
|
except Exception as _moa_trace_exc: # pragma: no cover - defensive
|
|
logger.debug("MoA trace flush failed: %s", _moa_trace_exc)
|
|
return _moa_client, canonical_usage, _moa_ref_cost
|
|
|
|
|
|
def record_response_usage(
|
|
agent: Any, response: Any, *, messages: List[Dict[str, Any]], api_call_count: int,
|
|
api_duration: float, compression_attempts: int, max_compression_attempts: int,
|
|
) -> ResponseUsageOutcome:
|
|
"""Fold ``response.usage`` into compressor, anchors, session counters, state.db
|
|
and the API-call log line (see module docstring). No-usage responses only
|
|
consume a pending compaction verdict. Returns the loop-visible outcome."""
|
|
rearmed = False
|
|
compressor = agent.context_compressor
|
|
# Count every completed provider attempt, including providers that omit usage.
|
|
# Token/cost accounting below stays gated on real usage, but the request itself
|
|
# must remain observable.
|
|
agent.session_api_calls += 1
|
|
if not (hasattr(response, 'usage') and response.usage):
|
|
if getattr(compressor, "awaiting_real_usage_after_compression", False):
|
|
# No usage -> cannot adjudicate the prior compaction; consume the
|
|
# pending verdict so later readings aren't charged to it and
|
|
# preflight deferral isn't latched indefinitely.
|
|
compressor.update_from_response({})
|
|
_note_usage_less = getattr(compressor, "note_usage_less_response", None)
|
|
if callable(_note_usage_less):
|
|
_note_usage_less()
|
|
logger.info(
|
|
"API call #%d: model=%s provider=%s in=? out=? total=? latency=%.1fs usage=unavailable",
|
|
agent.session_api_calls, agent.model, agent.provider or "unknown", api_duration,
|
|
)
|
|
return ResponseUsageOutcome(compression_attempts=compression_attempts, rearmed=rearmed)
|
|
|
|
canonical_usage = normalize_usage(response.usage, provider=agent.provider, api_mode=agent.api_mode)
|
|
# Aggregator-only usage kept for pricing: advisor tokens are priced at each advisor's
|
|
# OWN model rate and added as dollars below.
|
|
aggregator_usage = canonical_usage
|
|
_moa_client, canonical_usage, _moa_ref_cost = _fold_moa_usage(agent, canonical_usage)
|
|
prompt_tokens = canonical_usage.prompt_tokens
|
|
completion_tokens = canonical_usage.output_tokens
|
|
total_tokens = canonical_usage.total_tokens
|
|
# Canonical token + cache buckets for context engines; legacy keys stay for back-compat.
|
|
usage_dict = {
|
|
"prompt_tokens": prompt_tokens,
|
|
"completion_tokens": completion_tokens,
|
|
"total_tokens": total_tokens,
|
|
"input_tokens": canonical_usage.input_tokens,
|
|
"output_tokens": canonical_usage.output_tokens,
|
|
"cache_read_tokens": canonical_usage.cache_read_tokens,
|
|
"cache_write_tokens": canonical_usage.cache_write_tokens,
|
|
"reasoning_tokens": canonical_usage.reasoning_tokens,
|
|
}
|
|
# Capture the boundary latch before update_from_response() consumes it: only the real
|
|
# prompt count right after a compaction rearms the budget.
|
|
_completed_compaction_pending = bool(
|
|
getattr(compressor, "_verify_compaction_cleared_threshold", False)
|
|
)
|
|
compressor.update_from_response(usage_dict)
|
|
# Usage-anchored accounting: snapshot exact provider usage against the durable
|
|
# transcript (main-loop ONLY; MoA uses pre-fold aggregator usage). The display meter
|
|
# anchors on the turn's FIRST response: later same-turn responses inflate
|
|
# prompt_tokens with replayed thinking. Display-only; compression math uses real usage.
|
|
# The provider just priced this request exactly: if the delta since the previous anchor
|
|
# introduced images, the residual is their real per-image cost (learned before re-anchoring).
|
|
calibrate_from_usage(agent, messages, aggregator_usage.prompt_tokens)
|
|
_new_anchor = capture_usage_anchor(
|
|
aggregator_usage.prompt_tokens, aggregator_usage.output_tokens, messages
|
|
)
|
|
if _new_anchor is not None:
|
|
set_usage_anchor(agent, _new_anchor, turn_base=api_call_count == 1)
|
|
_compression_threshold = int(getattr(compressor, "threshold_tokens", 0) or 0)
|
|
if _loop_mod()._should_rearm_compression_budget(
|
|
compression_attempts, completed_compaction_pending=_completed_compaction_pending,
|
|
prompt_tokens=prompt_tokens, threshold_tokens=_compression_threshold,
|
|
):
|
|
logger.info(
|
|
"Compression budget rearmed after provider-confirmed "
|
|
"recovery: prompt=%s < threshold=%s (attempts were %s/%s)",
|
|
f"{prompt_tokens:,}",
|
|
f"{_compression_threshold:,}",
|
|
compression_attempts,
|
|
max_compression_attempts,
|
|
)
|
|
compression_attempts = 0
|
|
# Confirmed recovery also clears the loop's stale insufficient-progress verdict
|
|
# (``_preflight_compression_blocked``), else a later pressure spike grows unchecked.
|
|
rearmed = True
|
|
|
|
# Stash canonical usage for on_turn_complete(); keep the latest call's.
|
|
agent._last_turn_usage = dict(usage_dict)
|
|
# The parent's CURRENT prompt size for headroom math (delegate summary budgets): the
|
|
# aggregator's own prompt, never the MoA-folded total (advisor prompts are not in this context).
|
|
agent._last_prompt_size_tokens = int(aggregator_usage.prompt_tokens or 0)
|
|
|
|
# Persist only provider-confirmed context lengths, not probe tiers.
|
|
if getattr(compressor, "_context_probed", False):
|
|
ctx = compressor.context_length
|
|
if getattr(compressor, "_context_probe_persistable", False):
|
|
from agent.model_metadata import save_provider_context_length
|
|
|
|
save_provider_context_length(agent.model, agent.base_url, ctx, agent.provider)
|
|
agent._safe_print(f"{agent.log_prefix}💾 Cached context length: {ctx:,} tokens for {agent.model}")
|
|
compressor._context_probed = False
|
|
compressor._context_probe_persistable = False
|
|
|
|
agent.session_prompt_tokens += prompt_tokens
|
|
agent.session_completion_tokens += completion_tokens
|
|
agent.session_total_tokens += total_tokens
|
|
agent.session_input_tokens += canonical_usage.input_tokens
|
|
agent.session_output_tokens += canonical_usage.output_tokens
|
|
agent.session_cache_read_tokens += canonical_usage.cache_read_tokens
|
|
agent.session_cache_write_tokens += canonical_usage.cache_write_tokens
|
|
agent.session_reasoning_tokens += canonical_usage.reasoning_tokens
|
|
# Rolling history for status-bar averages (last 10).
|
|
with suppress(Exception):
|
|
hist = getattr(agent, "_api_latency_history", None)
|
|
if hist is not None:
|
|
hist.append(float(api_duration))
|
|
ohist = getattr(agent, "_api_output_history", None)
|
|
if ohist is not None:
|
|
ohist.append(int(canonical_usage.output_tokens or 0))
|
|
|
|
_cache_pct = ""
|
|
if canonical_usage.cache_read_tokens and prompt_tokens:
|
|
_cache_pct = f" cache={canonical_usage.cache_read_tokens}/{prompt_tokens} ({100*canonical_usage.cache_read_tokens/prompt_tokens:.0f}%)"
|
|
# write= is the money (cache writes cost 50x a read); id= is what a provider needs to look the
|
|
# request up; upstream= is who actually served it when the route reports that (OpenRouter's
|
|
# `provider`). Diagnosing the 1,393-agent run's cache misses took a DB join and a live probe
|
|
# because none of the three were on this line.
|
|
if canonical_usage.cache_write_tokens:
|
|
_cache_pct += f" write={canonical_usage.cache_write_tokens}"
|
|
_rid = getattr(response, "id", None)
|
|
_ident = f" id={_rid}" if isinstance(_rid, str) and _rid else ""
|
|
_upstream = getattr(response, "provider", None)
|
|
if isinstance(_upstream, str) and _upstream:
|
|
_ident += f" upstream={_upstream}"
|
|
logger.info(
|
|
"API call #%d: model=%s provider=%s in=%d out=%d total=%d latency=%.1fs%s%s",
|
|
agent.session_api_calls, agent.model, agent.provider or "unknown",
|
|
prompt_tokens, completion_tokens, total_tokens,
|
|
api_duration, _cache_pct, _ident,
|
|
)
|
|
# nous.anthropic_wire=auto: the session's wire is decided once, from this first response.
|
|
if agent.session_api_calls == 1 and (agent.provider or "") == "nous":
|
|
with suppress(Exception):
|
|
from agent.nous_wire import maybe_switch_wire_after_first_response
|
|
maybe_switch_wire_after_first_response(agent, response, agent.session_api_calls)
|
|
|
|
# MoA: agent.model/provider are the virtual preset/"moa" with no pricing entry, silently
|
|
# dropping aggregator spend. Price at the REAL model/provider from the aggregator slot.
|
|
_agg_cost_model, _agg_cost_provider, _agg_cost_base_url = agent.model, agent.provider, agent.base_url
|
|
_agg_slot = getattr(_moa_client, "last_aggregator_slot", None) if _moa_client is not None else None
|
|
if _agg_slot and _agg_slot.get("model"):
|
|
_agg_cost_model = _agg_slot["model"]
|
|
_agg_cost_provider = _agg_slot.get("provider") or agent.provider
|
|
_agg_cost_base_url = _agg_slot.get("base_url") or agent.base_url
|
|
cost_result = estimate_usage_cost(
|
|
_agg_cost_model, aggregator_usage, provider=_agg_cost_provider,
|
|
base_url=_agg_cost_base_url, api_key=getattr(agent, "api_key", ""),
|
|
)
|
|
# Cost delta = aggregator + MoA advisor cost (already priced per-advisor at each
|
|
# advisor's own model rate), so state.db's estimated_cost_usd matches the folded
|
|
# token counts.
|
|
_cost_delta = None
|
|
if cost_result.amount_usd is not None:
|
|
_cost_delta = float(cost_result.amount_usd)
|
|
agent.session_estimated_cost_usd += _cost_delta
|
|
if _moa_ref_cost is not None:
|
|
try:
|
|
_moa_cost = float(_moa_ref_cost)
|
|
except (TypeError, ValueError): # pragma: no cover - defensive
|
|
_moa_cost = None
|
|
if _moa_cost is not None:
|
|
agent.session_estimated_cost_usd += _moa_cost
|
|
_cost_delta = (_cost_delta or 0.0) + _moa_cost
|
|
agent.session_cost_status = cost_result.status
|
|
agent.session_cost_source = cost_result.source
|
|
|
|
# Persist per-call token deltas for any session_id so non-CLI runs can't lose
|
|
# accounting; gateway/session-store writes use absolute totals and safely overwrite
|
|
# these deltas. Enqueued, not written (a cold state.db UPDATE here stalled the tool
|
|
# loop); drained at finalize via _persist_session.
|
|
if agent._session_db and agent.session_id:
|
|
try:
|
|
# Ensure the row exists: under concurrent SQLite load the initial
|
|
# _ensure_db_session() may fail, and UPDATE on a missing row affects 0 rows.
|
|
if not agent._session_db_created:
|
|
agent._ensure_db_session()
|
|
agent._session_db.queue_token_counts(
|
|
agent.session_id,
|
|
source=_agent_session_source(agent),
|
|
input_tokens=canonical_usage.input_tokens,
|
|
output_tokens=canonical_usage.output_tokens,
|
|
cache_read_tokens=canonical_usage.cache_read_tokens,
|
|
cache_write_tokens=canonical_usage.cache_write_tokens,
|
|
reasoning_tokens=canonical_usage.reasoning_tokens,
|
|
estimated_cost_usd=_cost_delta,
|
|
cost_status=cost_result.status,
|
|
cost_source=cost_result.source,
|
|
billing_provider=agent.provider,
|
|
billing_base_url=agent.base_url,
|
|
billing_mode="subscription_included"
|
|
if cost_result.status == "included" else None,
|
|
model=agent.model,
|
|
api_call_count=1,
|
|
)
|
|
except Exception as e: # silent loss here undercounts analytics
|
|
logger.debug(
|
|
"Token persistence failed (session=%s, tokens=%d): %s",
|
|
agent.session_id, total_tokens, e,
|
|
)
|
|
|
|
if agent.verbose_logging:
|
|
logging.debug(f"Token usage: prompt={usage_dict['prompt_tokens']:,}, completion={usage_dict['completion_tokens']:,}, total={usage_dict['total_tokens']:,}")
|
|
|
|
# Report cache stats for any provider that returns ``prompt_tokens_details.cached_tokens``,
|
|
# not only when we inject cache_control markers.
|
|
cached = canonical_usage.cache_read_tokens
|
|
written = canonical_usage.cache_write_tokens
|
|
prompt = usage_dict["prompt_tokens"]
|
|
if (cached or written) and not agent.quiet_mode:
|
|
hit_pct = (cached / prompt * 100) if prompt > 0 else 0
|
|
agent._vprint(
|
|
f"{agent.log_prefix} 💾 Cache: "
|
|
f"{cached:,}/{prompt:,} tokens "
|
|
f"({hit_pct:.0f}% hit, {written:,} written)"
|
|
)
|
|
return ResponseUsageOutcome(compression_attempts=compression_attempts, rearmed=rearmed)
|