semblance-dev/backend/scripts/backfill_usage.py
Vadym Samoilenko d0ad8e67be Fix backfill: use accumulated conversation context for prompt estimation
Old logic used output text length as a proxy for prompt tokens — completely
wrong. Real Gemini calls send the full conversation history as context, so
prompt grows with every turn.

New logic:
- completion_tokens = len(response_text) / 3.8 (what was generated)
- prompt_tokens = base_template + sum(all_prior_messages_in_fg) / 3.8
  - persona_response base: 1500 tok (template + persona details + topic)
  - moderator base: 1200 tok (moderator template + fg context)
  - persona_generate base: 2500 tok (persona-detailed-generation.md template)

Also:
- Sorts messages chronologically per focus group before processing
- Accumulates context correctly so turn N includes turns 0..N-1 as context
- Idempotency via pre-fetched set instead of per-doc find_one queries
- cost_usd breakdown now has correct input/output split (not 40/60 guess)
- Dry-run prints per-focus-group cost estimates for sanity checking

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-24 19:11:01 +01:00

399 lines
17 KiB
Python

#!/usr/bin/env python3
"""
Backfill usage_events from existing focus-group messages and personas.
Creates estimated usage_event docs (is_estimated=True) so the admin dashboard
can show historical cost data for sessions that pre-date the usage tracking system.
Token estimation approach:
- completion = actual output text length / 3.8 chars-per-token
- prompt = base template size + ALL prior messages in conversation (accumulated context)
This mirrors the real LLM call: each turn sends the full conversation history.
Usage:
cd backend
python scripts/backfill_usage.py [--dry-run] [--delete-existing-estimates]
Environment:
MONGO_URI — connection string (falls back to localhost:27017 without auth)
DB_NAME — database name (default: semblance_db)
"""
import argparse
import os
import sys
from collections import defaultdict
from datetime import datetime, timezone
from pymongo import MongoClient
# ─────────────────────────────────────────────────────────────────────────────
# Prompt template size constants (measured from actual files in backend/prompts/)
# These are the BASE tokens before any dynamic content is added.
# ─────────────────────────────────────────────────────────────────────────────
# focus-group-response.md (~941 tok) + persona details (~350 tok) + topic/instructions (~200 tok)
BASE_PROMPT_PERSONA_RESPONSE = 1_500
# ai-moderator-system.md (~738 tok) + focus group context (~500 tok)
BASE_PROMPT_MODERATOR = 1_200
# persona-detailed-generation.md (~2307 tok) + focus group brief (~200 tok)
BASE_PROMPT_PERSONA_GENERATE = 2_500
CHARS_PER_TOKEN = 3.8 # Gemini approximation
# ─────────────────────────────────────────────────────────────────────────────
# Token helpers
# ─────────────────────────────────────────────────────────────────────────────
def _chars_to_tokens(chars: int) -> int:
return max(1, int(chars / CHARS_PER_TOKEN))
def _to_str(v) -> str:
if isinstance(v, list):
return " ".join(str(i) for i in v if i)
return str(v) if v else ""
# ─────────────────────────────────────────────────────────────────────────────
# Pricing
# ─────────────────────────────────────────────────────────────────────────────
_pricing_cache: dict = {}
def _load_pricing(db) -> None:
for row in db.model_pricing.find({"effective_until": None}):
model = row.get("model", "")
tiers = row.get("tiers") or []
if tiers:
t = tiers[0]
_pricing_cache[model] = (
t.get("input_per_mtok", 2.0),
t.get("output_per_mtok", 12.0),
)
def _estimate_cost(prompt_tokens: int, completion_tokens: int, model: str) -> dict:
rates = _pricing_cache.get(model)
if not rates:
for key, val in _pricing_cache.items():
if model and key and (key in model or model in key):
rates = val
break
if not rates:
m = (model or "").lower()
if "gpt-5" in m or "gpt-4" in m:
rates = (2.50, 15.00)
else:
rates = (2.00, 12.00)
input_rate, output_rate = rates
cost_input = (prompt_tokens / 1_000_000) * input_rate
cost_output = (completion_tokens / 1_000_000) * output_rate
total = round(cost_input + cost_output, 8)
return {
"input": round(cost_input, 8),
"output": round(cost_output, 8),
"cached": 0.0,
"reasoning": 0.0,
"total": total,
}
# ─────────────────────────────────────────────────────────────────────────────
# DB connection (sync PyMongo)
# ─────────────────────────────────────────────────────────────────────────────
def connect():
mongo_uri = os.environ.get("MONGO_URI", "mongodb://localhost:27017")
db_name = os.environ.get("DB_NAME", "semblance_db")
try:
client = MongoClient(mongo_uri, serverSelectionTimeoutMS=5000)
client.admin.command("ping")
print(f"Connected to MongoDB: {db_name}")
return client[db_name]
except Exception as e:
print(f"ERROR: Could not connect to MongoDB: {e}")
sys.exit(1)
# ─────────────────────────────────────────────────────────────────────────────
# Backfill focus-group messages
#
# Real prompt structure per call:
# system prompt template (~1200-1500 tok) + all prior messages (accumulated)
# Real completion:
# the response text
#
# We sort messages per focus group by timestamp and accumulate context,
# so that message N has all N-1 prior messages as context — matching reality.
# ─────────────────────────────────────────────────────────────────────────────
def backfill_messages(db, dry_run: bool) -> int:
created = 0
# Build focus-group metadata lookup
fg_meta = {}
for fg in db.focus_groups.find({}, {"llm_model": 1, "created_by": 1}):
fg_meta[str(fg["_id"])] = {
"model": fg.get("llm_model") or "gemini-3.1-pro-preview",
"user_id": str(fg.get("created_by") or ""),
}
# Collect all messages, group by focus_group_id, sort chronologically
all_messages = list(db.focus_group_messages.find({}))
print(f"\n[messages] Found {len(all_messages)} messages across all focus groups")
# Bucket by focus group
by_fg: dict = defaultdict(list)
for msg in all_messages:
fg_id = str(msg.get("focus_group_id") or "")
msg_type = msg.get("type", "")
# Only AI-generated messages cost money
if msg_type not in ("response", "question", "moderator", "ai", ""):
continue
by_fg[fg_id].append(msg)
# Sort each group chronologically
def _ts(m):
t = m.get("created_at") or m.get("timestamp")
if isinstance(t, str):
try:
return datetime.fromisoformat(t)
except Exception:
pass
if isinstance(t, datetime):
return t
return datetime.min
for fg_id, msgs in by_fg.items():
msgs.sort(key=_ts)
# Already-estimated message IDs (for idempotency)
existing_ids = set(
str(e["source_message_id"])
for e in db.usage_events.find(
{"is_estimated": True, "source_message_id": {"$exists": True}},
{"source_message_id": 1}
)
)
for fg_id, msgs in by_fg.items():
meta = fg_meta.get(fg_id, {"model": "gemini-3.1-pro-preview", "user_id": ""})
fg_model = meta["model"]
user_id = meta["user_id"]
provider = "gemini" if "gemini" in fg_model.lower() else "openai"
accumulated_context_chars = 0 # sum of all prior message text lengths
for msg in msgs:
msg_id = str(msg["_id"])
if msg_id in existing_ids:
# Still accumulate context so subsequent messages are correct
text = msg.get("text") or msg.get("content") or ""
accumulated_context_chars += len(text)
continue
text = msg.get("text") or msg.get("content") or ""
msg_type = msg.get("type", "")
# completion = what the model actually generated
completion_tokens = _chars_to_tokens(len(text))
# prompt = base template + full conversation history up to this point
context_tokens = _chars_to_tokens(accumulated_context_chars)
if msg_type in ("question", "moderator"):
prompt_tokens = BASE_PROMPT_MODERATOR + context_tokens
else:
prompt_tokens = BASE_PROMPT_PERSONA_RESPONSE + context_tokens
cost = _estimate_cost(prompt_tokens, completion_tokens, fg_model)
ts = _ts(msg) or datetime.now(timezone.utc)
feature = "moderator" if msg_type in ("question", "moderator") else "persona_response"
event = {
"ts": ts,
"provider": provider,
"model": fg_model,
"feature": feature,
"user_id": user_id,
"focus_group_id": fg_id,
"persona_id": str(msg.get("senderId") or msg.get("persona_id") or ""),
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"cached_tokens": 0,
"reasoning_tokens": 0,
"total_tokens": prompt_tokens + completion_tokens,
"cost_usd": cost,
"duration_ms": 0,
"retry_count": 0,
"status": "success",
"is_estimated": True,
"estimate_method": "accumulated_context",
"source_message_id": msg_id,
}
if not dry_run:
db.usage_events.insert_one(event)
created += 1
# Add this message to the accumulated context for subsequent messages
accumulated_context_chars += len(text)
print(f"[messages] {'Would create' if dry_run else 'Created'} {created} estimated usage events")
return created
# ─────────────────────────────────────────────────────────────────────────────
# Backfill persona generation
#
# Real prompt: persona-detailed-generation.md template (~2307 tok) + fg brief (~200 tok)
# Real completion: the generated persona profile text
# ─────────────────────────────────────────────────────────────────────────────
def backfill_personas(db, dry_run: bool) -> int:
created = 0
personas = list(db.personas.find({}))
print(f"\n[personas] Found {len(personas)} personas to process")
existing_persona_ids = set(
str(e["source_persona_id"])
for e in db.usage_events.find(
{"is_estimated": True, "source_persona_id": {"$exists": True}, "feature": "persona_generate"},
{"source_persona_id": 1}
)
)
for persona in personas:
persona_id = str(persona["_id"])
if persona_id in existing_persona_ids:
continue
# The generated output is the persona profile text
text = " ".join(filter(None, [
_to_str(persona.get("background")),
_to_str(persona.get("description")),
_to_str(persona.get("goals")),
_to_str(persona.get("name")),
])).strip()
if not text:
continue
model = "gemini-3.1-pro-preview"
# completion = the generated persona text
completion_tokens = _chars_to_tokens(len(text))
# prompt = template + focus group brief (fixed base)
prompt_tokens = BASE_PROMPT_PERSONA_GENERATE
cost = _estimate_cost(prompt_tokens, completion_tokens, model)
ts = persona.get("created_at") or persona.get("updatedAt") or datetime.now(timezone.utc)
if isinstance(ts, str):
try:
ts = datetime.fromisoformat(ts)
except Exception:
ts = datetime.now(timezone.utc)
event = {
"ts": ts,
"provider": "gemini",
"model": model,
"feature": "persona_generate",
"user_id": str(persona.get("created_by") or persona.get("user_id") or ""),
"focus_group_id": str(persona.get("focus_group_id") or ""),
"persona_id": persona_id,
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"cached_tokens": 0,
"reasoning_tokens": 0,
"total_tokens": prompt_tokens + completion_tokens,
"cost_usd": cost,
"duration_ms": 0,
"retry_count": 0,
"status": "success",
"is_estimated": True,
"estimate_method": "accumulated_context",
"source_persona_id": persona_id,
}
if not dry_run:
db.usage_events.insert_one(event)
created += 1
print(f"[personas] {'Would create' if dry_run else 'Created'} {created} estimated usage events")
return created
# ─────────────────────────────────────────────────────────────────────────────
# Main
# ─────────────────────────────────────────────────────────────────────────────
def main():
parser = argparse.ArgumentParser(description="Backfill usage_events from existing data")
parser.add_argument("--dry-run", action="store_true", help="Preview without writing")
parser.add_argument("--delete-existing-estimates", action="store_true",
help="Delete previously created estimated events before backfilling")
args = parser.parse_args()
if args.dry_run:
print("=== DRY RUN — no data will be written ===\n")
db = connect()
if args.delete_existing_estimates and not args.dry_run:
result = db.usage_events.delete_many({"is_estimated": True})
print(f"Deleted {result.deleted_count} existing estimated events\n")
_load_pricing(db)
print(f"Loaded {len(_pricing_cache)} pricing rows: {list(_pricing_cache.keys())}")
# Dry-run: show a sample of what the cost distribution looks like
if args.dry_run:
_dry_run_sample(db)
total = 0
total += backfill_messages(db, args.dry_run)
total += backfill_personas(db, args.dry_run)
print(f"\n{'[DRY RUN] ' if args.dry_run else ''}Backfill complete — {total} events total")
def _dry_run_sample(db):
"""Print a sample of estimated costs to sanity-check before real run."""
from collections import defaultdict
by_fg: dict = defaultdict(list)
for msg in db.focus_group_messages.find({}):
fg_id = str(msg.get("focus_group_id") or "")
if msg.get("type", "") in ("response", "question", "moderator", "ai", ""):
by_fg[fg_id].append(msg)
print("\n[dry-run sample] Estimated cost per focus group (top 5 by message count):")
fg_meta = {str(fg["_id"]): fg.get("llm_model") or "gemini-3.1-pro-preview"
for fg in db.focus_groups.find({}, {"llm_model": 1})}
rows = []
for fg_id, msgs in by_fg.items():
model = fg_meta.get(fg_id, "gemini-3.1-pro-preview")
accumulated = 0
total_cost = 0
for msg in sorted(msgs, key=lambda m: m.get("created_at") or datetime.min):
text = msg.get("text") or msg.get("content") or ""
completion = _chars_to_tokens(len(text))
prompt = BASE_PROMPT_PERSONA_RESPONSE + _chars_to_tokens(accumulated)
cost = _estimate_cost(prompt, completion, model)
total_cost += cost["total"]
accumulated += len(text)
rows.append((fg_id, len(msgs), total_cost))
for fg_id, count, cost in sorted(rows, key=lambda r: -r[1])[:5]:
fg = db.focus_groups.find_one({"_id": __import__("bson").ObjectId(fg_id)}, {"name": 1}) if fg_id else None
name = (fg or {}).get("name", fg_id[:8])
print(f" {name}: {count} messages → estimated ${cost:.4f}")
if __name__ == "__main__":
main()