app.py

# ── Cell 2: Imports ────────────────────────────────────────────────────────────
import os, re, json, time, random, shutil, unicodedata, numpy as np, pandas as pd
from getpass import getpass
from pymilvus import MilvusClient
from groq import Groq
from openai import OpenAI
from sentence_transformers import SentenceTransformer, CrossEncoder
from rank_bm25 import BM25Okapi
from sklearn.metrics import roc_auc_score
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from huggingface_hub import hf_hub_download, list_repo_files, HfFileSystem, login
from datasets import load_dataset
import gradio as gr


# ── Cell 3: API Keys ───────────────────────────────────────────────────────────
# Choose your provider: "groq" or "openrouter"
LLM_PROVIDER = "openrouter"   # ← change to "groq" if preferred

import os
from getpass import getpass
from huggingface_hub import login

def get_secret_or_prompt(secret_name, prompt_text=None):
    """
    Try to read secret from Google Colab Secrets.
    If not available, ask user securely using getpass().
    """

    value = None

    # Try Colab Secrets first
    try:
        #from google.colab import userdata
        value = os.environ.get(secret_name)
    except Exception:
        value = None

    # Fallback to environment variable
    if not value:
        value = os.environ.get(secret_name)

    # Fallback to manual secure input
    if not value:
        prompt_text = prompt_text or f"Enter {secret_name}: "
        value = getpass(prompt_text)

    return value


# ── HuggingFace Token ─────────────────────────────────────────────────────────

HF_TOKEN = get_secret_or_prompt(
    "HF_TOKEN",
    "Enter HuggingFace Token: "
)

login(token=HF_TOKEN)
os.environ["HF_TOKEN"] = HF_TOKEN

print("✅ HuggingFace token loaded and login completed")


# ── LLM Provider API Key ──────────────────────────────────────────────────────

if LLM_PROVIDER == "groq":

    GROQ_API_KEY = get_secret_or_prompt(
        "GROQ_API_KEY",
        "Enter GROQ API Key: "
    )

    OPENROUTER_API_KEY = None

    os.environ["GROQ_API_KEY"] = GROQ_API_KEY

    print("✅ GROQ API key loaded")

elif LLM_PROVIDER == "openrouter":

    OPENROUTER_API_KEY = get_secret_or_prompt(
        "OPENROUTER_API_KEY",
        "Enter OpenRouter API Key: "
    )

    GROQ_API_KEY = None

    os.environ["OPENROUTER_API_KEY"] = OPENROUTER_API_KEY

    print("✅ OpenRouter API key loaded")

else:
    raise ValueError(f"Unknown LLM_PROVIDER: {LLM_PROVIDER}")


# ── Cell 4: Global configuration ──────────────────────────────────────────────
BUCKET_ID      = "Phani555/IIITH-Cohort26-RAG-Batch37-storage"
BUCKET_PREFIX  = f"hf://buckets/{BUCKET_ID}/milvus_dbs"
MILVUS_DIR     = "/content/milvus_store/milvus_dbs"
HF_REPO_ID     = "Phani555/IIITH-Cohort26-RAG-Batch37-storage"
HF_REPO_TYPE   = "dataset"
HF_FOLDER      = "ablations"

# ── Download mode ─────────────────────────────────────────────────────────────
# "chunk_v5_domain_aware" : new advanced domain-aware chunk_v5 indexes (recommended)
# "llm_embedder_default"  : legacy llm_embedder default indexes
DOWNLOAD_MODE = "chunk_v5_domain_aware"

if DOWNLOAD_MODE == "chunk_v5_domain_aware":
    INDEX_VERSION = "chunk_v5_domain_aware"
    DOMAIN_EMBEDDING_RECOMMENDATION = {
        "Customer_Support": "qwen3_embedding_0_6b",
        "Bio_Medical":      "bge_m3",
        "General_Knowledge":"qwen3_embedding_0_6b",
        "Legal_Contracts":  "bge_m3",
        "Finance":          "bge_m3",
    }
    EMBEDDING_TYPE = "bge_m3"
else:  # llm_embedder_default
    INDEX_VERSION  = "default"
    DOMAIN_EMBEDDING_RECOMMENDATION = None
    EMBEDDING_TYPE = "llm_embedder"

# ── Embedding models ───────────────────────────────────────────────────────────
EMBED_MODELS = {
    "bge_small":             "BAAI/bge-small-en-v1.5",
    "llm_embedder":          "BAAI/llm-embedder",
    "bge_m3":                "BAAI/bge-m3",
    "qwen3_embedding_0_6b":  "Qwen/Qwen3-Embedding-0.6B",
}
EMBEDDING_CHOICES = list(EMBED_MODELS.keys())

# ── Model lists per LLM provider ──────────────────────────────────────────────
GROQ_LLM_CHOICES = [
    "llama-3.1-8b-instant",
    "gemma2-9b-it",
    "llama-3.3-70b-versatile",
    "mixtral-8x7b-32768",
    "qwen/qwen3-32b",
    "qwen-qwq-32b",
    "deepseek-r1-distill-llama-70b",
]
OPENROUTER_LLM_CHOICES = [
    "meta-llama/llama-3.1-8b-instruct",
    "meta-llama/llama-3.3-70b-instruct",
    "openai/gpt-oss-20b",
    "openai/gpt-oss-120b",
    "qwen/qwen3-32b",
    "deepseek/deepseek-r1",
    "moonshotai/kimi-k2-instruct",
    "openai/gpt-oss-safeguard-20b",
]
LLM_CHOICES = OPENROUTER_LLM_CHOICES if LLM_PROVIDER == "openrouter" else GROQ_LLM_CHOICES

# ── Runtime globals ────────────────────────────────────────────────────────────
MODEL_NAME     = LLM_CHOICES[0]
MODEL_NAME_BIG = LLM_CHOICES[4] if len(LLM_CHOICES) > 4 else LLM_CHOICES[-1]

# ── Feature flags ──────────────────────────────────────────────────────────────
ENABLE_HYBRID        = True
ENABLE_HYDE          = False
ENABLE_RERANKING     = False
RERANKER_TYPE        = "monot5"    # monot5 | tilde
ENABLE_RRF           = True        # Reciprocal Rank Fusion inside hybrid search
RRF_K                = 60          # standard RRF constant
PROMPT_STRATEGY      = "short"     # short | long | long_cot
ENABLE_REPACKING     = False
REPACK_STRATEGY      = "sides"     # forward | reverse | sides
ENABLE_SUMMARIZATION = False
SUMMARIZATION_TYPE   = "recomp"    # recomp | longllmlingua
ENABLE_QUERY_REWRITING     = False
ENABLE_QUERY_DECOMPOSITION = False
ENABLE_QUERY_CLASSIFICATION= False
MAX_SUBQUERIES             = 3
QUERY_REWRITE_MODEL        = None
QUERY_DECOMPOSE_MODEL      = None
RETRIEVE_DEBUG             = False

# ── Tunable knobs ─────────────────────────────────────────────────────────────
RETRIEVE_TOP_K         = 10
RERANK_TOP_K           = 3
HYBRID_ALPHA           = 0.5
MONOT5_MODEL           = "castorini/monot5-base-msmarco-10k"
TILDE_MODEL            = "BAAI/bge-reranker-base"
RECOMP_TOP_K_SENTS     = 6
RECOMP_MIN_SCORE       = 0.00
RECOMP_GROUNDING_BOOST = 0.15
RECOMP_MIN_KEEP_RATIO  = 0.30
RECOMP_KEEP_CRITICAL   = True
LLMLINGUA_RATE         = 0.5

# ── Runtime state ─────────────────────────────────────────────────────────────
milvus_clients              = {}
bm25_indexes                = {}
loaded_embedding_models     = {}   # keyed by embedding_type string
embed_model                 = None # single fallback embed model
llm_client                  = None
monot5_reranker             = None
tilde_reranker              = None
llmlingua_compressor        = None
ragbench_by_domain          = {}
LEGAL_SAMPLE_TO_CONTRACT_ID = {}

DOMAIN_NAMES = [
    "Bio_Medical",
    "General_Knowledge",
    "Customer_Support",
    "Finance",
    "Legal_Contracts",
]

GROUNDING_PATTERNS = [
    r"\b(?:must|should|shall|cannot|can't|never|always|only|except|unless|required|recommended)\b",
    r"\b(?:warning|caution|note|important|attention)\b",
    r"\b(?:do not|don't|does not|did not|not allowed|not recommended|never)\b",
    r"\b\d+(?:\.\d+)?\s*(?:%|percent|seconds?|minutes?|hours?|days?|weeks?|months?|years?)\b",
    r"\b\d+(?:\.\d+)?\s*(?:GB|MB|KB|TB|kg|g|mg|mm|cm|m|km|degrees?|°C|°F)\b",
    r"[$€£¥]\s*\d+(?:,\d{3})*(?:\.\d+)?",
    r"\b\d+(?:,\d{3})*(?:\.\d+)?\s*(?:dollars?|rupees?|crores?|lakhs?|million|billion)\b",
    r"\b(?:19|20)\d{2}\b",
    r"\b\d{2,}\b",
    r"\b[A-Z]{2,}[-_]?\d+[A-Z0-9-]*\b",
    r"\b[A-Z0-9]{3,}[-_][A-Z0-9]{2,}\b",
    r"\b[A-Z]{3,}\b",
]
GROUNDING_REGEX = re.compile("|".join(GROUNDING_PATTERNS), re.IGNORECASE)

print(f"Config loaded.  Mode: {DOWNLOAD_MODE}  |  Provider: {LLM_PROVIDER}  |  Models: {len(LLM_CHOICES)}")


# ── Cell 5: Pipeline functions (Advanced – chunk_v5 + RRF + Legal contract filtering) ──
# NOTE: _hf_fs is initialized in Cell 7. hf_path_exists() uses globals() so it
# safely resolves _hf_fs at call-time, not at definition-time.

# ── Utilities ──────────────────────────────────────────────────────────────────

def _safe_message_content(response):
    try:
        msg = response.choices[0].message
        content = getattr(msg, "content", None)
        return str(content).strip() if content else ""
    except Exception:
        return ""

def _sanitize(text):
    if not text: return text
    text = unicodedata.normalize("NFC", str(text))
    return text.encode("ascii", errors="replace").decode("ascii")

def get_domain(dataset):
    if dataset in ("covidqa","pubmedqa"):             return "Bio_Medical"
    elif dataset in ("expertqa","hagrid","hotpotqa","msmarco"): return "General_Knowledge"
    elif dataset in ("delucionqa","emanual","techqa"): return "Customer_Support"
    elif dataset in ("finqa","tatqa"):                return "Finance"
    else:                                             return "Legal_Contracts"

