"""Token counts per RAG stage on the FastAPI docs, then cost under listed prices.
Run from experiments/: python rag-cost-per-query-calculator/cost.py
Tokenizer: tiktoken o200k_base. Retrieval: BM25 over token-window chunks (50-token overlap).
Queries: 200 h2/h3 headings (evenly spaced) turned into 'How do I: <heading>?'."""
import re, glob, math, collections, statistics as st, tiktoken
enc = tiktoken.get_encoding('o200k_base')
T = re.compile(r"[a-z0-9_]+")
words = lambda s: T.findall(s.lower())
docs = {p: open(p, encoding='utf-8').read() for p in sorted(glob.glob('corpus/*.md'))}
heads = []
for p, t in docs.items():
    for m in re.finditer(r"^#{2,3} (.+?)\s*(\{.*\})?$", t, re.M): heads.append(m.group(1).strip())
step = len(heads) // 200
queries = ["How do I: %s?" % h for h in heads[::step][:200]]
SYSTEM = ("You are a support assistant for the FastAPI documentation. Answer using only the provided context. "
          "Cite sources as [n]. If the context does not contain the answer, say you don't know.")
def chunks(size):
    out = []
    for p, t in docs.items():
        ids = enc.encode(t)
        for i in range(0, len(ids), size - 50):
            out.append(enc.decode(ids[i:i + size]))
            if i + size >= len(ids): break
    return out
class BM25:
    def __init__(s, cs):
        s.tf = [collections.Counter(words(c)) for c in cs]; s.len = [sum(t.values()) for t in s.tf]
        s.avg = st.mean(s.len); s.df = collections.Counter(w for t in s.tf for w in t); s.n = len(cs)
    def top(s, q, k):
        sc = []
        for i, t in enumerate(s.tf):
            v = 0
            for w in set(words(q)):
                if w in t:
                    idf = math.log(1 + (s.n - s.df[w] + .5) / (s.df[w] + .5))
                    v += idf * t[w] * 2.5 / (t[w] + 1.5 * (.25 + .75 * s.len[i] / s.avg))
            sc.append(v)
        return sorted(range(s.n), key=lambda i: -sc[i])[:k]
ntok = lambda s: len(enc.encode(s))
PRICES = {  # USD per 1M tokens (input, output), OpenAI pricing page, 2026-11-11
    'gpt-5-nano': (0.05, 0.40), 'gpt-5-mini': (0.25, 2.00), 'gpt-5.4': (2.50, 15.00)}
EMB = 0.02
OUT = 250  # ASSUMED answer length in tokens (not measured)
print("queries:", len(queries), " mean query tokens: %.1f" % st.mean(ntok(q) for q in queries))
print("system prompt tokens:", ntok(SYSTEM))
rows = []
for size in (256, 512, 1024):
    cs = chunks(size); bm = BM25(cs)
    for k in (3, 5, 10):
        ctx = []; qt = []
        for q in queries:
            sel = bm.top(q, k)
            c = "\n\n".join("[%d] %s" % (j + 1, cs[i]) for j, i in enumerate(sel))
            ctx.append(ntok(c)); qt.append(ntok(q))
        prompt = [ntok(SYSTEM) + a + b + 6 for a, b in zip(ctx, qt)]  # +6 ~ template overhead (assumed)
        rows.append((size, k, st.mean(ctx), st.mean(prompt)))
        print("chunk=%4d k=%2d chunks=%d context=%7.1f prompt_in=%7.1f" % (size, k, len(cs), st.mean(ctx), st.mean(prompt)))
print("\nCost per 1,000 queries (USD), output=%d tokens assumed; query embedding (10 tokens, approx) included" % OUT)
print("chunk  k  prompt_in " + " ".join("%12s" % m for m in PRICES) + "  ctx_share_of_gpt-5-mini")
for size, k, c, p in rows:
    cost = [(p * pi + OUT * po) / 1e6 * 1000 + 10 * EMB / 1e6 * 1000 for pi, po in PRICES.values()]
    mini = PRICES['gpt-5-mini']; ctxcost = c * mini[0] / 1e6 * 1000
    print("%5d %2d %9.0f " % (size, k, p) + " ".join("%12.4f" % x for x in cost) + "  %.1f%%" % (100 * ctxcost / cost[1]))
# one-off indexing cost
cs = chunks(512); tot = sum(ntok(c) for c in cs)
print("\nIndexing: 512-token chunks=%d total tokens=%d embed cost text-embedding-3-small=$%.4f" % (len(cs), tot, tot * EMB / 1e6))