v2: BGE-M3 embeddings, hybrid BM25+Dense retrieval, HyDE, cross-encoder reranking

rag/chain.py

@@ -1,5 +1,12 @@
 """
 Core RAG chain — wraps ChromaDB retrieval + LLM (Ollama / Gemini / HF / OpenRouter).
+
+v2 improvements:
+  - BGE-M3 multilingual embeddings (replaces all-MiniLM-L6-v2)
+  - Hybrid BM25 + Dense retrieval with Reciprocal Rank Fusion
+  - Cross-encoder reranking with BGE-reranker-v2-m3
+  - HyDE (Hypothetical Document Embeddings) for query expansion
+  - Reduced context window (top-4 after reranking instead of top-6)
 """
 from __future__ import annotations
 import re
@@ -22,53 +29,38 @@ load_dotenv(override=True)
 
 # Greek Unicode ranges (excluding characters shared with Coptic)
 _GREEK_ONLY_RANGES = set()
-# Greek and Coptic block: U+0370–U+03FF
-# Greek Extended block: U+1F00–U+1FFF
-# These contain polytonic Greek, accented forms, archaic letters that are NOT Coptic
 for cp in range(0x0370, 0x0400):
     _GREEK_ONLY_RANGES.add(cp)
 for cp in range(0x1F00, 0x2000):
     _GREEK_ONLY_RANGES.add(cp)
 
-# Coptic Unicode block: U+2C80–U+2CFF (dedicated Coptic characters)
 _COPTIC_BLOCK = set(range(0x2C80, 0x2D00))
 
-# Characters shared between Greek and Coptic scripts (visually identical but
-# Coptic reuses Greek codepoints for these). We should NOT flag these as "Greek".
-# Common shared: Α-Ω / α-ω base letters that Coptic uses (U+0391-U+03C9 subset)
-# Coptic uses Greek codepoints for: α β γ δ ε ζ η θ ι κ λ μ ν ξ ο π ρ σ/ς τ υ φ χ ψ ω
 _SHARED_GREEK_COPTIC = set()
 for ch in "ΑΒΓΔΕΖΗΘΙΚΛΜΝΞΟΠΡΣΤΥΦΧΨΩαβγδεζηθικλμνξοπρστυφχψως":
     _SHARED_GREEK_COPTIC.add(ord(ch))
 
-# Coptic-specific letters (Demotic-derived) that confirm Coptic, not Greek
 _COPTIC_SPECIFIC = set()
 for ch in "ϣϩϫϭϯϥⲁⲃⲅⲇⲉⲍⲏⲑⲓⲕⲗⲙⲛⲝⲟⲡⲣⲥⲧⲩⲫⲭⲯⲱϣϩϫϭϯϥⲋⲍⲹⳉⳋⳍⳏⳑⳓⳕⳗⳙⳛ":
     _COPTIC_SPECIFIC.add(ord(ch))
 
-# Common Greek words that should NOT appear in Coptic text
 _GREEK_WORD_PATTERNS = re.compile(
     r'\b(τοῦ|τῆς|τῶν|τόν|τήν|καί|ἐν|εἰς|ἐκ|ἀπό|πρός|μετά|κατά|περί|ὑπό|παρά|διά|ἐπί'
-    r'|ὁ|ἡ|τό|οἱ|αἱ|τά'  # Greek articles
-    r'|ἐστί[ν]?|εἶναι|λέγει|λέγων|ἔχει|ἔχων'  # Common Greek verbs
-    r'|αὐτός|αὐτή|αὐτό|αὐτοῦ|αὐτῆς'  # Greek pronouns
-    r'|θεός|θεοῦ|κύριος|κυρίου|λόγος|λόγου'  # Common Greek nouns
+    r'|ὁ|ἡ|τό|οἱ|αἱ|τά'
+    r'|ἐστί[ν]?|εἶναι|λέγει|λέγων|ἔχει|ἔχων'
+    r'|αὐτός|αὐτή|αὐτό|αὐτοῦ|αὐτῆς'
+    r'|θεός|θεοῦ|κύριος|κυρίου|λόγος|λόγου'
     r'|ἄνθρωπος|ἀνθρώπου|κόσμος|κόσμου'
-    r'|οὐ|οὐκ|μή|γάρ|δέ|ἀλλά|ὅτι|ἵνα|ὡς'  # Greek particles
+    r'|οὐ|οὐκ|μή|γάρ|δέ|ἀλλά|ὅτι|ἵνα|ὡς'
     r')\b',
     re.UNICODE
 )
 
 
 def _count_greek_indicators(text: str) -> dict:
