superadmin / scripts /benchmark_uncached_rag_llm.py
suhailshafi007's picture
feat(enterprise): fullstack performance optimizations, cloudflare image pipeline, and edge hardening
e91bc61
Raw History Blame Contribute Delete
13.2 kB
"""
Sahih Islam - Uncached Real RAG Pipeline Benchmark (Full Neural Vector Search + Reranking + LLM Synthesis)
Measures:
1. Dense BGE-M3 1024-dim Embedding Latency (OpenRouter)
2. MongoDB Atlas Vector Search ($vectorSearch) Retrieval Latency
3. Cloud Cross-Encoder Reranking Latency (Cohere / OpenRouter)
4. LLM Generation Latency (DeepSeek / OpenRouter)
5. Exact Token Consumption: Prompt Tokens, Completion Tokens, Total Tokens
6. Financial Cost Breakdown in USD ($)
"""
import os
import sys
import time
import asyncio
import statistics
from typing import List, Dict, Any
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if project_root not in sys.path:
sys.path.insert(0, project_root)
for _s in (sys.stdout, sys.stderr):
try:
_s.reconfigure(encoding="utf-8", errors="replace")
except (AttributeError, OSError):
pass
import numpy as np
from app.core.config import settings
from app.infrastructure.mongodb import hadith_collection, quran_verses_collection
from app.inference_client import get_embedding as async_get_embedding, get_rerank_scores as async_get_rerank_scores
from app.rag.generate import generate_chat_completion_with_usage
from app.rag.prompts import SYSTEM_HADITH_PROMPT
# Realistic, uncached queries covering distinct fiqh and theological topics across English, Roman Urdu, and Arabic
UNCACHED_QUERIES = [
{
"id": "Q01",
"topic": "Salah / Footwear",
"language": "English",
"query": "Is it permissible to offer prayer while wearing clean shoes according to authentic Sunnah?",
"k": 3
},
{
"id": "Q02",
"topic": "Salah / Forgetfulness",
"language": "English",
"query": "How did Prophet Muhammad peace be upon him perform Sujood as-Sahw when adding or missing a rakah?",
"k": 3
},
{
"id": "Q03",
"topic": "Wudu / Leather Socks",
"language": "English",
"query": "What are the rules and time limits for wiping over leather socks khuffain during ablution?",
"k": 3
},
{
"id": "Q04",
"topic": "Taharah / Miswak",
"language": "English",
"query": "What did the Prophet say regarding the virtue and emphasis on using the Siwak before prayer?",
"k": 3
},
{
"id": "Q05",
"topic": "Ethics / Anger Management",
"language": "Roman Urdu / Hinglish",
"query": "Ghusse par qabu pane ke bare me Rasool Allah SAW ne kya hidayat aur dua sikhayi hai?",
"k": 3
},
{
"id": "Q06",
"topic": "Social / Visiting the Sick",
"language": "English",
"query": "What immense reward and angels prayer is promised for visiting a sick Muslim brother or sister?",
"k": 3
},
{
"id": "Q07",
"topic": "Fasting / Friday Fasting",
"language": "English",
"query": "Is it prohibited or disliked to single out Friday alone for voluntary fasting without fasting Thursday or Saturday?",
"k": 3
},
{
"id": "Q08",
"topic": "Zakat / Wealth Protection",
"language": "Roman Urdu / Hinglish",
"query": "Zakat ada karne se maal me barkat aur hifazat ke bare me Sahih Hadith me kya aya hai?",
"k": 3
},
{
"id": "Q09",
"topic": "Solar Eclipse / Kusuf",
"language": "Arabic",
"query": "ูƒูŠู ุตู„ู‰ ุฑุณูˆู„ ุงู„ู„ู‡ ุตู„ู‰ ุงู„ู„ู‡ ุนู„ูŠู‡ ูˆุณู„ู… ุตู„ุงุฉ ุงู„ูƒุณูˆู ูˆุงู„ุฎุณูˆู ูˆูƒู… ุฑูƒูˆุนุงู‹ ุฑูƒุน ููŠ ูƒู„ ุฑูƒุนุฉุŸ",
"k": 3
},
{
"id": "Q10",
"topic": "Repentance / Istighfar",
"language": "English",
"query": "What is Sayyid al-Istighfar and what is the virtue of reciting it in the morning and evening?",
"k": 3
}
]
# Pricing model (USD per 1,000,000 tokens)
# DeepSeek-V3: $0.14 input / $0.28 output
# OpenRouter Llama-3.1-8B: $0.055 input / $0.055 output
INPUT_TOKEN_RATE = 0.14 / 1_000_000.0
OUTPUT_TOKEN_RATE = 0.28 / 1_000_000.0
async def run_single_uncached_rag(q_spec: Dict[str, Any]) -> Dict[str, Any]:
query_text = q_spec["query"]
top_k = q_spec["k"]
t_start = time.perf_counter()
# -------------------------------------------------------------
# Step 1: OpenRouter BGE-M3 Dense Embedding (1024-dim)
# -------------------------------------------------------------
t_emb_0 = time.perf_counter()
emb_vec = await async_get_embedding(query_text)
t_emb_ms = (time.perf_counter() - t_emb_0) * 1000
query_vector = emb_vec.tolist()
# -------------------------------------------------------------
# Step 2: MongoDB Atlas Vector Search ($vectorSearch)
# -------------------------------------------------------------
t_ret_0 = time.perf_counter()
num_candidates = 20
pipeline = [
{
"$vectorSearch": {
"index": "vector_index",
"path": "embedding_1024",
"queryVector": query_vector,
"numCandidates": num_candidates * 10,
"limit": num_candidates
}
},
{
"$addFields": {"score": {"$meta": "vectorSearchScore"}}
},
{
"$project": {
"embedding_1024": 0
}
}
]
raw_candidates = await hadith_collection.aggregate(pipeline).to_list(length=num_candidates)
t_ret_ms = (time.perf_counter() - t_ret_0) * 1000
# -------------------------------------------------------------
# Step 3: Cloud Cross-Encoder Reranking
# -------------------------------------------------------------
t_rerank_0 = time.perf_counter()
if raw_candidates:
pairs = []
for doc in raw_candidates:
doc_text = doc.get("hadith_english") or doc.get("text_english") or doc.get("hadith_urdu") or ""
pairs.append([query_text, doc_text[:500]])
try:
scores = await async_get_rerank_scores(pairs)
for idx, doc in enumerate(raw_candidates):
doc["rerank_score"] = float(scores[idx]) if idx < len(scores) else float(doc.get("score", 0.0))
ranked_candidates = sorted(raw_candidates, key=lambda x: x.get("rerank_score", 0.0), reverse=True)
except Exception:
ranked_candidates = raw_candidates
else:
ranked_candidates = []
t_rerank_ms = (time.perf_counter() - t_rerank_0) * 1000
top_docs = ranked_candidates[:top_k]
# -------------------------------------------------------------
# Step 4: Context Assembly
# -------------------------------------------------------------
context_blocks = []
for i, doc in enumerate(top_docs, 1):
col = doc.get("collection", "Hadith")
h_no = doc.get("hadith_no") or doc.get("hadith_id")
narrator = doc.get("narrator", "Companion")
grade = doc.get("grade", "Sahih")
txt = doc.get("hadith_english") or doc.get("text_english") or doc.get("hadith_urdu") or ""
context_blocks.append(
f"--- SOURCE {i} ---\n"
f"Book: {col} | Hadith #{h_no} | Grade: {grade}\n"
f"Narrated by: {narrator}\n"
f"Text: {txt}\n"
)
combined_context = "\n".join(context_blocks)
# -------------------------------------------------------------
# Step 5: LLM Synthesis with Token Tracking
# -------------------------------------------------------------
user_prompt = (
f"User Question: {query_text}\n\n"
f"Grounded Hadith Context:\n{combined_context}\n\n"
f"Please provide an authentic, concise Islamic answer citing the exact Book and Hadith number provided above."
)
messages = [
{"role": "system", "content": SYSTEM_HADITH_PROMPT},
{"role": "user", "content": user_prompt}
]
t_llm_0 = time.perf_counter()
llm_resp = await generate_chat_completion_with_usage(messages, max_tokens=600)
t_llm_ms = (time.perf_counter() - t_llm_0) * 1000
t_total_ms = (time.perf_counter() - t_start) * 1000
usage = llm_resp.get("usage", {})
p_tokens = usage.get("prompt_tokens", 0)
c_tokens = usage.get("completion_tokens", 0)
tot_tokens = usage.get("total_tokens", p_tokens + c_tokens)
cost_usd = (p_tokens * INPUT_TOKEN_RATE) + (c_tokens * OUTPUT_TOKEN_RATE)
return {
"id": q_spec["id"],
"topic": q_spec["topic"],
"language": q_spec["language"],
"query": query_text,
"k": top_k,
"docs_retrieved": len(raw_candidates),
"docs_ranked": len(top_docs),
"t_emb_ms": t_emb_ms,
"t_ret_ms": t_ret_ms,
"t_rerank_ms": t_rerank_ms,
"t_llm_ms": t_llm_ms,
"t_total_ms": t_total_ms,
"model": llm_resp.get("model", "unknown"),
"prompt_tokens": p_tokens,
"completion_tokens": c_tokens,
"total_tokens": tot_tokens,
"cost_usd": cost_usd,
"answer": llm_resp.get("content", ""),
"citations": [f"{d.get('collection')} #{d.get('hadith_no') or d.get('hadith_id')}" for d in top_docs],
"error": llm_resp.get("error")
}
async def main():
print("=" * 115)
print("SAHIH ISLAM - UNCACHED RAG END-TO-END BENCHMARK (REAL LLM CALLS & TOKEN COSTS)")
print(f"Total Uncached Queries : {len(UNCACHED_QUERIES)}")
print("Pipeline Stages Tested : Dense Embedding -> Atlas Vector Search -> Cross-Encoder Rerank -> LLM Generation")
print(f"Pricing Model (DeepSeek-V3) : Input: ${INPUT_TOKEN_RATE*1e6:.2f}/M tokens | Output: ${OUTPUT_TOKEN_RATE*1e6:.2f}/M tokens")
print("=" * 115)
results = []
for q_spec in UNCACHED_QUERIES:
print(f"\nโšก Processing [{q_spec['id']}] {q_spec['topic']} ({q_spec['language']})...")
print(f" Query: \"{q_spec['query']}\"")
res = await run_single_uncached_rag(q_spec)
results.append(res)
print(f" โœ… Done: Total={res['t_total_ms']:.1f}ms | Emb={res['t_emb_ms']:.1f}ms | Atlas={res['t_ret_ms']:.1f}ms | Rerank={res['t_rerank_ms']:.1f}ms | LLM={res['t_llm_ms']:.1f}ms")
print(f" ๐Ÿ“Š Tokens: {res['prompt_tokens']} prompt + {res['completion_tokens']} comp = {res['total_tokens']} total | Cost: ${res['cost_usd']:.6f}")
print(f" ๐Ÿ“– Citations Found: {', '.join(res['citations'])}")
# Print first 140 chars of answer
clean_ans = res['answer'].replace('\n', ' ')[:140]
print(f" ๐Ÿ’ฌ Synthesis Preview: \"{clean_ans}...\"")
# Global Aggregations
total_time_ms = [r["t_total_ms"] for r in results]
emb_times = [r["t_emb_ms"] for r in results]
ret_times = [r["t_ret_ms"] for r in results]
rerank_times = [r["t_rerank_ms"] for r in results]
llm_times = [r["t_llm_ms"] for r in results]
total_prompt_tokens = sum(r["prompt_tokens"] for r in results)
total_comp_tokens = sum(r["completion_tokens"] for r in results)
total_tokens = sum(r["total_tokens"] for r in results)
total_cost = sum(r["cost_usd"] for r in results)
print("\n" + "=" * 115)
print("UNCACHED RAG BENCHMARK PERFORMANCE & TOKEN COST SUMMARY")
print("=" * 115)
print(f"Total Uncached Queries Tested : {len(results)}")
print(f"Total Tokens Consumed : {total_tokens:,} tokens ({total_prompt_tokens:,} prompt + {total_comp_tokens:,} completion)")
print(f"Average Tokens per Query : {total_tokens / len(results):.1f} tokens/query")
print(f"Total Financial Cost (USD) : ${total_cost:.5f} USD")
print(f"Average Cost per Query (USD) : ${total_cost / len(results):.5f} USD (approx {total_cost/len(results)*100:.3f} cents)")
print("-" * 115)
print("PIPELINE LATENCY BREAKDOWN (Averages across 10 Uncached Queries):")
print(f" 1. Dense BGE-M3 Embedding : {statistics.mean(emb_times):7.2f} ms (Min: {min(emb_times):.1f}ms, Max: {max(emb_times):.1f}ms)")
print(f" 2. Atlas Vector Search : {statistics.mean(ret_times):7.2f} ms (Min: {min(ret_times):.1f}ms, Max: {max(ret_times):.1f}ms)")
print(f" 3. Cross-Encoder Reranker : {statistics.mean(rerank_times):7.2f} ms (Min: {min(rerank_times):.1f}ms, Max: {max(rerank_times):.1f}ms)")
print(f" 4. LLM Generation (DeepSeek) : {statistics.mean(llm_times):7.2f} ms (Min: {min(llm_times):.1f}ms, Max: {max(llm_times):.1f}ms)")
print(f" ๐Ÿš€ Total End-to-End Latency : {statistics.mean(total_time_ms):7.2f} ms (Min: {min(total_time_ms):.1f}ms, Max: {max(total_time_ms):.1f}ms)")
print("=" * 115)
print("\nDETAILED PER-QUERY PERFORMANCE MATRIX:")
print(f"{'ID':<4} | {'TOPIC':<22} | {'EMBED':<8} | {'ATLAS':<8} | {'RERANK':<8} | {'LLM GEN':<9} | {'TOTAL':<9} | {'TOKENS':<10} | {'COST ($)'}")
print("-" * 115)
for r in results:
tok_str = f"{r['total_tokens']} ({r['completion_tokens']}c)"
print(f"{r['id']:<4} | {r['topic']:<22} | {r['t_emb_ms']:6.1f}ms | {r['t_ret_ms']:6.1f}ms | {r['t_rerank_ms']:6.1f}ms | {r['t_llm_ms']:7.1f}ms | {r['t_total_ms']:7.1f}ms | {tok_str:<10} | ${r['cost_usd']:.5f}")
print("=" * 115)
if __name__ == "__main__":
asyncio.run(main())