feat(CG3e): add llm_token_usage BigQuery logger + /billing/tokens/summary endpoint
This commit is contained in:
parent
6dbaa49a9c
commit
33482aee2c
|
|
@ -17,6 +17,7 @@ Gjeldende modell-tilgjengelighet (mai 2026):
|
|||
import asyncio
|
||||
import os
|
||||
import logging
|
||||
import uuid
|
||||
from typing import Literal
|
||||
|
||||
from google.adk.agents import Agent
|
||||
|
|
@ -48,6 +49,13 @@ HEAVY_MODE_ALLOWED_USERS = ["opax", "admin"]
|
|||
Mode = Literal["light", "heavy"]
|
||||
APP_NAME = "opax"
|
||||
|
||||
# Token logger — feiler stille, stopper aldri agent
|
||||
try:
|
||||
from token_logger import log_token_usage
|
||||
except ImportError:
|
||||
def log_token_usage(*args, **kwargs):
|
||||
pass
|
||||
|
||||
|
||||
def _normalize_mode(mode: str) -> str:
|
||||
mapping = {"A": "light", "A+": "heavy", "light": "light", "heavy": "heavy"}
|
||||
|
|
@ -138,6 +146,7 @@ root_agent = build_agent(mode="light")
|
|||
async def _run_async(message: str, user_id: str, session_id: str, mode: str) -> str:
|
||||
mode = _normalize_mode(mode)
|
||||
agent = build_agent(mode=mode)
|
||||
models = get_models_for_mode(mode)
|
||||
session_service = InMemorySessionService()
|
||||
session = await session_service.create_session(
|
||||
app_name=APP_NAME, user_id=user_id, session_id=session_id,
|
||||
|
|
@ -145,11 +154,31 @@ async def _run_async(message: str, user_id: str, session_id: str, mode: str) ->
|
|||
runner = Runner(agent=agent, app_name=APP_NAME, session_service=session_service)
|
||||
new_message = types.Content(role="user", parts=[types.Part(text=message)])
|
||||
final_text = ""
|
||||
input_tokens = 0
|
||||
output_tokens = 0
|
||||
request_id = str(uuid.uuid4())
|
||||
|
||||
async for event in runner.run_async(
|
||||
user_id=user_id, session_id=session.id, new_message=new_message,
|
||||
):
|
||||
if event.is_final_response() and event.content and event.content.parts:
|
||||
final_text = event.content.parts[0].text or ""
|
||||
# Hent token-metadata fra event hvis tilgjengelig
|
||||
if hasattr(event, "usage_metadata") and event.usage_metadata:
|
||||
um = event.usage_metadata
|
||||
input_tokens += getattr(um, "prompt_token_count", 0) or 0
|
||||
output_tokens += getattr(um, "candidates_token_count", 0) or 0
|
||||
|
||||
# Logg token-bruk — feiler stille
|
||||
if input_tokens > 0 or output_tokens > 0:
|
||||
log_token_usage(
|
||||
agent_name=user_id,
|
||||
model_name=models["orchestrator"],
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
request_id=request_id,
|
||||
)
|
||||
|
||||
return final_text
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -23,6 +23,8 @@ except ImportError as e:
|
|||
|
||||
session_service = InMemorySessionService()
|
||||
APP_NAME = os.environ.get("CLOUD_RUN_SERVICE", "gcp-orchestrator")
|
||||
PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT", "propane-will-491900-m5")
|
||||
BQ_BILLING_DATASET = os.environ.get("BQ_BILLING_DATASET", "billing_data")
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
|
|
@ -116,3 +118,50 @@ async def run(req: RunRequest):
|
|||
except Exception as e:
|
||||
logger.error(f"Agent run failed: {e}", exc_info=True)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@app.get("/billing/tokens/summary")
|
||||
async def billing_tokens_summary():
|
||||
"""
|
||||
CG3e — Aggreger LLM token-bruk og estimert kostnad per agent siste 30 dager.
|
||||
Returnerer JSON-liste sortert etter total_cost DESC.
|
||||
"""
|
||||
try:
|
||||
from google.cloud import bigquery
|
||||
client = bigquery.Client(project=PROJECT_ID)
|
||||
|
||||
query = f"""
|
||||
SELECT
|
||||
agent_name,
|
||||
model_name,
|
||||
SUM(input_tokens) AS total_input_tokens,
|
||||
SUM(output_tokens) AS total_output_tokens,
|
||||
SUM(total_tokens) AS total_tokens,
|
||||
SUM(estimated_cost_usd) AS total_cost_usd
|
||||
FROM
|
||||
`{PROJECT_ID}.{BQ_BILLING_DATASET}.llm_token_usage`
|
||||
WHERE
|
||||
timestamp >= TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL 30 DAY)
|
||||
GROUP BY
|
||||
agent_name, model_name
|
||||
ORDER BY
|
||||
total_cost_usd DESC
|
||||
"""
|
||||
|
||||
results = client.query(query).result()
|
||||
rows = [
|
||||
{
|
||||
"agent_name": row.agent_name,
|
||||
"model_name": row.model_name,
|
||||
"total_input_tokens": row.total_input_tokens,
|
||||
"total_output_tokens": row.total_output_tokens,
|
||||
"total_tokens": row.total_tokens,
|
||||
"total_cost_usd": round(float(row.total_cost_usd), 6),
|
||||
}
|
||||
for row in results
|
||||
]
|
||||
return JSONResponse({"period_days": 30, "rows": rows})
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"[/billing/tokens/summary] Failed: {e}", exc_info=True)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
|
|
|||
82
agents/core-logic/token_logger.py
Normal file
82
agents/core-logic/token_logger.py
Normal file
|
|
@ -0,0 +1,82 @@
|
|||
# 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"},
|
||||
]
|
||||
18
scripts/create_llm_token_usage_table.sh
Normal file
18
scripts/create_llm_token_usage_table.sh
Normal file
|
|
@ -0,0 +1,18 @@
|
|||
#!/bin/bash
|
||||
# CG3e — Opprett BigQuery-tabell llm_token_usage
|
||||
# Kjøres én gang manuelt: bash scripts/create_llm_token_usage_table.sh
|
||||
|
||||
set -e
|
||||
|
||||
PROJECT="propane-will-491900-m5"
|
||||
DATASET="billing_data"
|
||||
TABLE="llm_token_usage"
|
||||
|
||||
echo "Oppretter BigQuery-tabell ${PROJECT}:${DATASET}.${TABLE} ..."
|
||||
|
||||
bq mk --table \
|
||||
--description "Logg over all LLM token-bruk for CostGuard (CG3e)" \
|
||||
"${PROJECT}:${DATASET}.${TABLE}" \
|
||||
timestamp:TIMESTAMP,agent_name:STRING,model_name:STRING,input_tokens:INTEGER,output_tokens:INTEGER,total_tokens:INTEGER,estimated_cost_usd:FLOAT,request_id:STRING
|
||||
|
||||
echo "✅ Tabell ${FULL_TABLE} opprettet."
|
||||
Loading…
Reference in New Issue
Block a user