fix: 503 retry + token tracking + cost calc per response

acra.py

@@ -1,4 +1,4 @@
-import os, httpx
+import os, httpx, time, re
 from google import genai
 from google.genai import types
 from db import supabase
@@ -10,6 +10,44 @@ EMBED_MODEL = "gemini-embedding-001"
 GEN_MODEL   = "gemini-3.1-flash-lite-preview"
 DEPTH       = {0: 3, 1: 3, 2: 6, 3: 10}
 
+# ── Pricing ────────────────────────────────────────────────────
+PRICE_IN  = 0.075 / 1_000_000
+PRICE_OUT = 0.30  / 1_000_000
+GPT4O_IN  = 2.50  / 1_000_000
+GPT4O_OUT = 10.00 / 1_000_000
+
+_total_input_tokens  = 0
+_total_output_tokens = 0
+
+def _generate(contents, retries=3):
+    global _total_input_tokens, _total_output_tokens
+    for attempt in range(retries):
+        try:
+            r = client.models.generate_content(model=GEN_MODEL, contents=contents)
+            if hasattr(r, "usage_metadata") and r.usage_metadata:
+                _total_input_tokens  += r.usage_metadata.prompt_token_count or 0
+                _total_output_tokens += r.usage_metadata.candidates_token_count or 0
+            return r
+        except Exception as e:
+            if ("503" in str(e) or "UNAVAILABLE" in str(e)) and attempt < retries - 1:
+                wait = 2 ** attempt
+                print(f"Gemini 503 — retrying in {wait}s (attempt {attempt+1})")
+                time.sleep(wait)
+            else:
+                raise
+
+def calc_cost(in_tok, out_tok):
+    acra_cost  = in_tok * PRICE_IN  + out_tok * PRICE_OUT
+    gpt4o_cost = in_tok * GPT4O_IN  + out_tok * GPT4O_OUT
+    savings_pct = round((1 - acra_cost / gpt4o_cost) * 100, 1) if gpt4o_cost else 0
+    return {
+        "input_tokens":   in_tok,
+        "output_tokens":  out_tok,
+        "acra_cost_usd":  round(acra_cost,  6),
+        "gpt4o_cost_usd": round(gpt4o_cost, 6),
+        "savings_pct":    savings_pct,
+    }
+
 def embed_texts(texts):
     r = client.models.embed_content(model=EMBED_MODEL, contents=texts,
         config=types.EmbedContentConfig(task_type="RETRIEVAL_DOCUMENT", output_dimensionality=768))
@@ -33,10 +71,9 @@ def adaptive_chunk(text, max_tok=512):
     return chunks or [text]
 
 def web_search(query: str, max_results: int = 5) -> List[dict]:
-    """Web search via Jina AI — requires JINA_API_KEY secret in Space."""
     try:
         import urllib.parse
-        encoded = urllib.parse.quote(query)
+        encoded  = urllib.parse.quote(query)
         jina_key = os.environ.get("JINA_API_KEY", "")
         r = httpx.get(
             f"https://s.jina.ai/?q={encoded}",
@@ -46,7 +83,7 @@ def web_search(query: str, max_results: int = 5) -> List[dict]:
                 "X-Retain-Images": "none",
                 "X-Engine": "direct",
             },
-            timeout=30.0,
+            timeout=20.0,
             follow_redirects=True
         )
         if r.status_code != 200:
@@ -69,20 +106,18 @@ def web_search(query: str, max_results: int = 5) -> List[dict]:
         return []
 
 def decompose(query):
-    r = client.models.generate_content(model=GEN_MODEL,
-        contents=f"Decompose into 2-4 simpler sub-queries. Numbered list only.\n\nQuery: {query}")
+    r = _generate(f"Decompose into 2-4 simpler sub-queries. Numbered list only.\n\nQuery: {query}")
     lines = [l.strip().lstrip("1234567890.). ") for l in r.text.strip().split("\n") if l.strip()]
     return lines[:4] or [query]
 
 def compress(query, chunks):
     if not chunks: return []
     numbered = "\n\n".join(f"[{i+1}]\n{c}" for i, c in enumerate(chunks))
-    r = client.models.generate_content(model=GEN_MODEL, contents=(
+    r = _generate(
         f"You have {len(chunks)} text chunks and a query.\n"
         f"For each chunk extract ONLY sentences relevant to the query.\n"
         f"Reply as [1] <text or EMPTY> [2] <text or EMPTY> etc.\n\n"
-        f"Query: {query}\n\nChunks:\n{numbered}"))
-    import re
+        f"Query: {query}\n\nChunks:\n{numbered}")
     out = []
     for m in re.finditer(r"\[(\d+)\]\s*(.*?)(?=\[\d+\]|$)", r.text, re.DOTALL):
         t = m.group(2).strip()
@@ -123,20 +158,22 @@ async def ingest_pipeline(texts, metadata, namespace, user_id):
     return len(chunks)
 
 async def query_pipeline(query, namespace, top_k, rerank, user_id, use_web=False):
