DavidL72Code/UMB_Sustainable_Chatbot

1"""Per-request telemetry: stage latency, pipeline path and token cost.2 3Split out of Chatbot.py unchanged. Nothing here touches RetrievalChatbot or4ChatbotConfig, which is what makes it safe to import from anywhere — the5dependency runs one way, from the pipeline into this module.6"""7from __future__ import annotations8 9import json10import os11import sys12import time13from contextvars import ContextVar14from typing import Optional15 16 17# ---------------------------------------------------------------------------18# Per-request telemetry: stage latency, pipeline path, token cost.19# ---------------------------------------------------------------------------20_ACTIVE_TELEMETRY: ContextVar[Optional[dict]] = ContextVar("active_telemetry", default=None)21 22# USD per 1M tokens, paid tier, from the published Gemini API pricing23# (ai.google.dev/gemini-api/docs/pricing, checked 2026-09-05). Override per24# deployment with LLM_PRICE_TABLE_JSON, e.g.25# {"gemini-3.1-flash-lite": {"input": 0.25, "output": 1.5}}26#27# These were previously a flat 0.10/0.40 for every lite model, which28# understated output on gemini-3.5-flash-lite by 6.25x and made the dashboard's29# cost figure meaningless. Thinking tokens bill at the output rate, which30# estimate_call_cost already does.31_DEFAULT_LLM_PRICES: dict[str, dict[str, float]] = {32    "gemini-3.5-flash-lite": {"input": 0.30, "output": 2.50, "cached": 0.03},33    "gemini-3.5-flash": {"input": 1.50, "output": 9.00, "cached": 0.15},34    # 3.1-flash promotional rates run through 2026-12-31, then double.35    "gemini-3.1-flash-lite": {"input": 0.25, "output": 1.50, "cached": 0.025},36    "gemini-3.1-flash": {"input": 0.75, "output": 3.75, "cached": 0.075},37    # 3.1-pro: the higher tier applies to prompts over 200k tokens.38    "gemini-3.1-pro": {"input": 2.00, "output": 12.00, "cached": 0.20},39    "gemma": {"input": 0.0, "output": 0.0},40}41 42 43def llm_price_table() -> dict[str, dict[str, float]]:44    table = {name: dict(prices) for name, prices in _DEFAULT_LLM_PRICES.items()}45    raw = os.getenv("LLM_PRICE_TABLE_JSON", "").strip()46    if raw:47        try:48            overrides = json.loads(raw)49        except json.JSONDecodeError:50            overrides = {}51        for model, prices in (overrides or {}).items():52            if isinstance(prices, dict):53                entry = {54                    "input": float(prices.get("input", 0.0) or 0.0),55                    "output": float(prices.get("output", 0.0) or 0.0),56                }57                if prices.get("cached") is not None:58                    entry["cached"] = float(prices.get("cached") or 0.0)59                table[str(model)] = entry60    return table61 62 63def price_for_model(model: str) -> Optional[dict[str, float]]:64    table = llm_price_table()65    if model in table:66        return table[model]67    for name, prices in table.items():68        if model.startswith(name) or name in model:69            return prices70    return None71 72 73def reset_request_telemetry() -> dict:74    """Start a fresh telemetry record for the current request/answer call."""75    telemetry = {"steps": [], "started_at": time.perf_counter()}76    _ACTIVE_TELEMETRY.set(telemetry)77    return telemetry78 79 80def record_pipeline_step(step: str, kind: str, latency_ms: float, **detail) -> None:81    telemetry = _ACTIVE_TELEMETRY.get()82    if telemetry is None:83        return84    record = {"step": step, "kind": kind, "latency_ms": round(float(latency_ms), 2)}85    record.update({key: value for key, value in detail.items() if value is not None})86    telemetry["steps"].append(record)87 88 89def _usage_counts(usage: object) -> dict[str, int]:90    def count(*names: str) -> int:91        for name in names:92            value = getattr(usage, name, None)93            if value is None and isinstance(usage, dict):94                value = usage.get(name)95            if value:96                try:97                    return int(value)98                except (TypeError, ValueError):99                    continue100        return 0101 102    input_tokens = count("prompt_token_count", "input_tokens")103    output_tokens = count("candidates_token_count", "output_tokens")104    thinking_tokens = count("thoughts_token_count", "thinking_tokens")105    cached_tokens = count("cached_content_token_count", "cached_tokens")106    total_tokens = count("total_token_count", "total_tokens") or (107        input_tokens + output_tokens + thinking_tokens108    )109    return {110        "input_tokens": input_tokens,111        "output_tokens": output_tokens,112        "thinking_tokens": thinking_tokens,113        "cached_tokens": cached_tokens,114        "total_tokens": total_tokens,115    }116 117 118def estimate_call_cost(model: str, counts: dict[str, int]) -> Optional[float]:119    prices = price_for_model(model)120    if prices is None:121        return None122    billed_output = counts.get("output_tokens", 0) + counts.get("thinking_tokens", 0)123    # prompt_token_count already includes cached tokens, which bill at a lower124    # rate. Default the cached rate to a quarter of the input rate unless the125    # deployment configures one explicitly.126    cached_tokens = min(counts.get("cached_tokens", 0), counts.get("input_tokens", 0))127    fresh_input = counts.get("input_tokens", 0) - cached_tokens128    cached_rate = prices.get("cached", prices.get("input", 0.0) * 0.25)129    cost = (130        fresh_input * prices.get("input", 0.0)131        + cached_tokens * cached_rate132        + billed_output * prices.get("output", 0.0)133    ) / 1_000_000134    return round(cost, 8)135 136 137def record_llm_call(*, model: str, stage: str, usage: object, latency_ms: float, streamed: bool) -> None:138    counts = _usage_counts(usage)139    record_pipeline_step(140        stage,141        "llm",142        latency_ms,143        model=model,144        streamed=streamed,145        cost_usd=estimate_call_cost(model, counts),146        priced=price_for_model(model) is not None,147        **counts,148    )149 150 151def _caller_stage(default: str = "llm_call") -> str:152    """Label an LLM call by the pipeline function that issued it."""153    frame = sys._getframe(1)154    for _ in range(6):155        frame = frame.f_back156        if frame is None:157            return default158        name = frame.f_code.co_name159        if name.startswith("_") or name in {160            "call_gemini", "call_gemini_stream", "<lambda>", "record_llm_call",161        }:162            continue163        return name164    return default165 166 167def summarize_request_telemetry(response_mode: str = "", query_route: Optional[dict] = None) -> dict:168    """Roll the recorded steps into dashboard-ready latency/path/cost fields."""169    telemetry = _ACTIVE_TELEMETRY.get() or {"steps": []}170    steps = list(telemetry.get("steps", []))171 172    totals = {173        "input_tokens": 0,174        "output_tokens": 0,175        "thinking_tokens": 0,176        "cached_tokens": 0,177        "total_tokens": 0,178        "cost_usd": 0.0,179        "call_count": 0,180    }181    by_model: dict[str, dict] = {}182    fully_priced = True183    for step in steps:184        if step.get("kind") != "llm":185            continue186        totals["call_count"] += 1187        for key in ("input_tokens", "output_tokens", "thinking_tokens", "cached_tokens", "total_tokens"):188            totals[key] += int(step.get(key, 0) or 0)189        cost = step.get("cost_usd")190        if cost is None:191            fully_priced = False192        else:193            totals["cost_usd"] += float(cost)194        model_bucket = by_model.setdefault(195            str(step.get("model", "unknown")),196            {"calls": 0, "total_tokens": 0, "cost_usd": 0.0},197        )198        model_bucket["calls"] += 1199        model_bucket["total_tokens"] += int(step.get("total_tokens", 0) or 0)200        model_bucket["cost_usd"] = round(model_bucket["cost_usd"] + float(cost or 0.0), 8)201 202    totals["cost_usd"] = round(totals["cost_usd"], 8)203    totals["fully_priced"] = fully_priced204    totals["by_model"] = by_model205 206    latency_breakdown = {"retrieval_ms": 0.0, "llm_ms": 0.0, "other_ms": 0.0}207    for step in steps:208        bucket = {209            "retrieval": "retrieval_ms",210            "llm": "llm_ms",211        }.get(str(step.get("kind")), "other_ms")212        latency_breakdown[bucket] += float(step.get("latency_ms", 0.0) or 0.0)213    latency_breakdown = {key: round(value, 2) for key, value in latency_breakdown.items()}214 215    route = query_route or {}216    path_steps = [217        {218            "step": step.get("step", ""),219            "kind": step.get("kind", ""),220            "latency_ms": step.get("latency_ms", 0.0),221            "total_tokens": step.get("total_tokens"),222            "model": step.get("model"),223        }224        for step in steps225    ]226    label_parts = [str(step.get("step", "")) for step in steps if step.get("step")]227    if response_mode:228        label_parts.append(str(response_mode))229    return {230        "token_usage": totals,231        "llm_calls": [step for step in steps if step.get("kind") == "llm"],232        "stage_timings": steps,233        "latency_breakdown": latency_breakdown,234        "path": path_steps,235        "path_label": " -> ".join(dict.fromkeys(label_parts)) or (response_mode or "direct"),236        "route_summary": {237            "response_mode": response_mode,238            "routing_mode": route.get("routing_mode", ""),239            "question_type": route.get("question_type", ""),240            "prefer_summary": route.get("prefer_summary"),241        },242    }243