def split_into_sentences(text):
    return [s.strip() for s in re.split(r'(?<=[.!?])\s+', str(text).strip()) if s.strip()]

def _tokenize(text):
    return re.findall(r'\w+', str(text).lower())

def _normalize(scores):
    arr = np.array(scores, dtype=float)
    if len(arr) == 0 or arr.max() == arr.min(): return np.zeros_like(arr)
    return (arr - arr.min()) / (arr.max() - arr.min())

def _count_grounding_signals(sentence):
    return len(GROUNDING_REGEX.findall(str(sentence)))

def _is_critical_sentence(sentence):
    pat = re.compile(
        r"\b(?:warning|caution|important|must|must not|cannot|can't|do not|don't|never|only|except|unless|required)\b",
        re.IGNORECASE)
    return bool(pat.search(str(sentence)))


# ── LLM client factory ─────────────────────────────────────────────────────────

def get_llm_client():
    if LLM_PROVIDER == "groq":
        return Groq(api_key=GROQ_API_KEY)
    elif LLM_PROVIDER == "openrouter":
        return OpenAI(api_key=OPENROUTER_API_KEY, base_url="https://openrouter.ai/api/v1")
    raise ValueError(f"Unknown LLM_PROVIDER: {LLM_PROVIDER}")


# ── DB path helpers ────────────────────────────────────────────────────────────

def get_index_folder(embedding_type=None, index_version=None):
    embedding_type = embedding_type or EMBEDDING_TYPE
    index_version  = index_version  or INDEX_VERSION
    return embedding_type if index_version == "default" else f"{embedding_type}_{index_version}"

def get_db_path(domain_name, embedding_type=None, index_version=None):
    folder = get_index_folder(embedding_type, index_version)
    db_dir = os.path.join(MILVUS_DIR, folder)
    os.makedirs(db_dir, exist_ok=True)
    return os.path.join(db_dir, f"{domain_name}.db")

def get_embedding_type_for_domain(domain_name):
    rec = globals().get("DOMAIN_EMBEDDING_RECOMMENDATION")
    if rec: return rec.get(domain_name, EMBEDDING_TYPE)
    return EMBEDDING_TYPE

def get_embed_model_for_domain(domain_name):
    emb_type = get_embedding_type_for_domain(domain_name)
    models   = globals().get("loaded_embedding_models", {})
    if emb_type not in models:
        raise ValueError(f"Embedding type '{emb_type}' not in loaded_embedding_models. Run Cell 7 first.")
    return models[emb_type]


# ── HF filesystem helper ───────────────────────────────────────────────────────
# Uses globals() so _hf_fs is resolved at call-time (Cell 7), not import-time (Cell 5).

def hf_path_exists(path):
    fs = globals().get("_hf_fs")
    if fs is None:
        raise RuntimeError("_hf_fs not initialised — run Cell 7 before Cell 8.")
    try:
        fs.ls(path); return True
    except Exception:
        return False


# ── Legal contract helpers ─────────────────────────────────────────────────────

def build_contract_filter_expr(contract_id):
    contract_id = str(contract_id).replace('"', '\\"')
    return f'contract_id == "{contract_id}"'

def load_legal_sample_to_contract_mapping(local_path=None):
    if local_path is None:
        legal_folder = get_index_folder(get_embedding_type_for_domain("Legal_Contracts"), INDEX_VERSION)
        local_path   = os.path.join(MILVUS_DIR, legal_folder, "legal_sample_to_contract_id.json")
    if not os.path.exists(local_path):
        print(f"  WARNING: Legal mapping not found: {local_path}")
        return {}
    with open(local_path, "r") as f:
        mapping = json.load(f)
    print(f"  Legal sample->contract mapping loaded: {len(mapping):,} entries")
    return mapping

def get_contract_id_for_legal_sample(sample_id):
    mapping = globals().get("LEGAL_SAMPLE_TO_CONTRACT_ID", {})
    sid = str(sample_id)
    if sid not in mapping:
        raise ValueError(f"sample_id '{sid}' not found in Legal mapping.")
    return mapping[sid]


# ── Query Classification ───────────────────────────────────────────────────────

def classify_query(query, domain_name=None):
    if not ENABLE_QUERY_CLASSIFICATION: return "RAG"
    rag_domains = {"Bio_Medical","General_Knowledge","Customer_Support","Finance","Legal_Contracts"}
    if domain_name in rag_domains: return "RAG"
    llm_keywords = ["who is","what is","when was","where is","define","explain",
                    "tell me about","what are","why is","how does","what does"]
    if any(kw in str(query).lower() for kw in llm_keywords): return "LLM"
    return "RAG"


# ── Query Rewriting ────────────────────────────────────────────────────────────

def rewrite_query(query, domain_name, llm_client, model_name=None):
    if not ENABLE_QUERY_REWRITING: return query
    model  = model_name or QUERY_REWRITE_MODEL or MODEL_NAME
    prompt = f"""Rewrite the question to improve document retrieval. Apply only when needed.
Domain: {domain_name}
Rules: Preserve meaning. Fix grammar. Expand abbreviations. Preserve all names/numbers/terms.
Do not answer. Return ONLY the rewritten query.
Original question: {query}""".strip()
    try:
        resp = llm_client.chat.completions.create(
            model=model,
            messages=[{"role":"system","content":"You rewrite questions to improve semantic document retrieval. Return only the rewritten question."},
                      {"role":"user","content":_sanitize(prompt)}],
            temperature=0.0, max_tokens=150,
        )
        rewritten = _safe_message_content(resp).strip()
        return rewritten if rewritten and len(rewritten) < 600 else query
    except Exception as e:
        print(f"Query rewriting failed: {e}"); return query


# ── Query Decomposition helpers ────────────────────────────────────────────────

def _clean_subquery_text(text):
    if text is None: return ""
    text = str(text).strip().replace("```json","").replace("```","").strip()
    text = text.rstrip(",").strip('"').strip("'").strip()
    text = re.sub(r"^\s*[-*]\s*","",text); text = re.sub(r"^\s*\d+[\).\:\-]\s*","",text)
    return text.strip()

def _looks_like_explanation_line(text):
    if not text: return True
    tl = text.lower().strip()
    bad = ["here are","here is","decomposed","search queries","the decomposed",
           "queries:","subqueries:","output:","json:","answer:"]
    if any(tl.startswith(p) for p in bad): return True
    if tl in {"queries","subqueries","search queries","decomposed search queries"}: return True
    return False

def _parse_json_object_line(line):
    line = _clean_subquery_text(line)
    if not line: return None
    try:
        obj = json.loads(line)
        if isinstance(obj, dict):
            for k in ["query","question","subquery","search_query"]:
                if k in obj and str(obj[k]).strip(): return str(obj[k]).strip()
        if isinstance(obj, str): return obj.strip()
    except Exception: pass
    m = re.search(r'"(?:query|question|subquery|search_query)"\s*:\s*"([^"]+)"', line)
    if m: return m.group(1).strip()
    return None

def _split_multi_question_locally(query, max_subqueries=None):
    max_subqueries = max_subqueries or MAX_SUBQUERIES
    parts = [p.strip() for p in re.split(r"\?\s*", str(query).strip()) if p.strip()]
    if len(parts) <= 1: return None
    return [(p+"?" if not p.endswith("?") else p) for p in parts[:max_subqueries]]

def _parse_subqueries(raw_text, original_query, max_subqueries=None):
    max_subqueries = max_subqueries or MAX_SUBQUERIES
    if not raw_text: return [original_query]
    text = str(raw_text).strip().replace("```json","").replace("```","").strip()
    try:
        parsed = json.loads(text)
        if isinstance(parsed, list):
            subs = []
            for item in parsed:
                if isinstance(item, dict):
                    for k in ["query","question","subquery","search_query"]:
                        if k in item and str(item[k]).strip(): subs.append(str(item[k]).strip()); break
                elif isinstance(item, str): subs.append(item.strip())
            subs = [_clean_subquery_text(q) for q in subs if _clean_subquery_text(q)]
            return subs[:max_subqueries] or [original_query]
        elif isinstance(parsed, dict):
            raw_list = parsed.get("subqueries") or parsed.get("queries") or parsed.get("questions") or []
            if isinstance(raw_list, list):
                subs = [_clean_subquery_text(q) for q in raw_list if _clean_subquery_text(q)]
                return subs[:max_subqueries] or [original_query]
    except Exception: pass
    subqueries = []
    for raw_line in text.splitlines():
        line = _clean_subquery_text(raw_line)
        if not line or _looks_like_explanation_line(line): continue
        obj_q = _parse_json_object_line(line)
        if obj_q:
            obj_q = _clean_subquery_text(obj_q)
            if obj_q and not _looks_like_explanation_line(obj_q): subqueries.append(obj_q)
            continue
        if line.startswith("{") or line.endswith("}") or line in {"[","]","{","}"}: continue
        subqueries.append(line)
    deduped = []
    for q in subqueries:
        q = _clean_subquery_text(q)
        if q and q not in deduped: deduped.append(q)
    return deduped[:max_subqueries] or [original_query]

def decompose_query(query, llm_client, domain=None, model=None, max_subqueries=None):
    if not ENABLE_QUERY_DECOMPOSITION: return [query]
    max_subqueries = max_subqueries or MAX_SUBQUERIES
    local_split = _split_multi_question_locally(query, max_subqueries)
    if local_split: return local_split
    model = model or QUERY_DECOMPOSE_MODEL or MODEL_NAME
    if not model: return [query]
    prompt = f"""Decompose the question into at most {max_subqueries} retrieval-focused search queries.
Return ONLY a valid JSON list of strings. No explanations. No markdown.
Example: ["What caused the 2008 crisis?", "Which banks failed in 2008?"]
Rules: If already simple return list with original. Preserve all technical terms. Do not answer.
Domain: {domain}
Question: {query}""".strip()
    try:
        resp = llm_client.chat.completions.create(
            model=model,
            messages=[{"role":"system","content":"You decompose complex questions into retrieval subqueries and return only a JSON list of strings."},
                      {"role":"user","content":_sanitize(prompt)}],
            temperature=0.0, max_tokens=300,
        )
        return _parse_subqueries(_safe_message_content(resp), original_query=query, max_subqueries=max_subqueries)
    except Exception as e:
        print(f"Query decomposition failed: {e}"); return [query]


