diff --git a/backend/agents/document.py b/backend/agents/document.py index 49dfe16b..9dccad42 100644 --- a/backend/agents/document.py +++ b/backend/agents/document.py @@ -52,6 +52,7 @@ from agents.concept_extraction import concept_extraction_agent, ConceptList from agents.syllabus_extraction import syllabus_extraction_agent, SyllabusAssignments from agents.tools.graph import apply_concepts_to_graph +from agents.usage import record_agent_usage from services.durable import workflow as durable_workflow, step as durable_step @@ -86,25 +87,37 @@ class _WorkerResults: @durable_step async def _step_classify(text: str, deps: SaplingDeps) -> DocumentClassification: - result = await classifier_agent.run(text, deps=deps, usage_limits=WORKER_LIMITS) + result = record_agent_usage( + await classifier_agent.run(text, deps=deps, usage_limits=WORKER_LIMITS), + feature="document", task="classifier", user_id=deps.user_id, + ) return result.output @durable_step async def _step_summary(text: str, deps: SaplingDeps) -> Summary: - result = await summary_agent.run(text, deps=deps, usage_limits=WORKER_LIMITS) + result = record_agent_usage( + await summary_agent.run(text, deps=deps, usage_limits=WORKER_LIMITS), + feature="document", task="summary", user_id=deps.user_id, + ) return result.output @durable_step async def _step_concepts(text: str, deps: SaplingDeps) -> ConceptList: - result = await concept_extraction_agent.run(text, deps=deps, usage_limits=WORKER_LIMITS) + result = record_agent_usage( + await concept_extraction_agent.run(text, deps=deps, usage_limits=WORKER_LIMITS), + feature="document", task="concepts", user_id=deps.user_id, + ) return result.output @durable_step async def _step_syllabus(text: str, deps: SaplingDeps) -> SyllabusAssignments: - result = await syllabus_extraction_agent.run(text, deps=deps, usage_limits=WORKER_LIMITS) + result = record_agent_usage( + await syllabus_extraction_agent.run(text, deps=deps, usage_limits=WORKER_LIMITS), + feature="document", task="syllabus", user_id=deps.user_id, + ) return result.output diff --git a/backend/agents/usage.py b/backend/agents/usage.py new file mode 100644 index 00000000..0403ab17 --- /dev/null +++ b/backend/agents/usage.py @@ -0,0 +1,82 @@ +"""One-line usage capture for Pydantic AI agent runs (issue #118). + +Every agent call site wraps its result in ``record_agent_usage(result, +feature=..., task=...)``. The helper reads ``result.usage()`` and the model +actually used, then hands them to ``events_service.log_llm_usage`` (which +normalizes tokens, computes cost, and enqueues off the request thread). + +Two properties matter: + +* **One line per call site.** Because it returns ``result`` unchanged, a call + site can wrap inline (``result = record_agent_usage(await agent.run(...), + feature=..., task=...)``) or add it as a trailing statement. +* **Never raises.** Instrumentation must not break an agent run, so every + failure — a result shape we don't recognize, a usage extraction slip — is + swallowed and logged at debug level. + +The ``model`` is read from the result's final ``ModelResponse`` (the model the +provider actually served); if that's unavailable it falls back to the task's +configured model via ``_providers.model_for``. +""" + +from __future__ import annotations + +import logging +from typing import Any + +from agents._providers import AgentTask, model_for +from services import events_service + +logger = logging.getLogger("sapling.agents.usage") + + +def _model_name(result: Any, task: AgentTask | None) -> str: + """Best-effort model id for the run, resilient to Pydantic AI churn.""" + # Preferred: the final ModelResponse carries the served model name. + try: + name = getattr(result.response, "model_name", None) + if name: + return name + except Exception: + pass + # Fallback: scan the message history for the last response with a model. + try: + for msg in reversed(result.all_messages()): + name = getattr(msg, "model_name", None) + if name: + return name + except Exception: + pass + # Last resort: the task's configured default model. + if task is not None: + try: + return model_for(task).model_name + except Exception: + pass + return "unknown" + + +def record_agent_usage( + result: Any, + *, + feature: str, + task: AgentTask | None = None, + user_id: str | None = None, +) -> Any: + """Record token usage for an agent run and return ``result`` unchanged. + + ``user_id`` is optional: pass it where the actor is in scope (routes with a + ``deps.user_id`` / request body) for per-user rollups; omit it and the + request_id from the contextvar still attributes the row. + """ + try: + events_service.log_llm_usage( + feature=feature, + task=task, + model=_model_name(result, task), + usage=result.usage(), + user_id=user_id, + ) + except Exception: + logger.debug("record_agent_usage: could not capture usage", exc_info=True) + return result diff --git a/backend/db/migrations/0035_observability.sql b/backend/db/migrations/0035_observability.sql new file mode 100644 index 00000000..121231f2 --- /dev/null +++ b/backend/db/migrations/0035_observability.sql @@ -0,0 +1,57 @@ +-- 0035: observability tables (issue #115) — the two tables backing the site +-- logging + usage-tracking system. +-- +-- events — flexible, high-volume analytics / audit / error events. +-- llm_usage — structured per-call LLM token + cost rows, queried with +-- SUM/GROUP BY for billing rollups. +-- +-- Notes / deviations from the original issue spec: +-- * Ids and user_id are TEXT, not uuid, to match the rest of this schema +-- (users.id is TEXT, e.g. 'user_andres'); a uuid user_id column could not +-- hold the existing text ids. Same gen_random_uuid()::text PK style as +-- 0026_ops.sql. +-- * No FK on user_id. These are append-only, high-write analytics tables and +-- user_id is intentionally nullable (system/anonymous actors); we don't +-- want a per-row FK check on the write path or cascade coupling to the +-- users lifecycle. Kept as a plain nullable column. +-- * NO raw-content columns. Sensitive content is represented only by +-- content_fp (a 16-hex sha256 fingerprint). Never store message text, +-- document text, or names here. +-- * admin_audit_log is deliberately untouched — events complements it. +-- +-- Idempotent (IF NOT EXISTS throughout) so it is safe to re-run. + +-- ── events ────────────────────────────────────────────────────────────────── +CREATE TABLE IF NOT EXISTS events ( + id TEXT PRIMARY KEY DEFAULT gen_random_uuid()::text, + event_type TEXT NOT NULL, -- dotted taxonomy: document.upload, auth.login, error.5xx + category TEXT NOT NULL, -- usage | audit | error + user_id TEXT, -- actor; NULL for anonymous/system + request_id TEXT, -- correlates to RequestIDMiddleware + Logfire + payload JSONB NOT NULL DEFAULT '{}', -- type-specific metadata (counts, ids, status_code, duration_ms…) + content_fp TEXT, -- sha256 fingerprint (16 hex) — never raw content + created_at TIMESTAMPTZ NOT NULL DEFAULT now() +); + +CREATE INDEX IF NOT EXISTS idx_events_user_created ON events (user_id, created_at DESC); +CREATE INDEX IF NOT EXISTS idx_events_type_created ON events (event_type, created_at DESC); +CREATE INDEX IF NOT EXISTS idx_events_cat_created ON events (category, created_at DESC); + +-- ── llm_usage ──────────────────────────────────────────────────────────────── +CREATE TABLE IF NOT EXISTS llm_usage ( + id TEXT PRIMARY KEY DEFAULT gen_random_uuid()::text, + user_id TEXT, -- per-user rollups; NULL for system + request_id TEXT, + feature TEXT NOT NULL, -- quiz | chat_tutor | document | notes … + task TEXT, -- matches agents/_providers.py task slots + model TEXT NOT NULL, -- e.g. gemini-2.5-flash + provider TEXT NOT NULL DEFAULT 'gemini', + prompt_tokens INTEGER NOT NULL DEFAULT 0, + completion_tokens INTEGER NOT NULL DEFAULT 0, + total_tokens INTEGER NOT NULL DEFAULT 0, + cost_usd NUMERIC(12,6), -- NULL when the model isn't priced + created_at TIMESTAMPTZ NOT NULL DEFAULT now() +); + +CREATE INDEX IF NOT EXISTS idx_llm_usage_user_created ON llm_usage (user_id, created_at DESC); +CREATE INDEX IF NOT EXISTS idx_llm_usage_feature_created ON llm_usage (feature, created_at DESC); diff --git a/backend/main.py b/backend/main.py index bb8a0427..8c7fc61f 100644 --- a/backend/main.py +++ b/backend/main.py @@ -24,6 +24,7 @@ from routes import graph, learn, quiz, calendar, social, extract, auth, documents, flashcards, study_guide, feedback, careers, onboarding, gradebook, gradescope, notes, academics from routes.profile import router as profile_router from routes.admin import router as admin_router +from routes.admin_analytics import router as admin_analytics_router from routes.newsletter import router as newsletter_router from services.logfire_scrubber import EXTRA_PATTERNS, scrub_value from services.request_context import RequestIDMiddleware, current_request_id @@ -80,8 +81,14 @@ async def _lifespan(_app: FastAPI): file_size_limit=MAX_AVATAR_SIZE, allowed_mime_types=sorted(ALLOWED_CONTENT_TYPES), ) + # #116/#118: start the fire-and-forget observability drain thread so LLM + # usage + event rows flush off the request path. + from services import events_service + events_service.start_worker() yield - # No shutdown hooks today. + # Stop the drain thread and flush anything still queued so the last batch + # of usage rows isn't lost on shutdown. + events_service.shutdown() def _drop_request_arguments(_request, _attributes): @@ -199,6 +206,7 @@ async def unhandled_exception_handler(request: Request, exc: Exception): app.include_router(onboarding.router, prefix="/api/onboarding") app.include_router(profile_router, prefix="/api/profile") app.include_router(admin_router, prefix="/api/admin") +app.include_router(admin_analytics_router, prefix="/api/admin/analytics") app.include_router(newsletter_router, prefix="/api/newsletter") app.include_router(gradebook.router, prefix="/api/gradebook") app.include_router(gradescope.router, prefix="/api/gradescope") @@ -267,9 +275,13 @@ def gemini_test(request: Request): require_admin(request) # 403 unless the session belongs to an admin; 401 if unauthenticated from agents._run import run_agent_sync from agents.health import health_probe_agent + from agents.usage import record_agent_usage try: - result = run_agent_sync( - health_probe_agent.run('Reply with exactly the text: Gemini OK') + result = record_agent_usage( + run_agent_sync( + health_probe_agent.run('Reply with exactly the text: Gemini OK') + ), + feature="health", ) return {"ok": True, "reply": result.output.strip()} except Exception as e: diff --git a/backend/routes/admin_analytics.py b/backend/routes/admin_analytics.py new file mode 100644 index 00000000..1c3690f3 --- /dev/null +++ b/backend/routes/admin_analytics.py @@ -0,0 +1,395 @@ +"""Admin-only, read-only analytics + cost-rollup API (issue #120). + +Turns the `events` and `llm_usage` tables (written by services/events_service.py, +#116/#118) into usage summaries, per-user rollups, LLM cost breakdowns, and an +error feed. Mounted at `/api/admin/analytics`; every endpoint is gated by +`require_admin`. Read-only — no mutation endpoints. + +Aggregation strategy: PostgREST (via `db/connection.py::table()`) has no +GROUP BY, so the grouped endpoints scan the (date-bounded) rows and aggregate +in Python. Scans page through `select_with_count` and use its exact count both +to know when to stop and to detect the rare truncation case (logged and +surfaced as `truncated: true` in the response, never silent). The `/errors` +feed needs no aggregation, so it paginates server-side. +""" + +from __future__ import annotations + +import logging +from collections import defaultdict +from datetime import datetime, timedelta, timezone +from typing import Literal + +from fastapi import APIRouter, HTTPException, Query, Request, Response +from pydantic import BaseModel + +from db.connection import table +from services.auth_guard import require_admin + +logger = logging.getLogger("sapling.admin_analytics") + +router = APIRouter() + +# Page size for range scans, and a hard ceiling so a pathological range can't +# pull unbounded rows into memory. Hitting the cap is logged and surfaced via +# the response's `truncated` flag, never silent. +_PAGE = 1000 +_SCAN_CAP = 100_000 + +GroupBy = Literal["user", "feature", "model"] +_GROUP_COLUMN = {"user": "user_id", "feature": "feature", "model": "model"} + + +# ── Response models ────────────────────────────────────────────────────────── + + +class Range(BaseModel): + from_: str + to: str + + +class EventTypeCount(BaseModel): + event_type: str + count: int + + +class UsageSummary(BaseModel): + range: Range + total_events: int + distinct_active_users: int + by_event_type: list[EventTypeCount] + truncated: bool = False + + +class UserUsage(BaseModel): + user_id: str + event_count: int + by_category: dict[str, int] + llm_cost_usd: float + total_tokens: int + + +class UsageByUser(BaseModel): + range: Range + total_users: int + limit: int + offset: int + users: list[UserUsage] + truncated: bool = False + + +class CostRow(BaseModel): + key: str + calls: int + prompt_tokens: int + completion_tokens: int + total_tokens: int + cost_usd: float + + +class CostTotals(BaseModel): + calls: int + prompt_tokens: int + completion_tokens: int + total_tokens: int + cost_usd: float + + +class LLMCost(BaseModel): + range: Range + group_by: GroupBy + rows: list[CostRow] + totals: CostTotals + truncated: bool = False + + +class ErrorEvent(BaseModel): + created_at: str | None + event_type: str + request_id: str | None + user_id: str | None + path: str | None + method: str | None + status_code: int | None + duration_ms: float | None + + +class ErrorsPage(BaseModel): + range: Range + total: int + limit: int + offset: int + errors: list[ErrorEvent] + + +# ── Date range helpers ─────────────────────────────────────────────────────── + + +def _resolve_range(from_: str | None, to: str | None) -> tuple[str, str]: + """Default to the last 30 days; echo caller-supplied ISO bounds otherwise. + + Bounds are validated before use: each must parse as ISO 8601 (422 naming + the bad param otherwise) and `from` must not be after `to`. The strings are + returned as supplied — validation never reformats them. + """ + now = datetime.now(timezone.utc) + to_iso = to or now.isoformat() + from_iso = from_ or (now - timedelta(days=30)).isoformat() + + def _parse(param: str, value: str) -> datetime: + try: + parsed = datetime.fromisoformat(value) + except ValueError: + raise HTTPException( + status_code=422, + detail=f"Invalid {param!r} datetime: {value!r} is not ISO 8601", + ) + # Treat naive bounds as UTC so mixed naive/aware bounds stay comparable. + return parsed if parsed.tzinfo is not None else parsed.replace(tzinfo=timezone.utc) + + from_dt = _parse("from", from_iso) + to_dt = _parse("to", to_iso) + if from_dt > to_dt: + raise HTTPException(status_code=422, detail="'from' must not be after 'to'") + return from_iso, to_iso + + +def _scan_range( + table_name: str, columns: str, from_iso: str, to_iso: str, + extra_filters: dict | None = None, +) -> tuple[list[dict], bool]: + """Fetch every row in [from, to] for a table, paging via select_with_count. + + Aggregation endpoints need the full (date-bounded) set — PostgREST won't + GROUP BY for us — so we page to completion rather than relying on the + server's default row cap. Returns ``(rows, truncated)``: a range large + enough to hit _SCAN_CAP stops the scan early, logs a warning, and sets + ``truncated=True`` so callers can surface the partial aggregation. + """ + out: list[dict] = [] + offset = 0 + truncated = False + while True: + filters: dict = {"created_at": [f"gte.{from_iso}", f"lte.{to_iso}"]} + if extra_filters: + filters.update(extra_filters) + rows, total = table(table_name).select_with_count( + columns, filters=filters, order="created_at.asc", limit=_PAGE, offset=offset, + ) + out.extend(rows) + if len(out) >= total or not rows: + break + if len(out) >= _SCAN_CAP: + truncated = True + logger.warning( + "admin_analytics scan hit cap %d on %r (total=%d); results truncated", + _SCAN_CAP, table_name, total, + ) + break + offset += _PAGE + return out, truncated + + +def _as_float(value) -> float: + try: + return float(value) if value is not None else 0.0 + except (TypeError, ValueError): + return 0.0 + + +def _as_int(value) -> int: + try: + return int(value) if value is not None else 0 + except (TypeError, ValueError): + return 0 + + +# ── Endpoints ──────────────────────────────────────────────────────────────── + + +@router.get("/usage/summary", response_model=UsageSummary) +def usage_summary( + request: Request, + response: Response, + from_: str | None = Query(None, alias="from"), + to: str | None = Query(None), +) -> UsageSummary: + """Event totals for the range; `truncated: true` means the scan cap cut the aggregation short.""" + require_admin(request) + response.headers["Cache-Control"] = "private" + from_iso, to_iso = _resolve_range(from_, to) + rows, truncated = _scan_range("events", "event_type,user_id,created_at", from_iso, to_iso) + + by_type: dict[str, int] = defaultdict(int) + users: set[str] = set() + for r in rows: + by_type[r.get("event_type") or ""] += 1 + if r.get("user_id"): + users.add(r["user_id"]) + + return UsageSummary( + range=Range(from_=from_iso, to=to_iso), + total_events=len(rows), + distinct_active_users=len(users), + by_event_type=sorted( + (EventTypeCount(event_type=k, count=v) for k, v in by_type.items()), + key=lambda e: e.count, reverse=True, + ), + truncated=truncated, + ) + + +@router.get("/usage/by-user", response_model=UsageByUser) +def usage_by_user( + request: Request, + response: Response, + from_: str | None = Query(None, alias="from"), + to: str | None = Query(None), + limit: int = Query(50, ge=1, le=500), + offset: int = Query(0, ge=0), +) -> UsageByUser: + """Per-user event/cost rollup; `truncated: true` means a scan cap cut the aggregation short.""" + require_admin(request) + response.headers["Cache-Control"] = "private" + from_iso, to_iso = _resolve_range(from_, to) + + event_rows, events_truncated = _scan_range( + "events", "user_id,category,created_at", from_iso, to_iso, + ) + usage_rows, usage_truncated = _scan_range( + "llm_usage", "user_id,cost_usd,total_tokens,created_at", from_iso, to_iso, + ) + + agg: dict[str, dict] = defaultdict( + lambda: {"event_count": 0, "by_category": defaultdict(int), "llm_cost_usd": 0.0, "total_tokens": 0}, + ) + for r in event_rows: + uid = r.get("user_id") + if not uid: + continue + a = agg[uid] + a["event_count"] += 1 + a["by_category"][r.get("category") or ""] += 1 + for r in usage_rows: + uid = r.get("user_id") + if not uid: + continue + a = agg[uid] + a["llm_cost_usd"] += _as_float(r.get("cost_usd")) + a["total_tokens"] += _as_int(r.get("total_tokens")) + + # Sort by spend then activity so the most expensive users surface first. + ordered = sorted( + agg.items(), + key=lambda kv: (kv[1]["llm_cost_usd"], kv[1]["event_count"]), + reverse=True, + ) + page = ordered[offset:offset + limit] + users = [ + UserUsage( + user_id=uid, + event_count=a["event_count"], + by_category=dict(a["by_category"]), + llm_cost_usd=round(a["llm_cost_usd"], 6), + total_tokens=a["total_tokens"], + ) + for uid, a in page + ] + return UsageByUser( + range=Range(from_=from_iso, to=to_iso), + total_users=len(ordered), limit=limit, offset=offset, users=users, + truncated=events_truncated or usage_truncated, + ) + + +@router.get("/llm/cost", response_model=LLMCost) +def llm_cost( + request: Request, + response: Response, + from_: str | None = Query(None, alias="from"), + to: str | None = Query(None), + group_by: GroupBy = Query("feature"), +) -> LLMCost: + """LLM token/cost rollup; `truncated: true` means the scan cap cut the aggregation short.""" + require_admin(request) + response.headers["Cache-Control"] = "private" + from_iso, to_iso = _resolve_range(from_, to) + column = _GROUP_COLUMN[group_by] + rows, truncated = _scan_range( + "llm_usage", + f"{column},prompt_tokens,completion_tokens,total_tokens,cost_usd,created_at", + from_iso, to_iso, + ) + + buckets: dict[str, dict] = defaultdict( + lambda: {"calls": 0, "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0, "cost_usd": 0.0}, + ) + for r in rows: + key = r.get(column) + b = buckets["" if key is None else str(key)] + b["calls"] += 1 + b["prompt_tokens"] += _as_int(r.get("prompt_tokens")) + b["completion_tokens"] += _as_int(r.get("completion_tokens")) + b["total_tokens"] += _as_int(r.get("total_tokens")) + b["cost_usd"] += _as_float(r.get("cost_usd")) + + cost_rows = sorted( + ( + CostRow( + key=k, calls=b["calls"], + prompt_tokens=b["prompt_tokens"], completion_tokens=b["completion_tokens"], + total_tokens=b["total_tokens"], cost_usd=round(b["cost_usd"], 6), + ) + for k, b in buckets.items() + ), + key=lambda c: c.cost_usd, reverse=True, + ) + totals = CostTotals( + calls=sum(c.calls for c in cost_rows), + prompt_tokens=sum(c.prompt_tokens for c in cost_rows), + completion_tokens=sum(c.completion_tokens for c in cost_rows), + total_tokens=sum(c.total_tokens for c in cost_rows), + cost_usd=round(sum(c.cost_usd for c in cost_rows), 6), + ) + return LLMCost( + range=Range(from_=from_iso, to=to_iso), group_by=group_by, + rows=cost_rows, totals=totals, truncated=truncated, + ) + + +@router.get("/errors", response_model=ErrorsPage) +def errors( + request: Request, + response: Response, + from_: str | None = Query(None, alias="from"), + to: str | None = Query(None), + limit: int = Query(50, ge=1, le=500), + offset: int = Query(0, ge=0), +) -> ErrorsPage: + """Paginated error.* event feed (server-side pagination — no range scan, so no cap).""" + require_admin(request) + response.headers["Cache-Control"] = "private" + from_iso, to_iso = _resolve_range(from_, to) + # error.* events, newest first — paginated server-side (no aggregation). + rows, total = table("events").select_with_count( + "created_at,event_type,request_id,user_id,payload", + filters={"created_at": [f"gte.{from_iso}", f"lte.{to_iso}"], "event_type": "like.error.*"}, + order="created_at.desc", limit=limit, offset=offset, + ) + items = [] + for r in rows: + payload = r.get("payload") or {} + items.append(ErrorEvent( + created_at=r.get("created_at"), + event_type=r.get("event_type") or "", + request_id=r.get("request_id"), + user_id=r.get("user_id"), + path=payload.get("path"), + method=payload.get("method"), + status_code=payload.get("status_code"), + duration_ms=payload.get("duration_ms"), + )) + return ErrorsPage( + range=Range(from_=from_iso, to=to_iso), + total=total, limit=limit, offset=offset, errors=items, + ) diff --git a/backend/routes/documents.py b/backend/routes/documents.py index 79cd7553..451d2d3c 100644 --- a/backend/routes/documents.py +++ b/backend/routes/documents.py @@ -52,6 +52,7 @@ from agents.tools.graph import apply_concepts_to_graph from agents._run import run_agent_sync from agents.concept_scan import concept_scan_agent +from agents.usage import record_agent_usage logger = logging.getLogger(__name__) @@ -177,7 +178,7 @@ def _extend_course_concepts( "administrative items.\n" "- concepts must be a JSON array of strings." ) - raw = call_gemini_json(prompt, model=MODEL_LITE) + raw = call_gemini_json(prompt, model=MODEL_LITE, feature="document") if not isinstance(raw, dict): return [] return _coerce_str_list(raw.get("concepts")) @@ -246,8 +247,11 @@ async def _extend_via_agent( doc_summary=doc_summary, doc_concept_notes=doc_concept_notes, ) - result = await concept_scan_agent.run( - message, deps=deps, usage_limits=WORKER_LIMITS, + result = record_agent_usage( + await concept_scan_agent.run( + message, deps=deps, usage_limits=WORKER_LIMITS, + ), + feature="document", task="concept_scan", ) return list(result.output.concepts) @@ -362,7 +366,7 @@ def _process_document(filename: str, extracted_text: str) -> dict: "\n" '"concept_notes" must be a JSON array of {"name": str, "description": str} objects.' ) - raw = call_gemini_json(prompt) + raw = call_gemini_json(prompt, feature="document") if not isinstance(raw, dict): raw = {} @@ -915,8 +919,11 @@ async def event_stream(): type="progress", step="classify", message="Classifying document...", )) - cls_run = await classifier_agent.run( - extracted_text, deps=deps, usage_limits=WORKER_LIMITS, + cls_run = record_agent_usage( + await classifier_agent.run( + extracted_text, deps=deps, usage_limits=WORKER_LIMITS, + ), + feature="document", task="classifier", ) classification = cls_run.output yield sapling_event_to_sse(SaplingEvent( @@ -948,6 +955,7 @@ async def event_stream(): summary_r, concepts_r, syllabus_r = await asyncio.gather( summary_task, concepts_task, syllabus_task, ) + record_agent_usage(syllabus_r, feature="document", task="syllabus") summary = summary_r.output concepts = concepts_r.output syllabus = syllabus_r.output @@ -956,6 +964,8 @@ async def event_stream(): summary = summary_r.output concepts = concepts_r.output syllabus = None + record_agent_usage(summary_r, feature="document", task="summary") + record_agent_usage(concepts_r, feature="document", task="concepts") yield sapling_event_to_sse(SaplingEvent( type="progress", step="extracted", message=f"Extracted {len(concepts.concepts)} concept(s).", diff --git a/backend/routes/graph.py b/backend/routes/graph.py index d019e1f6..eb434f36 100644 --- a/backend/routes/graph.py +++ b/backend/routes/graph.py @@ -18,6 +18,7 @@ from agents._run import run_agent_sync from agents.deps import SaplingDeps from agents.concept_describe import concept_describe_agent, build_message +from agents.usage import record_agent_usage router = APIRouter() @@ -130,12 +131,15 @@ def describe_concept(user_id: str, body: ConceptDescriptionBody, request: Reques request_id=current_request_id() or str(uuid.uuid4()), ) try: - result = run_agent_sync( - concept_describe_agent.run( - build_message(concept, course_label), - deps=deps, - usage_limits=WORKER_LIMITS, - ) + result = record_agent_usage( + run_agent_sync( + concept_describe_agent.run( + build_message(concept, course_label), + deps=deps, + usage_limits=WORKER_LIMITS, + ) + ), + feature="graph", task="concept_describe", user_id=user_id, ) except (AgentRunError, httpx.HTTPError, ValidationError, UnregisteredHandlerError) as e: # Model / transport / output-validation failures are upstream problems — diff --git a/backend/routes/learn.py b/backend/routes/learn.py index 21d269e9..e382b779 100644 --- a/backend/routes/learn.py +++ b/backend/routes/learn.py @@ -15,6 +15,7 @@ from agents import ORCHESTRATOR_LIMITS from agents.chat_tutor import agent_for_mode from agents.deps import SaplingDeps +from agents.usage import record_agent_usage from db.connection import table from services.academics import offering_course_id, resolve_offering from models import StartSessionBody, ChatBody, EndSessionBody, ActionBody, ModeSwitchBody, RenameSessionBody @@ -498,7 +499,8 @@ def _start_session_legacy(body: StartSessionBody, session_id: str | None = None) try: raw = call_gemini_multiturn( - system_prompt, [], user_message, model=_resolve_legacy_model(body.model_pref) + system_prompt, [], user_message, model=_resolve_legacy_model(body.model_pref), + feature="chat_tutor", ) except Exception as e: raise HTTPException(status_code=502, detail=f"Gemini error: {e}") @@ -657,7 +659,10 @@ async def _chat_via_agent( model_pref=model_pref, ) - result = await agent.run(user_message, **run_kwargs) + result = record_agent_usage( + await agent.run(user_message, **run_kwargs), + feature="chat_tutor", task="chat_tutor", user_id=deps.user_id, + ) reply = result.output # str — chat_tutor agents return plain Markdown. # Merge all graph update payloads accumulated by tools during this run @@ -716,7 +721,8 @@ async def _legacy_chat(body: ChatBody, request: Request) -> dict: try: raw = call_gemini_multiturn( - system_prompt, history, body.message, model=_resolve_legacy_model(body.model_pref) + system_prompt, history, body.message, model=_resolve_legacy_model(body.model_pref), + feature="chat_tutor", ) except Exception as e: raise HTTPException(status_code=502, detail=f"Gemini error: {e}") @@ -842,6 +848,13 @@ async def _legacy() -> dict: # stream_agent_turn enforces that exclusivity. return await _legacy_chat(body, request) + def _usage(run_result) -> None: + # #118: streaming turns report usage via the final AgentRunResultEvent, + # surfaced by stream_agent_turn's on_usage hook after the run completes. + record_agent_usage( + run_result, feature="chat_tutor", task="chat_tutor", user_id=body.user_id, + ) + async def event_stream(): async for ev in stream_agent_turn( agent=agent, @@ -850,6 +863,7 @@ async def event_stream(): deps=deps, on_complete=_persist, legacy_fallback=_legacy, + on_usage=_usage, request_id=request_id, ): yield sapling_event_to_sse(ev) @@ -940,6 +954,13 @@ async def _legacy() -> dict: # per turn — see stream_agent_turn's docstring). return _start_session_legacy(body, session_id) + def _usage(run_result) -> None: + # #118: same hook as /chat/stream — the opener runs the same + # chat_tutor agent, so it rolls up under the same feature/task. + record_agent_usage( + run_result, feature="chat_tutor", task="chat_tutor", user_id=body.user_id, + ) + async def event_stream(): async for ev in stream_agent_turn( agent=agent, @@ -948,6 +969,7 @@ async def event_stream(): deps=deps, on_complete=_stash, legacy_fallback=_legacy, + on_usage=_usage, request_id=request_id, ): yield sapling_event_to_sse(ev) @@ -1211,7 +1233,8 @@ def action(body: ActionBody, request: Request): try: raw = call_gemini_multiturn( - system_prompt, history, action_message, model=_resolve_legacy_model(body.model_pref) + system_prompt, history, action_message, model=_resolve_legacy_model(body.model_pref), + feature="chat_tutor", ) except Exception as e: raise HTTPException(status_code=502, detail=f"Gemini error: {e}") diff --git a/backend/routes/notes.py b/backend/routes/notes.py index afcddca8..50789010 100644 --- a/backend/routes/notes.py +++ b/backend/routes/notes.py @@ -17,6 +17,7 @@ from agents.note_concepts import note_concepts_agent from agents.note_summary import note_summary_agent from agents.tools.graph import apply_concepts_to_graph +from agents.usage import record_agent_usage from db.connection import table from services.academics import offering_course_id, resolve_offering from services.auth_guard import get_session_user_id, require_self @@ -267,8 +268,11 @@ async def summarize(note_id: str, body: AgentActionBody, request: Request): # The graph keys on the abstract course; the note carries the offering. course_id = offering_course_id(note.get("offering_id")) deps = _deps_for(body.user_id, course_id, note_id) - result = await _run_note_worker( - note_summary_agent, user_prompt, deps, action="summarization" + result = record_agent_usage( + await _run_note_worker( + note_summary_agent, user_prompt, deps, action="summarization" + ), + feature="notes", task="note_summary", user_id=body.user_id, ) summary_text = result.output.summary await save_summary(note_id=note_id, user_id=body.user_id, summary=summary_text) @@ -290,8 +294,11 @@ async def extract_concepts( # Concepts land in the abstract-course graph; resolve offering → course. course_id = offering_course_id(note.get("offering_id")) deps = _deps_for(body.user_id, course_id, note_id) - result = await _run_note_worker( - note_concepts_agent, user_prompt, deps, action="concept extraction" + result = record_agent_usage( + await _run_note_worker( + note_concepts_agent, user_prompt, deps, action="concept extraction" + ), + feature="notes", task="note_concepts", user_id=body.user_id, ) names = [n.strip() for n in (result.output.concepts or []) if n and n.strip()] await apply_concepts_to_graph( @@ -322,8 +329,11 @@ async def note_chat(note_id: str, body: NoteChatBody, request: Request): course_id = offering_course_id(note.get("offering_id")) deps = _deps_for(body.user_id, course_id, note_id) try: - result = await note_chat_agent.run( - body.message, deps=deps, usage_limits=ORCHESTRATOR_LIMITS + result = record_agent_usage( + await note_chat_agent.run( + body.message, deps=deps, usage_limits=ORCHESTRATOR_LIMITS + ), + feature="notes", task="note_chat", user_id=body.user_id, ) except UsageLimitExceeded as e: # Only a real budget trip reaches the in-band degrade, so the budget diff --git a/backend/routes/quiz.py b/backend/routes/quiz.py index 9b67e844..1e677116 100644 --- a/backend/routes/quiz.py +++ b/backend/routes/quiz.py @@ -14,6 +14,7 @@ from agents.deps import SaplingDeps from agents._run import run_agent_sync from agents.quiz_context import quiz_context_agent +from agents.usage import record_agent_usage from db.connection import table from models import GenerateQuizBody, SubmitQuizBody from routes.learn import _get_catalog_chunk @@ -264,7 +265,10 @@ async def _quiz_via_agent( run_kwargs: dict = {"deps": deps, "usage_limits": ORCHESTRATOR_LIMITS} if model_override is not None: run_kwargs["model"] = model_override - result = await quiz_agent.run(user_message, **run_kwargs) + result = record_agent_usage( + await quiz_agent.run(user_message, **run_kwargs), + feature="quiz", task="quiz", user_id=deps.user_id, + ) quiz: Quiz = result.output # Filter out questions where the agent's correct_answer didn't match # any option verbatim — _agent_question_to_wire returns None for those. @@ -465,7 +469,10 @@ def submit_quiz(body: SubmitQuizBody, background_tasks: BackgroundTasks, request def _update_context(prompt: str, uid: str, node_id: str): try: - result = run_agent_sync(quiz_context_agent.run(prompt)) + result = record_agent_usage( + run_agent_sync(quiz_context_agent.run(prompt)), + feature="quiz", task="quiz_context", user_id=uid, + ) save_quiz_context(uid, node_id, result.output.model_dump()) except Exception: pass diff --git a/backend/routes/social.py b/backend/routes/social.py index 358da0c9..d4811f1e 100644 --- a/backend/routes/social.py +++ b/backend/routes/social.py @@ -8,6 +8,7 @@ from agents._run import run_agent_sync from agents.deps import SaplingDeps from agents.social_summary import social_summary_agent +from agents.usage import record_agent_usage from db.connection import table from models import CreateRoomBody, JoinRoomBody, MatchBody, SendMessageBody, EditMessageBody, ToggleReactionBody, LeaveRoomBody from services.auth_guard import require_self, get_session_user_id @@ -140,7 +141,10 @@ def room_overview(room_id: str, request: Request): "Summarize this study group's collective knowledge:\n" + "\n".join(member_summaries) ) - result = run_agent_sync(social_summary_agent.run(user_message, deps=deps)) + result = record_agent_usage( + run_agent_sync(social_summary_agent.run(user_message, deps=deps)), + feature="social", task="social_summary", + ) ai_summary = result.output.summary save_summary(room_id, member_summaries, ai_summary) except Exception as e: diff --git a/backend/routes/study_guide.py b/backend/routes/study_guide.py index 14dc5652..c7244b52 100644 --- a/backend/routes/study_guide.py +++ b/backend/routes/study_guide.py @@ -12,6 +12,7 @@ from agents._run import run_agent_sync from agents.deps import SaplingDeps from agents.study_guide import study_guide_agent +from agents.usage import record_agent_usage from db.connection import table from services.academics import offering_course_id, resolve_offering from services.auth_guard import require_self @@ -106,7 +107,10 @@ def _generate_and_insert(user_id: str, offering_id: str, exam_id: str) -> dict: request_id=current_request_id() or "", ) try: - result = run_agent_sync(study_guide_agent.run(user_message, deps=deps)) + result = record_agent_usage( + run_agent_sync(study_guide_agent.run(user_message, deps=deps)), + feature="study_guide", task="study_guide", user_id=user_id, + ) except Exception as e: # generation/transport failure → 502, not a raw 500 raise HTTPException( status_code=502, detail="Study guide generation failed." diff --git a/backend/services/calendar_service.py b/backend/services/calendar_service.py index 2f5cd2ac..9c18d641 100644 --- a/backend/services/calendar_service.py +++ b/backend/services/calendar_service.py @@ -10,6 +10,7 @@ from agents.deps import SaplingDeps from agents.syllabus_extraction import syllabus_extraction_agent from agents.tools.syllabus_adapter import syllabus_to_wire_dict +from agents.usage import record_agent_usage from services.extraction_service import extract_text_from_file from services.assignment_dedupe import assignment_dedupe_key from services.encryption import encrypt_if_present @@ -112,8 +113,11 @@ async def _extract_via_agent( supabase=None, request_id=request_id or "", ) - result = await syllabus_extraction_agent.run( - extracted_text, deps=deps, usage_limits=WORKER_LIMITS + result = record_agent_usage( + await syllabus_extraction_agent.run( + extracted_text, deps=deps, usage_limits=WORKER_LIMITS + ), + feature="document", task="syllabus", user_id=user_id, ) return syllabus_to_wire_dict(result.output, raw_text=extracted_text) diff --git a/backend/services/chat_stream.py b/backend/services/chat_stream.py index f29f40e7..0a61550d 100644 --- a/backend/services/chat_stream.py +++ b/backend/services/chat_stream.py @@ -109,6 +109,7 @@ async def stream_agent_turn( deps: Any, on_complete: Callable[[str, dict, list], dict | None], legacy_fallback: Callable[[], Awaitable[dict]] | None = None, + on_usage: Callable[[Any], None] | None = None, request_id: str = "", ) -> AsyncIterator[SaplingEvent]: """Stream one agent turn as SaplingEvents. @@ -118,6 +119,14 @@ async def stream_agent_turn( completes and BEFORE `done` is yielded, so a mid-generation disconnect (which cancels this generator at its current yield) persists nothing. + on_usage(run_result) -> observability hook (#118): called once with the + final `AgentRunResult` after the stream completes, BEFORE on_complete — + tokens were spent even if persistence subsequently fails. Routes pass + `agents.usage.record_agent_usage` here. Not called on the error rungs + (no result event was seen) nor on the legacy fallback, whose usage is + captured inside `call_gemini_multiturn` (feature=). A hook failure is + swallowed: usage capture must never break the stream. + legacy_fallback() -> awaitable returning the route's pre-agent result, used ONLY when the agent fails before emitting any text (Rung 1). It is async because the routes' legacy paths are (`_legacy_chat`). It owns its @@ -134,6 +143,7 @@ async def stream_agent_turn( chunks: list[str] = [] final_output: str | None = None + run_result: Any = None # High-water marks: how much of deps.* we have already emitted. graph_hw = 0 mastery_hw = 0 @@ -173,7 +183,8 @@ async def stream_agent_turn( continue if cls_name == "AgentRunResultEvent": - output = getattr(getattr(event, "result", None), "output", None) + run_result = getattr(event, "result", None) + output = getattr(run_result, "output", None) if isinstance(output, str): final_output = output @@ -229,6 +240,15 @@ async def stream_agent_turn( merged = merge_graph_updates(deps.graph_updates) mastery = list(deps.mastery_changes) + # Usage first, persistence second: the tokens were spent regardless of + # whether on_complete manages to persist. Guarded — instrumentation must + # never turn a fully-streamed reply into an error event. + if on_usage is not None and run_result is not None: + try: + on_usage(run_result) + except Exception: + logger.debug("on_usage hook failed; usage row dropped", exc_info=True) + try: extra = on_complete(reply, merged, mastery) or {} except Exception as exc: diff --git a/backend/services/course_context_service.py b/backend/services/course_context_service.py index bae36839..401aab31 100644 --- a/backend/services/course_context_service.py +++ b/backend/services/course_context_service.py @@ -20,6 +20,7 @@ from db.connection import table from agents._run import run_agent_sync from agents.course_summary import course_summary_agent +from agents.usage import record_agent_usage def _generate_data_hash(stats_rows: list) -> str: @@ -52,7 +53,10 @@ def _generate_summary_with_gemini( ) try: - result = run_agent_sync(course_summary_agent.run(user_message)) + result = record_agent_usage( + run_agent_sync(course_summary_agent.run(user_message)), + feature="course_summary", task="course_summary", + ) return result.output.summary except Exception: # Fallback summary if the agent fails diff --git a/backend/services/events_service.py b/backend/services/events_service.py new file mode 100644 index 00000000..261e4a8f --- /dev/null +++ b/backend/services/events_service.py @@ -0,0 +1,292 @@ +"""Fire-and-forget write path for observability events (issue #116). + +Two public helpers — ``log_event`` (analytics / audit / error) and +``log_llm_usage`` (per-call LLM token + cost) — enqueue a row onto a bounded +in-process queue and return immediately. A single daemon worker thread drains +the queue in small batches and inserts through the sanctioned +``db/connection.py::table()`` seam. + +Design guarantees: + +* **Never adds latency.** The calling thread only builds a dict and does a + non-blocking ``put``. No DB call happens on the request thread. +* **Never raises into request handling.** Every failure — a full queue, a + serialization slip, a PostgREST error inside the worker — is caught and + logged via the stdlib logger. Logging observability data must not break the + thing being observed. +* **No raw content.** ``log_event(content=...)`` is hashed to a 16-hex + fingerprint (``content_fp``); the raw string is never enqueued or persisted. +* **Kill switch.** ``EVENTS_LOGGING_ENABLED=false`` turns both helpers into + no-ops. + +Cost computation and token-field normalization live in +``services/llm_pricing.py``; this module just persists what it's given. +""" + +from __future__ import annotations + +import logging +import os +import queue +import threading +from typing import Any, Optional + +from db.connection import table +from services import llm_pricing +from services.fingerprint import fingerprint_text +from services.request_context import current_request_id + +logger = logging.getLogger("sapling.events") + +# Tunables (env-driven). Read at queue-construction time so tests can shrink +# the queue via reset_for_tests(maxsize=...). +_DEFAULT_QUEUE_MAX = int(os.getenv("EVENTS_QUEUE_MAX", "10000")) + +# Batch flush parameters: drain up to _BATCH_MAX rows or wait _FLUSH_INTERVAL +# seconds, whichever comes first. +_BATCH_MAX = 50 +_FLUSH_INTERVAL = 1.0 + +# ── Module state ──────────────────────────────────────────────────────────── + +_queue: "queue.Queue[dict]" = queue.Queue(maxsize=_DEFAULT_QUEUE_MAX) +_dropped = 0 +_dropped_lock = threading.Lock() + +_worker: Optional[threading.Thread] = None +_worker_lock = threading.Lock() +_stop = threading.Event() + + +def _logging_enabled() -> bool: + """Read the kill switch each call so it can be toggled at runtime / in tests.""" + return os.getenv("EVENTS_LOGGING_ENABLED", "true").strip().lower() not in { + "false", "0", "no", "off", + } + + +# ── Public API ────────────────────────────────────────────────────────────── + + +def log_event( + event_type: str, + *, + category: str, + user_id: str | None = None, + request_id: str | None = None, + payload: dict | None = None, + content: str | None = None, +) -> None: + """Enqueue a row for the ``events`` table. Never raises, never blocks. + + ``content`` is fingerprinted to ``content_fp`` (16 hex chars); the raw + string is hashed here and immediately dropped — it never enters the queue. + """ + if not _logging_enabled(): + return + try: + row = { + "event_type": event_type, + "category": category, + "user_id": user_id, + "request_id": request_id if request_id is not None else current_request_id(), + "payload": payload or {}, + "content_fp": fingerprint_text(content, length=16) if content else None, + } + _enqueue("events", row) + except Exception: # pragma: no cover - defensive; enqueue is already guarded + logger.exception("log_event failed; event dropped") + + +def log_llm_usage( + *, + feature: str, + task: str | None, + model: str, + usage: Any, + provider: str = "gemini", + user_id: str | None = None, + request_id: str | None = None, +) -> None: + """Enqueue a row for the ``llm_usage`` table. Never raises, never blocks. + + ``usage`` is any Pydantic AI / Gemini usage object (or dict); it is + normalized here and the cost computed from ``llm_pricing.MODEL_PRICING`` + (``cost_usd = NULL`` for unpriced models). + """ + if not _logging_enabled(): + return + try: + tokens = llm_pricing.normalize_usage(usage) + row = { + "user_id": user_id, + "request_id": request_id if request_id is not None else current_request_id(), + "feature": feature, + "task": task, + "model": model, + "provider": provider, + "prompt_tokens": tokens["prompt_tokens"], + "completion_tokens": tokens["completion_tokens"], + "total_tokens": tokens["total_tokens"], + "cost_usd": llm_pricing.cost_usd( + model, tokens["prompt_tokens"], tokens["completion_tokens"], + ), + } + _enqueue("llm_usage", row) + except Exception: # pragma: no cover - defensive + logger.exception("log_llm_usage failed; row dropped") + + +# ── Enqueue + drop accounting ─────────────────────────────────────────────── + + +def _enqueue(table_name: str, row: dict) -> None: + global _dropped + try: + _queue.put_nowait({"table": table_name, "row": row}) + except queue.Full: + with _dropped_lock: + _dropped += 1 + n = _dropped + # Throttle: warn on the first drop and every 1000th thereafter so a + # sustained overflow doesn't flood the logs. + if n == 1 or n % 1000 == 0: + logger.warning( + "events queue full (max=%d); dropped %d event(s) so far", + _queue.maxsize, n, + ) + + +def dropped_count() -> int: + """Number of events dropped due to queue overflow (test/metrics hook).""" + return _dropped + + +# ── Flush ─────────────────────────────────────────────────────────────────── + + +def _flush_batch(items: list[dict]) -> None: + """Insert a batch grouped by table. Errors are swallowed + logged. + + A failed bulk insert falls back to inserting that table's rows one at a + time, so a single poison row can't take its whole batch down with it. + """ + if not items: + return + grouped: dict[str, list[dict]] = {} + for item in items: + grouped.setdefault(item["table"], []).append(item["row"]) + for table_name, rows in grouped.items(): + try: + table(table_name).insert(rows) + except Exception: + logger.info( + "events bulk insert failed for table %r (%d row(s)); " + "retrying rows individually", + table_name, len(rows), exc_info=True, + ) + _flush_rows_individually(table_name, rows) + + +def _flush_rows_individually(table_name: str, rows: list[dict]) -> None: + """Per-row salvage after a failed bulk insert: only the rows that + individually fail are dropped (logged per-row at debug, plus one warning + with the drop count). Never raises — same contract as _flush_batch.""" + dropped = 0 + for row in rows: + try: + table(table_name).insert([row]) + except Exception: + dropped += 1 + logger.debug( + "events row insert failed for table %r; row dropped", + table_name, exc_info=True, + ) + if dropped: + logger.warning( + "events flush dropped %d of %d row(s) for table %r after per-row retry", + dropped, len(rows), table_name, + ) + + +def flush_now() -> None: + """Synchronously drain the queue on the calling thread. + + Test-only determinism hook (and used by ``shutdown()``): pulls everything + currently queued and inserts it, so tests never race the worker thread. + """ + items: list[dict] = [] + while True: + try: + items.append(_queue.get_nowait()) + except queue.Empty: + break + _flush_batch(items) + + +# ── Background worker lifecycle ───────────────────────────────────────────── + + +def _worker_loop() -> None: + """Drain-and-insert loop: batch up to _BATCH_MAX rows or _FLUSH_INTERVAL.""" + while not _stop.is_set(): + batch: list[dict] = [] + try: + batch.append(_queue.get(timeout=_FLUSH_INTERVAL)) + except queue.Empty: + continue + while len(batch) < _BATCH_MAX: + try: + batch.append(_queue.get_nowait()) + except queue.Empty: + break + _flush_batch(batch) + + +def start_worker() -> None: + """Start the daemon drain thread (idempotent). + + Called from the FastAPI lifespan in production. Unit tests deliberately do + not call this — they use ``flush_now()`` for deterministic draining. + """ + global _worker + with _worker_lock: + if _worker is not None and _worker.is_alive(): + return + _stop.clear() + _worker = threading.Thread( + target=_worker_loop, name="sapling-events-worker", daemon=True, + ) + _worker.start() + + +def shutdown() -> None: + """Stop the worker and flush anything still queued (called on app shutdown).""" + global _worker + _stop.set() + worker = _worker + if worker is not None: + worker.join(timeout=2.0) + _worker = None + flush_now() + + +# ── Test hook ──────────────────────────────────────────────────────────────── + + +def reset_for_tests(maxsize: int | None = None) -> None: + """Reset queue, counters, and one-time-warning state between tests. + + Wired into an autouse fixture in ``tests/conftest.py`` so no test leaks + queued rows or a tripped drop-counter into the next. + """ + global _queue, _dropped, _worker + _stop.set() + if _worker is not None: + _worker.join(timeout=1.0) + _worker = None + _stop.clear() + _queue = queue.Queue(maxsize=maxsize if maxsize is not None else _DEFAULT_QUEUE_MAX) + with _dropped_lock: + _dropped = 0 + llm_pricing._warned_models.clear() diff --git a/backend/services/extraction_backends/gemini_vision_backend.py b/backend/services/extraction_backends/gemini_vision_backend.py index 7a7d246d..f9438fdd 100644 --- a/backend/services/extraction_backends/gemini_vision_backend.py +++ b/backend/services/extraction_backends/gemini_vision_backend.py @@ -42,6 +42,7 @@ TRANSCRIBE_PROMPT, ocr_vision_agent, ) +from agents.usage import record_agent_usage class GeminiVisionUnavailableError(RuntimeError): @@ -112,22 +113,30 @@ def extract_page_with_gemini_vision(image_bytes: bytes) -> str: raise GeminiVisionUnavailableError("GEMINI_API_KEY is not set") try: - result = _run_from_anywhere( - ocr_vision_agent.run( - # Image first, then the instruction — the wire shape this was - # measured against. Moving the instruction to a system prompt - # makes the model emit a whole LaTeX document; see ocr_vision. - [ - BinaryContent(data=image_bytes, media_type="image/png"), - TRANSCRIBE_PROMPT, - ], - # No per-run `model=` override needed: `ocr_vision_agent`'s own - # default model (`model_for("ocr_vision")`) is loop-safe on its - # own now (#354/#436 — see agents/_providers.py), so it - # survives this per-page loop's alternating asyncio.run() - # loops without help from this call site. - usage_limits=WORKER_LIMITS, - ) + # #118: one usage row per transcribed page — this per-page loop is the + # largest per-document LLM spend in the app, so it must show up in + # llm_usage. No user_id here: extraction is content-addressed and + # user-agnostic; the request-id contextvar still attributes the row. + result = record_agent_usage( + _run_from_anywhere( + ocr_vision_agent.run( + # Image first, then the instruction — the wire shape this + # was measured against. Moving the instruction to a system + # prompt makes the model emit a whole LaTeX document; see + # ocr_vision. + [ + BinaryContent(data=image_bytes, media_type="image/png"), + TRANSCRIBE_PROMPT, + ], + # No per-run `model=` override needed: `ocr_vision_agent`'s + # own default model (`model_for("ocr_vision")`) is loop-safe + # on its own now (#354/#436 — see agents/_providers.py), so + # it survives this per-page loop's alternating asyncio.run() + # loops without help from this call site. + usage_limits=WORKER_LIMITS, + ) + ), + feature="document", task="ocr_vision", ) except UnexpectedModelBehavior: # The model produced no usable text for this page (empty candidate, a diff --git a/backend/services/flashcard_import_service.py b/backend/services/flashcard_import_service.py index 6c28f50f..5bf6a2b3 100644 --- a/backend/services/flashcard_import_service.py +++ b/backend/services/flashcard_import_service.py @@ -27,6 +27,7 @@ from db.connection import table from agents._run import run_agent_sync from agents.flashcard import flashcard_agent +from agents.usage import record_agent_usage from services import extraction_service logger = logging.getLogger(__name__) @@ -253,7 +254,10 @@ def _run_flashcard_agent(prompt: str) -> list[Card]: result = None for attempt in range(_AGENT_RETRIES + 1): try: - result = run_agent_sync(flashcard_agent.run(prompt)) + result = record_agent_usage( + run_agent_sync(flashcard_agent.run(prompt)), + feature="flashcard", task="flashcard", + ) break except UnexpectedModelBehavior: logger.exception("flashcard agent produced unusable output; degrading to []") diff --git a/backend/services/gemini_service.py b/backend/services/gemini_service.py index 685d32a4..28158dfa 100644 --- a/backend/services/gemini_service.py +++ b/backend/services/gemini_service.py @@ -23,6 +23,29 @@ MODEL_SMART = "gemini-2.5-pro" +def _log_gemini_usage(response, *, feature: str, model: str) -> None: + """Record token usage for a direct Gemini call (#118). + + Reads ``response.usage_metadata`` and hands it to the fire-and-forget + events writer. Import is local to avoid a module-load cycle + (events_service → db/connection) and any import cost when Gemini is unused; + events_service.log_llm_usage is itself failure-isolated, but we still guard + here so a missing/odd ``usage_metadata`` can never break a real call. + """ + try: + from services.events_service import log_llm_usage + + log_llm_usage( + feature=feature, + task=None, + model=model, + usage=getattr(response, "usage_metadata", None), + provider="gemini", + ) + except Exception: + pass + + def _thinking_budget_for(model: str) -> int: """Thinking-token budget for a model, shared by every call path here. @@ -75,13 +98,22 @@ def _extract_json(text: str) -> str: return text -def call_gemini(prompt: str, retries: int = 1, json_mode: bool = False, model: str = MODEL_DEFAULT) -> str: +def call_gemini( + prompt: str, + retries: int = 1, + json_mode: bool = False, + model: str = MODEL_DEFAULT, + feature: str = "misc", +) -> str: """Single-turn call to Gemini with a plain string prompt. Pro's thinking cap (`_thinking_budget_for`) also applies here — that fix originally lived only in call_gemini_multiturn (PR #74), so any caller passing model="gemini-2.5-pro" here (e.g. an LLM-judge model override) used to 400 on thinking_budget=0. + + ``feature`` tags the resulting llm_usage row (#118); it defaults to + ``"misc"`` so uninstrumented callers still attribute somewhere. """ thinking_budget = _thinking_budget_for(model) for attempt in range(retries + 1): @@ -97,6 +129,7 @@ def call_gemini(prompt: str, retries: int = 1, json_mode: bool = False, model: s contents=prompt, config=config, ) + _log_gemini_usage(response, feature=feature, model=model) if not response.text: raise ValueError("Gemini returned empty response (content may have been filtered)") return response.text @@ -108,13 +141,15 @@ def call_gemini(prompt: str, retries: int = 1, json_mode: bool = False, model: s raise -def call_gemini_multiturn(system_prompt: str, history: list[dict], user_message: str, retries: int = 1, model: str = MODEL_DEFAULT) -> str: +def call_gemini_multiturn(system_prompt: str, history: list[dict], user_message: str, retries: int = 1, model: str = MODEL_DEFAULT, feature: str = "misc") -> str: """ Multi-turn call to Gemini using native chat history. history: list of {"role": "user"|"model", "content": "..."} dicts from the DB (role "assistant" is remapped to "model"). Returns the assistant reply as a plain string. + + ``feature`` tags the resulting llm_usage row (#118). """ # Gemini expects role to be "user" or "model" (not "assistant") def _normalise_role(role: str) -> str: @@ -141,6 +176,7 @@ def _normalise_role(role: str) -> str: ) chat = _client.chats.create(model=model, config=config, history=gemini_history) response = chat.send_message(user_message) + _log_gemini_usage(response, feature=feature, model=model) if not response.text: raise ValueError("Gemini returned empty response (content may have been filtered)") return response.text @@ -152,8 +188,10 @@ def _normalise_role(role: str) -> str: raise -def call_gemini_json(prompt: str, model: str = MODEL_DEFAULT): - raw = call_gemini(prompt, json_mode=True, model=model) +def call_gemini_json(prompt: str, model: str = MODEL_DEFAULT, feature: str = "misc"): + # Delegates to call_gemini, which logs usage — so JSON-mode calls are + # captured exactly once (no double-count here). + raw = call_gemini(prompt, json_mode=True, model=model, feature=feature) try: return json.loads(raw) except json.JSONDecodeError: diff --git a/backend/services/llm_pricing.py b/backend/services/llm_pricing.py new file mode 100644 index 00000000..745ba2e9 --- /dev/null +++ b/backend/services/llm_pricing.py @@ -0,0 +1,132 @@ +"""LLM token-usage normalization and cost computation (issue #118). + +Two concerns live here, kept separate from the write path +(``services/events_service.py``) so the persistence layer stays provider- +agnostic: + +1. ``normalize_usage`` — different SDKs name their token fields differently + (Pydantic AI's ``RunUsage`` uses ``input_tokens``/``output_tokens``; older + builds used ``request_tokens``/``response_tokens``; Gemini's + ``usage_metadata`` uses ``prompt_token_count``/``candidates_token_count``). + This reduces any of them to a single ``prompt``/``completion``/``total`` + dict before the row is persisted. + +2. ``cost_usd`` — a small, editable per-1K-token price map. Known models get a + computed cost; unknown models return ``None`` (persisted as SQL NULL) and + emit a one-time warning so an un-priced model surfaces in the logs without + spamming one line per call. +""" + +from __future__ import annotations + +import logging +from decimal import ROUND_HALF_UP, Decimal +from typing import Any + +logger = logging.getLogger("sapling.llm_pricing") + + +# Per-1,000-token USD prices as ``(input_rate, output_rate)``. Sourced from +# Google Gemini API list pricing; kept deliberately small and editable. A model +# missing here is not an error — its usage is still recorded, just with +# ``cost_usd = NULL``. Update this map (not the call sites) when prices change +# or a new model ships. +MODEL_PRICING: dict[str, tuple[float, float]] = { + "gemini-2.5-pro": (0.00125, 0.010), + "gemini-2.5-flash": (0.0003, 0.0025), + "gemini-2.5-flash-lite": (0.0001, 0.0004), + "gemini-2.0-flash": (0.0001, 0.0004), + "gemini-2.0-flash-lite": (0.000075, 0.0003), +} + +# Models we've already warned about — so an un-priced model logs once, not +# once per call. Module-level (per-process); tests reset entries as needed. +_warned_models: set[str] = set() + +# Token-field aliases across the SDKs we touch, in priority order. First +# non-None hit wins. +_PROMPT_FIELDS = ("prompt_tokens", "input_tokens", "request_tokens", "prompt_token_count") +_COMPLETION_FIELDS = ( + "completion_tokens", "output_tokens", "response_tokens", "candidates_token_count", +) +_TOTAL_FIELDS = ("total_tokens", "total_token_count") + + +def _read_int(usage: Any, names: tuple[str, ...]) -> int: + """Return the first present, int-coercible field in ``names`` (attr or key). + + Missing / None / non-numeric fields are skipped; nothing matches → 0. + """ + for name in names: + value = usage.get(name) if isinstance(usage, dict) else getattr(usage, name, None) + if value is None: + continue + try: + return int(value) + except (TypeError, ValueError): + continue + return 0 + + +def normalize_usage(usage: Any) -> dict[str, int]: + """Reduce any supported usage object/dict to prompt/completion/total ints. + + Handles Pydantic AI ``RunUsage`` (both current input/output and legacy + request/response naming), Gemini ``usage_metadata``, and a plain dict that + is already normalized. ``total`` is derived from prompt + completion when + the source reports it as zero or omits it. + """ + if usage is None: + return {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + + prompt = _read_int(usage, _PROMPT_FIELDS) + completion = _read_int(usage, _COMPLETION_FIELDS) + total = _read_int(usage, _TOTAL_FIELDS) + if total <= 0: + total = prompt + completion + return { + "prompt_tokens": prompt, + "completion_tokens": completion, + "total_tokens": total, + } + + +def _canonical_model(model: str) -> str: + """Strip a provider qualifier so 'google-gla:gemini-2.5-flash' matches. + + Pydantic AI may report a provider-prefixed model name; the price map is + keyed on the bare Gemini model id. + """ + return model.rsplit(":", 1)[-1].strip() if model else model + + +def cost_usd(model: str, prompt_tokens: int, completion_tokens: int) -> float | None: + """Compute USD cost for a call, or ``None`` if the model isn't priced. + + Rounds to 6 decimal places (half-up) to fit ``llm_usage.cost_usd + numeric(12,6)``. An unknown model returns ``None`` and warns once. + """ + rates = MODEL_PRICING.get(model) or MODEL_PRICING.get(_canonical_model(model)) + if rates is None: + # SAPLING_MODEL_MODE=function runs (the e2e/CI seam) report the model + # as 'function:' — a deliberate free stand-in, not an unpriced + # real model. Record the row with cost NULL and stay silent: warning + # here would emit one line per task on every e2e run. Real unknown + # models keep the one-time warning below. + if model and model.startswith("function:"): + return None + if model not in _warned_models: + _warned_models.add(model) + logger.warning( + "No pricing entry for model %r; llm_usage.cost_usd will be NULL. " + "Add it to services/llm_pricing.MODEL_PRICING to enable cost rollups.", + model, + ) + return None + + in_rate, out_rate = rates + cost = ( + Decimal(str(in_rate)) * Decimal(int(prompt_tokens)) + + Decimal(str(out_rate)) * Decimal(int(completion_tokens)) + ) / Decimal(1000) + return float(cost.quantize(Decimal("0.000001"), rounding=ROUND_HALF_UP)) diff --git a/backend/tests/conftest.py b/backend/tests/conftest.py index 92060cc9..86f149f5 100644 --- a/backend/tests/conftest.py +++ b/backend/tests/conftest.py @@ -59,6 +59,17 @@ def _clear_lru_caches(): course_context_service.clear_course_context_cache() +@pytest.fixture(autouse=True) +def _reset_events_service(): + """#118/#116: reset the observability queue, drop-counter, and one-time + pricing-warning state around every test, so a queued row, a tripped + overflow counter, or a shrunk test queue can't leak into the next test.""" + from services import events_service + events_service.reset_for_tests() + yield + events_service.reset_for_tests() + + @pytest.fixture(autouse=True) def _hermetic_supabase_client(request, monkeypatch): """Hermetic safety net (#210): no test may make a real Supabase call. diff --git a/backend/tests/test_admin_analytics_routes.py b/backend/tests/test_admin_analytics_routes.py new file mode 100644 index 00000000..1b255267 --- /dev/null +++ b/backend/tests/test_admin_analytics_routes.py @@ -0,0 +1,284 @@ +"""Tests for routes/admin_analytics.py — the admin cost-rollup / analytics API +(issue #120). + +A small but faithful fake of the PostgREST `table()` seam interprets the +`eq.`/`gte.`/`lte.`/`like.` filters, ordering, and limit/offset that the route +builds, so date-range filtering, aggregation, and pagination are exercised for +real (not mocked away). +""" +from __future__ import annotations + +import fnmatch +from datetime import datetime + +import pytest +from fastapi.testclient import TestClient + +from main import app +import routes.admin_analytics as analytics + +client = TestClient(app) + +BASE = "/api/admin/analytics" + +# In-range = July 2026; OUT = 2020. The default-range test freezes the module +# clock at 2026-07-21 so its 30-day window covers the July rows, not the 2020 one. +IN1 = "2026-07-10T09:00:00+00:00" +IN2 = "2026-07-12T09:00:00+00:00" +IN3 = "2026-07-15T09:00:00+00:00" +OUT = "2020-01-01T00:00:00+00:00" + + +def _seed(): + events = [ + {"event_type": "document.upload", "category": "usage", "user_id": "u1", "request_id": "r1", "payload": {}, "created_at": IN1}, + {"event_type": "quiz.completed", "category": "usage", "user_id": "u1", "request_id": "r2", "payload": {}, "created_at": IN2}, + {"event_type": "quiz.completed", "category": "usage", "user_id": "u2", "request_id": "r3", "payload": {}, "created_at": IN3}, + {"event_type": "error.5xx", "category": "error", "user_id": "u2", "request_id": "r4", + "payload": {"path": "/api/quiz", "method": "POST", "status_code": 500, "duration_ms": 12.3}, "created_at": IN3}, + {"event_type": "auth.login", "category": "audit", "user_id": "u1", "request_id": "r0", "payload": {}, "created_at": OUT}, + ] + llm = [ + {"user_id": "u1", "feature": "quiz", "task": "quiz", "model": "gemini-2.5-flash", "provider": "gemini", + "prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150, "cost_usd": 0.01, "created_at": IN1}, + {"user_id": "u1", "feature": "chat_tutor", "task": "chat_tutor", "model": "gemini-2.5-pro", "provider": "gemini", + "prompt_tokens": 200, "completion_tokens": 100, "total_tokens": 300, "cost_usd": 0.05, "created_at": IN2}, + {"user_id": "u2", "feature": "quiz", "task": "quiz", "model": "gemini-2.5-flash", "provider": "gemini", + "prompt_tokens": 80, "completion_tokens": 40, "total_tokens": 120, "cost_usd": 0.02, "created_at": IN3}, + {"user_id": "u1", "feature": "quiz", "task": "quiz", "model": "gemini-2.5-flash", "provider": "gemini", + "prompt_tokens": 999, "completion_tokens": 999, "total_tokens": 1998, "cost_usd": 9.99, "created_at": OUT}, + ] + return {"events": events, "llm_usage": llm} + + +class _FakeTable: + def __init__(self, rows): + self.rows = rows + + @staticmethod + def _match(value, cond: str) -> bool: + op, _, target = cond.partition(".") + sval = "" if value is None else str(value) + if op == "eq": + return sval == target + if op == "gte": + return sval >= target + if op == "lte": + return sval <= target + if op == "like": + return fnmatch.fnmatch(sval, target) + raise AssertionError(f"unsupported op in test fake: {op}") + + def _filtered(self, filters): + rows = list(self.rows) + for col, cond in (filters or {}).items(): + conds = cond if isinstance(cond, list) else [cond] + for c in conds: + rows = [r for r in rows if self._match(r.get(col), c)] + return rows + + @staticmethod + def _ordered(rows, order): + if not order: + return rows + col, _, direction = order.partition(".") + return sorted(rows, key=lambda r: (r.get(col) is None, r.get(col)), reverse=direction == "desc") + + def select_with_count(self, columns="*", filters=None, order=None, limit=None, offset=None): + rows = self._ordered(self._filtered(filters), order) + total = len(rows) + if offset: + rows = rows[offset:] + if limit is not None: + rows = rows[:limit] + return rows, total + + def select(self, columns="*", filters=None, order=None, limit=None): + rows, _ = self.select_with_count(columns, filters, order, limit, None) + return rows + + +@pytest.fixture +def seeded(monkeypatch): + store = _seed() + monkeypatch.setattr(analytics, "table", lambda name: _FakeTable(store.get(name, []))) + return store + + +# Explicit window covering the July rows (and excluding the 2020 rows). +RANGE = {"from": "2026-07-01T00:00:00+00:00", "to": "2026-08-01T00:00:00+00:00"} + + +# ── /usage/summary ─────────────────────────────────────────────────────────── + + +def test_usage_summary_counts(seeded): + r = client.get(f"{BASE}/usage/summary", params=RANGE) + assert r.status_code == 200 + body = r.json() + assert body["total_events"] == 4 # the 2020 auth.login is excluded + assert body["distinct_active_users"] == 2 + assert body["truncated"] is False # nowhere near the scan cap + by_type = {row["event_type"]: row["count"] for row in body["by_event_type"]} + assert by_type["quiz.completed"] == 2 + assert by_type["document.upload"] == 1 + assert "auth.login" not in by_type + + +# ── /usage/by-user ─────────────────────────────────────────────────────────── + + +def test_usage_by_user_totals(seeded): + r = client.get(f"{BASE}/usage/by-user", params=RANGE) + assert r.status_code == 200 + body = r.json() + users = {u["user_id"]: u for u in body["users"]} + assert body["total_users"] == 2 + assert users["u1"]["event_count"] == 2 + assert users["u1"]["llm_cost_usd"] == pytest.approx(0.06) + assert users["u1"]["total_tokens"] == 450 + assert users["u2"]["llm_cost_usd"] == pytest.approx(0.02) + + +def test_usage_by_user_pagination(seeded): + r = client.get(f"{BASE}/usage/by-user", params={**RANGE, "limit": 1, "offset": 0}) + body = r.json() + assert body["total_users"] == 2 + assert len(body["users"]) == 1 + + +# ── /llm/cost ──────────────────────────────────────────────────────────────── + + +def test_llm_cost_group_by_feature(seeded): + r = client.get(f"{BASE}/llm/cost", params={**RANGE, "group_by": "feature"}) + assert r.status_code == 200 + body = r.json() + rows = {row["key"]: row for row in body["rows"]} + assert rows["quiz"]["cost_usd"] == pytest.approx(0.03) + assert rows["quiz"]["calls"] == 2 + assert rows["chat_tutor"]["cost_usd"] == pytest.approx(0.05) + assert body["totals"]["cost_usd"] == pytest.approx(0.08) + + +def test_llm_cost_group_by_model(seeded): + r = client.get(f"{BASE}/llm/cost", params={**RANGE, "group_by": "model"}) + rows = {row["key"]: row for row in r.json()["rows"]} + assert rows["gemini-2.5-flash"]["cost_usd"] == pytest.approx(0.03) + assert rows["gemini-2.5-pro"]["cost_usd"] == pytest.approx(0.05) + + +def test_llm_cost_group_by_user(seeded): + r = client.get(f"{BASE}/llm/cost", params={**RANGE, "group_by": "user"}) + rows = {row["key"]: row for row in r.json()["rows"]} + assert rows["u1"]["cost_usd"] == pytest.approx(0.06) + assert rows["u2"]["total_tokens"] == 120 + + +def test_llm_cost_rejects_bad_group_by(seeded): + r = client.get(f"{BASE}/llm/cost", params={**RANGE, "group_by": "banana"}) + assert r.status_code == 422 + + +# ── /errors ────────────────────────────────────────────────────────────────── + + +def test_errors_returns_error_events_with_payload_fields(seeded): + r = client.get(f"{BASE}/errors", params=RANGE) + assert r.status_code == 200 + body = r.json() + assert body["total"] == 1 + err = body["errors"][0] + assert err["event_type"] == "error.5xx" + assert err["path"] == "/api/quiz" + assert err["method"] == "POST" + assert err["status_code"] == 500 + assert err["duration_ms"] == pytest.approx(12.3) + + +# ── date filtering + defaults ──────────────────────────────────────────────── + + +def test_narrow_range_excludes_rows(seeded): + # A window that only covers IN1 (2026-07-10). + r = client.get(f"{BASE}/usage/summary", params={ + "from": "2026-07-09T00:00:00+00:00", "to": "2026-07-11T00:00:00+00:00", + }) + assert r.json()["total_events"] == 1 + + +def test_default_range_is_last_30_days(seeded, monkeypatch): + # Freeze the module clock at 2026-07-21 so the default 30-day window covers + # the July fixture rows but not the 2020 one — whatever today's date is. + class _FrozenDatetime(datetime): + @classmethod + def now(cls, tz=None): + return cls(2026, 7, 21, 12, 0, 0, tzinfo=tz) + + monkeypatch.setattr(analytics, "datetime", _FrozenDatetime) + r = client.get(f"{BASE}/usage/summary") + assert r.status_code == 200 + assert r.json()["total_events"] == 4 + + +def test_rejects_malformed_from(seeded): + r = client.get(f"{BASE}/usage/summary", params={"from": "not-a-date", "to": RANGE["to"]}) + assert r.status_code == 422 + assert "'from'" in r.json()["detail"] + + +def test_rejects_malformed_to(seeded): + r = client.get(f"{BASE}/usage/summary", params={"from": RANGE["from"], "to": "2026-13-45"}) + assert r.status_code == 422 + assert "'to'" in r.json()["detail"] + + +def test_rejects_from_after_to(seeded): + r = client.get(f"{BASE}/usage/summary", params={"from": RANGE["to"], "to": RANGE["from"]}) + assert r.status_code == 422 + + +# ── scan-cap truncation + response headers ─────────────────────────────────── + + +def test_scan_cap_truncation_is_surfaced(seeded, monkeypatch): + # Shrink the paging + cap so the 4 in-range event rows overflow the scan. + monkeypatch.setattr(analytics, "_PAGE", 1) + monkeypatch.setattr(analytics, "_SCAN_CAP", 2) + r = client.get(f"{BASE}/usage/summary", params=RANGE) + assert r.status_code == 200 + body = r.json() + assert body["truncated"] is True + assert body["total_events"] == 2 # capped, and the response says so + + +def test_responses_are_cache_control_private(seeded): + for path, params in [ + ("/usage/summary", RANGE), + ("/usage/by-user", RANGE), + ("/llm/cost", {**RANGE, "group_by": "feature"}), + ("/errors", RANGE), + ]: + r = client.get(f"{BASE}{path}", params=params) + assert r.status_code == 200 + assert r.headers.get("Cache-Control") == "private", path + + +# ── admin gating ───────────────────────────────────────────────────────────── + + +def test_endpoints_reject_non_admin(seeded, monkeypatch): + from fastapi import HTTPException + + def _deny(request): + raise HTTPException(status_code=403, detail="Admin access required") + + monkeypatch.setattr(analytics, "require_admin", _deny) + for path, params in [ + ("/usage/summary", RANGE), + ("/usage/by-user", RANGE), + ("/llm/cost", {**RANGE, "group_by": "feature"}), + ("/errors", RANGE), + ]: + resp = client.get(f"{BASE}{path}", params=params) + assert resp.status_code == 403, f"{path} should be admin-only" diff --git a/backend/tests/test_agent_usage.py b/backend/tests/test_agent_usage.py new file mode 100644 index 00000000..b8c2ba56 --- /dev/null +++ b/backend/tests/test_agent_usage.py @@ -0,0 +1,99 @@ +"""Unit tests for agents/usage.py::record_agent_usage (issue #118). + +The helper is the single, one-line-per-call-site seam that reads a Pydantic AI +run result's usage + model and forwards it to events_service.log_llm_usage. It +must never raise (instrumentation can't break the agent run) and must return +the result so it can be used inline. +""" +from __future__ import annotations + +import pytest + +from agents.usage import record_agent_usage +from agents._providers import model_for +from services import events_service + + +class _FakeUsage: + input_tokens = 320 + output_tokens = 80 + total_tokens = 400 + + +class _FakeResponse: + model_name = "gemini-2.5-pro" + + +class _FakeResult: + output = "hello" + + def usage(self): + return _FakeUsage() + + @property + def response(self): + return _FakeResponse() + + +@pytest.fixture +def sink(monkeypatch): + rows: list = [] + + class _FakeTable: + def __init__(self, name): + self.name = name + + def insert(self, r): + rows.append((self.name, r)) + return r + + monkeypatch.setattr(events_service, "table", lambda name: _FakeTable(name)) + return rows + + +def test_records_usage_from_result(sink): + result = record_agent_usage(_FakeResult(), feature="chat_tutor", task="chat_tutor") + events_service.flush_now() + + assert result.output == "hello", "must return the original result for inline use" + name, rows = sink[0] + assert name == "llm_usage" + row = rows[0] + assert row["feature"] == "chat_tutor" + assert row["task"] == "chat_tutor" + assert row["model"] == "gemini-2.5-pro" + assert row["prompt_tokens"] == 320 + assert row["completion_tokens"] == 80 + assert row["total_tokens"] == 400 + + +def test_falls_back_to_task_model_when_result_has_no_model(sink): + class _NoModelResult: + output = "x" + + def usage(self): + return _FakeUsage() + + @property + def response(self): + raise AttributeError("no response") + + def all_messages(self): + return [] + + record_agent_usage(_NoModelResult(), feature="quiz", task="quiz") + events_service.flush_now() + row = sink[0][1][0] + assert row["model"] == model_for("quiz").model_name + + +def test_never_raises_on_broken_result(sink): + class _Broken: + def usage(self): + raise RuntimeError("usage exploded") + + # Must swallow: instrumentation cannot break the agent run. + out = record_agent_usage(_Broken(), feature="notes", task="note_chat") + events_service.flush_now() + assert out is not None + assert sink == [] # nothing logged, but no exception either diff --git a/backend/tests/test_chat_stream.py b/backend/tests/test_chat_stream.py index 11c4079a..862c8a23 100644 --- a/backend/tests/test_chat_stream.py +++ b/backend/tests/test_chat_stream.py @@ -110,12 +110,12 @@ def make_deps(): ) -async def collect(agent, deps, on_complete, legacy_fallback=None): +async def collect(agent, deps, on_complete, legacy_fallback=None, on_usage=None): events = [] async for ev in stream_agent_turn( agent=agent, user_message="hi", run_kwargs={}, deps=deps, on_complete=on_complete, legacy_fallback=legacy_fallback, - request_id="r1", + on_usage=on_usage, request_id="r1", ): events.append(ev) return events @@ -340,3 +340,79 @@ async def fake_legacy(): assert legacy_calls == [], "reply already streamed — never re-run legacy" asyncio.run(run()) + + +# ── on_usage hook (#118) ────────────────────────────────────────────────── + +def test_on_usage_called_once_with_run_result_before_done(): + """The success path hands the final AgentRunResult to on_usage exactly + once — the seam the routes use to record streamed-tutor token usage.""" + async def run(): + agent = FakeAgent([ + PartStartEvent("Hi"), + AgentRunResultEvent("Hi"), + ]) + usage_calls = [] + events = await collect( + agent, make_deps(), lambda r, g, m: {}, + on_usage=lambda res: usage_calls.append(res), + ) + assert len(usage_calls) == 1, "on_usage must fire exactly once" + assert usage_calls[0].output == "Hi", "hook receives the run result itself" + assert events[-1].type == "done" + + asyncio.run(run()) + + +def test_on_usage_failure_never_breaks_the_stream(): + """Instrumentation must not turn a fully-streamed reply into an error: + a raising on_usage is swallowed and the turn still persists + dones.""" + async def run(): + agent = FakeAgent([ + PartStartEvent("Hi"), + AgentRunResultEvent("Hi"), + ]) + persisted = [] + + def bad_usage(res): + raise RuntimeError("usage capture blew up") + + events = await collect( + agent, make_deps(), lambda r, g, m: persisted.append(r) or {}, + on_usage=bad_usage, + ) + assert persisted == ["Hi"], "persistence still runs after a usage slip" + assert events[-1].type == "done" + + asyncio.run(run()) + + +def test_on_usage_not_called_on_error_rungs_or_legacy_fallback(): + """No result event was seen on Rung 1/2, and the legacy fallback's usage + is captured inside call_gemini_multiturn — the hook must stay silent.""" + async def run(): + usage_calls = [] + + async def fake_legacy(): + return {"reply": "legacy"} + + # Rung 1: failure before any token → legacy fallback. + events = await collect( + FakeAgent([AgentRunResultEvent("x")], raise_after=0), + make_deps(), lambda r, g, m: {}, + legacy_fallback=fake_legacy, + on_usage=lambda res: usage_calls.append(res), + ) + assert events[-1].type == "done" and events[-1].data["reply"] == "legacy" + assert usage_calls == [], "legacy fallback must not trigger on_usage" + + # Rung 2: failure after tokens streamed → terminal error. + events = await collect( + FakeAgent([PartStartEvent("Hi"), AgentRunResultEvent("x")], raise_after=1), + make_deps(), lambda r, g, m: {}, + on_usage=lambda res: usage_calls.append(res), + ) + assert events[-1].type == "error" + assert usage_calls == [], "an aborted run has no result to record" + + asyncio.run(run()) diff --git a/backend/tests/test_events_service.py b/backend/tests/test_events_service.py new file mode 100644 index 00000000..bd2eb861 --- /dev/null +++ b/backend/tests/test_events_service.py @@ -0,0 +1,248 @@ +"""Unit tests for services/events_service.py — the fire-and-forget write +path for `events` + `llm_usage` (issue #116). + +The worker thread is never started in these tests; `flush_now()` drains the +queue synchronously on the calling thread, which makes assertions +deterministic (no races with a background drainer). +""" +from __future__ import annotations + +import pytest + +from services import events_service + + +# ── Test doubles for the db seam ──────────────────────────────────────────── + + +class _FakeTable: + def __init__(self, name: str, sink: list): + self.name = name + self.sink = sink + + def insert(self, rows): + # PostgREST accepts a list for a batch insert; record what was sent. + self.sink.append((self.name, rows)) + return rows if isinstance(rows, list) else [rows] + + +def _fake_table_factory(sink: list): + def factory(name: str): + return _FakeTable(name, sink) + + return factory + + +def _raising_table_factory(): + def factory(name: str): + class _T: + def insert(self, rows): + raise RuntimeError("simulated PostgREST failure") + + return _T() + + return factory + + +@pytest.fixture +def sink(monkeypatch): + rows: list = [] + monkeypatch.setattr(events_service, "table", _fake_table_factory(rows)) + return rows + + +# ── log_llm_usage ─────────────────────────────────────────────────────────── + + +def test_log_llm_usage_enqueues_normalized_row_with_cost(sink, monkeypatch): + monkeypatch.setenv("EVENTS_LOGGING_ENABLED", "true") + + class FakeUsage: + input_tokens = 1000 + output_tokens = 1000 + total_tokens = 2000 + + events_service.log_llm_usage( + feature="quiz", task="quiz", model="gemini-2.5-flash", usage=FakeUsage(), + user_id="user_andres", request_id="req-1", + ) + events_service.flush_now() + + assert len(sink) == 1 + name, rows = sink[0] + assert name == "llm_usage" + row = rows[0] + assert row["feature"] == "quiz" + assert row["task"] == "quiz" + assert row["model"] == "gemini-2.5-flash" + assert row["provider"] == "gemini" + assert row["prompt_tokens"] == 1000 + assert row["completion_tokens"] == 1000 + assert row["total_tokens"] == 2000 + assert row["user_id"] == "user_andres" + assert row["request_id"] == "req-1" + assert row["cost_usd"] == pytest.approx(0.0028) + + +def test_log_llm_usage_unknown_model_persists_null_cost(sink): + class FakeUsage: + input_tokens = 5 + output_tokens = 5 + total_tokens = 10 + + events_service.log_llm_usage( + feature="notes", task="note_chat", model="mystery-model-9000", usage=FakeUsage(), + ) + events_service.flush_now() + + row = sink[0][1][0] + assert row["cost_usd"] is None + + +def test_log_llm_usage_gemini_metadata_shape(sink): + class FakeMeta: + prompt_token_count = 30 + candidates_token_count = 12 + total_token_count = 42 + + events_service.log_llm_usage( + feature="document", task=None, model="gemini-2.5-flash-lite", + usage=FakeMeta(), provider="gemini", + ) + events_service.flush_now() + row = sink[0][1][0] + assert (row["prompt_tokens"], row["completion_tokens"], row["total_tokens"]) == (30, 12, 42) + assert row["task"] is None + + +# ── log_event + content fingerprinting ────────────────────────────────────── + + +def test_log_event_hashes_content_and_never_stores_raw(sink): + secret = "student's private essay body that must never be persisted" + events_service.log_event( + "document.upload", category="usage", user_id="u1", content=secret, + ) + events_service.flush_now() + + name, rows = sink[0] + assert name == "events" + row = rows[0] + assert row["event_type"] == "document.upload" + assert row["category"] == "usage" + # content_fp is a 16-hex fingerprint; raw content is absent everywhere. + assert row["content_fp"] is not None + assert len(row["content_fp"]) == 16 + assert "content" not in row + assert secret not in str(row) + + +def test_log_event_defaults_request_id_from_contextvar(sink): + from services import request_context + + token = request_context._REQUEST_ID_CTX.set("ctx-req-42") + try: + events_service.log_event("auth.login", category="audit") + events_service.flush_now() + finally: + request_context._REQUEST_ID_CTX.reset(token) + + assert sink[0][1][0]["request_id"] == "ctx-req-42" + + +# ── Non-blocking + failure isolation ──────────────────────────────────────── + + +def test_calling_thread_never_hits_db(monkeypatch): + """log_* must enqueue only — no DB call happens on the caller's thread, + so even a table() that raises on insert can't affect the caller.""" + monkeypatch.setattr(events_service, "table", _raising_table_factory()) + + class FakeUsage: + input_tokens = 1 + output_tokens = 1 + total_tokens = 2 + + # Neither call raises, because the (raising) insert only runs at flush time. + events_service.log_llm_usage(feature="quiz", task="quiz", model="gemini-2.5-flash", usage=FakeUsage()) + events_service.log_event("quiz.completed", category="usage") + + +def test_worker_insert_error_is_swallowed(monkeypatch, caplog): + monkeypatch.setattr(events_service, "table", _raising_table_factory()) + events_service.log_event("error.5xx", category="error") + # flush_now performs the insert; the error must be caught and logged, + # never propagated. + with caplog.at_level("WARNING"): + events_service.flush_now() # must not raise + assert any("simulated PostgREST failure" in r.getMessage() or "flush" in r.getMessage().lower() + for r in caplog.records) + + +def test_bulk_insert_failure_falls_back_to_per_row(monkeypatch, caplog): + """A poison row must only take down itself: the failed bulk insert is + retried row by row, the good rows land, and the one drop is summarized in + a single warning.""" + landed: list = [] + + class _PoisonTable: + def insert(self, rows): + if any(r.get("event_type") == "poison.row" for r in rows): + raise RuntimeError("simulated poison row") + landed.extend(rows) + return rows + + monkeypatch.setattr(events_service, "table", lambda name: _PoisonTable()) + + events_service.log_event("ok.one", category="usage") + events_service.log_event("poison.row", category="usage") + events_service.log_event("ok.two", category="usage") + + with caplog.at_level("WARNING", logger="sapling.events"): + events_service.flush_now() # must not raise + + assert [r["event_type"] for r in landed] == ["ok.one", "ok.two"] + warnings = [r for r in caplog.records if r.levelno >= 30] + assert len(warnings) == 1 + assert "dropped 1 of 3" in warnings[0].getMessage() + + +def test_queue_overflow_drops_and_increments_counter(monkeypatch, sink): + events_service.reset_for_tests(maxsize=1) + monkeypatch.setattr(events_service, "table", _fake_table_factory(sink)) + + # First enqueues; the rest overflow a size-1 queue (nothing drains yet). + for i in range(4): + events_service.log_event(f"usage.tick.{i}", category="usage") + + assert events_service.dropped_count() == 3 + events_service.flush_now() + # Only the single row that fit is inserted. + assert sum(len(rows) for _, rows in sink) == 1 + + +def test_kill_switch_makes_helpers_noops(monkeypatch, sink): + monkeypatch.setenv("EVENTS_LOGGING_ENABLED", "false") + + class FakeUsage: + input_tokens = 1 + output_tokens = 1 + total_tokens = 2 + + events_service.log_event("usage.tick", category="usage") + events_service.log_llm_usage(feature="quiz", task="quiz", model="gemini-2.5-flash", usage=FakeUsage()) + events_service.flush_now() + + assert sink == [] + + +def test_flush_now_drains_queue(sink): + events_service.log_event("a.b", category="usage") + events_service.log_event("c.d", category="audit") + events_service.flush_now() + total = sum(len(rows) for _, rows in sink) + assert total == 2 + # Queue is empty afterwards — a second flush inserts nothing more. + sink.clear() + events_service.flush_now() + assert sink == [] diff --git a/backend/tests/test_gemini_usage_logging.py b/backend/tests/test_gemini_usage_logging.py new file mode 100644 index 00000000..bfff30a3 --- /dev/null +++ b/backend/tests/test_gemini_usage_logging.py @@ -0,0 +1,104 @@ +"""Gemini direct-call instrumentation (issue #118). + +call_gemini / call_gemini_multiturn / call_gemini_json must each emit exactly +one llm_usage row from ``response.usage_metadata``, tagged with the threaded +``feature``. call_gemini_json must NOT double-count (it delegates to +call_gemini). +""" +from __future__ import annotations + +import types + +import pytest + +from services import gemini_service, events_service + + +def _usage_metadata(p, c, t): + m = types.SimpleNamespace() + m.prompt_token_count = p + m.candidates_token_count = c + m.total_token_count = t + return m + + +@pytest.fixture +def sink(monkeypatch): + rows: list = [] + + class _FakeTable: + def __init__(self, name): + self.name = name + + def insert(self, r): + rows.append((self.name, r)) + return r + + monkeypatch.setattr(events_service, "table", lambda name: _FakeTable(name)) + return rows + + +@pytest.fixture +def fake_single_turn(monkeypatch): + """Stub the genai client's single-shot generate_content. + + The real genai Client exposes ``.models``/``.chats`` as read-only + properties, so we replace the whole client object rather than a sub-attr. + """ + resp = types.SimpleNamespace(text='{"ok": true}', usage_metadata=_usage_metadata(100, 20, 120)) + fake_client = types.SimpleNamespace( + models=types.SimpleNamespace(generate_content=lambda **kw: resp), + ) + monkeypatch.setattr(gemini_service, "_client", fake_client) + return resp + + +def test_call_gemini_logs_usage(sink, fake_single_turn): + gemini_service.call_gemini("hi", feature="quiz", model="gemini-2.5-flash") + events_service.flush_now() + + assert len(sink) == 1 + name, rows = sink[0] + assert name == "llm_usage" + row = rows[0] + assert row["provider"] == "gemini" + assert row["feature"] == "quiz" + assert row["model"] == "gemini-2.5-flash" + assert (row["prompt_tokens"], row["completion_tokens"], row["total_tokens"]) == (100, 20, 120) + + +def test_call_gemini_default_feature_is_misc(sink, fake_single_turn): + gemini_service.call_gemini("hi", model="gemini-2.5-flash") + events_service.flush_now() + assert sink[0][1][0]["feature"] == "misc" + + +def test_call_gemini_json_does_not_double_count(sink, fake_single_turn): + gemini_service.call_gemini_json("hi", feature="document", model="gemini-2.5-flash") + events_service.flush_now() + # Exactly one row — call_gemini_json delegates to call_gemini, which logs. + total = sum(len(rows) for _, rows in sink) + assert total == 1 + assert sink[0][1][0]["feature"] == "document" + + +def test_call_gemini_multiturn_logs_usage(sink, monkeypatch): + resp = types.SimpleNamespace(text="hello", usage_metadata=_usage_metadata(200, 50, 250)) + + class _Chat: + def send_message(self, msg): + return resp + + fake_client = types.SimpleNamespace( + chats=types.SimpleNamespace(create=lambda **kw: _Chat()), + ) + monkeypatch.setattr(gemini_service, "_client", fake_client) + + gemini_service.call_gemini_multiturn( + "sys", [], "hey", feature="chat_tutor", model="gemini-2.5-pro", + ) + events_service.flush_now() + row = sink[0][1][0] + assert row["feature"] == "chat_tutor" + assert row["model"] == "gemini-2.5-pro" + assert (row["prompt_tokens"], row["completion_tokens"], row["total_tokens"]) == (200, 50, 250) diff --git a/backend/tests/test_llm_pricing.py b/backend/tests/test_llm_pricing.py new file mode 100644 index 00000000..eb26d53d --- /dev/null +++ b/backend/tests/test_llm_pricing.py @@ -0,0 +1,150 @@ +"""Unit tests for services/llm_pricing.py — token-field normalization and +cost computation (issue #118). + +Normalization is exercised against a *fake Pydantic AI result usage* object +and a *fake Gemini usage_metadata* object, per the success criteria: the two +SDKs name their token fields differently and both must reduce to the same +prompt/completion/total triple before persistence. +""" +from __future__ import annotations + +from dataclasses import dataclass + +import pytest + +from services import llm_pricing + + +# ── Fakes mirroring the two real SDK shapes ───────────────────────────────── + + +@dataclass +class FakePydanticUsage: + """Mirror of pydantic_ai.usage.RunUsage (v2.x): input/output/total.""" + + input_tokens: int + output_tokens: int + total_tokens: int + + +@dataclass +class FakeLegacyPydanticUsage: + """Older Pydantic AI naming the issue references: request/response.""" + + request_tokens: int + response_tokens: int + total_tokens: int + + +@dataclass +class FakeGeminiUsageMetadata: + """Mirror of google.genai response.usage_metadata.""" + + prompt_token_count: int + candidates_token_count: int + total_token_count: int + + +# ── normalize_usage ───────────────────────────────────────────────────────── + + +def test_normalize_pydantic_ai_usage(): + usage = FakePydanticUsage(input_tokens=100, output_tokens=40, total_tokens=140) + assert llm_pricing.normalize_usage(usage) == { + "prompt_tokens": 100, + "completion_tokens": 40, + "total_tokens": 140, + } + + +def test_normalize_legacy_pydantic_ai_usage(): + usage = FakeLegacyPydanticUsage(request_tokens=7, response_tokens=3, total_tokens=10) + assert llm_pricing.normalize_usage(usage) == { + "prompt_tokens": 7, + "completion_tokens": 3, + "total_tokens": 10, + } + + +def test_normalize_gemini_usage_metadata(): + usage = FakeGeminiUsageMetadata( + prompt_token_count=200, candidates_token_count=55, total_token_count=255, + ) + assert llm_pricing.normalize_usage(usage) == { + "prompt_tokens": 200, + "completion_tokens": 55, + "total_tokens": 255, + } + + +def test_normalize_derives_total_when_missing_or_zero(): + """A provider may omit/zero the total; we derive prompt + completion.""" + usage = FakeGeminiUsageMetadata( + prompt_token_count=10, candidates_token_count=5, total_token_count=0, + ) + assert llm_pricing.normalize_usage(usage)["total_tokens"] == 15 + + +def test_normalize_handles_none_and_missing_fields(): + assert llm_pricing.normalize_usage(None) == { + "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0, + } + # A gemini metadata whose candidates count came back None (filtered reply). + usage = FakeGeminiUsageMetadata( + prompt_token_count=12, candidates_token_count=None, total_token_count=12, + ) + assert llm_pricing.normalize_usage(usage) == { + "prompt_tokens": 12, "completion_tokens": 0, "total_tokens": 12, + } + + +def test_normalize_accepts_plain_dict(): + usage = {"prompt_tokens": 3, "completion_tokens": 4, "total_tokens": 7} + assert llm_pricing.normalize_usage(usage) == usage + + +# ── cost_usd ──────────────────────────────────────────────────────────────── + + +def test_cost_for_known_model(): + # gemini-2.5-flash: (0.0003 in, 0.0025 out) per 1K tokens. + # 1000 prompt -> 0.0003 ; 1000 completion -> 0.0025 ; total 0.0028. + cost = llm_pricing.cost_usd("gemini-2.5-flash", 1000, 1000) + assert cost == pytest.approx(0.0028) + + +def test_cost_for_known_model_rounds_to_six_dp(): + cost = llm_pricing.cost_usd("gemini-2.5-flash-lite", 1, 1) + # (0.0001 + 0.0004)/1000 = 0.0000005 -> quantized to 6dp = 0.000001 (half-up) + assert cost == pytest.approx(0.000001) + + +def test_cost_for_unknown_model_returns_none_and_warns_once(caplog): + llm_pricing._warned_models.discard("totally-made-up-model") + with caplog.at_level("WARNING"): + assert llm_pricing.cost_usd("totally-made-up-model", 100, 100) is None + assert llm_pricing.cost_usd("totally-made-up-model", 100, 100) is None + warnings = [r for r in caplog.records if "totally-made-up-model" in r.getMessage()] + assert len(warnings) == 1, "unknown model must warn exactly once" + + +def test_cost_strips_provider_prefix(): + """Pydantic AI may report a provider-qualified name like 'google-gla:...'.""" + plain = llm_pricing.cost_usd("gemini-2.5-flash", 1000, 0) + prefixed = llm_pricing.cost_usd("google-gla:gemini-2.5-flash", 1000, 0) + assert prefixed == plain + + +def test_cost_for_function_mode_model_is_none_and_silent(caplog): + """SAPLING_MODEL_MODE=function runs report 'function:' models. + + They are deliberate free stand-ins (the e2e/CI seam), not unpriced real + models: cost is NULL and NO warning fires — a warning here would spam + every e2e run, once per task. + """ + llm_pricing._warned_models.clear() + with caplog.at_level("WARNING"): + assert llm_pricing.cost_usd("function:chat_tutor", 100, 100) is None + assert llm_pricing.cost_usd("function:quiz", 100, 100) is None + assert not caplog.records, "function:* models must not warn" + assert "function:chat_tutor" not in llm_pricing._warned_models diff --git a/backend/tests/test_usage_instrumentation_coverage.py b/backend/tests/test_usage_instrumentation_coverage.py new file mode 100644 index 00000000..34564daa --- /dev/null +++ b/backend/tests/test_usage_instrumentation_coverage.py @@ -0,0 +1,156 @@ +"""Guard: no LLM seam is left uninstrumented (issue #118 success criterion). + +Two invariants, enforced by static analysis so a *future* call site that skips +usage capture fails CI rather than silently dropping billing data: + +1. Every ``*_agent.run(...)`` / ``.run_sync(...)`` call site in a production + module must be usage-recorded **per call site**: a function enclosing the + call must invoke ``record_agent_usage``. One documented escape hatch: a + helper that directly ``return``s the run result (a "pass-through runner", + e.g. ``notes._run_note_worker``) defers recording to its callers — every + module-local call to such a helper is then itself checked under the same + rule. Cross-module runner indirection is NOT tracked: keep the run and its + ``record_agent_usage`` wrap in the same module, or this guard goes blind. +2. Every direct-Gemini helper in ``gemini_service`` (``call_gemini`` and + ``call_gemini_multiturn``) must call ``_log_gemini_usage``; ``call_gemini_json`` + delegates to ``call_gemini`` and is exempt (it would double-count otherwise). +""" +from __future__ import annotations + +import ast +from pathlib import Path + +BACKEND = Path(__file__).resolve().parents[1] + +# Production trees that may run agents. Tests and one-off scripts are excluded: +# scripts are dev tooling, not the served app. +_SCAN_ROOTS = ("routes", "services", "agents") + +_FUNC_DEFS = (ast.FunctionDef, ast.AsyncFunctionDef) + + +def _is_agent_run_call(node: ast.AST, runner_names: set[str] = frozenset()) -> bool: + """True for ``agent.run(...)`` / ``.run_sync(...)`` calls, and + for calls to a known module-local pass-through runner helper.""" + if not isinstance(node, ast.Call): + return False + func = node.func + if isinstance(func, ast.Attribute) and func.attr in {"run", "run_sync"}: + # Receiver is a bare name like `quiz_agent`, `agent`, `health_probe_agent`. + recv = func.value + if isinstance(recv, ast.Name): + return recv.id == "agent" or recv.id.endswith("_agent") + return False + if isinstance(func, ast.Name): + return func.id in runner_names + return False + + +def _calls_record_usage(fn: ast.AST) -> bool: + """True if the function body contains a ``record_agent_usage(...)`` call.""" + for n in ast.walk(fn): + if isinstance(n, ast.Call): + f = n.func + if isinstance(f, ast.Name) and f.id == "record_agent_usage": + return True + if isinstance(f, ast.Attribute) and f.attr == "record_agent_usage": + return True + return False + + +def _returned_call(stmt: ast.AST) -> ast.Call | None: + """The Call a ``return``/``return await`` statement passes through, if any.""" + if not isinstance(stmt, ast.Return): + return None + value = stmt.value + if isinstance(value, ast.Await): + value = value.value + return value if isinstance(value, ast.Call) else None + + +def _module_offenders(py: Path) -> list[str]: + """Uninstrumented agent-run call sites in one module, as ``path:line``.""" + tree = ast.parse(py.read_text(encoding="utf-8"), filename=str(py)) + functions = [n for n in ast.walk(tree) if isinstance(n, _FUNC_DEFS)] + + # Pass 1 (fixpoint): find pass-through runners — functions that hand an + # agent-run result straight back via ``return`` — so their callers can be + # held to the recording rule instead. + runners: set[str] = set() + changed = True + while changed: + changed = False + for fn in functions: + if fn.name in runners: + continue + for stmt in ast.walk(fn): + call = _returned_call(stmt) + if call is not None and _is_agent_run_call(call, runners): + runners.add(fn.name) + changed = True + break + + # Pass 2: every run site (including calls to pass-through runners) must sit + # inside an enclosing function that calls record_agent_usage — unless the + # site is itself a runner's returned expression (its callers are checked). + offenders: list[str] = [] + + def walk(node: ast.AST, stack: list) -> None: + if isinstance(node, _FUNC_DEFS): + stack = stack + [node] + if _is_agent_run_call(node, runners): + innermost = stack[-1] if stack else None + deferred_to_callers = ( + innermost is not None + and innermost.name in runners + and any(_returned_call(s) is node for s in ast.walk(innermost)) + ) + if not deferred_to_callers and not any(_calls_record_usage(f) for f in stack): + offenders.append(f"{py.relative_to(BACKEND)}:{node.lineno}") + for child in ast.iter_child_nodes(node): + walk(child, stack) + + walk(tree, []) + return offenders + + +def _production_py_files(): + files = [BACKEND / "main.py"] + for root in _SCAN_ROOTS: + for py in (BACKEND / root).rglob("*.py"): + if py.name.startswith("test_") or "tests" in py.parts: + continue + files.append(py) + return files + + +def test_every_agent_run_site_records_usage(): + offenders: list[str] = [] + for py in _production_py_files(): + offenders.extend(_module_offenders(py)) + assert not offenders, ( + "These agent-run call sites have no enclosing record_agent_usage call " + f"(uninstrumented LLM spend): {offenders}" + ) + + +def test_gemini_helpers_log_usage(): + src = (BACKEND / "services" / "gemini_service.py").read_text(encoding="utf-8") + tree = ast.parse(src) + + logged: dict[str, bool] = {} + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef) and node.name in { + "call_gemini", "call_gemini_multiturn", + }: + logged[node.name] = any( + isinstance(c, ast.Call) + and isinstance(c.func, ast.Name) + and c.func.id == "_log_gemini_usage" + for c in ast.walk(node) + ) + + assert logged.get("call_gemini"), "call_gemini must call _log_gemini_usage" + assert logged.get("call_gemini_multiturn"), ( + "call_gemini_multiturn must call _log_gemini_usage" + )