83 lines
3.1 KiB
Python
83 lines
3.1 KiB
Python
# token_logger.py — Logg LLM token-bruk til BigQuery (CG3e)
|
|
|
|
import os
|
|
import logging
|
|
import uuid
|
|
from datetime import datetime, timezone
|
|
from typing import Optional
|
|
|
|
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}"
|
|
|
|
# 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,
|
|
model_name: str,
|
|
input_tokens: int,
|
|
output_tokens: int,
|
|
request_id: Optional[str] = None,
|
|
) -> None:
|
|
"""
|
|
Logg ett LLM-kall til BigQuery-tabellen llm_token_usage.
|
|
Feiler stille slik at applikasjonen aldri krasjer pga logging.
|
|
"""
|
|
try:
|
|
from google.cloud import bigquery
|
|
client = bigquery.Client(project=PROJECT_ID)
|
|
|
|
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": agent_name,
|
|
"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: agent={agent_name} 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 eller Terraform
|
|
BQ_SCHEMA = [
|
|
{"name": "timestamp", "type": "TIMESTAMP", "mode": "REQUIRED"},
|
|
{"name": "agent_name", "type": "STRING", "mode": "REQUIRED"},
|
|
{"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"},
|
|
]
|