# ── Reranking ──────────────────────────────────────────────────────────────────

class MonoT5Reranker:
    def __init__(self, model_name=None):
        model_name = model_name or MONOT5_MODEL
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model     = AutoModelForSeq2SeqLM.from_pretrained(model_name)
        self.device    = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device); self.model.eval()
        self.true_id  = self.tokenizer.convert_tokens_to_ids("▁true")
        self.false_id = self.tokenizer.convert_tokens_to_ids("▁false")
        if not self.true_id  or self.true_id  < 0: self.true_id  = self.tokenizer.encode("true",  add_special_tokens=False)[0]
        if not self.false_id or self.false_id < 0: self.false_id = self.tokenizer.encode("false", add_special_tokens=False)[0]
        print(f"MonoT5 loaded: {model_name} on {self.device}")

    def score(self, query, document):
        text = f"Query: {query} Document: {document} Relevant:"
        enc  = self.tokenizer(text, return_tensors="pt", max_length=512, truncation=True).to(self.device)
        with torch.no_grad():
            out = self.model.generate(**enc, max_new_tokens=1, return_dict_in_generate=True, output_scores=True)
        logits = out.scores[0][0]
        probs  = torch.softmax(torch.stack([logits[self.false_id], logits[self.true_id]]), dim=0)
        return float(probs[1].item())

    def compute_scores(self, query, texts):
        return np.array([self.score(query, t) for t in texts], dtype=float)

def get_monot5_reranker():
    global monot5_reranker
    if monot5_reranker is None: monot5_reranker = MonoT5Reranker(MONOT5_MODEL)
    return monot5_reranker

def get_tilde_reranker():
    global tilde_reranker
    if tilde_reranker is None:
        dev = "cuda" if torch.cuda.is_available() else "cpu"
        tilde_reranker = CrossEncoder(TILDE_MODEL, device=dev)
        print(f"TILDE reranker loaded on {dev}")
    return tilde_reranker

def rerank_documents(query, documents, top_k=3):
    """Rerank documents; preserves all metadata fields including Legal contract_id."""
    if not documents: return []
    texts  = [d["text"] if isinstance(d, dict) else d for d in documents]
    rtype  = RERANKER_TYPE.lower().strip()
    scores = get_monot5_reranker().compute_scores(query, texts) if rtype == "monot5" \
             else np.asarray(get_tilde_reranker().predict([(query, t) for t in texts], show_progress_bar=False), dtype=float).reshape(-1)
    ranked_idx = np.argsort(scores)[::-1][:top_k]
    reranked = []
    for i in ranked_idx:
        # dict() shallow-copies ALL fields (dense_score, bm25_score, contract_id, etc.)
        item = dict(documents[i]) if isinstance(documents[i], dict) else {"text": documents[i]}
        item["base_score"]    = item.get("score")   # preserve original retrieval score
        item["score"]         = float(scores[i])
        item["rerank_score"]  = float(scores[i])
        item["reranker_type"] = rtype
        reranked.append(item)
    return reranked


# ── BM25 (stores contract_ids for Legal to enable per-contract filtering) ─────

def build_bm25_index(domain_name, clients):
    client = clients[domain_name]; col = domain_name.lower()
    try:
        if "Loaded" not in str(client.get_load_state(col)): client.load_collection(col)
    except Exception: pass
    try:    n = int(client.get_collection_stats(col).get("row_count", 0))
    except Exception: return
    if n == 0: return
    output_fields = ["text"]
    if domain_name == "Legal_Contracts": output_fields.append("contract_id")
    rows = client.query(collection_name=col, filter="", limit=n, output_fields=output_fields)
    if not rows: return
    texts = []; contract_ids = []
    for r in rows:
        t = r.get("text","")
        if not t: continue
        texts.append(t)
        if domain_name == "Legal_Contracts": contract_ids.append(r.get("contract_id"))
    if not texts: return
    bm25_indexes[domain_name] = {
        "bm25":         BM25Okapi([_tokenize(t) for t in texts]),
        "texts":        texts,
        "contract_ids": contract_ids if domain_name == "Legal_Contracts" else None,
    }
    print(f"  BM25 built: {len(texts)} docs  [{domain_name}]")
    if domain_name == "Legal_Contracts":
        valid = sum(1 for c in contract_ids if c is not None)
        print(f"    Legal contract_ids: {valid:,}/{len(texts):,}")

def build_all_bm25_indexes(clients):
    bm25_indexes.clear()
    for d in clients: build_bm25_index(d, clients)


# ── HyDE ───────────────────────────────────────────────────────────────────────

def generate_hyde(query, llm_client, model_name=None):
    model  = model_name or MODEL_NAME
    prompt = f"Write a brief factual passage answering this question (under 4 sentences).\nQuestion: {query}\nPassage:"
    try:
        resp = llm_client.chat.completions.create(
            model=model,
            messages=[{"role":"system","content":"You write hypothetical answer passages for retrieval."},
                      {"role":"user","content":_sanitize(prompt)}],
            temperature=0.2, max_tokens=300,
        )
        return _safe_message_content(resp)
    except Exception as e:
        print(f"HyDE failed: {e}"); return ""


# ── Reciprocal Rank Fusion ─────────────────────────────────────────────────────

def reciprocal_rank_fusion(dense_results, bm25_results, top_k=20, rrf_k=60,
                           dense_meta=None, bm25_meta=None):
    """
    Fuse dense and BM25 ranked lists using RRF.
    score(doc) = 1/(k + rank_dense) + 1/(k + rank_bm25)
    dense_meta / bm25_meta: optional dicts of extra fields per text (e.g. Legal metadata).
    """
    dense_meta = dense_meta or {}; bm25_meta = bm25_meta or {}
    rrf_scores = {}
    for rank, (text, score) in enumerate(
            sorted(dense_results.items(), key=lambda x: x[1], reverse=True), start=1):
        rrf_scores.setdefault(text, {"text":text,"dense_score":float(score),"bm25_score":0.0,"score":0.0})
        rrf_scores[text]["dense_score"] = float(score)
        rrf_scores[text]["score"] += 1.0 / (rrf_k + rank)
        if text in dense_meta: rrf_scores[text].update(dense_meta[text])
    for rank, (text, score) in enumerate(
            sorted(bm25_results.items(), key=lambda x: x[1], reverse=True), start=1):
        rrf_scores.setdefault(text, {"text":text,"dense_score":0.0,"bm25_score":float(score),"score":0.0})
        rrf_scores[text]["bm25_score"] = float(score)
        rrf_scores[text]["score"] += 1.0 / (rrf_k + rank)
        # dense_meta takes priority over bm25_meta for Legal metadata consistency
        if text not in dense_meta and text in bm25_meta:
            rrf_scores[text].update(bm25_meta[text])
    fused = sorted(rrf_scores.values(), key=lambda x: x["score"], reverse=True)
    return fused[:top_k]


# ── Hybrid search (dense + BM25, Legal contract filtering, RRF or alpha fusion) ─

def hybrid_search(query, domain_name, embed_model, top_k=20, alpha=0.5, contract_id=None):
    client = milvus_clients[domain_name]; col = domain_name.lower()
    try:
        if "Loaded" not in str(client.get_load_state(col)): client.load_collection(col)
    except Exception: pass

    # ── Dense search ─────────────────────────────────────────────────────────
    q_emb = embed_model.encode([query], normalize_embeddings=True, convert_to_numpy=True).astype("float32")
    output_fields = ["text"]
    if domain_name == "Legal_Contracts": output_fields += ["contract_id","source_doc_id","source_hash"]
    search_kwargs = dict(collection_name=col, data=q_emb.tolist(), limit=top_k,
                         output_fields=output_fields, search_params={"metric_type":"IP","params":{}})
    if domain_name == "Legal_Contracts" and contract_id is not None:
        search_kwargs["filter"] = build_contract_filter_expr(contract_id)
    dense_hits = client.search(**search_kwargs)
    dense_results = {}; dense_meta = {}
    hit_list = dense_hits[0] if (dense_hits and isinstance(dense_hits[0], (list,tuple))) else dense_hits
    for hit in hit_list:
        entity   = (hit.get("entity",{}) or hit) if isinstance(hit,dict) else (getattr(hit,"entity",{}) or {})
        distance = hit.get("distance",0.0)        if isinstance(hit,dict) else getattr(hit,"distance",0.0)
        text = entity.get("text","")
        if not text: continue
        dense_results[text] = float(distance)
        if domain_name == "Legal_Contracts":
            dense_meta[text] = {"contract_id":entity.get("contract_id"),
                                 "source_doc_id":entity.get("source_doc_id"),
                                 "source_hash":entity.get("source_hash")}

    # ── BM25 search ───────────────────────────────────────────────────────────
    bm25_obj = bm25_indexes.get(domain_name)
    if not bm25_obj:  # fallback to dense-only
        rows = [{"text":t,"score":s,"dense_score":s,"bm25_score":0.0,"hybrid_fallback":True}
                for t,s in sorted(dense_results.items(),key=lambda x:-x[1])[:top_k]]
        if domain_name == "Legal_Contracts":
            for r in rows: r.update(dense_meta.get(r["text"],{}))
        return rows

    bm25_scores = bm25_obj["bm25"].get_scores(_tokenize(query))
    bm25_texts  = bm25_obj["texts"]
    bm25_cids   = bm25_obj.get("contract_ids")

    # For Legal: filter BM25 candidates to the same contract before ranking
    if domain_name == "Legal_Contracts" and contract_id is not None and bm25_cids:
        candidate_idx = [i for i,cid in enumerate(bm25_cids) if str(cid)==str(contract_id)]
    else:
        candidate_idx = list(range(len(bm25_texts)))
    top_bm25_idx = sorted(candidate_idx, key=lambda i: bm25_scores[i], reverse=True)[:top_k]

    bm25_results = {}; bm25_meta = {}
    for i in top_bm25_idx:
        text = bm25_texts[i]
        bm25_results[text] = float(bm25_scores[i])
        if domain_name == "Legal_Contracts":
            bm25_meta[text] = {"contract_id": str(contract_id) if contract_id else (bm25_cids[i] if bm25_cids else None)}

    # ── Fusion ────────────────────────────────────────────────────────────────
    if ENABLE_RRF:
        return reciprocal_rank_fusion(dense_results, bm25_results, top_k=top_k, rrf_k=RRF_K,
                                      dense_meta=dense_meta, bm25_meta=bm25_meta)

    # Alpha-weighted min-max fusion (fallback when RRF disabled)
    all_texts = sorted(set(dense_results)|set(bm25_results))
    d_vals    = [dense_results.get(t,0.0) for t in all_texts]
    b_vals    = [bm25_results.get(t,0.0) for t in all_texts]
    d_norm, b_norm = _normalize(d_vals), _normalize(b_vals)
    combined = []
    for i, text in enumerate(all_texts):
        row = {"text":text,"score":float(alpha*d_norm[i]+(1-alpha)*b_norm[i]),
               "dense_score":float(d_vals[i]),"bm25_score":float(b_vals[i]),"alpha":alpha}
        if domain_name == "Legal_Contracts":
            row.update(dense_meta.get(text, bm25_meta.get(text,{})))
            if "contract_id" not in row and contract_id is not None:
                row["contract_id"] = str(contract_id)
        combined.append(row)
    combined.sort(key=lambda x: -x["score"])
    return combined[:top_k]


