""" EpiRAG - server.py ------------------ Flask server (v2) with: - /api/query - standard JSON response (with result cache) - /api/query/stream - SSE streaming (live debate events) - /api/stats - corpus stats - /api/metrics - session performance metrics - /api/feedback - thumbs up/down user feedback - /api/rewrite - query improvement endpoint - /api/cache-stats - result cache statistics - /performance - performance dashboard page """ import os import time import json import queue import threading import chromadb from flask import Flask, jsonify, request, send_from_directory, Response, stream_with_context from flask_cors import CORS from query import rag_query, set_components, improve_query from dotenv import load_dotenv load_dotenv() app = Flask(__name__, static_folder="static") CORS(app) COLLECTION_NAME = "epirag" IS_CLOUD = os.environ.get("EPIRAG_ENV", "").lower() == "cloud" # -- Session metrics ------------------------------------------------------------ SESSION_METRICS = { "queries_total": 0, "queries_local": 0, "queries_web": 0, "queries_hybrid": 0, "queries_debate": 0, "queries_cached": 0, "latencies_ms": [], "avg_sims": [], # per-query avg similarity "local_chunk_counts": [], # per-query local chunk count "web_chunk_counts": [], # per-query web chunk count "started_at": time.time(), } # -- User feedback log ---------------------------------------------------------- FEEDBACK_LOG = [] # list of {"query": str, "rating": int, "ts": float} def record_metric(result, elapsed_ms): SESSION_METRICS["queries_total"] += 1 SESSION_METRICS["latencies_ms"].append(elapsed_ms) mode = result.get("mode", "") if mode == "local": SESSION_METRICS["queries_local"] += 1 if mode == "web": SESSION_METRICS["queries_web"] += 1 if mode == "hybrid": SESSION_METRICS["queries_hybrid"] += 1 if result.get("is_debate"): SESSION_METRICS["queries_debate"] += 1 if result.get("from_cache"): SESSION_METRICS["queries_cached"] += 1 # Extended metrics SESSION_METRICS["avg_sims"].append(result.get("avg_sim", 0.0)) SESSION_METRICS["local_chunk_counts"].append(result.get("local_chunks", 0)) SESSION_METRICS["web_chunk_counts"].append(result.get("web_chunks", 0)) # Keep rolling window of last 200 entries to avoid unbounded memory growth for key in ("latencies_ms", "avg_sims", "local_chunk_counts", "web_chunk_counts"): if len(SESSION_METRICS[key]) > 200: SESSION_METRICS[key] = SESSION_METRICS[key][-200:] # -- Corpus startup------------------------------------------------------------ _collection = None _embedder = None CORPUS_STATS = {} def init_corpus(): global _collection, _embedder, CORPUS_STATS if IS_CLOUD: print("Cloud mode - building in-memory corpus", flush=True) from ingest import build_collection_in_memory _collection, _embedder = build_collection_in_memory() else: print("Local mode - loading from ./chroma_db/", flush=True) from sentence_transformers import SentenceTransformer client = chromadb.PersistentClient(path="./chroma_db") _collection = client.get_collection(COLLECTION_NAME) _embedder = SentenceTransformer("sentence-transformers/all-mpnet-base-v2") # Build BM25 index from stored docs try: count = _collection.count() results = _collection.get(limit=count, include=["documents"]) docs = results.get("documents", []) if docs: from hybrid_retrieval import build_bm25_index build_bm25_index(docs) except Exception as e: print(f" [BM25] Local index build skipped: {e}", flush=True) set_components(_embedder, _collection) count = _collection.count() results = _collection.get(limit=count, include=["metadatas"]) papers = sorted(set( m.get("paper_name", m.get("source", "Unknown")) for m in results["metadatas"] )) CORPUS_STATS.update({ "chunks": count, "papers": len(papers), "paperList": papers, "status": "online", "mode": "cloud (in-memory)" if IS_CLOUD else "local (persistent)" }) print(f"Corpus ready: {count} chunks / {len(papers)} papers", flush=True) # Pre-warm the cross-encoder in the background so the first query doesn't pay # the model-load cost (~3-5s). Runs in a daemon thread, does not block startup. def _warm_reranker(): try: from rerank import _get_cross_encoder _get_cross_encoder() print(" [Rerank] Pre-warmed cross-encoder ready.", flush=True) except Exception as e: print(f" [Rerank] Pre-warm skipped: {e}", flush=True) threading.Thread(target=_warm_reranker, daemon=True).start() init_corpus() # -- Routes ------------------------------------------------------------ @app.route("/") def index(): return send_from_directory("static", "index.html") @app.route("/performance") def performance(): return send_from_directory("static", "performance.html") @app.route("/api/stats") def stats(): return jsonify(CORPUS_STATS) @app.route("/api/metrics") def metrics(): lats = SESSION_METRICS["latencies_ms"] avg = int(sum(lats) / len(lats)) if lats else 0 sims = SESSION_METRICS["avg_sims"] avg_sim = round(sum(sims) / len(sims), 4) if sims else 0.0 return jsonify({ **SESSION_METRICS, "avg_latency_ms": avg, "avg_similarity": avg_sim, "uptime_seconds": int(time.time() - SESSION_METRICS["started_at"]), "latencies_ms": lats[-50:], "avg_sims": sims[-50:], "local_chunk_counts": SESSION_METRICS["local_chunk_counts"][-50:], "web_chunk_counts": SESSION_METRICS["web_chunk_counts"][-50:], }) @app.route("/api/feedback", methods=["POST"]) def feedback(): """ Record a thumbs-up (rating=1) or thumbs-down (rating=-1) from the user. Body: { "query": str, "rating": 1 | -1 } """ data = request.json or {} query = (data.get("query") or "").strip() rating = data.get("rating") if not query or rating not in (1, -1): return jsonify({"error": "Provide 'query' and 'rating' (1 or -1)"}), 400 entry = {"query": query, "rating": rating, "ts": time.time()} FEEDBACK_LOG.append(entry) # Persist to a local file for later analysis try: with open("feedback.jsonl", "a", encoding="utf-8") as f: f.write(json.dumps(entry) + "\n") except Exception: pass return jsonify({"status": "recorded"}) @app.route("/api/rewrite", methods=["POST"]) def rewrite(): """ Rewrite a user query into a more precise domain-specific form. Body: { "query": str } Returns: { "improved": str } """ data = request.json or {} question = (data.get("query") or "").strip() if not question: return jsonify({"error": "No query provided"}), 400 groq_key = os.environ.get("GROQ_API_KEY") if not groq_key: return jsonify({"improved": question}) # graceful fallback improved = improve_query(question, groq_key) return jsonify({"improved": improved, "original": question}) @app.route("/api/cache-stats") def cache_stats_route(): try: from cache import cache_stats return jsonify(cache_stats()) except Exception as e: return jsonify({"error": str(e)}), 500 @app.route("/api/query", methods=["POST"]) def query(): data = request.json or {} question = (data.get("question") or "").strip() if not question: return jsonify({"error": "No question provided"}), 400 groq_key = os.environ.get("GROQ_API_KEY") tavily_key = os.environ.get("TAVILY_API_KEY") hf_token = os.environ.get("HF_TOKEN") if not groq_key: return jsonify({"error": "GROQ_API_KEY not set on server"}), 500 start = time.time() result = rag_query( question, groq_api_key = groq_key, tavily_api_key = tavily_key, hf_token = hf_token, use_debate = bool(hf_token), auto_improve = True ) elapsed_ms = int((time.time() - start) * 1000) record_metric(result, elapsed_ms) return jsonify({ "answer": result["answer"], "sources": result["sources"], "mode": result["mode"], "avg_sim": result["avg_sim"], "latency_ms": elapsed_ms, "tokens": len(result["answer"]) // 4, "question": question, "improved_query": result.get("improved_query"), "is_debate": result.get("is_debate", False), "debate_rounds": result.get("debate_rounds", []), "consensus": result.get("consensus", False), "rounds_run": result.get("rounds_run", 0), "from_cache": result.get("from_cache", False), }) @app.route("/api/query/stream", methods=["POST"]) def query_stream(): """ SSE endpoint. Streams debate events in real time, then sends final result. Event types sent to browser: data: {"type": "status", "text": "..."} data: {"type": "round_start", "round": N} data: {"type": "agent_done", "round": N, "name": "...", "color": "...", "text": "..."} data: {"type": "synthesizing"} data: {"type": "factchecking"} data: {"type": "result", ...full result payload...} data: {"type": "error", "text": "..."} """ data = request.json or {} question = (data.get("question") or "").strip() if not question: return jsonify({"error": "No question provided"}), 400 groq_key = os.environ.get("GROQ_API_KEY") tavily_key = os.environ.get("TAVILY_API_KEY") hf_token = os.environ.get("HF_TOKEN") event_queue = queue.Queue() def callback(event): event_queue.put(event) def run_in_thread(): try: start = time.time() result = rag_query( question, groq_api_key = groq_key, tavily_api_key = tavily_key, hf_token = hf_token, use_debate = bool(hf_token), sse_callback = callback, auto_improve = True ) elapsed_ms = int((time.time() - start) * 1000) record_metric(result, elapsed_ms) event_queue.put({ "type": "result", "answer": result["answer"], "sources": result["sources"], "mode": result["mode"], "avg_sim": result["avg_sim"], "latency_ms": elapsed_ms, "tokens": len(result["answer"]) // 4, "is_debate": result.get("is_debate", False), "debate_rounds": result.get("debate_rounds", []), "consensus": result.get("consensus", False), "rounds_run": result.get("rounds_run", 0), "improved_query": result.get("improved_query"), "from_cache": result.get("from_cache", False), }) except Exception as e: event_queue.put({"type": "error", "text": str(e)}) finally: event_queue.put(None) # sentinel thread = threading.Thread(target=run_in_thread, daemon=True) thread.start() def generate(): yield "data: " + json.dumps({"type": "status", "text": "Retrieving context..."}) + "\n\n" while True: try: event = event_queue.get(timeout=60) except queue.Empty: yield "data: " + json.dumps({"type": "error", "text": "Timeout"}) + "\n\n" break if event is None: break yield "data: " + json.dumps(event) + "\n\n" return Response( stream_with_context(generate()), mimetype="text/event-stream", headers={ "Cache-Control": "no-cache", "X-Accel-Buffering": "no", } ) @app.route("/api/health") def health(): return jsonify({"status": "ok", "corpus": CORPUS_STATS.get("status", "unknown")}) if __name__ == "__main__": port = int(os.environ.get("PORT", 7860)) app.run(debug=False, host="0.0.0.0", port=port, threaded=True)