""" Token usage tracker for AI file processing. Records per-model input/output token counts during a processing session and logs a record per model when the session ends. (Posting usage to the backend ``{BASE_URL}/api/v1/compliance/token`` is disabled — usage is logged only.) Design — avoid double counting ------------------------------ LLM calls in this repo flow through several wrappers (GeminiBedrockClient, think_tool, RLM engine). If every layer recorded usage, each API call would be counted multiple times (the 5x bug we hit before). Rule: record at the **terminal caller** — the function that issues a single network call to the model. Shim layers (e.g. ``GeminiBedrockClient.invoke_model`` which only repackages a response) must NOT record. Usage ----- from utils.token_tracker import token_session, record_usage with token_session(task_name="doc.pdf", entity_id=eid, reference_id=event_id): # any LLM call inside here will feed record_usage(model, in, out) ... Thread-safety ------------- Uses ``contextvars.ContextVar`` so each request/thread gets an isolated session. The SQS processor's ThreadPoolExecutor calls ``handler()`` per thread; each handler opens its own ``token_session``. """ from __future__ import annotations import contextlib import contextvars import json import logging import os import threading from collections import defaultdict from datetime import datetime, timezone from typing import Optional import requests logger = logging.getLogger(__name__) # USD per token. Source: Anthropic/Google published list prices. _M = 1_000_000 # (input_price_per_token, output_price_per_token) # Fund-upload flow standardizes on Claude Sonnet 4.6 — Bedrock and Vertex # both use the same list price ($3 in / $15 out per 1M). Gemini entries are # kept because image_template_analysis (document-studio flow) uses Gemini. PRICING: dict[str, tuple[float, float]] = { "claude-sonnet-4-6": (3.0 / _M, 15.0 / _M), "anthropic.claude-sonnet-4-6": (3.0 / _M, 15.0 / _M), "us.anthropic.claude-sonnet-4-6": (3.0 / _M, 15.0 / _M), # Gemini 3 Pro family (image processor, separate flow from fund uploads) "gemini-3-pro-preview": (2.0 / _M, 12.0 / _M), "gemini-3-pro-image-preview": (2.0 / _M, 12.0 / _M), # Gemini 2.5 Pro — default for LLM_BACKEND=gemini. List price for ≤ 200k # prompt tokens (per https://ai.google.dev/gemini-api/docs/pricing — # tiered: $1.25/$10 ≤200k, $2.50/$15 >200k incl. thinking). We bill the # lower tier; ops can override via TOKEN_PRICING_OVERRIDES for >200k flows. "gemini-2.5-pro": (1.25 / _M, 10.0 / _M), } # Fallback for unrecognized models — Sonnet 4.6 list price. This is a safety # net only; unknown models log a WARNING so Ops can add them to PRICING. _DEFAULT_PRICE = (3.0 / _M, 15.0 / _M) _COMPLIANCE_PATH = "/api/v1/compliance/token" # Ops can override pricing without a deploy by setting TOKEN_PRICING_OVERRIDES # to a JSON object of {"model_id": [input_per_token, output_per_token]} or # {"model_id": {"input_per_M": 3.0, "output_per_M": 15.0}}. def _load_price_overrides() -> dict[str, tuple[float, float]]: raw = os.getenv("TOKEN_PRICING_OVERRIDES", "").strip() if not raw: return {} try: data = json.loads(raw) except Exception as e: logger.warning("TOKEN_PRICING_OVERRIDES is not valid JSON: %s", e) return {} out: dict[str, tuple[float, float]] = {} for model, val in data.items(): try: if isinstance(val, (list, tuple)) and len(val) == 2: out[model] = (float(val[0]), float(val[1])) elif isinstance(val, dict): pin = float(val.get("input_per_M", 0)) / _M pout = float(val.get("output_per_M", 0)) / _M out[model] = (pin, pout) except Exception as e: logger.warning("TOKEN_PRICING_OVERRIDES bad entry for %r: %s", model, e) return out PRICING.update(_load_price_overrides()) _unknown_models_warned: set[str] = set() def _price_for(model: str) -> tuple[float, float]: if not model: return _DEFAULT_PRICE if model in PRICING: return PRICING[model] # Prefix match: "us.anthropic.claude-sonnet-4-6:v1" etc. for key, price in PRICING.items(): if model.startswith(key): return price if model not in _unknown_models_warned: _unknown_models_warned.add(model) logger.warning( "token_tracker: no PRICING entry for %r — using default $3/$15 per 1M. " "Add to PRICING or set TOKEN_PRICING_OVERRIDES to correct billing.", model, ) return _DEFAULT_PRICE # Status values the backend AiUsageLog schema accepts. Anything else is # rejected with a ValidatorError (served as a 500, so it would retry to # exhaustion and then drop the record). _BACKEND_STATUS_ENUM = {"in_progress", "completed", "failed"} class _Session: def __init__( self, task_name: str, type_: str, entity_id: str, reference_id: Optional[str] = None, reference_model: Optional[str] = None, meta: Optional[dict] = None, role_id: Optional[str] = None, ): self.task_name = task_name self.type = type_ self.entity_id = entity_id self.reference_id = reference_id self.reference_model = reference_model self.meta = meta self.role_id = role_id self.started_at = datetime.now(timezone.utc).isoformat() self.ended_at: Optional[str] = None self.status = "in_progress" self._lock = threading.Lock() # {model: {"input": int, "output": int}} self._usage: dict[str, dict[str, int]] = defaultdict( lambda: {"input": 0, "output": 0} ) def record(self, model: str, input_tokens: int, output_tokens: int) -> None: if not (input_tokens or output_tokens): return model_key = model or "unknown" with self._lock: u = self._usage[model_key] u["input"] += int(input_tokens or 0) u["output"] += int(output_tokens or 0) def set_status(self, status: str) -> None: self.status = status def set_meta(self, meta: Optional[dict]) -> None: self.meta = meta def to_records(self) -> list[dict]: # The backend AiUsageLog schema only accepts _BACKEND_STATUS_ENUM. # Custom statuses set via set_status (e.g. "duplicate", # "classified_other") 500 with a ValidatorError on EVERY retry, then # drop the record. Post a valid status and keep the semantic one in # meta.statusDetail. status = self.status meta = self.meta if status not in _BACKEND_STATUS_ENUM: meta = {**(meta or {}), "statusDetail": status} status = "completed" records: list[dict] = [] with self._lock: items = list(self._usage.items()) for model, u in items: p_in, p_out = _price_for(model) price = u["input"] * p_in + u["output"] * p_out records.append({ "taskName": self.task_name, "type": self.type, "inputTokens": u["input"], "outputTokens": u["output"], "price": round(price, 6), "startedAt": self.started_at, "endedAt": self.ended_at, "model": model, "status": status, "referenceId": self.reference_id, "referenceModel": self.reference_model, "meta": meta, "entityId": self.entity_id, }) if not records: # Post a zero record so the backend still sees the session. records.append({ "taskName": self.task_name, "type": self.type, "inputTokens": 0, "outputTokens": 0, "price": 0, "startedAt": self.started_at, "endedAt": self.ended_at, "model": None, "status": status, "referenceId": self.reference_id, "referenceModel": self.reference_model, "meta": meta, "entityId": self.entity_id, }) return records _current: contextvars.ContextVar[Optional[_Session]] = contextvars.ContextVar( "_token_tracker_session", default=None, ) def current_session() -> Optional[_Session]: return _current.get() _no_session_warned = False def record_usage(model: str, input_tokens: int, output_tokens: int) -> None: """Terminal call sites invoke this after parsing a single API response. If there is no active token_session we log a WARNING (once per process) so Ops can catch regressions where an entry point forgot to open a session. """ sess = _current.get() if sess is None: if input_tokens or output_tokens: global _no_session_warned if not _no_session_warned: _no_session_warned = True logger.warning( "record_usage(%s, in=%s, out=%s) called with NO active " "token_session — usage NOT billed. Check that handler() / " "process_single_document() opened a session around this path.", model, input_tokens, output_tokens, ) return try: sess.record(model, input_tokens, output_tokens) except Exception as e: # never let tracking break the caller logger.warning("record_usage failed: %s", e) def _get_auth_token() -> Optional[str]: # Reuse lambda_function.get_token() so we share the same auth flow. try: from lambda_function import get_token # type: ignore return get_token() except Exception: pass try: from task.lambda_function import get_token # type: ignore return get_token() except Exception: return None _POST_MAX_ATTEMPTS = int(os.getenv("TOKEN_POST_MAX_ATTEMPTS", "5")) _POST_TIMEOUT = float(os.getenv("TOKEN_POST_TIMEOUT", "15")) # HTTP status codes where a retry is worthwhile (transient). _RETRYABLE_STATUS = {408, 425, 429, 500, 502, 503, 504} def _post_record(record: dict, role_id: Optional[str] = None) -> bool: """Log one compliance/token usage record. Posting usage to the backend ({BASE_URL}/api/v1/compliance/token) has been disabled for this service — token usage is recorded to the logs only, never sent over the network. The ``role_id`` argument is kept for call-site compatibility. Always returns True (the record was handled). """ logger.info( "compliance/token usage model=%s in=%s out=%s price=%s role=%s " "(log-only; not POSTed)", record.get("model"), record.get("inputTokens"), record.get("outputTokens"), record.get("price"), role_id, ) return True @contextlib.contextmanager def token_session( task_name: str, entity_id: str, type_: str = "ai_upload", reference_id: Optional[str] = None, reference_model: Optional[str] = None, meta: Optional[dict] = None, role_id: Optional[str] = None, ): """Open a tracking session. On exit, POST one record per model used. Status is "completed" on clean exit and "failed" if an exception escapes the block. Use ``session.set_status(...)`` inside the block to override (e.g. to mark "duplicate" or "classified_other"). Custom statuses outside the backend enum are POSTed as "completed" with the real value carried in ``meta.statusDetail`` — the backend rejects unknown enum values. """ sess = _Session( task_name=task_name, type_=type_, entity_id=entity_id, reference_id=reference_id, reference_model=reference_model, meta=meta, role_id=role_id, ) reset_token = _current.set(sess) exc: Optional[BaseException] = None try: yield sess except BaseException as e: exc = e raise finally: # An explicit set_status (e.g. "duplicate" before a control-flow # raise) outranks the generic exception default. if exc is not None and sess.status == "in_progress": sess.status = "failed" elif sess.status == "in_progress": sess.status = "completed" sess.ended_at = datetime.now(timezone.utc).isoformat() # Emit an audit log line BEFORE POSTing. This gives Ops a # guaranteed CloudWatch record even if every POST retry fails, # so billing can be reconstructed from logs if needed. records = sess.to_records() try: totals_in = sum(r.get("inputTokens", 0) for r in records) totals_out = sum(r.get("outputTokens", 0) for r in records) totals_price = sum(r.get("price", 0) for r in records) logger.info( "[TOKEN_AUDIT] task=%s entity=%s ref=%s type=%s status=%s " "total_in=%d total_out=%d total_price=$%.6f models=%s " "started=%s ended=%s", sess.task_name, sess.entity_id, sess.reference_id, sess.type, sess.status, totals_in, totals_out, totals_price, [r.get("model") for r in records], sess.started_at, sess.ended_at, ) except Exception as audit_err: logger.warning("token audit log failed: %s", audit_err) try: for rec in records: _post_record(rec, role_id=sess.role_id) except Exception as post_err: logger.error( "token_session flush raised unexpectedly: %s — records were: %s", post_err, json.dumps(records), ) _current.reset(reset_token)