# ── Repacking ──────────────────────────────────────────────────────────────────

def repack_documents(docs, strategy="sides"):
    """
    Reorder retrieved documents for LLM attention bias.
    forward: most-relevant first (no change)
    reverse: most-relevant last (benefits models that attend to end of context)
    sides:   U-shape — highest-relevance at both ends, lowest in the middle
    """
    if not docs: return []
    if strategy == "forward": return docs
    if strategy == "reverse": return docs[::-1]
    if strategy == "sides":
        n, result, left, right = len(docs), [None]*len(docs), 0, len(docs)-1
        for i, doc in enumerate(docs):
            if i % 2 == 0: result[left]  = doc; left  += 1
            else:           result[right] = doc; right -= 1
        return result
    raise ValueError(f"Unknown REPACK_STRATEGY: {strategy}")


# ── Summarization ──────────────────────────────────────────────────────────────

def recomp_summarize(query, docs, em, top_k=6, min_score=0.0,
                     grounding_boost=0.15, min_keep_ratio=0.30, keep_critical=True):
    texts     = [d.get("text","") if isinstance(d,dict) else d for d in docs]
    sentences = [s.strip() for doc in texts for s in split_into_sentences(doc) if s.strip()]
    if not sentences: return ""
    q_emb  = em.encode([query], normalize_embeddings=True)
    s_emb  = em.encode(sentences, normalize_embeddings=True)
    scores = (q_emb @ s_emb.T).flatten() + np.array([_count_grounding_signals(s)*grounding_boost for s in sentences])
    crits  = {i for i,s in enumerate(sentences) if keep_critical and _is_critical_sentence(s)}
    valid  = np.where(scores >= min_score)[0]
    if len(valid) == 0: valid = np.array([int(np.argmax(scores))])
    keep   = min(max(top_k, int(np.ceil(len(sentences)*min_keep_ratio))), len(sentences))
    chosen = sorted(set(list(valid[np.argsort(scores[valid])[::-1][:keep]])) | crits)
    return " ".join(sentences[i] for i in chosen)

def _get_llmlingua():
    global llmlingua_compressor
    if llmlingua_compressor is None:
        from llmlingua import PromptCompressor
        dev = "cuda" if torch.cuda.is_available() else "cpu"
        llmlingua_compressor = PromptCompressor(
            model_name="microsoft/llmlingua-2-bert-base-multilingual-cased-meetingbank",
            use_llmlingua2=True, device_map=dev)
        print(f"LLMLingua loaded on {dev}")
    return llmlingua_compressor

def llmlingua_compress(query, docs, rate=0.5):
    texts = [d.get("text","") if isinstance(d,dict) else d for d in docs]
    comp  = _get_llmlingua()
    parts = []
    for t in texts:
        if not t or not t.strip(): continue
        try:    parts.append(comp.compress_prompt(t, question=query, rate=rate)["compressed_prompt"])
        except Exception as e: print(f"LLMLingua chunk failed: {e}"); parts.append(t)
    return "\n\n".join(parts)

def summarize_docs(query, docs, em=None, llm_client=None):
    """
    FIX: explicit None guard on em before calling encode().
    Falls back to global embed_model, then raises a clear error.
    """
    if not ENABLE_SUMMARIZATION:
        return [d.get("text","") if isinstance(d,dict) else d for d in docs]

    # Resolve embed model — must be non-None before encode()
    if em is None:
        em = globals().get("embed_model")
    if em is None:
        raise RuntimeError("summarize_docs: no embed_model available. Run Cell 7 first.")

    if SUMMARIZATION_TYPE == "recomp":
        s = recomp_summarize(query, docs, em, RECOMP_TOP_K_SENTS, RECOMP_MIN_SCORE,
                             RECOMP_GROUNDING_BOOST, RECOMP_MIN_KEEP_RATIO, RECOMP_KEEP_CRITICAL)
        return [s] if s else []
    elif SUMMARIZATION_TYPE == "longllmlingua":
        c = llmlingua_compress(query, docs, LLMLINGUA_RATE)
        return [c] if c else []
    raise ValueError(f"Unknown SUMMARIZATION_TYPE: {SUMMARIZATION_TYPE}")


# ── Main retrieve ──────────────────────────────────────────────────────────────
# Pipeline order: HyDE → Retrieve → Rerank → Repack → Summarize

def retrieve(query, domain_name, embed_model=None, llm_client=None, top_k=None,
             rewritten_query=None, sample_id=None, contract_id=None):
    """
    FIX: top_k now defaults to RETRIEVE_TOP_K (10), not RERANK_TOP_K (3).
    The fetch_k logic already enlarges the initial pool; top_k is the final
    count after reranking.
    """
    if domain_name not in milvus_clients:
        raise ValueError(f"Domain '{domain_name}' not loaded.")

    # FIX: default to RETRIEVE_TOP_K for initial fetch, not RERANK_TOP_K
    if top_k is None:
        top_k = globals().get("RETRIEVE_TOP_K", 10)
    top_k = int(top_k)

    llm = llm_client or globals().get("llm_client")

    # Resolve domain-specific embed model
    em = embed_model
    if em is None:
        try:   em = get_embed_model_for_domain(domain_name)
        except Exception: em = globals().get("embed_model")
    if em is None:
        raise ValueError(f"No embed_model available for domain '{domain_name}'. Run Cell 7 first.")

    # Legal contract filtering (section 1.5)
    legal_contract_id = None
    if domain_name == "Legal_Contracts":
        if contract_id is not None:
            legal_contract_id = str(contract_id)
        elif sample_id is not None:
            try:   legal_contract_id = get_contract_id_for_legal_sample(sample_id)
            except Exception as e:
                print(f"  Could not resolve Legal contract_id for sample_id={sample_id}: {e}")
        if legal_contract_id is None:
            print("  WARNING: Legal_Contracts retrieval without contract filter — cross-contract contamination possible")

    # Fetch more candidates if downstream processing will reduce count
    fetch_k = max(int(RETRIEVE_TOP_K if (ENABLE_HYBRID or ENABLE_RERANKING or ENABLE_SUMMARIZATION) else top_k), top_k)

    eff_q = rewritten_query or query; search_q = eff_q

    # HyDE query expansion
    if ENABLE_HYDE and llm:
        try:
            hyde = generate_hyde(eff_q, llm)
            if hyde: search_q = f"{eff_q} {hyde}"
        except Exception as e: print(f"HyDE failed: {e}")

    # Retrieve
    if ENABLE_HYBRID:
        retrieved = hybrid_search(search_q, domain_name, em, top_k=fetch_k,
                                  alpha=HYBRID_ALPHA, contract_id=legal_contract_id) or []
        for d in retrieved:
            if isinstance(d, dict): d.setdefault("retrieval_type","hybrid")
    else:
        client = milvus_clients[domain_name]; col = domain_name.lower()
        try:
            if "Loaded" not in str(client.get_load_state(col)): client.load_collection(col)
        except Exception: pass
        q_emb = em.encode([search_q], normalize_embeddings=True, convert_to_numpy=True).astype("float32")
        output_fields = ["text"]
        if domain_name == "Legal_Contracts": output_fields += ["contract_id","source_doc_id","source_hash"]
        skw = dict(collection_name=col, data=q_emb.tolist(), limit=fetch_k,
                   output_fields=output_fields, search_params={"metric_type":"IP","params":{}})
        if domain_name == "Legal_Contracts" and legal_contract_id:
            skw["filter"] = build_contract_filter_expr(legal_contract_id)
        hits     = client.search(**skw)
        hit_list = hits[0] if (hits and isinstance(hits[0],(list,tuple))) else hits
        seen, retrieved = set(), []
        for hit in hit_list:
            entity   = (hit.get("entity",{}) or hit) if isinstance(hit,dict) else (getattr(hit,"entity",{}) or {})
            distance = hit.get("distance",0.0)        if isinstance(hit,dict) else getattr(hit,"distance",0.0)
            text = entity.get("text","")
            if text and text not in seen:
                item = {"text":text,"score":float(distance),"retrieval_type":"dense"}
                if domain_name == "Legal_Contracts":
                    item.update({"contract_id":entity.get("contract_id"),
                                 "source_doc_id":entity.get("source_doc_id"),
                                 "source_hash":entity.get("source_hash")})
                retrieved.append(item); seen.add(text)

    if not retrieved: return []

    # Rerank → trim to top_k
    if ENABLE_RERANKING: retrieved = rerank_documents(eff_q, retrieved, top_k)
    else:                retrieved = retrieved[:top_k]

    # Repack (reorder for LLM attention)
    if ENABLE_REPACKING: retrieved = repack_documents(retrieved, REPACK_STRATEGY)

    # Summarize / compress context
    if ENABLE_SUMMARIZATION:
        orig       = retrieved
        summarized = summarize_docs(eff_q, retrieved, em=em, llm_client=llm)
        if not summarized: return orig
        scores_list = [d.get("rerank_score",d.get("score",0.0)) for d in retrieved if isinstance(d,dict)]
        avg = float(np.mean(scores_list)) if scores_list else 1.0
        mx  = float(max(scores_list))     if scores_list else avg
        rtype   = retrieved[0].get("reranker_type")  if retrieved and isinstance(retrieved[0],dict) else None
        rettype = retrieved[0].get("retrieval_type") if retrieved and isinstance(retrieved[0],dict) else None
        # Preserve Legal metadata from the first (highest-relevance) source chunk
        smeta = {}
        if domain_name == "Legal_Contracts" and retrieved and isinstance(retrieved[0],dict):
            smeta = {k: retrieved[0].get(k) for k in ("contract_id","source_doc_id","source_hash")}
        retrieved = [{"text":s,"score":avg,"rerank_score":mx,"reranker_type":rtype,
                      "retrieval_type":rettype,"summarized":True,"summary_type":SUMMARIZATION_TYPE,**smeta}
                     for s in summarized if s and str(s).strip()]
        if not retrieved: return orig

    return retrieved


