mirror of
https://github.com/x1xhlol/system-prompts-and-models-of-ai-tools.git
synced 2026-08-16 21:24:09 +00:00
374 lines
14 KiB
Python
374 lines
14 KiB
Python
"""
|
|
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)
|