fix(copilot): baseline cost tracking and cache token display

autogpt_platform/backend/backend/copilot/baseline/service.py

@@ -14,12 +14,14 @@
 import re
 import shutil
 import tempfile
+import time
 import uuid
 from collections.abc import AsyncGenerator, Sequence
 from dataclasses import dataclass, field
 from functools import partial
 from typing import TYPE_CHECKING, Any, cast
 
+import httpx
 import orjson
 from langfuse import propagate_attributes
 from openai.types.chat import ChatCompletionMessageParam, ChatCompletionToolParam
@@ -103,6 +105,108 @@
 # MIME types that can be embedded as vision content blocks (OpenAI format).
 _VISION_MIME_TYPES = frozenset({"image/png", "image/jpeg", "image/gif", "image/webp"})
 
+# OpenRouter model pricing cache:
+# maps model id -> (input_rate, output_rate, cache_read_rate | None) per token.
+# cache_read_rate is None when OpenRouter doesn't publish it for that model;
+# those tokens fall back to the full input rate (safe over-estimate).
+# Populated lazily by _fetch_openrouter_pricing() and refreshed every hour.
+_OPENROUTER_PRICING_CACHE: dict[str, tuple[float, float, float | None]] = {}
+_OPENROUTER_PRICING_CACHE_TTL = 3600  # seconds
+_OPENROUTER_PRICING_CACHE_FETCHED_AT: float = float("-inf")
+
+
+async def _fetch_openrouter_pricing() -> dict[str, tuple[float, float, float | None]]:
+    """Return per-token pricing for all OpenRouter models.
+
+    Fetches https://openrouter.ai/api/v1/models and caches the result for
+    ``_OPENROUTER_PRICING_CACHE_TTL`` seconds (1 hour).  On any network or
+    parse error the previous cache value (possibly empty dict) is returned so
+    that cost estimation gracefully degrades to ``None`` rather than crashing.
+    """
+    global _OPENROUTER_PRICING_CACHE, _OPENROUTER_PRICING_CACHE_FETCHED_AT
+
+    now = time.monotonic()
+    if now - _OPENROUTER_PRICING_CACHE_FETCHED_AT < _OPENROUTER_PRICING_CACHE_TTL:
+        return _OPENROUTER_PRICING_CACHE
+
+    try:
+        async with httpx.AsyncClient(timeout=10.0) as client:
+            response = await client.get("https://openrouter.ai/api/v1/models")
+            response.raise_for_status()
+            data = response.json()
+
+        pricing: dict[str, tuple[float, float, float | None]] = {}
+        for model in data.get("data", []):
+            model_id = model.get("id")
+            model_pricing = model.get("pricing") or {}
+            prompt_str = model_pricing.get("prompt")
+            completion_str = model_pricing.get("completion")
+            if model_id and prompt_str is not None and completion_str is not None:
+                try:
+                    cache_read_str = model_pricing.get("cache_read")
+                    cache_read_rate = (
+                        float(cache_read_str) if cache_read_str is not None else None
+                    )
+                    pricing[model_id] = (
+                        float(prompt_str),
+                        float(completion_str),
+                        cache_read_rate,
+                    )
+                except (ValueError, TypeError):
+                    continue
+
+        _OPENROUTER_PRICING_CACHE = pricing
+        _OPENROUTER_PRICING_CACHE_FETCHED_AT = now
+    except Exception:
+        logger.warning(
+            "Failed to fetch OpenRouter model pricing; using cached data (%d models)",
+            len(_OPENROUTER_PRICING_CACHE),
+        )
+
+    return _OPENROUTER_PRICING_CACHE
+
+
+async def _estimate_cost_from_tokens(
+    model: str,
+    prompt_tokens: int,
+    completion_tokens: int,
+    cache_read_tokens: int = 0,
+) -> float | None:
+    """Estimate USD cost from token counts using live OpenRouter model pricing.
+
+    ``prompt_tokens`` should be the *total* prompt token count as reported by
+    the API (includes both regular and cache-read tokens as a subset).
+    ``cache_read_tokens`` is the subset of prompt tokens served from cache.
+    When OpenRouter publishes a ``cache_read`` rate for the model it is used
+    directly; otherwise cache-read tokens fall back to the full input rate
+    (a safe over-estimate).
+
+    Cache writes (creation) are intentionally ignored as an acceptable
+    approximation for a fallback cost estimate.
+
+    Returns None if the model is not found in the OpenRouter pricing response.
+    """
+    pricing_table = await _fetch_openrouter_pricing()
+    pricing = pricing_table.get(model)
+    if pricing is None:
+        return None
+    input_rate, output_rate, cache_read_rate = pricing
+
+    # Regular (non-cached) input tokens billed at full price;
+    # cache-read tokens billed at the OpenRouter-published rate when available,
+    # or at the full input rate when not (safe over-estimate).
+    effective_cache_rate = (
+        cache_read_rate if cache_read_rate is not None else input_rate
+    )
+    regular_prompt = max(0, prompt_tokens - cache_read_tokens)
+    cost = (
+        regular_prompt * input_rate
+        + cache_read_tokens * effective_cache_rate
+        + completion_tokens * output_rate
+    )
+    return cost
+
+
 # Max size for embedding images directly in the user message (20 MiB raw).
 _MAX_INLINE_IMAGE_BYTES = 20 * 1024 * 1024
 