# ── Prompt / generation ────────────────────────────────────────────────────────

def _build_prompt(context, question, strategy="short"):
    context = _sanitize(context); question = _sanitize(question)
    if strategy == "short":
        return f"Answer the question using the provided context.\n\nContext:\n{context}\n\nQuestion:\n{question}".strip()
    elif strategy == "long":
        return ("You are a chatbot providing answers to user queries. Use the context documents to answer the question.\n"
                'If the documents do not provide enough information, say "The documents are missing some of the information required to answer the question."\n'
                f"Do not use external knowledge. Do not make up an answer.\n\nContext Documents:\n{context}\n\nQuestion: {question}").strip()
    elif strategy == "long_cot":
        return ("You are a chatbot providing answers to user queries. Use the context documents to answer the question.\n"
                'If the documents do not provide enough information, say "The documents are missing some of the information required to answer the question."\n'
                f"Do not use external knowledge. Think step by step and quote documents when necessary.\n\nContext Documents:\n{context}\n\nQuestion: {question}").strip()
    raise ValueError(f"Unknown PROMPT_STRATEGY: {strategy}")

def ask_rag(context, question, llm_client, strategy=None):
    strategy = strategy or PROMPT_STRATEGY
    resp = llm_client.chat.completions.create(
        model=MODEL_NAME,
        messages=[{"role":"system","content":"You are a helpful RAG assistant"},
                  {"role":"user","content":_build_prompt(context, question, strategy)}],
        temperature=0.3,
    )
    return _safe_message_content(resp)


print("Pipeline functions defined.")

# ── Cell 6: Initialize LLM client ─────────────────────────────────────────────
llm_client = get_llm_client()
print(f"LLM client ready.  Provider: {LLM_PROVIDER}")


# ── Cell 7: HF filesystem + embedding model loading ───────────────────────────
# NOTE: _hf_fs must be initialized here before Cell 8 calls hf_path_exists()
from sentence_transformers import SentenceTransformer
import torch
from huggingface_hub import HfFileSystem

device   = "cuda" if torch.cuda.is_available() else "cpu"
_hf_fs   = HfFileSystem()
print(f"Device: {device}  |  HfFileSystem ready")

# Determine which embedding types to load
if DOMAIN_EMBEDDING_RECOMMENDATION:
    models_to_load = sorted(set(DOMAIN_EMBEDDING_RECOMMENDATION.values()))
    print(f"Domain-aware mode — loading: {models_to_load}")
else:
    # Legacy single-model mode: always load EMBEDDING_TYPE
    models_to_load = [EMBEDDING_TYPE]
    print(f"Single-model mode — loading: {models_to_load}")

loaded_embedding_models = {}
for emb_type in models_to_load:
    if emb_type not in EMBED_MODELS:
        raise ValueError(f"Unknown embedding type: {emb_type}.  Available: {list(EMBED_MODELS.keys())}")
    model_name = EMBED_MODELS[emb_type]
    print(f"  Loading {emb_type}: {model_name}")
    m   = SentenceTransformer(model_name, device=device)
    dim = m.get_sentence_embedding_dimension() if hasattr(m, "get_sentence_embedding_dimension") else getattr(m, "get_embedding_dimension", lambda: "?")()
    loaded_embedding_models[emb_type] = m
    print(f"    OK — dim={dim}")

if not loaded_embedding_models:
    raise RuntimeError("No embedding models were loaded. Check DOWNLOAD_MODE and EMBED_MODELS.")

# Fallback single embed_model used by RECOMP summarization
embed_model = loaded_embedding_models.get(EMBEDDING_TYPE) or next(iter(loaded_embedding_models.values()))
print(f"\nAll embedding models ready.  Fallback embed_model: {EMBEDDING_TYPE}")

# ── Cell 8: Download Milvus DBs + Legal mapping from HuggingFace ──────────────
# Uses download_indexes() from Cell 5 which mirrors the Advanced notebook logic:
# preferred path: BUCKET_PREFIX/{embedding_type}_{index_version}/{domain}.db
# fallback path:  BUCKET_PREFIX/{embedding_type}/{domain}.db
os.makedirs(MILVUS_DIR, exist_ok=True)

def download_indexes():
    """Download all domain DBs using domain-aware embedding types."""
    report = []
    for domain_name in DOMAIN_NAMES:
        embedding_type = get_embedding_type_for_domain(domain_name)
        local_target   = get_db_path(domain_name, embedding_type=embedding_type, index_version=INDEX_VERSION)
        os.makedirs(os.path.dirname(local_target), exist_ok=True)

        preferred_remote = f"{BUCKET_PREFIX}/{get_index_folder(embedding_type, INDEX_VERSION)}/{domain_name}.db"
        fallback_remote  = f"{BUCKET_PREFIX}/{embedding_type}/{domain_name}.db"

        # Remove stale file before re-download
        if os.path.exists(local_target):
            if os.path.isdir(local_target): shutil.rmtree(local_target)
            else:                            os.remove(local_target)

        selected_remote, source_type = None, None
        if hf_path_exists(preferred_remote):
            selected_remote = preferred_remote
            source_type     = get_index_folder(embedding_type, INDEX_VERSION)
        elif hf_path_exists(fallback_remote):
            selected_remote = fallback_remote
            source_type     = embedding_type

        if selected_remote is None:
            print(f"  MISSING: {domain_name} ({embedding_type})")
            report.append({"domain":domain_name,"status":"missing","source_type":None,"local_target":local_target})
            continue

        print(f"  Downloading: {domain_name}  [{source_type}]")
        try:
            _hf_fs.get(selected_remote, local_target, recursive=True)
            ok = os.path.exists(local_target) and os.path.getsize(local_target) > 0
            status = "downloaded" if ok else "empty"
            print(f"    {'OK' if ok else 'EMPTY'}: {local_target}")
            report.append({"domain":domain_name,"status":status,"source_type":source_type,"local_target":local_target})
        except Exception as e:
            print(f"    FAILED: {e}")
            report.append({"domain":domain_name,"status":"failed","source_type":source_type,"local_target":local_target,"error":str(e)})
    return report

print("Downloading domain DBs...")
dl_report = download_indexes()

# ── Download Legal sample→contract mapping ────────────────────────────────────
legal_emb    = get_embedding_type_for_domain("Legal_Contracts")
legal_folder = get_index_folder(legal_emb, INDEX_VERSION)
remote_mapping = f"{BUCKET_PREFIX}/{legal_folder}/legal_sample_to_contract_id.json"
local_mapping  = os.path.join(MILVUS_DIR, legal_folder, "legal_sample_to_contract_id.json")
os.makedirs(os.path.dirname(local_mapping), exist_ok=True)
print(f"\nDownloading Legal mapping: {remote_mapping}")
try:
    _hf_fs.get(remote_mapping, local_mapping)
    if os.path.exists(local_mapping) and os.path.getsize(local_mapping) > 0:
        print(f"  OK: {local_mapping}")
    else:
        print("  WARNING: Legal mapping download failed or empty")
except Exception as e:
    print(f"  WARNING: Could not download Legal mapping: {e}")

# ── Sanity check ──────────────────────────────────────────────────────────────
print("\nSanity check:")
for domain in DOMAIN_NAMES:
    p = get_db_path(domain, get_embedding_type_for_domain(domain), INDEX_VERSION)
    print(f"  {'OK' if os.path.exists(p) else 'MISSING'}: {p}")


    # ── Cell 9: Open Milvus clients + BM25 indexes + Legal mapping ────────────────
def load_milvus_clients():
    global milvus_clients, LEGAL_SAMPLE_TO_CONTRACT_ID
    milvus_clients = {}
    for domain_name in DOMAIN_NAMES:
        embedding_type = get_embedding_type_for_domain(domain_name)
        db_path        = get_db_path(domain_name, embedding_type=embedding_type, index_version=INDEX_VERSION)
        col            = domain_name.lower()
        if not os.path.exists(db_path):
            print(f"  DB not found, skipping: {db_path}")
            continue
        try:
            client = MilvusClient(db_path)
            if not client.has_collection(col):
                print(f"  Collection missing in {db_path}, skipping")
                continue
            client.load_collection(col)
            stats = client.get_collection_stats(col)
            rows  = int(stats.get("row_count", 0))
            milvus_clients[domain_name] = client
            print(f"  {domain_name}: {rows:,} rows  [{embedding_type}]")
        except Exception as e:
            print(f"  Failed to open {domain_name}: {e}")
    print(f"\nLoaded {len(milvus_clients)} domain clients: {list(milvus_clients.keys())}")

    # Legal sample→contract mapping
    LEGAL_SAMPLE_TO_CONTRACT_ID = load_legal_sample_to_contract_mapping()

load_milvus_clients()

# Build BM25 indexes (stores contract_ids for Legal)
build_all_bm25_indexes(milvus_clients)
print("BM25 indexes ready.")

# ── Cell 10: Load RAGBench (test split only) + sample catalogue ───────────────
DATASET_BY_DOMAIN = {
    "Bio_Medical":       ["covidqa", "pubmedqa"],
    "General_Knowledge": ["expertqa", "hagrid", "hotpotqa", "msmarco"],
    "Customer_Support":  ["delucionqa", "emanual", "techqa"],
    "Finance":           ["finqa", "tatqa"],
    "Legal_Contracts":   ["cuad"],
}

# sample_store[domain][dataset] = list of row dicts from the test split
sample_store = {}