+    global _total_input_tokens, _total_output_tokens
+    _total_input_tokens = _total_output_tokens = 0
+
     cls   = classify_query(query)
     level = cls["level"]
     k     = DEPTH[level]
 
-    # use_web=True: skip ALL local retrieval, pure Jina web search
     if use_web:
         hits = web_search(query, max_results=6)
         if not hits:
             return {"answer": "No web results found.", "sources": [],
-                    "complexity": cls, "retrieval_source": "none"}
+                    "complexity": cls, "retrieval_source": "none",
+                    "cost": calc_cost(0, 0)}
         ctx = "\n\n---\n\n".join(
             f"Source: {h['title']}\nURL: {h['url']}\n{h['snippet']}" for h in hits)
-        r = client.models.generate_content(model=GEN_MODEL,
-            contents=WEB_PROMPT.format(ctx=ctx, q=query))
+        r = _generate(WEB_PROMPT.format(ctx=ctx, q=query))
         return {
             "answer": r.text.strip(),
             "sources": [{"content": h["snippet"][:200],
@@ -144,17 +181,17 @@ async def query_pipeline(query, namespace, top_k, rerank, user_id, use_web=False
                          "score": 1.0, "source": "web"} for h in hits],
             "complexity": cls,
             "retrieval_source": "web",
+            "cost": calc_cost(_total_input_tokens, _total_output_tokens),
         }
 
-    # L0: check docs first, fall back to model knowledge
     if level == 0:
         doc_hits = vsearch(query, namespace, user_id, 2)
         if doc_hits:
             ctx = "\n\n---\n\n".join(h["content"] for h in doc_hits)
-            r   = client.models.generate_content(model=GEN_MODEL, contents=(
+            r   = _generate(
                 f"Use the context if it contains a relevant answer. "
                 f"Otherwise answer from your own knowledge.\n\n"
-                f"Context:\n{ctx}\n\nQuestion: {query}\nAnswer:"))
+                f"Context:\n{ctx}\n\nQuestion: {query}\nAnswer:")
             top_score = doc_hits[0].get("similarity", 0)
             return {
                 "answer": r.text.strip(),
@@ -165,13 +202,13 @@ async def query_pipeline(query, namespace, top_k, rerank, user_id, use_web=False
                             for h in doc_hits if h.get("similarity", 0) > 0.5],
                 "complexity": cls,
                 "retrieval_source": "local" if top_score > 0.5 else "model_knowledge",
+                "cost": calc_cost(_total_input_tokens, _total_output_tokens),
             }
-        r = client.models.generate_content(model=GEN_MODEL,
-            contents=f"Answer from your knowledge:\n\n{query}")
+        r = _generate(f"Answer from your knowledge:\n\n{query}")
         return {"answer": r.text.strip(), "sources": [],
-                "complexity": cls, "retrieval_source": "model_knowledge"}
+                "complexity": cls, "retrieval_source": "model_knowledge",
+                "cost": calc_cost(_total_input_tokens, _total_output_tokens)}
 
-    # L1-L3: local vector retrieval
     hits = []
     if level == 3:
         seen = set()
@@ -185,28 +222,29 @@ async def query_pipeline(query, namespace, top_k, rerank, user_id, use_web=False
         web_hits = web_search(query, max_results=k)
         if not web_hits:
             return {"answer": "Nothing found locally or on the web.",
-                    "sources": [], "complexity": cls, "retrieval_source": "none"}
+                    "sources": [], "complexity": cls, "retrieval_source": "none",
+                    "cost": calc_cost(_total_input_tokens, _total_output_tokens)}
         ctx = "\n\n---\n\n".join(
             f"Source: {h['title']}\n{h['snippet']}" for h in web_hits)
-        r = client.models.generate_content(model=GEN_MODEL,
-            contents=WEB_PROMPT.format(ctx=ctx, q=query))
+        r = _generate(WEB_PROMPT.format(ctx=ctx, q=query))
         return {"answer": r.text.strip(),
                 "sources": [{"content": h["snippet"][:200],
                              "metadata": {"title": h["title"], "url": h["url"]},
                              "score": 1.0, "source": "web"} for h in web_hits],
-                "complexity": cls, "retrieval_source": "web"}
+                "complexity": cls, "retrieval_source": "web",
+                "cost": calc_cost(_total_input_tokens, _total_output_tokens)}
 
     lc = [h["content"] for h in hits]
     if rerank and level >= 2: lc = [c for c in compress(query, lc) if c.strip()]
     ctx = "\n\n---\n\n".join(lc[:k])
-    r   = client.models.generate_content(model=GEN_MODEL,
-        contents=PROMPTS[level].format(ctx=ctx, q=query))
+    r   = _generate(PROMPTS[level].format(ctx=ctx, q=query))
     return {
         "answer": r.text.strip(),
         "sources": [{"content": h["content"][:200], "metadata": h.get("metadata", {}),
                      "score": h.get("similarity", 0), "source": "local"}
                     for h in hits[:len(lc)]],
         "complexity": cls, "retrieval_source": "local",
+        "cost": calc_cost(_total_input_tokens, _total_output_tokens),
     }
 
 async def run_acra_pipeline(mode, **kw):