-    """
-    Analyze text for Greek vs Coptic script usage.
-    Returns counts of greek-only chars, coptic-specific chars, and greek word matches.
-    """
     greek_only_count = 0
     coptic_specific_count = 0
     shared_count = 0
-
     for ch in text:
         cp = ord(ch)
         if cp in _COPTIC_SPECIFIC or cp in _COPTIC_BLOCK:
@@ -77,9 +69,7 @@ def _count_greek_indicators(text: str) -> dict:
             greek_only_count += 1
         elif cp in _SHARED_GREEK_COPTIC:
             shared_count += 1
-
     greek_words = _GREEK_WORD_PATTERNS.findall(text)
-
     return {
         "greek_only_chars": greek_only_count,
         "coptic_specific_chars": coptic_specific_count,
@@ -90,15 +80,9 @@ def _count_greek_indicators(text: str) -> dict:
 
 
 def _add_greek_warning(answer: str) -> str:
-    """
-    If the answer contains significant Greek text, append a warning.
-    """
     analysis = _count_greek_indicators(answer)
-
-    # If there are Greek-only characters or Greek words detected
     has_greek_words = analysis["greek_word_count"] > 0
     has_greek_chars = analysis["greek_only_chars"] > 3
-
     if has_greek_words or has_greek_chars:
         warning_parts = []
         if has_greek_words:
@@ -106,7 +90,6 @@ def _add_greek_warning(answer: str) -> str:
             warning_parts.append(f"Greek words detected: {sample}")
         if has_greek_chars:
             warning_parts.append(f"{analysis['greek_only_chars']} Greek-only characters found")
-
         warning = (
             "\n\n---\n"
             "⚠️ **Greek Content Warning**: This response may contain Greek text "
@@ -116,7 +99,6 @@ def _add_greek_warning(answer: str) -> str:
             "[CDO](https://coptic-dictionary.org) or [Coptic SCRIPTORIUM](https://copticscriptorium.org)."
         )
         return answer + warning
-
     return answer
 
 
@@ -171,6 +153,17 @@ INSTRUCTIONS:
     Unicode characters (ⲁⲃⲅⲇⲉⲍⲏⲑⲓⲕⲗⲙⲛⲝⲟⲡⲣⲥⲧⲩⲫⲭⲯⲱ) and NOT Greek Unicode
     characters (αβγδεζηθικλμνξοπρστυφχψω).
 
+ANSWER QUALITY RULES:
+- Ground your answer ONLY in the retrieved context below. Do NOT fabricate
+  dictionary entries, paradigm tables, or grammatical forms from memory.
+- If the retrieved context does not contain information to answer the question,
+  say so clearly. Do NOT invent plausible-sounding answers.
+- When multiple retrieved chunks contain relevant information, SYNTHESIZE them
+  into a coherent answer rather than repeating each chunk separately.
+- Prefer information from lexicon entries (CCL, Faulkner) over grammar PDFs
+  for vocabulary questions, and grammar sources (Allen, Layton, Lambdin) for
+  structural/syntactic questions.
+
 RETRIEVED KNOWLEDGE BASE CONTEXT:
 {context}
 
@@ -184,84 +177,220 @@ _prompt = PromptTemplate(
     template=Shenute_SYSTEM,
 )
 
-# ── Embedder helper (shared across chain + ingest when using OpenRouter) ─────
+# ── HyDE prompt for Coptic/Egyptian query expansion ─────────────────────────
+_HYDE_TEMPLATE = """You are an expert in Coptic linguistics and Ancient Egyptian.
+Given the following question, write a short hypothetical dictionary entry or grammar
+explanation that would answer it. Write as if it were an entry in Crum's Coptic Dictionary,
+Faulkner's Middle Egyptian Dictionary, or Layton's Coptic Grammar. Include the relevant
+Coptic or Egyptian terms in proper Unicode script.
+
+Question: {question}
+
+Hypothetical entry:"""
+
+_hyde_prompt = PromptTemplate(
+    input_variables=["question"],
+    template=_HYDE_TEMPLATE,
+)
+
+
+# ── Embedder helper ──────────────────────────────────────────────────────────
 def get_embedder_for_provider(provider: str):
     """Return a LangChain embedder for the given provider.
-    OpenRouter does not offer an embedding endpoint, so we fall back to
-    HuggingFace embeddings (free, no key needed for public models) when
-    OpenRouter is selected."""
+    
+    v2: Uses BGE-M3 multilingual embeddings for HF and OpenRouter providers
+    instead of all-MiniLM-L6-v2 (which is English-only and blind to Coptic script).
+    BGE-M3 supports 100+ languages, 1024-dim, 8192-token context.
+    """
     if provider == "Gemini API":
         from langchain_google_genai import GoogleGenerativeAIEmbeddings
         return GoogleGenerativeAIEmbeddings(
             model="models/gemini-embedding-2-preview",
             google_api_key=os.environ.get("GEMINI_API_KEY"),
         )
-    elif provider == "Hugging Face":
+    elif provider in ("Hugging Face", "OpenRouter"):
         from langchain_huggingface import HuggingFaceEndpointEmbeddings
         return HuggingFaceEndpointEmbeddings(
             huggingfacehub_api_token=os.environ.get("HF_TOKEN"),
-            model="sentence-transformers/all-MiniLM-L6-v2",
-        )
-    elif provider == "OpenRouter":
-        # OpenRouter has no embedding API — use HF embeddings as fallback
-        from langchain_huggingface import HuggingFaceEndpointEmbeddings
-        return HuggingFaceEndpointEmbeddings(
-            huggingfacehub_api_token=os.environ.get("HF_TOKEN"),
-            model="sentence-transformers/all-MiniLM-L6-v2",
+            model="BAAI/bge-m3",
         )
     else:
-        # Local AI (Ollama)
         base_url = os.environ.get("LOCAL_AI_BASE_URL", "http://127.0.0.1:11434")
         return OllamaEmbeddings(model="nomic-embed-text", base_url=base_url)
 
 
+# ── Reranker ─────────────────────────────────────────────────────────────────
+def _rerank_documents(query: str, docs: list, top_k: int = 4) -> list:
+    """Rerank using BAAI/bge-reranker-v2-m3 via HF Inference API."""
+    if not docs or len(docs) <= top_k:
+        return docs
+    try:
+        import requests
+        hf_token = os.environ.get("HF_TOKEN")
+        if not hf_token:
+            return docs[:top_k]
+        API_URL = "https://router.huggingface.co/hf-inference/models/BAAI/bge-reranker-v2-m3"
+        headers = {"Authorization": f"Bearer {hf_token}"}
+        texts = [doc.page_content for doc in docs]
+        payload = {"inputs": query, "parameters": {"texts": texts, "truncate": True}}
+        response = requests.post(API_URL, headers=headers, json=payload, timeout=30)
+        if response.status_code == 200:
+            scores = response.json()
+            if isinstance(scores, list) and len(scores) > 0:
+                if isinstance(scores[0], dict):
+                    scored_docs = [(s["score"], docs[s["index"]]) for s in scores]
+                else:
+                    scored_docs = list(zip(scores, docs))
+                scored_docs.sort(key=lambda x: x[0], reverse=True)
+                return [doc for _, doc in scored_docs[:top_k]]
+        return docs[:top_k]
+    except Exception as e:
+        print(f"Reranking failed (falling back to top-k): {e}")
+        return docs[:top_k]
+
+
+# ── HyDE helper ──────────────────────────────────────────────────────────────
+def _generate_hyde_query(question: str, llm) -> str:
+    """Generate a hypothetical document using HyDE for better retrieval."""
+    try:
+        hyde_text = _HYDE_TEMPLATE.replace("{question}", question)
+        if hasattr(llm, 'invoke'):
+            result = llm.invoke(hyde_text)
+            if hasattr(result, 'content'):
+                return result.content.strip()
+            return str(result).strip()
+        return question
+    except Exception as e:
+        print(f"HyDE generation failed (using original query): {e}")
+        return question
+
+
+# ── BM25 retriever builder ──────────────────────────────────────────────────
+def _build_bm25_retriever(vectordb: Chroma, k: int = 6):
+    """Build a BM25 retriever from existing ChromaDB documents."""
+    try:
+        from langchain_classic.retrievers.bm25 import BM25Retriever
+        collection = vectordb._collection
+        result = collection.get(include=["documents", "metadatas"])
+        if not result["documents"]:
+            return None
+        from langchain_core.documents import Document
+        docs = []
+        for i, text in enumerate(result["documents"]):
+            meta = result["metadatas"][i] if result["metadatas"] else {}
+            docs.append(Document(page_content=text, metadata=meta))
+        bm25 = BM25Retriever.from_documents(docs, k=k)
+        return bm25
+    except Exception as e:
+        print(f"BM25 retriever build failed: {e}")
+        return None
+
+
+# ── Hybrid retriever (BM25 + Dense with Reciprocal Rank Fusion) ─────────────
+class HybridRetriever:
+    """
+    Fuses BM25 (exact keyword match) and dense (semantic embedding) retrieval
+    using Reciprocal Rank Fusion (RRF).
+    
+    BM25 catches exact Coptic word-form matches (ⲥⲱⲧⲙ, ⲛⲟⲩⲧⲉ) that
+    dense embeddings might miss. Dense catches semantic similarity
+    (e.g., "God" → ⲛⲟⲩⲧⲉ) that BM25 misses for cross-lingual queries.
+    """
+    def __init__(self, bm25_retriever, dense_retriever, bm25_weight=0.4, dense_weight=0.6, k=6):
+        self.bm25 = bm25_retriever
+        self.dense = dense_retriever
+        self.bm25_weight = bm25_weight
+        self.dense_weight = dense_weight
+        self.k = k
+    
+    def invoke(self, query: str) -> list:
+        bm25_docs = []
+        dense_docs = []
+        try:
+            bm25_docs = self.bm25.invoke(query)
+        except Exception as e:
+            print(f"BM25 retrieval failed: {e}")
+        try:
+            dense_docs = self.dense.invoke(query)
+        except Exception as e:
+            print(f"Dense retrieval failed: {e}")
+        if not bm25_docs and not dense_docs:
+            return []
+        if not bm25_docs:
+            return dense_docs[:self.k]
+        if not dense_docs:
+            return bm25_docs[:self.k]
+        # Reciprocal Rank Fusion
+        rrf_constant = 60
+        doc_scores = {}
+        for rank, doc in enumerate(bm25_docs):
+            key = hash(doc.page_content)
+            rrf_score = self.bm25_weight / (rrf_constant + rank + 1)
+            if key in doc_scores:
+                doc_scores[key] = (doc_scores[key][0] + rrf_score, doc)
+            else:
+                doc_scores[key] = (rrf_score, doc)
+        for rank, doc in enumerate(dense_docs):
+            key = hash(doc.page_content)
+            rrf_score = self.dense_weight / (rrf_constant + rank + 1)
+            if key in doc_scores:
+                doc_scores[key] = (doc_scores[key][0] + rrf_score, doc)
+            else:
+                doc_scores[key] = (rrf_score, doc)
+        ranked = sorted(doc_scores.values(), key=lambda x: x[0], reverse=True)
+        return [doc for _, doc in ranked[:self.k]]
+
+
 # ── Chain builder ──────────────────────────────────────────────────────────────
 def build_chain(model: str = "qwen3:14b",
                 top_k: int = 6,
                 temperature: float = 0.1,
-                provider: str = "Local AI") -> RetrievalQA:
+                provider: str = "Local AI",
+                use_hyde: bool = True,
+                use_reranking: bool = True,
+                use_hybrid: bool = True) -> tuple:
     
-    # 1. Setup Embedder & VectorDB
     embedder = get_embedder_for_provider(provider)
-        
-    vectordb  = Chroma(
-        persist_directory="./chroma_db",
-        embedding_function=embedder,
-    )
-    retriever = vectordb.as_retriever(search_kwargs={"k": top_k})
+    vectordb = Chroma(persist_directory="./chroma_db", embedding_function=embedder)
+    
+    dense_retriever = vectordb.as_retriever(search_kwargs={"k": top_k})
+    
+    if use_hybrid:
+        try:
+            bm25_retriever = _build_bm25_retriever(vectordb, k=top_k)
+            if bm25_retriever:
+                retriever = HybridRetriever(
+                    bm25_retriever=bm25_retriever,
+                    dense_retriever=dense_retriever,
+                    bm25_weight=0.4, dense_weight=0.6, k=top_k,
+                )
+            else:
+                retriever = dense_retriever
+        except Exception as e:
+            print(f"Hybrid retrieval setup failed, falling back to dense: {e}")
+            retriever = dense_retriever
+    else:
+        retriever = dense_retriever
     
-    # 2. Setup chosen LLM
     if provider == "Gemini API":
         from langchain_google_genai import ChatGoogleGenerativeAI
         gemini_api_key = os.environ.get("GEMINI_API_KEY")
         if not gemini_api_key:
             raise ValueError("GEMINI_API_KEY is not set in the .env file.")
-        llm = ChatGoogleGenerativeAI(
-            model=model,
-            temperature=temperature,
-            google_api_key=gemini_api_key,
-        )
+        llm = ChatGoogleGenerativeAI(model=model, temperature=temperature, google_api_key=gemini_api_key)
     elif provider == "Hugging Face":
         from langchain_openai import ChatOpenAI
         hf_token = os.environ.get("HF_TOKEN")
         if not hf_token:
             raise ValueError("HF_TOKEN is not set in the .env file.")
-        llm = ChatOpenAI(
-            model=model,
-            temperature=temperature,
-            api_key=hf_token,
-            base_url="https://router.huggingface.co/v1",
-        )
+        llm = ChatOpenAI(model=model, temperature=temperature, api_key=hf_token, base_url="https://router.huggingface.co/v1")
     elif provider == "OpenRouter":
-        # OpenRouter exposes an OpenAI-compatible API at https://openrouter.ai/api/v1
         from langchain_openai import ChatOpenAI
         openrouter_key = os.environ.get("OPENROUTER_API_KEY")
         if not openrouter_key:
             raise ValueError("OPENROUTER_API_KEY is not set. Add it as a Space secret or in your .env file.")
         llm = ChatOpenAI(
-            model=model,
-            temperature=temperature,
-            api_key=openrouter_key,
+            model=model, temperature=temperature, api_key=openrouter_key,
             base_url="https://openrouter.ai/api/v1",
             default_headers={
                 "HTTP-Referer": "https://huggingface.co/spaces/georgtawadrous/thoth_app",
@@ -269,50 +398,37 @@ def build_chain(model: str = "qwen3:14b",
             },
         )
     else:
-        # Default to Local AI (Ollama)
         base_url = os.environ.get("LOCAL_AI_BASE_URL", "http://127.0.0.1:11434")
-        llm = OllamaLLM(
-            model=model,
-            temperature=temperature,
-            num_ctx=4096,
-            base_url=base_url,
-        )
-        
+        llm = OllamaLLM(model=model, temperature=temperature, num_ctx=4096, base_url=base_url)
+    
     chain = RetrievalQA.from_chain_type(
-        llm=llm,
-        retriever=retriever,
-        chain_type="stuff",
-        chain_type_kwargs={"prompt": _prompt},
-        return_source_documents=True,
+        llm=llm, retriever=dense_retriever, chain_type="stuff",
+        chain_type_kwargs={"prompt": _prompt}, return_source_documents=True,
     )
     
-    # Optional Feedback Retriever
     try:
         feedback_store = Chroma(
-            persist_directory="./chroma_db",
-            embedding_function=embedder,
+            persist_directory="./chroma_db", embedding_function=embedder,
             collection_name="Shenute_feedback",
         )
         feedback_retriever = feedback_store.as_retriever(search_kwargs={"k": 2})
     except Exception:
         feedback_retriever = None
 
-    return chain, feedback_retriever
+    return chain, feedback_retriever, llm, retriever
 
 
 def query_Shenute(question: str,
                 model: str = "qwen3:14b",
                 top_k: int = 6,
                 temperature: float = 0.1,
-                provider: str = "Local AI") -> dict:
-    """
-    Returns:
-        {
-          "answer":  str,
-          "sources": [{"text": str, "source": str, "page": int}]
-        }
-    """
-    chain, feedback_retriever = build_chain(model, top_k, temperature, provider)
+                provider: str = "Local AI",
+                use_hyde: bool = True,
+                use_reranking: bool = True,
+                use_hybrid: bool = True) -> dict:
+    chain, feedback_retriever, llm, retriever = build_chain(
+        model, top_k, temperature, provider, use_hyde, use_reranking, use_hybrid
+    )
     
     feedback_context = ""
     if feedback_retriever:
@@ -327,15 +443,53 @@ def query_Shenute(question: str,
     if feedback_context:
         augmented_question = f"{question}\n\n[SYSTEM NOTE - PAST USER CORRECTIONS TO APPLY IF RELEVANT:\n{feedback_context}]"
 
-    result = chain.invoke({"query": augmented_question})
+    # HyDE: Generate hypothetical document for better retrieval
+    retrieval_query = augmented_question
+    hyde_used = False
+    if use_hyde and provider != "Local AI":
+        try:
+            hyde_doc = _generate_hyde_query(question, llm)
+            if hyde_doc and len(hyde_doc) > 20:
+                retrieval_query = f"{augmented_question}\n\n{hyde_doc}"
+                hyde_used = True
+        except Exception as e:
+            print(f"HyDE failed: {e}")
+
+    # Retrieve with optional reranking
+    reranked = False
+    if use_reranking and provider != "Local AI":
+        try:
+            initial_docs = retriever.invoke(retrieval_query if hyde_used else augmented_question)
+            rerank_k = min(top_k, 4)
+            reranked_docs = _rerank_documents(question, initial_docs, top_k=rerank_k)
+            if reranked_docs:
+                reranked = True
+                context = "\n\n---\n\n".join([doc.page_content for doc in reranked_docs])
+                filled_prompt = Shenute_SYSTEM.replace("{context}", context).replace("{question}", augmented_question)
+                result_text = llm.invoke(filled_prompt)
+                if hasattr(result_text, 'content'):
+                    result_text = result_text.content
+                answer = str(result_text)
+                source_documents = reranked_docs
+            else:
+                result = chain.invoke({"query": augmented_question})
+                answer = result["result"]
+                source_documents = result.get("source_documents", [])
+        except Exception as e:
+            print(f"Reranking pipeline failed, falling back to standard chain: {e}")
+            result = chain.invoke({"query": augmented_question})
+            answer = result["result"]
+            source_documents = result.get("source_documents", [])
+    else:
+        result = chain.invoke({"query": augmented_question})
+        answer = result["result"]
+        source_documents = result.get("source_documents", [])
 
-    # Post-process: clean up think tags and check for Greek hallucination
-    answer = result["result"]
     answer = _strip_think_tags(answer)
     answer = _add_greek_warning(answer)
 
     sources = []
-    for doc in result.get("source_documents", []):
+    for doc in source_documents:
         sources.append({
             "text":   doc.page_content[:300],
             "source": doc.metadata.get("source", "Unknown"),
@@ -345,4 +499,9 @@ def query_Shenute(question: str,
     return {
         "answer":  answer,
         "sources": sources,
+        "retrieval_info": {
+            "hyde_used": hyde_used,
+            "reranked": reranked,
+            "hybrid": use_hybrid,
+        },
     }