def load_ragbench(domains=None):
    global ragbench_by_domain, sample_store
    domains = domains or list(DATASET_BY_DOMAIN.keys())
    for domain in domains:
        ragbench_by_domain[domain] = {}
        sample_store[domain]       = {}
        for ds_name in DATASET_BY_DOMAIN.get(domain, []):
            try:
                ds = load_dataset("rungalileo/ragbench", ds_name)
                ragbench_by_domain[domain][ds_name] = ds
                if "test" not in ds:
                    print(f"  WARNING: no 'test' split for {domain}/{ds_name}, skipping")
                    continue
                rows = []
                for idx, row in enumerate(ds["test"]):
                    # For Legal_Contracts resolve contract_id from mapping
                    contract_id = None
                    if domain == "Legal_Contracts":
                        try:   contract_id = get_contract_id_for_legal_sample(idx)
                        except Exception: pass
                    rows.append({
                        "idx":         idx,
                        "question":    row.get("question", ""),
                        "response":    row.get("response", ""),
                        "documents":   row.get("documents", []),
                        "contract_id": contract_id,
                        "gold_relevance":    row.get("relevance_score"),
                        "gold_utilization":  row.get("utilization_score"),
                        "gold_completeness": row.get("completeness_score"),
                        "gold_adherence":    row.get("adherence_score"),
                    })
                sample_store[domain][ds_name] = rows
                print(f"  Loaded: {domain}/{ds_name}  test rows={len(rows)}")
            except Exception as e:
                print(f"  Failed: {domain}/{ds_name}: {e}")
    print(f"\nRAGBench loaded (test only). Domains: {list(sample_store.keys())}")

load_ragbench()


# ── Helpers for cascading dropdowns ───────────────────────────────────────────

def get_datasets_for_domain(domain):
    return list(sample_store.get(domain, {}).keys())

def get_sample_ids_for_dataset(domain, dataset):
    """
    Returns label strings for the Sample ID dropdown.
    For Legal_Contracts uses 'Contract ID' wording and shows contract hash.
    """
    rows   = sample_store.get(domain, {}).get(dataset, [])
    is_legal = (domain == "Legal_Contracts")
    labels = []
    for r in rows:
        q   = r["question"]
        cid = r.get("contract_id")
        if is_legal and cid:
            prefix = f"Contract {str(cid)[:8]}…  |  idx={r['idx']}  –  "
        else:
            prefix = f"{r['idx']}  –  "
        labels.append(f"{prefix}{q[:70]}{'…' if len(q)>70 else ''}")
    return labels

def get_row_by_label(domain, dataset, label):
    """Retrieve a stored row dict from a label string."""
    if not label: return None
    rows     = sample_store.get(domain, {}).get(dataset, [])
    is_legal = (domain == "Legal_Contracts")
    # Legal labels: "Contract <hash8>…  |  idx=N  –  ..."
    # Regular labels: "N  –  ..."
    if is_legal:
        m = re.search(r"idx=(\d+)", label)
        try:   idx = int(m.group(1)) if m else int(label.split("–")[0].strip())
        except ValueError: return None
    else:
        try:   idx = int(label.split("–")[0].strip())
        except ValueError: return None
    return next((r for r in rows if r["idx"] == idx), None)

print("Sample catalogue ready.")

# ── Cell 11: Evaluation helpers + all Gradio handlers ─────────────────────────

# ── Judge / evaluation ─────────────────────────────────────────────────────────

def build_keyed_response(answer):
    return {f"r_{i}": s for i, s in enumerate(split_into_sentences(answer))}

def build_sentence_keyed_docs(retrieved_docs):
    keyed = {}
    for di, doc in enumerate(retrieved_docs):
        text = doc.get("text","") if isinstance(doc, dict) else doc
        for si, s in enumerate(split_into_sentences(text)):
            keyed[f"{di}_{si}"] = s
    return keyed

def build_evaluation_prompt(documents_text, question, answer_text):
    return f"""Evaluate the RAG response using the provided documents.

Documents (sentence-keyed):
{documents_text}

Question:
{question}

Response (sentence-keyed):
{answer_text}

Return ONLY valid JSON:
{{
  "overall_supported": true,
  "all_relevant_sentence_keys": ["0_0"],
  "all_utilized_sentence_keys": ["0_0"],
  "sentence_support_information": [
    {{"response_sentence_key": "r_0", "supporting_sentence_keys": ["0_0"], "fully_supported": true}}
  ]
}}
Rules: document keys look like 0_0; response keys like r_0. Return only JSON.""".strip()

def ask_judge(prompt, llm_client, judge_model, max_retries=5):
    last_error = None
    for attempt in range(max_retries):
        try:
            resp = llm_client.chat.completions.create(
                model=judge_model,
                messages=[
                    {"role":"system","content":"You are a strict RAG evaluation judge. Return ONLY valid JSON. No markdown. No <think> tags."},
                    {"role":"user","content":_sanitize(prompt)},
                ],
                temperature=0.0, max_tokens=3000,
            )
            return _safe_message_content(resp)
        except Exception as e:
            last_error = e; msg = str(e)
            wait = 2**attempt
            if "429" in msg or "rate_limit" in msg:
                m = re.search(r"try again in ([\\d.]+)s", msg)
                if m: wait = float(m.group(1))
            elif not any(x in msg for x in ["503","502","504","over capacity","gateway"]):
                raise
            time.sleep(wait + random.uniform(0.1, 0.5))
    raise RuntimeError(f"Judge failed after {max_retries} retries: {last_error}")

def parse_judge_json(raw):
    if not raw: raise ValueError("Judge output empty")
    cleaned = re.sub(r"<think>.*?</think>","",str(raw),flags=re.DOTALL).strip()
    cleaned = cleaned.replace("```json","").replace("```","").strip()
    s, e = cleaned.find("{"), cleaned.rfind("}")
    if s == -1 or e == -1: raise ValueError(f"No JSON: {cleaned[:300]}")
    cleaned = cleaned[s:e+1]
    cleaned = re.sub(r"}\s*{","}, {",cleaned)
    cleaned = re.sub(r",\s*([}\]])",r"\1",cleaned)
    return json.loads(cleaned)

def evaluate_ragbench_json(judge_json, keyed_docs):
    vk  = set(keyed_docs.keys())
    rel = set(judge_json.get("all_relevant_sentence_keys", [])) & vk
    utl = set(judge_json.get("all_utilized_sentence_keys", [])) & vk
    ovl = rel & utl; n = len(vk)
    return {
        "adherence_score":    int(bool(judge_json.get("overall_supported", False))),
        "hallucination_flag": 1 - int(bool(judge_json.get("overall_supported", False))),
        "relevance_score":    float(np.clip(len(rel)/n        if n   else 0, 0, 1)),
        "utilization_score":  float(np.clip(len(utl)/n        if n   else 0, 0, 1)),
        "completeness_score": float(np.clip(len(ovl)/len(rel) if rel else 0, 0, 1)),
    }


# ── Source badge helper ────────────────────────────────────────────────────────

def _source_badge(source, model, extra=None):
    parts = [f"[Source: {source}  |  model: {model}"]
    if extra: parts += [f"  |  {k}: {v}" for k, v in extra.items()]
    parts.append("]")
    return "".join(parts)


# ── DB status helper ───────────────────────────────────────────────────────────

def db_status_md():
    if not milvus_clients:
        return ("> **No vector DBs loaded.**  Re-run Cell 8 (download) then Cell 9 (open), then re-run Cell 12.")
    rows = []
    for d in sorted(milvus_clients.keys()):
        emb = get_embedding_type_for_domain(d)
        rows.append(f"`{d}` ({emb})")
    return f"> **Loaded domains ({len(milvus_clients)}):** {', '.join(rows)}"


# ── Config applier ─────────────────────────────────────────────────────────────

def apply_config(llm_choice, embed_choice,
                 enable_hybrid, enable_hyde, enable_reranking, reranker_type,
                 enable_rrf, rrf_k,
                 enable_repacking, repack_strategy,
                 enable_summarization, summarization_type,
                 prompt_strategy, hybrid_alpha, top_k,
                 enable_query_classification, enable_query_rewriting, enable_query_decomp):
    global MODEL_NAME, EMBEDDING_TYPE, embed_model
    global ENABLE_HYBRID, ENABLE_HYDE, ENABLE_RERANKING, RERANKER_TYPE
    global ENABLE_RRF, RRF_K
    global ENABLE_REPACKING, REPACK_STRATEGY, ENABLE_SUMMARIZATION, SUMMARIZATION_TYPE
    global PROMPT_STRATEGY, HYBRID_ALPHA
    global ENABLE_QUERY_CLASSIFICATION, ENABLE_QUERY_REWRITING, ENABLE_QUERY_DECOMPOSITION

    MODEL_NAME                  = llm_choice
    ENABLE_HYBRID               = enable_hybrid
    ENABLE_HYDE                 = enable_hyde
    ENABLE_RERANKING            = enable_reranking
    RERANKER_TYPE               = reranker_type
    ENABLE_RRF                  = enable_rrf
    RRF_K                       = int(rrf_k)
    ENABLE_REPACKING            = enable_repacking
    REPACK_STRATEGY             = repack_strategy
    ENABLE_SUMMARIZATION        = enable_summarization
    SUMMARIZATION_TYPE          = summarization_type
    PROMPT_STRATEGY             = prompt_strategy
    HYBRID_ALPHA                = float(hybrid_alpha)
    ENABLE_QUERY_CLASSIFICATION = enable_query_classification
    ENABLE_QUERY_REWRITING      = enable_query_rewriting
    ENABLE_QUERY_DECOMPOSITION  = enable_query_decomp

    # Update single fallback embed_model if user changes embedding choice
    if embed_choice != EMBEDDING_TYPE:
        EMBEDDING_TYPE = embed_choice
        if embed_choice in loaded_embedding_models:
            embed_model = loaded_embedding_models[embed_choice]
        else:
            print(f"Embedding type '{embed_choice}' not preloaded; loading now...")
            embed_model = SentenceTransformer(EMBED_MODELS[embed_choice], device=device)
            loaded_embedding_models[embed_choice] = embed_model


# ── Cascading dropdown callbacks ───────────────────────────────────────────────

_NONE_DOMAIN = "None (direct LLM, no retrieval)"

def on_domain_change(domain):
    if domain == _NONE_DOMAIN:
        return gr.update(choices=[], value=None), gr.update(choices=[], value=None), gr.update()
    datasets   = get_datasets_for_domain(domain)
    ds         = datasets[0] if datasets else None
    sample_ids = get_sample_ids_for_dataset(domain, ds) if ds else []
    label_text = "Contract ID  (contract hash | idx – question preview)" if domain == "Legal_Contracts" else "Sample ID  (idx – question preview)"
    return (
        gr.update(choices=datasets,   value=ds),
        gr.update(choices=sample_ids, value=None, label=label_text),
        gr.update(value=""),
    )

def on_dataset_change(domain, dataset):
    if domain == _NONE_DOMAIN or not dataset:
        return gr.update(choices=[], value=None), gr.update(value="")
    sample_ids = get_sample_ids_for_dataset(domain, dataset)
    label_text = "Contract ID  (contract hash | idx – question preview)" if domain == "Legal_Contracts" else "Sample ID  (idx – question preview)"
    return gr.update(choices=sample_ids, value=None, label=label_text), gr.update(value="")