@@ -247,6 +351,8 @@ class _BaselineStreamState:
     text_started: bool = False
     turn_prompt_tokens: int = 0
     turn_completion_tokens: int = 0
+    turn_cache_read_tokens: int = 0
+    turn_cache_creation_tokens: int = 0
     cost_usd: float | None = None
     thinking_stripper: _ThinkingStripper = field(default_factory=_ThinkingStripper)
     session_messages: list[ChatMessage] = field(default_factory=list)
@@ -269,6 +375,12 @@ async def _baseline_llm_caller(
 
     round_text = ""
     response = None  # initialized before try so finally block can access it
+    # Snapshot token counts before this call so we can compute the delta used
+    # for fallback cost estimation.  Must be set before the try so the finally
+    # block can always reference them even when the API call raises.
+    prompt_tokens_before = state.turn_prompt_tokens
+    completion_tokens_before = state.turn_completion_tokens
+    cache_read_tokens_before = state.turn_cache_read_tokens
     try:
         client = _get_openai_client()
         typed_messages = cast(list[ChatCompletionMessageParam], messages)
@@ -294,6 +406,18 @@ async def _baseline_llm_caller(
             if chunk.usage:
                 state.turn_prompt_tokens += chunk.usage.prompt_tokens or 0
                 state.turn_completion_tokens += chunk.usage.completion_tokens or 0
+                # Extract cache token details when available (OpenAI /
+                # OpenRouter include these in prompt_tokens_details).
+                ptd = getattr(chunk.usage, "prompt_tokens_details", None)
+                if ptd:
+                    state.turn_cache_read_tokens += (
+                        getattr(ptd, "cached_tokens", 0) or 0
+                    )
+                    # cache_creation_input_tokens is reported by some providers
+                    # (e.g. Anthropic native) but not standard OpenAI streaming.
+                    state.turn_cache_creation_tokens += (
+                        getattr(ptd, "cache_creation_input_tokens", 0) or 0
+                    )
 
             delta = chunk.choices[0].delta if chunk.choices else None
             if not delta:
@@ -348,6 +472,7 @@ async def _baseline_llm_caller(
         # Extract OpenRouter cost from response headers (in finally so we
         # capture cost even when the stream errors mid-way — we already paid).
         # Accumulate across multi-round tool-calling turns.
+        got_header_cost = False
         try:
             # Access undocumented _response attribute — same pattern as
             # extract_openrouter_cost() in blocks/llm.py.
@@ -356,9 +481,38 @@ async def _baseline_llm_caller(
                 cost = float(cost_header)
                 if math.isfinite(cost) and cost >= 0:
                     state.cost_usd = (state.cost_usd or 0.0) + cost
+                    got_header_cost = True
         except (AttributeError, ValueError):
             pass
 
+        # Fallback: estimate cost from token counts when x-total-cost is
+        # missing (e.g. some OpenRouter models don't report it).
+        # Use the delta for this call only -- the state accumulators grow across
+        # all tool-call turns, so passing the cumulative total would
+        # compound-overestimate costs on the 2nd+ turn.
+        # Separate out cached reads so we can apply provider-specific discounts
+        # (Anthropic: 10 %, OpenAI: 50 %) instead of billing them at full price.
+        call_prompt_tokens = state.turn_prompt_tokens - prompt_tokens_before
+        call_completion_tokens = state.turn_completion_tokens - completion_tokens_before
+        call_cache_read_tokens = state.turn_cache_read_tokens - cache_read_tokens_before
+        if not got_header_cost and (
+            call_prompt_tokens > 0 or call_completion_tokens > 0
+        ):
+            estimated = await _estimate_cost_from_tokens(
+                state.model,
+                call_prompt_tokens,
+                call_completion_tokens,
+                cache_read_tokens=call_cache_read_tokens,
+            )
+            if estimated is not None:
+                state.cost_usd = (state.cost_usd or 0.0) + estimated
+                logger.info(
+                    "[Baseline] x-total-cost header missing; estimated cost "
+                    "from token pricing: $%.6f (model=%s)",
+                    estimated,
+                    state.model,
+                )
+
         # Always persist partial text so the session history stays consistent,
         # even when the stream is interrupted by an exception.
         state.assistant_text += round_text
@@ -1190,16 +1344,39 @@ async def stream_chat_completion_baseline(
                 state.turn_prompt_tokens,
                 state.turn_completion_tokens,
             )
-
+            # Attempt a cost estimate from the tiktoken-derived counts so that
+            # persist_and_record_usage never receives cost_usd=None after this
+            # fallback fires.  Only fill in if no cost was already recorded.
+            if state.cost_usd is None:
+                tiktoken_estimated = await _estimate_cost_from_tokens(
+                    active_model,
+                    state.turn_prompt_tokens,
+                    state.turn_completion_tokens,
+                )
+                if tiktoken_estimated is not None:
+                    state.cost_usd = tiktoken_estimated
+                    logger.info(
+                        "[Baseline] Estimated cost from tiktoken counts: "
+                        "$%.6f (model=%s)",
+                        tiktoken_estimated,
+                        active_model,
+                    )
         # Persist token usage to session and record for rate limiting.
-        # NOTE: OpenRouter folds cached tokens into prompt_tokens, so we
-        # cannot break out cache_read/cache_creation weights. Users on the
-        # baseline path may be slightly over-counted vs the SDK path.
+        # When prompt_tokens_details.cached_tokens is reported, subtract
+        # them from prompt_tokens to get the uncached count so the cost
+        # breakdown stays accurate.
+        uncached_prompt = state.turn_prompt_tokens
+        if state.turn_cache_read_tokens > 0:
+            uncached_prompt = max(
+                0, state.turn_prompt_tokens - state.turn_cache_read_tokens
+            )
         await persist_and_record_usage(
             session=session,
             user_id=user_id,
-            prompt_tokens=state.turn_prompt_tokens,
+            prompt_tokens=uncached_prompt,
             completion_tokens=state.turn_completion_tokens,
+            cache_read_tokens=state.turn_cache_read_tokens,
+            cache_creation_tokens=state.turn_cache_creation_tokens,
             log_prefix="[Baseline]",
             cost_usd=state.cost_usd,
             model=active_model,
@@ -1269,10 +1446,13 @@ async def stream_chat_completion_baseline(
     # On GeneratorExit the client is already gone, so unreachable yields
     # are harmless; on normal completion they reach the SSE stream.
     if state.turn_prompt_tokens > 0 or state.turn_completion_tokens > 0:
+        # Report uncached prompt tokens to match what was billed — cached tokens
+        # are excluded so the frontend display is consistent with cost_usd.
+        billed_prompt = max(0, state.turn_prompt_tokens - state.turn_cache_read_tokens)
         yield StreamUsage(
-            prompt_tokens=state.turn_prompt_tokens,
+            prompt_tokens=billed_prompt,
             completion_tokens=state.turn_completion_tokens,
-            total_tokens=state.turn_prompt_tokens + state.turn_completion_tokens,
+            total_tokens=billed_prompt + state.turn_completion_tokens,
         )
 
     yield StreamFinish()