Chris Christiansen
9fcb9c354a
Some checks are pending
Check Python Version Consistency / Check Python Version (push) Waiting to run
143 lines
6.2 KiB
Python
143 lines
6.2 KiB
Python
# token_logger.py — Logg LLM token-bruk til BigQuery (CG3e + CG4)
|
|
|
|
import os
|
|
import logging
|
|
import uuid
|
|
from datetime import datetime, timezone
|
|
from typing import Optional, Literal
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT", "propane-will-491900-m5")
|
|
DATASET = os.environ.get("BQ_BILLING_DATASET", "billing_data")
|
|
TABLE = "llm_token_usage"
|
|
FULL_TABLE = f"{PROJECT_ID}.{DATASET}.{TABLE}"
|
|
|
|
# Gyldige caller_type-verdier
|
|
# agent = OPAX/Jason kjørt som agent via /run eller /run/dag
|
|
# browser = direkte kall fra browser-chat (opax.vauco.no)
|
|
# cli = Gemini TUI CLI eller lokal agent.py kjørt manuelt
|
|
# cron = bakgrunnsjobb (billing_agent, anomaly_detector, dag-runner)
|
|
# api = ekstern REST-kall til /run uten sesjon
|
|
CallerType = Literal["agent", "browser", "cli", "cron", "api"]
|
|
|
|
# Gyldige module_name-verdier (utvides etter hvert som moduler comes online)
|
|
# jason/light = OPAX-agenten, gemini-2.5-flash
|
|
# jason/heavy = OPAX-agenten, gemini-2.5-pro
|
|
# costguard = CostGuard-analyse og anomali-endepunkter
|
|
# billing_agent = ml/billing_agent.py bakgrunns-analyse
|
|
# threadstone = Threadstone-modul (fremtidig)
|
|
# cli/gemini-tui = lokal Gemini TUI (arkivert, men kan aktiveres)
|
|
ModuleName = str # ikke enum — for fremtidssikkerhet
|
|
|
|
# Prismodell (USD per 1M tokens) — oppdater ved modellbytte
|
|
MODEL_PRICING = {
|
|
"gemini-2.5-pro": {"input": 1.25, "output": 10.00},
|
|
"gemini-2.5-flash": {"input": 0.075, "output": 0.30},
|
|
}
|
|
DEFAULT_PRICING = {"input": 1.25, "output": 10.00}
|
|
|
|
|
|
def _estimate_cost(model_name: str, input_tokens: int, output_tokens: int) -> float:
|
|
pricing = MODEL_PRICING.get(model_name, DEFAULT_PRICING)
|
|
cost = (input_tokens / 1_000_000) * pricing["input"] + \
|
|
(output_tokens / 1_000_000) * pricing["output"]
|
|
return round(cost, 8)
|
|
|
|
|
|
def log_token_usage(
|
|
agent_name: str, # deprecated alias — bruk module_name
|
|
model_name: str,
|
|
input_tokens: int,
|
|
output_tokens: int,
|
|
request_id: Optional[str] = None,
|
|
module_name: Optional[str] = None, # CG4: f.eks. 'jason/light', 'costguard'
|
|
caller_type: CallerType = "agent", # CG4: hvem/hva som kalte
|
|
session_id: Optional[str] = None, # CG4: sesjon-ID for gruppering
|
|
) -> None:
|
|
"""
|
|
Logg ett LLM-kall til BigQuery-tabellen llm_token_usage.
|
|
Feiler stille slik at applikasjonen aldri krasjer pga logging.
|
|
|
|
Labeling-konvensjoner:
|
|
module_name — hvilken komponent (jason/light, costguard, billing_agent, …)
|
|
caller_type — hvordan kallet ble trigget (agent, browser, cli, cron, api)
|
|
agent_name — beholdt for bakoverkompatibilitet; settes til module_name om gitt
|
|
"""
|
|
try:
|
|
|
|
from google.cloud import bigquery
|
|
client = bigquery.Client(project=PROJECT_ID)
|
|
|
|
effective_module = module_name or agent_name
|
|
total_tokens = input_tokens + output_tokens
|
|
estimated_cost = _estimate_cost(model_name, input_tokens, output_tokens)
|
|
|
|
row = {
|
|
"timestamp": datetime.now(timezone.utc).isoformat(),
|
|
"agent_name": effective_module, # bakoverkompatibelt felt
|
|
"module_name": effective_module, # CG4: nytt felt
|
|
"caller_type": caller_type, # CG4: nytt felt
|
|
"session_id": session_id or "", # CG4: nytt felt
|
|
"model_name": model_name,
|
|
"input_tokens": input_tokens,
|
|
"output_tokens": output_tokens,
|
|
"total_tokens": total_tokens,
|
|
"estimated_cost_usd": estimated_cost,
|
|
"request_id": request_id or str(uuid.uuid4()),
|
|
}
|
|
|
|
errors = client.insert_rows_json(FULL_TABLE, [row])
|
|
if errors:
|
|
logger.warning(f"[token_logger] BQ insert errors: {errors}")
|
|
else:
|
|
logger.info(
|
|
f"[token_logger] Logged: module={effective_module} "
|
|
f"caller={caller_type} model={model_name} "
|
|
f"in={input_tokens} out={output_tokens} cost=${estimated_cost:.6f}"
|
|
)
|
|
except Exception as e:
|
|
logger.warning(f"[token_logger] Failed to log token usage (non-fatal): {e}")
|
|
|
|
|
|
# BQ table schema — brukes som referanse ved manuell oppretting, Terraform eller create_bq_table()
|
|
BQ_SCHEMA = [
|
|
{"name": "timestamp", "type": "TIMESTAMP", "mode": "REQUIRED"},
|
|
{"name": "agent_name", "type": "STRING", "mode": "REQUIRED"}, # bakoverkompatibelt
|
|
{"name": "module_name", "type": "STRING", "mode": "REQUIRED"}, # CG4
|
|
{"name": "caller_type", "type": "STRING", "mode": "REQUIRED"}, # CG4: agent|browser|cli|cron|api
|
|
{"name": "session_id", "type": "STRING", "mode": "NULLABLE"}, # CG4
|
|
{"name": "model_name", "type": "STRING", "mode": "REQUIRED"},
|
|
{"name": "input_tokens", "type": "INTEGER", "mode": "REQUIRED"},
|
|
{"name": "output_tokens", "type": "INTEGER", "mode": "REQUIRED"},
|
|
{"name": "total_tokens", "type": "INTEGER", "mode": "REQUIRED"},
|
|
{"name": "estimated_cost_usd", "type": "FLOAT", "mode": "REQUIRED"},
|
|
{"name": "request_id", "type": "STRING", "mode": "NULLABLE"},
|
|
]
|
|
|
|
|
|
def create_bq_table_if_not_exists() -> None:
|
|
"""
|
|
Opprett llm_token_usage-tabellen i BigQuery hvis den ikke finnes.
|
|
Trygt å kalle ved app-oppstart (CREATE TABLE IF NOT EXISTS-semantikk).
|
|
"""
|
|
try:
|
|
from google.cloud import bigquery
|
|
client = bigquery.Client(project=PROJECT_ID)
|
|
dataset_ref = client.dataset(DATASET)
|
|
table_ref = dataset_ref.table(TABLE)
|
|
|
|
try:
|
|
client.get_table(table_ref)
|
|
logger.info(f"[token_logger] Tabell {FULL_TABLE} finnes allerede.")
|
|
return
|
|
except Exception:
|
|
pass # tabell finnes ikke — opprett
|
|
|
|
schema = [bigquery.SchemaField(f["name"], f["type"], mode=f["mode"]) for f in BQ_SCHEMA]
|
|
table = bigquery.Table(table_ref, schema=schema)
|
|
client.create_table(table)
|
|
logger.info(f"[token_logger] Opprettet tabell {FULL_TABLE}.")
|
|
except Exception as e:
|
|
logger.warning(f"[token_logger] Kunne ikke opprette BQ-tabell (non-fatal): {e}")
|