def on_sample_select(domain, dataset, label):
    if domain == _NONE_DOMAIN or not label: return gr.update()
    row = get_row_by_label(domain, dataset, label)
    if row is None: return gr.update()
    return gr.update(value=row["question"])


# ── Chunk display helpers ──────────────────────────────────────────────────────

def _format_chunks(docs, title="Retrieved"):
    if not docs: return f"_No documents for {title}._"
    parts = []
    for i, doc in enumerate(docs):
        if isinstance(doc, dict):
            text  = doc.get("text", str(doc))
            score = doc.get("rerank_score", doc.get("score", 0.0))
            tags  = []
            if doc.get("summarized"):    tags.append(f"summarized/{doc.get('summary_type','')}")
            if doc.get("reranker_type"): tags.append(f"reranked/{doc.get('reranker_type','')}")
            if doc.get("retrieval_type"): tags.append(doc.get("retrieval_type",""))
            if ENABLE_HYBRID and ENABLE_RRF:
                tags.append(f"RRF score={score:.4f}")
            elif ENABLE_HYBRID:
                tags.append(f"d={doc.get('dense_score',0):.3f} b={doc.get('bm25_score',0):.3f}")
            if doc.get("contract_id"): tags.append(f"contract={str(doc.get('contract_id',''))[:8]}…")
            tag_str = f"  `{'  |  '.join(tags)}`" if tags else ""
        else:
            text, score, tag_str = str(doc), 0.0, ""
        parts.append(f"**{title} Chunk {i+1}** — score: `{score:.4f}`{tag_str}\n\n{text}")
    return "\n\n---\n\n".join(parts)

def _format_gt_docs(doc_list):
    if not doc_list: return "_No ground-truth documents stored for this sample._"
    parts = []
    for i, text in enumerate(doc_list):
        parts.append(f"**GT Doc {i+1}**\n\n{str(text)}")
    return "\n\n---\n\n".join(parts)


# ── Main run handler ───────────────────────────────────────────────────────────

def run_query(
    query, domain,
    dataset_sel, sample_label,
    llm_choice, judge_llm_choice, embed_choice,
    enable_hybrid, enable_hyde, enable_reranking, reranker_type,
    enable_rrf, rrf_k,
    enable_repacking, repack_strategy,
    enable_summarization, summarization_type,
    prompt_strategy, hybrid_alpha, top_k,
    enable_query_classification, enable_query_rewriting, enable_query_decomp,
    run_judge,
):
    query = _sanitize(query)
    if not query.strip():
        return ("Please enter a query.",) + ("",)*4
    if llm_client is None:
        return ("LLM client not initialised. Re-run Cell 6 then Cell 12.",) + ("",)*4

    apply_config(
        llm_choice, embed_choice,
        enable_hybrid, enable_hyde, enable_reranking, reranker_type,
        enable_rrf, rrf_k,
        enable_repacking, repack_strategy,
        enable_summarization, summarization_type,
        prompt_strategy, float(hybrid_alpha), int(top_k),
        enable_query_classification, enable_query_rewriting, enable_query_decomp,
    )

    # ── Domain = None → direct LLM ───────────────────────────────────────────
    if domain == _NONE_DOMAIN or not domain:
        try:
            direct_ans = _safe_message_content(llm_client.chat.completions.create(
                model=MODEL_NAME,
                messages=[{"role":"system","content":"You are a helpful assistant."},
                          {"role":"user","content":query}],
                temperature=0.3, max_tokens=800,
            ))
        except Exception as e: direct_ans = f"Direct LLM error: {e}"
        badge = _source_badge("Direct LLM (no retrieval)", MODEL_NAME)
        note  = "_[Domain = None — answered directly by LLM without vector DB retrieval]_"
        return f"{badge}\n\n{direct_ans}", note, note, note, note

    # ── Guard ──────────────────────────────────────────────────────────────────
    if not milvus_clients:
        return ("No vector DBs loaded. Re-run Cell 8 then Cell 9, then re-run Cell 12.",) + ("",)*4
    if domain not in milvus_clients:
        return (f"Domain '{domain}' not loaded. Loaded: {list(milvus_clients.keys())}",) + ("",)*4

    # ── Query Classification ──────────────────────────────────────────────────
    route = classify_query(query, domain_name=domain)
    if route == "LLM":
        try:
            direct_ans = _safe_message_content(llm_client.chat.completions.create(
                model=MODEL_NAME,
                messages=[{"role":"system","content":"You are a concise factual assistant."},
                          {"role":"user","content":query}],
                temperature=0.2, max_tokens=500,
            ))
        except Exception as e: direct_ans = f"Direct LLM error: {e}"
        badge = _source_badge("Direct LLM", MODEL_NAME)
        note  = "_[Query Classifier routed to direct LLM — no retrieval]_"
        return f"{badge}\n\n{direct_ans}", note, note, note, note

    # ── Query Rewriting + Decomposition ──────────────────────────────────────
    rewritten  = rewrite_query(query, domain, llm_client) if ENABLE_QUERY_REWRITING else query
    subqueries = decompose_query(rewritten, llm_client, domain=domain) if ENABLE_QUERY_DECOMPOSITION else [rewritten]

    # ── Resolve row + contract_id for Legal ───────────────────────────────────
    row = get_row_by_label(domain, dataset_sel, sample_label) if sample_label else None
    if row is None:
        for ds_name, rows in sample_store.get(domain, {}).items():
            match = next((r for r in rows if r["question"].strip().lower() == query.strip().lower()), None)
            if match: row = match; break

    legal_contract_id = None
    legal_sample_id   = None
    if domain == "Legal_Contracts" and row is not None:
        legal_contract_id = row.get("contract_id")
        legal_sample_id   = row.get("idx")

    # ── Retrieve + Generate ───────────────────────────────────────────────────
    all_retrieved, all_answers = [], []
    for sq in subqueries:
        try:
            docs = retrieve(sq, domain,
                            llm_client=llm_client, top_k=int(top_k),
                            sample_id=legal_sample_id, contract_id=legal_contract_id)
        except Exception as e:
            return (f"Retrieval error: {e}",) + ("",)*4
        if not docs: continue
        all_retrieved.extend(docs)
        ctx = _sanitize("\n\n".join(d.get("text","") if isinstance(d,dict) else d for d in docs))
        sq  = _sanitize(sq)
        try:
            all_answers.append(ask_rag(ctx, sq, llm_client, strategy=PROMPT_STRATEGY))
        except Exception as e:
            return (f"Generation error: {e}",) + ("",)*4

    if not all_retrieved:
        return ("No documents retrieved.",) + ("",)*4

    raw_answer = "\n\n".join(all_answers)

    # ── Source badge ──────────────────────────────────────────────────────────
    active = {"prompt": PROMPT_STRATEGY, "chunks": len(all_retrieved)}
    if ENABLE_HYBRID:
        active["hybrid"] = f"RRF(k={RRF_K})" if ENABLE_RRF else f"alpha={HYBRID_ALPHA}"
    if ENABLE_HYDE:          active["hyde"]   = "on"
    if ENABLE_RERANKING:     active["rerank"] = RERANKER_TYPE
    if ENABLE_SUMMARIZATION: active["summ"]   = SUMMARIZATION_TYPE
    if ENABLE_REPACKING:     active["repack"] = REPACK_STRATEGY
    if len(subqueries) > 1:  active["subq"]   = len(subqueries)
    if legal_contract_id:    active["contract"] = str(legal_contract_id)[:8] + "…"
    rag_response = f"{_source_badge('RAG', MODEL_NAME, extra=active)}\n\n{raw_answer}"

    ground_truth = row["response"]  if row else "_(no matching sample found)_"
    gt_docs_md   = _format_gt_docs(row["documents"] if row else [])
    rag_docs_md  = _format_chunks(all_retrieved, title="RAG")

    # ── Judge evaluation ──────────────────────────────────────────────────────
    metrics_md = "_Judge evaluation not requested._"
    if run_judge:
        try:
            keyed_docs   = build_sentence_keyed_docs(all_retrieved)
            keyed_answer = build_keyed_response(raw_answer)
            docs_text = _sanitize("\n".join(f"{k}: {v}" for k,v in keyed_docs.items()))
            ans_text  = _sanitize("\n".join(f"{k}: {v}" for k,v in keyed_answer.items()))
            raw       = ask_judge(build_evaluation_prompt(docs_text, query, ans_text), llm_client, judge_llm_choice)
            pred      = evaluate_ragbench_json(parse_judge_json(raw), keyed_docs)
            gold = {k: row.get(f"gold_{k}") for k in ("relevance","utilization","completeness","adherence")} if row else {}
            def _f(v): return f"{v:.3f}" if isinstance(v, float) else (str(v) if v is not None else "—")
            metrics_md = "\n".join([
                "| Metric | Predicted | Gold |",
                "|--------|-----------|------|",
                f"| Relevance    | {_f(pred['relevance_score'])}    | {_f(gold.get('relevance'))} |",
                f"| Utilization  | {_f(pred['utilization_score'])}  | {_f(gold.get('utilization'))} |",
                f"| Completeness | {_f(pred['completeness_score'])} | {_f(gold.get('completeness'))} |",
                f"| Adherence    | {_f(pred['adherence_score'])}    | {_f(gold.get('adherence'))} |",
                f"| Hallucination| {_f(pred['hallucination_flag'])} | — |",
            ])
        except Exception as e:
            metrics_md = f"Judge error: {e}"

    return ground_truth, rag_response, gt_docs_md, rag_docs_md, metrics_md


print("Handlers ready.")

# ── Cell 12: Gradio UI ────────────────────────────────────────────────────────

AVAILABLE_DOMAINS = list(milvus_clients.keys())
DEFAULT_DOMAIN    = AVAILABLE_DOMAINS[0] if AVAILABLE_DOMAINS else None

_init_datasets = get_datasets_for_domain(DEFAULT_DOMAIN) if DEFAULT_DOMAIN else []
_init_ds       = _init_datasets[0] if _init_datasets else None
_init_samples  = get_sample_ids_for_dataset(DEFAULT_DOMAIN, _init_ds) if _init_ds else []
_legal_first   = DEFAULT_DOMAIN == "Legal_Contracts"

CSS = """
footer { display: none !important; }
"""

