server.py

"""
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)