with gr.Blocks(title="RAG Capstone — Advanced Demo") as demo:

    # ── Header ────────────────────────────────────────────────────────────────
    gr.Markdown("# 🔍 RAG Capstone — Advanced Interactive Demo")
    gr.Markdown(
        f"**Provider:** {LLM_PROVIDER.upper()}  &nbsp;|&nbsp;  "
        f"**Index:** `{INDEX_VERSION}`  &nbsp;|&nbsp;  "
        "Type any question, or expand **Sample Selector** to load a test-split example."
    )
    gr.Markdown(db_status_md())

    # ══════════════════════════════════════════════════════════════════════════
    # SECTION 1 — Query + Domain  (always visible)
    # ══════════════════════════════════════════════════════════════════════════
    with gr.Row():
        query_input = gr.Textbox(
            lines=3,
            placeholder="Type any question here… or expand Sample Selector below to auto-fill.",
            label="Query",
            scale=4,
        )
        domain_dd = gr.Dropdown(
            choices=["None (direct LLM, no retrieval)"] + AVAILABLE_DOMAINS,
            value="None (direct LLM, no retrieval)" if not AVAILABLE_DOMAINS else DEFAULT_DOMAIN,
            label="Domain",
            info="None = direct LLM; pick a domain to run full RAG retrieval",
            scale=1,
        )

    # ══════════════════════════════════════════════════════════════════════════
    # SECTION 2 — Sample Selector  (collapsed, optional)
    # ══════════════════════════════════════════════════════════════════════════
    with gr.Accordion("📋  Sample Selector  (optional — expand to load a test-split example)", open=False):
        gr.Markdown(
            "_Select a sample to auto-fill Query above.  "
            "For **Legal_Contracts** the dropdown shows Contract ID (hash prefix) instead of plain Sample ID — "
            "retrieval is automatically scoped to that contract._"
        )
        with gr.Row():
            dataset_dd = gr.Dropdown(
                choices=_init_datasets, value=_init_ds, label="Dataset", scale=1)
            sample_dd  = gr.Dropdown(
                choices=_init_samples, value=None,
                label="Contract ID  (contract hash | idx – question preview)" if _legal_first else "Sample ID  (idx – question preview)",
                scale=4)

    # ══════════════════════════════════════════════════════════════════════════
    # SECTION 3 — Control Panel
    # ══════════════════════════════════════════════════════════════════════════
    with gr.Accordion("⚙️  Control Panel", open=False):
        with gr.Tabs():

            # ── Models ───────────────────────────────────────────────────────
            with gr.Tab("🤖 Models"):
                gr.Markdown(
                    f"**Domain embedding assignment** (chunk_v5_domain_aware):  \n"
                    + "  \n".join(
                        [f"- `{d}` → `{get_embedding_type_for_domain(d)}` ({EMBED_MODELS[get_embedding_type_for_domain(d)]})"
                         for d in DOMAIN_NAMES]
                    )
                )
                with gr.Row():
                    llm_choice = gr.Dropdown(
                        choices=LLM_CHOICES, value=LLM_CHOICES[0],
                        label="Generator LLM", info="Produces the RAG answer")
                    judge_llm_choice = gr.Dropdown(
                        choices=LLM_CHOICES,
                        value=LLM_CHOICES[4] if len(LLM_CHOICES) > 4 else LLM_CHOICES[-1],
                        label="Judge LLM", info="Used for evaluation scoring")
                    embed_choice = gr.Dropdown(
                        choices=EMBEDDING_CHOICES, value=EMBEDDING_TYPE,
                        label="Fallback Embedding Model",
                        info="Used only when domain-specific model is unavailable")

            # ── Query Processing ──────────────────────────────────────────────
            with gr.Tab("🔄 Query Processing"):
                gr.Markdown("Applied **before** retrieval: Classify → Rewrite → Decompose")
                with gr.Row():
                    enable_query_classification = gr.Checkbox(
                        label="Query Classification", value=False,
                        info="Route simple factual queries to LLM directly; benchmark domains always use RAG")
                with gr.Row():
                    enable_query_rewriting = gr.Checkbox(
                        label="Query Rewriting", value=False,
                        info="LLM rewrites the query for better retrieval")
                    enable_query_decomp = gr.Checkbox(
                        label="Query Decomposition", value=False,
                        info="Break multi-part queries into subqueries")

            # ── Retrieval ─────────────────────────────────────────────────────
            with gr.Tab("🔎 Retrieval"):
                with gr.Row():
                    top_k = gr.Slider(minimum=1, maximum=10, step=1, value=3,
                                      label="Top-K chunks returned")
                    hybrid_alpha = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.5,
                                             label="Hybrid Alpha  (1=dense, 0=BM25)  — used only when RRF is OFF")
                with gr.Row():
                    enable_hybrid = gr.Checkbox(label="Hybrid Search  (Dense + BM25)", value=True)
                    enable_hyde   = gr.Checkbox(label="HyDE  (query expansion)", value=False)

            # ── RRF ───────────────────────────────────────────────────────────
            with gr.Tab("🔀 RRF"):
                gr.Markdown(
                    "**Reciprocal Rank Fusion** replaces the weighted alpha fusion inside Hybrid Search.  \n"
                    "Score formula: `1/(k + rank_dense) + 1/(k + rank_bm25)`  \n"
                    "Standard literature value for k is **60** — lower k boosts top-ranked docs more aggressively."
                )
                with gr.Row():
                    enable_rrf = gr.Checkbox(
                        label="Enable RRF  (replaces alpha fusion inside Hybrid Search)",
                        value=True,
                        info="RRF is only active when Hybrid Search is also enabled")
                    rrf_k = gr.Slider(
                        minimum=1, maximum=200, step=1, value=60,
                        label="RRF k  (rank smoothing constant)")

            # ── Reranking ─────────────────────────────────────────────────────
            with gr.Tab("↕️ Reranking"):
                with gr.Row():
                    enable_reranking = gr.Checkbox(label="Enable Reranking", value=False)
                    reranker_type    = gr.Radio(choices=["monot5","tilde"], value="monot5",
                                               label="Reranker", info="MonoT5: seq2seq | TILDE: cross-encoder")

            # ── Repacking ─────────────────────────────────────────────────────
            with gr.Tab("📦 Repacking"):
                with gr.Row():
                    enable_repacking = gr.Checkbox(label="Enable Repacking", value=False)
                    repack_strategy  = gr.Radio(choices=["forward","reverse","sides"], value="sides",
                                               label="Strategy", info="forward | reverse | U-shape sides")

            # ── Summarization ─────────────────────────────────────────────────
            with gr.Tab("📝 Summarization"):
                with gr.Row():
                    enable_summarization = gr.Checkbox(label="Enable Summarization", value=False)
                    summarization_type   = gr.Radio(choices=["recomp","longllmlingua"], value="recomp",
                                                   label="Method", info="RECOMP: extractive | LLMLingua: token compression")

            # ── Prompt ────────────────────────────────────────────────────────
            with gr.Tab("💬 Prompt"):
                prompt_strategy = gr.Radio(
                    choices=["short","long","long_cot"], value="short",
                    label="Prompt Strategy",
                    info="short: minimal | long: strict no-hallucination | long_cot: step-by-step")

            # ── Judge ─────────────────────────────────────────────────────────
            with gr.Tab("⚖️ Judge"):
                run_judge = gr.Checkbox(
                    label="Run Judge evaluation after generation", value=False,
                    info="~1 extra LLM call. Gold scores shown only for preloaded samples.")
                gr.Markdown("_Judge LLM is configured in the **Models** tab._")

    # ── Run button ────────────────────────────────────────────────────────────
    run_btn = gr.Button("▶  Run Query", variant="primary", size="lg")

    # ══════════════════════════════════════════════════════════════════════════
    # SECTION 4 — Responses  (Ground Truth LEFT, RAG RIGHT)
    # ══════════════════════════════════════════════════════════════════════════
    gr.Markdown("## 💬 Responses")
    with gr.Row(equal_height=True):
        gt_out  = gr.Textbox(label="Ground Truth Response", lines=10, interactive=False, scale=1)
        rag_out = gr.Textbox(label="RAG Response",          lines=10, interactive=False, scale=1)

    # ══════════════════════════════════════════════════════════════════════════
    # SECTION 5 — Retrieved Documents  (GT LEFT, RAG RIGHT)
    # ══════════════════════════════════════════════════════════════════════════
    gr.Markdown("## 📄 Retrieved Documents")
    with gr.Row(equal_height=True):
        with gr.Column(scale=1):
            gr.Markdown("### Ground Truth Documents")
            gt_docs_out  = gr.Markdown(value="_Select a preloaded sample to see GT documents._")
        with gr.Column(scale=1):
            gr.Markdown("### RAG Retrieved Documents")
            rag_docs_out = gr.Markdown(value="_Run a query to see RAG retrieved chunks._")

    # ══════════════════════════════════════════════════════════════════════════
    # SECTION 6 — Metrics
    # ══════════════════════════════════════════════════════════════════════════
    with gr.Accordion("📊  Metrics  (Gold vs Predicted)", open=False):
        metrics_out = gr.Markdown(value="_Enable the Judge in the Control Panel and run a query._")

    # ── Cascading sample selector wiring ─────────────────────────────────────
    domain_dd.change(
        fn=on_domain_change, inputs=[domain_dd],
        outputs=[dataset_dd, sample_dd, query_input],
    )
    dataset_dd.change(
        fn=on_dataset_change, inputs=[domain_dd, dataset_dd],
        outputs=[sample_dd, query_input],
    )
    sample_dd.change(
        fn=on_sample_select, inputs=[domain_dd, dataset_dd, sample_dd],
        outputs=[query_input],
    )

    # ── Run wiring ────────────────────────────────────────────────────────────
    _config_inputs = [
        llm_choice, judge_llm_choice, embed_choice,
        enable_hybrid, enable_hyde, enable_reranking, reranker_type,
        enable_rrf, rrf_k,
        enable_repacking, repack_strategy,
        enable_summarization, summarization_type,
        prompt_strategy, hybrid_alpha, top_k,
        enable_query_classification, enable_query_rewriting, enable_query_decomp,
        run_judge,
    ]
    _all_inputs  = [query_input, domain_dd, dataset_dd, sample_dd] + _config_inputs
    _all_outputs = [gt_out, rag_out, gt_docs_out, rag_docs_out, metrics_out]

    run_btn.click(fn=run_query, inputs=_all_inputs, outputs=_all_outputs)
    query_input.submit(fn=run_query, inputs=_all_inputs, outputs=_all_outputs)

demo.launch(
    share=True,
    debug=True,
    theme=gr.themes.Soft(),
    css=CSS,
)