Spaces:

Lvvbaaal
/

legalgpt-api

Sleeping

legalgpt-api / backend /layer2_retrieval.py

Deploy SaulGPT backend with full ChromaDB, Supabase, and Gradio API server

d1be42f verified 12 days ago

23.7 kB

"""
SAULGPT — LAYER 2: HYBRID RETRIEVAL
======================================
Techniques:
1. Semantic Search      — ChromaDB (all-MiniLM-L6-v2)
2. BM25 Keyword Search  — rank_bm25
3. Metadata Citation    — exact section_number filter
4. CrossEncoder         — ms-marco-MiniLM reranker
5. Deduplication        — content hash merge
Input:  UnderstandingResult dict from Layer 1
Output: Top k most relevant law section dicts
Run standalone to test:
    python layer2_retrieval.py
When used by pipeline_orchestrator.py:
    from layer2_retrieval import retrieve_with_hybrid_logic
    results = retrieve_with_hybrid_logic(layer1_payload)
IMPORTANT — matches 02_chunk_and_embed.py exactly:
    Embedding model : all-MiniLM-L6-v2
    Vector DB path  : vector_db
    Collection name : saulgpt_indian_laws
    Metadata keys   : act_name, section_number,
                      is_repealed, law_type
"""

import os
import re

from rank_bm25 import BM25Okapi

# ─────────────────────────────────────────────────────────────
# CONFIG — must match 02_chunk_and_embed.py exactly
# Changing any of these will cause zero results
# ─────────────────────────────────────────────────────────────

# Same model used in chunk_and_embed
EMBEDDING_MODEL  = "all-MiniLM-L6-v2"

# Same absolute path logic used in 02_chunk_and_embed.py
BASE_DIR         = os.path.dirname(os.path.abspath(__file__))
REPO_DIR         = os.path.dirname(BASE_DIR)
VECTOR_DB_PATH = os.path.join(REPO_DIR, "data", "vector_db")

# Same collection name used in chunk_and_embed
COLLECTION_NAME  = "saulgpt_indian_laws"

# CrossEncoder reranker — scores results by true relevance
RERANKER_MODEL   = "cross-encoder/ms-marco-MiniLM-L-6-v2"

# Number of results to return after reranking
TOP_K            = 3

# ─────────────────────────────────────────────────────────────
# LAZY LOADING — heavy imports deferred until first query
# SentenceTransformer (~18s) + ChromaDB (~2s) loaded on demand
# ─────────────────────────────────────────────────────────────

_models_loaded = False
embedder = None
chroma_client = None
collection = None
reranker = None
_cached_bm25 = None
_cached_docs = None
_cached_metas = None


def _ensure_models():
    """Load SentenceTransformer + ChromaDB on first call only, with graceful fallbacks."""
    global _models_loaded, embedder, chroma_client, collection, reranker
    if _models_loaded:
        return

    print("Loading retrieval models (first query)...")
    try:
        import chromadb
        from chromadb.config import Settings
        chroma_client = chromadb.PersistentClient(
            path=VECTOR_DB_PATH,
            settings=Settings(anonymized_telemetry=False)
        )
        collection = chroma_client.get_or_create_collection(COLLECTION_NAME)
        print(f"ChromaDB loaded. Collection: {COLLECTION_NAME}")
        print(f"Total chunks in DB: {collection.count()}\n")
    except Exception as e:
        print(f"[Layer 2 Warning] ChromaDB load failed or unavailable: {e}")
        collection = None

    try:
        from sentence_transformers import SentenceTransformer
        embedder = SentenceTransformer(EMBEDDING_MODEL, device="cpu")
    except Exception as e:
        print(f"[Layer 2 Warning] SentenceTransformer load failed: {e}")
        embedder = None

    try:
        from sentence_transformers import CrossEncoder
        reranker = CrossEncoder(RERANKER_MODEL, device="cpu")
    except Exception as e:
        print(f"[Layer 2 Warning] CrossEncoder load skipped: {e}")
        reranker = None

    _models_loaded = True


# ─────────────────────────────────────────────────────────────
# BM25 INDEX BUILDER
# Loads all documents from ChromaDB and builds BM25 index (cached)
# ─────────────────────────────────────────────────────────────

def build_bm25_index(act_name_filter: str = None):
    """
    Builds BM25 keyword index from ChromaDB documents.
    Caches global index to avoid expensive re-indexing and RAM churn.
    """
    global _cached_bm25, _cached_docs, _cached_metas

    if not act_name_filter and _cached_bm25 is not None:
        return _cached_bm25, _cached_docs, _cached_metas

    if collection is None:
        return None, [], []

    try:
        if act_name_filter:
            result = collection.get(where={"act_name": act_name_filter})
        else:
            result = collection.get()
    except Exception as e:
        print(f"[BM25 Warning] collection.get failed: {e}")
        return None, [], []

    documents = result.get("documents", [])
    metadatas = result.get("metadatas", [])

    if not documents:
        return None, [], []

    tokenized_corpus = [doc.lower().split() for doc in documents]
    bm25 = BM25Okapi(tokenized_corpus)

    if not act_name_filter:
        _cached_bm25 = bm25
        _cached_docs = documents
        _cached_metas = metadatas

    return bm25, documents, metadatas


# ─────────────────────────────────────────────────────────────
# MAIN RETRIEVAL FUNCTION
# ─────────────────────────────────────────────────────────────

def retrieve_with_hybrid_logic(payload: dict, k: int = TOP_K) -> list:
    """
    Hybrid retrieval: semantic + BM25 + reranking.
    Takes Layer 1 output payload and returns
    top k most relevant law sections.
    Args:
        payload: dict from layer1_understanding.analyze_query()
                 must contain:
                 - search_optimized_query (str)
                 - explicit_citations (list)
                 - domain (str) [optional]
        k: number of results to return (default 8)
    Returns:
        list of result dicts each containing:
        content, act_name, section_number,
        is_repealed, law_type, relevance_score
    Example:
        results = retrieve_with_hybrid_logic({
            "search_optimized_query": "unpaid wages employer",
            "explicit_citations": ["SECTION 15"],
            "domain": "labour"
        })
    """

    _ensure_models()

    def _safe_top_k(value, default=TOP_K):
        """Guarantee a positive integer for ChromaDB n_results."""
        try:
            value = int(value)
        except (TypeError, ValueError):
            value = default
        return value if value > 0 else default

    def _has_word(w, text):
        return bool(re.search(rf"\b{re.escape(w)}\b", str(text), re.IGNORECASE))

    # Extract fields from Layer 1 payload
    semantic_query = payload.get("hyde_paragraph") or payload.get("semantic_query") or payload.get("search_optimized_query", "")
    keyword_query  = payload.get("keyword_synonyms") or payload.get("keywords") or payload.get("search_optimized_query", "")
    citations = payload.get("explicit_citations", [])
    top_k     = _safe_top_k(payload.get("top_k", k))
    temporal_status = payload.get("temporal_context", {}).get("temporal_status") or payload.get("temporal_status", "UNDATED")

    print(f"\n[Layer 2] Semantic  : {semantic_query[:70]}...")
    print(f"[Layer 2] Keywords  : {keyword_query[:70]}...")
    print(f"[Layer 2] Citations : {citations}")
    print(f"[Layer 2] Temporal  : {temporal_status}")
    print(f"[Layer 2] Top K     : {top_k}")

    # ── TECHNIQUE 1: Semantic Search via ChromaDB ──
    # Embeds query (uses HyDE paragraph) and finds nearest chunks by cosine similarity
    semantic_results_raw = []
    if embedder is not None and collection is not None:
        print("\n[1/3] Running semantic search (ChromaDB)...")
        try:
            query_embedding = embedder.encode(semantic_query).tolist()
            collection_count = collection.count()
        except Exception as e:
            print(f"[Layer 2 Warning] Embedding generation failed: {e}")
            query_embedding = None
            collection_count = 0
    else:
        print("\n[1/3] Semantic search skipped (fast keyword / low-memory mode).")
        query_embedding = None
        collection_count = 0
    target_sec_nums = set()
    if citations:
        for c in citations:
            m = re.search(r'\d+[A-Za-z]*', str(c))
            if m:
                target_sec_nums.add(m.group())

    # Detect implied statutory sections from core legal topics
    q_all = (semantic_query + " " + keyword_query).lower()
    if any(_has_word(w, q_all) for w in ["cheat", "cheated", "cheating", "fraud", "scam"]):
        target_sec_nums.update(["318", "420"])
    if any(_has_word(w, q_all) for w in ["fir", "zero fir", "e-fir"]):
        target_sec_nums.update(["173", "154"])
    if any(_has_word(w, q_all) for w in ["certificate", "electronic evidence"]):
        target_sec_nums.update(["63", "65B"])
    if any(_has_word(w, q_all) for w in ["community service"]):
        target_sec_nums.update(["4", "23"])

    if collection and target_sec_nums:
        print(f"[Layer 2] Citation/Topic filter: section_numbers = {list(target_sec_nums)}")
        for s_num in target_sec_nums:
            try:
                c_chunks = collection.get(where={"section_number": str(s_num)})
                if c_chunks and c_chunks.get("documents"):
                    for doc, meta in zip(c_chunks["documents"], c_chunks["metadatas"]):
                        semantic_results_raw.append((doc, meta))
            except Exception as e:
                pass

    # Always run broad semantic search too
    # Combines with citation results for broad coverage
    if collection and query_embedding and collection_count > 0:
        try:
            broad = collection.query(
                query_embeddings=[query_embedding],
                n_results=min(max(top_k * 2, 8), collection_count)
            )
            for doc, meta in zip(broad["documents"][0], broad["metadatas"][0]):
                semantic_results_raw.append((doc, meta))

            # ── PARTITION A: ACTIVE 2024 SANHITA RETRIEVAL (BNS / BNSS / BSA / Constitution) ──
            try:
                sanhita_broad = collection.query(
                    query_embeddings=[query_embedding],
                    where={"is_active_2024": True},
                    n_results=min(max(top_k, 5), collection_count)
                )
                if sanhita_broad and sanhita_broad.get("documents") and sanhita_broad["documents"][0]:
                    for doc, meta in zip(sanhita_broad["documents"][0], sanhita_broad["metadatas"][0]):
                        semantic_results_raw.append((doc, meta))
                    print(f"[Layer 2] Partition A (2024 Sanhitas): {len(sanhita_broad['documents'][0])} chunks retrieved.")
            except Exception as e_sanhita:
                print(f"[Layer 2] Partition A query error: {e_sanhita}")

            # ── PARTITION B: LEGACY COLONIAL CODE RETRIEVAL (IPC / CrPC / IEA) FOR UNDATED DUAL-TRACK ──
            if temporal_status == "UNDATED":
                try:
                    legacy_broad = collection.query(
                        query_embeddings=[query_embedding],
                        where={"is_repealed": True},
                        n_results=min(max(top_k, 5), collection_count)
                    )
                    if legacy_broad and legacy_broad.get("documents") and legacy_broad["documents"][0]:
                        for doc, meta in zip(legacy_broad["documents"][0], legacy_broad["metadatas"][0]):
                            semantic_results_raw.append((doc, meta))
                        print(f"[Layer 2] Partition B (Legacy Colonial): {len(legacy_broad['documents'][0])} chunks retrieved.")
                except Exception as e_legacy:
                    print(f"[Layer 2] Partition B query error: {e_legacy}")
        except Exception as e:
            print(f"[Layer 2] ChromaDB query failed: {e}")

    print(f"[Layer 2] Semantic raw results: {len(semantic_results_raw)}")

    # ── TECHNIQUE 2: BM25 Keyword Search ──
    print("[2/3] Running BM25 keyword search...")

    bm25, bm25_docs, bm25_metas = build_bm25_index()
    bm25_results_raw = []

    if bm25 and bm25_docs:
        tokenized_query = keyword_query.lower().split()
        scores = bm25.get_scores(tokenized_query)

        top_indices = sorted(
            range(len(scores)),
            key=lambda i: scores[i],
            reverse=True
        )[:max(top_k * 4, 25)]

        for idx in top_indices:
            if scores[idx] > 0:
                bm25_results_raw.append((
                    bm25_docs[idx],
                    bm25_metas[idx],
                    scores[idx]
                ))

    print(f"[Layer 2] BM25 raw results: {len(bm25_results_raw)}")

    # ── COMBINE AND DEDUPLICATE ──
    combined = {}

    for doc, meta in semantic_results_raw:
        if doc not in combined:
            act_n = meta.get("act_name") or meta.get("act_short") or meta.get("act") or ""
            sec_n = meta.get("section_number") or meta.get("section") or ""
            combined[doc] = {
                "content":        doc,
                "act_name":       act_n,
                "section_number": sec_n,
                "is_repealed":    meta.get("is_repealed", False),
                "law_type":       meta.get("law_type", "Statute"),
            }

    for doc, meta, score in bm25_results_raw:
        if doc not in combined:
            act_n = meta.get("act_name") or meta.get("act_short") or meta.get("act") or ""
            sec_n = meta.get("section_number") or meta.get("section") or ""
            combined[doc] = {
                "content":        doc,
                "act_name":       act_n,
                "section_number": sec_n,
                "is_repealed":    meta.get("is_repealed", False),
                "law_type":       meta.get("law_type", "Statute"),
                "bm25_score":     score,
            }

    all_results = list(combined.values())
    print(f"[Layer 2] Combined unique results: {len(all_results)}")

    if not all_results:
        print("[Layer 2] WARNING: No results found in DB.")
        return []

    # ── TECHNIQUE 3: CrossEncoder Reranker with Exact Citation Boosting ──
    if reranker is not None:
        print("[3/3] Reranking with CrossEncoder...")
        pairs  = [(semantic_query, r["content"]) for r in all_results]
        try:
            scores = reranker.predict(pairs)
        except Exception as e:
            print(f"[Layer 2 Warning] CrossEncoder inference failed ({e}), using default scores.")
            scores = [0.8] * len(all_results)
    else:
        print("[3/3] CrossEncoder skipped (low-memory / fast mode). Using relevance scoring.")
        scores = [0.8] * len(all_results)

    # Detect primary target act mentions in query to give strong affinity
    q_lower = (semantic_query + " " + keyword_query).lower()

    matched_target_acts = set()
    if any(_has_word(w, q_lower) for w in ["negotiable", "cheque", "nia", "promissory"]):
        matched_target_acts.add("nia")
    if any(_has_word(w, q_lower) for w in ["fir", "crpc", "bnss", "bail", "anticipatory", "cognizable", "complaint", "police"]):
        matched_target_acts.update(["crpc", "bnss"])
    if any(_has_word(w, q_lower) for w in ["marriage", "divorce", "hindu marriage", "hma", "restitution"]):
        matched_target_acts.add("hma")
    if any(_has_word(w, q_lower) for w in ["cpc", "civil procedure", "written statement", "injunction"]):
        matched_target_acts.add("cpc")
    if any(_has_word(w, q_lower) for w in ["motor", "mva", "vehicle", "traffic"]):
        matched_target_acts.add("mva")
    if any(_has_word(w, q_lower) for w in ["evidence", "iea", "bsa", "witness", "certificate"]):
        matched_target_acts.update(["iea", "bsa"])
    if any(_has_word(w, q_lower) for w in ["industrial", "workman", "ida", "retrenchment", "layoff"]):
        matched_target_acts.add("ida")
    if any(_has_word(w, q_lower) for w in [
        "ipc", "penal code", "bns", "cheated", "cheating", "fraud", "scam",
        "deception", "theft", "murder", "extortion", "dishonest", "misappropriation",
        "breach of trust", "community service", "420", "318"
    ]):
        matched_target_acts.update(["ipc", "bns"])
    if any(_has_word(w, q_lower) for w in ["constitution", "article", "fundamental right"]):
        matched_target_acts.add("constitution")

    is_modern_query = any(w in q_lower for w in ["today", "new law", "new", "2024", "bns", "bnss", "bsa"])

    boosted_ranked = []
    for score, r in zip(scores, all_results):
        adjusted_score = float(score)
        sec = str(r.get("section_number", "")).strip()
        act = str(r.get("act_name", "")).lower()
        content = r.get("content", "")

        # 1. Exact section citation boost
        if target_sec_nums and (sec in target_sec_nums or any(s in sec for s in target_sec_nums)):
            adjusted_score += 15.0
            # Prefer substantive chunks over short 1-line amendment notes
            if len(content) > 120:
                adjusted_score += 8.0

        # 2. Target act affinity boost / penalty
        if matched_target_acts:
            act_tokens = set(re.split(r'[^a-z0-9]+', act))
            if any(target in act_tokens or target == act for target in matched_target_acts):
                adjusted_score += 25.0
            else:
                adjusted_score -= 15.0  # Penalize chunks from completely unrelated acts

        # 3. Modern 2024 Sanhita preference for queries with modern markers
        if is_modern_query and any(s in act for s in ["bns", "bnss", "bsa"]):
            adjusted_score += 10.0

        boosted_ranked.append((adjusted_score, r))

    # Sort by adjusted score descending
    ranked = sorted(boosted_ranked, key=lambda x: x[0], reverse=True)

    # Build final top k results (default min 5 for complete legal context)
    effective_k = max(top_k, 5)

    if temporal_status == "UNDATED":
        def _is_legacy_chunk(item):
            a = (item.get("act_name") or item.get("act") or "").lower()
            return bool(item.get("is_repealed")) or "from_db" in a or a in ["ipc", "crpc", "iea"]

        def _is_active_sanhita(item):
            a = (item.get("act_name") or item.get("act") or "").lower()
            return bool(item.get("is_active_2024")) or a in ["bns", "bnss", "bsa", "constitution"]

        active_candidates = [r for score, r in ranked if _is_active_sanhita(r)]
        legacy_candidates = [r for score, r in ranked if _is_legacy_chunk(r)]

        balanced = []
        # Add up to 3 active 2024 chunks
        for c in active_candidates[:3]:
            if c not in balanced:
                balanced.append(c)
        # Add up to 3 legacy colonial chunks
        for c in legacy_candidates[:3]:
            if c not in balanced:
                balanced.append(c)
        # Fill remaining slots up to effective_k from highest ranked
        for score, r in ranked:
            if len(balanced) >= max(effective_k, 6):
                break
            if r not in balanced:
                balanced.append(r)
        final_list = balanced
    else:
        final_list = [r for score, r in ranked[:effective_k]]

    top_results = []
    for result in final_list:
        top_results.append({
            "content":         result["content"],
            "act_name":        result["act_name"],
            "section_number":  result["section_number"],
            "is_repealed":     result["is_repealed"],
            "law_type":        result["law_type"],
            "relevance_score": round(float(result.get("relevance_score", 0.9)), 4),
        })

    # ── TECHNIQUE 6: Shared Precedent Engine Integration (50,000+ Case Matrix) ──
    try:
        try:
            from backend.trial_studio.precedent_engine import get_precedent_engine
        except ImportError:
            from trial_studio.precedent_engine import get_precedent_engine
        pe = get_precedent_engine()
        precedent_matches = pe.search(query=semantic_query, top_k=2)
        for pm in precedent_matches:
            top_results.append({
                "content": f"[Ratio Decidendi]: {pm.ratio_decidendi} [Key Fact Anchor]: {pm.fact_anchor}",
                "act_name": f"{pm.court} Precedent — {pm.citation} ({pm.title})",
                "section_number": ", ".join(pm.sections_invoked),
                "is_repealed": False,
                "law_type": "CONTROLLING PRECEDENT RATIO",
                "relevance_score": 0.95,
            })
    except Exception as e:
        print(f"[Layer 2] Note: Precedent engine optional link skipped: {e}")

    print(f"[Layer 2] Final top {len(top_results)} sections & precedents ready for Layer 3 (Dual-Track balanced: {temporal_status == 'UNDATED'}).\n")
    return top_results


# ─────────────────────────────────────────────────────────────
# TEST RUNNER
# Only runs when executed directly:
#   python layer2_retrieval.py
#
# Uses a simulated Layer 1 payload
# In production this payload comes from analyze_query()
# Test query covers labour law + explicit section citation
# which exercises all 3 retrieval techniques at once
# ─────────────────────────────────────────────────────────────

if __name__ == "__main__":

    # Simulated Layer 1 output
    # In production: comes from layer1_understanding.analyze_query()
    test_payload = {
        "search_optimized_query": (
            "What are the legal conditions for a valid marriage regarding age and living spouses "
            "under Section 5 of the Hindu Marriage Act?"
        ),
        "domain":             "family",
        "explicit_citations": ["SECTION 5"],
        "is_ambiguous":       False,
    }

    print("-" * 55)
    print("SaulGPT — Layer 2 Hybrid Retrieval Test")
    print("-" * 55)

    results = retrieve_with_hybrid_logic(test_payload)

    print("\n" + "=" * 55)
    print("TOP RETRIEVED LEGAL SECTIONS")
    print("=" * 55)

    if not results:
        print("No results returned.")
        print("Troubleshooting:")
        print("  1. Run 02_chunk_and_embed.py first")
        print("  2. Check VECTOR_DB_PATH = 'vector_db'")
        print("  3. Check COLLECTION_NAME = 'saulgpt_indian_laws'")
    else:
        for i, result in enumerate(results, 1):
            print(f"\nRESULT {i}:")
            print(f"  ACT     : {result['act_name']}")
            print(f"  SECTION : {result['section_number']}")
            print(f"  SCORE   : {result['relevance_score']}")
            print(f"  REPEALED: {result['is_repealed']}")
            print(f"  CONTENT : {result['content'][:200]}...")

    print("\n" + "=" * 55)
    print("Layer 2 complete. Results ready for Layer 3.")
    print("=" * 55)

	"""
	SAULGPT — LAYER 2: HYBRID RETRIEVAL
	======================================
	Techniques:
	1. Semantic Search — ChromaDB (all-MiniLM-L6-v2)
	2. BM25 Keyword Search — rank_bm25
	3. Metadata Citation — exact section_number filter
	4. CrossEncoder — ms-marco-MiniLM reranker
	5. Deduplication — content hash merge

	Input: UnderstandingResult dict from Layer 1
	Output: Top k most relevant law section dicts

	Run standalone to test:
	python layer2_retrieval.py

	When used by pipeline_orchestrator.py:
	from layer2_retrieval import retrieve_with_hybrid_logic
	results = retrieve_with_hybrid_logic(layer1_payload)

	IMPORTANT — matches 02_chunk_and_embed.py exactly:
	Embedding model : all-MiniLM-L6-v2
	Vector DB path : vector_db
	Collection name : saulgpt_indian_laws
	Metadata keys : act_name, section_number,
	is_repealed, law_type
	"""

	import os
	import re

	from rank_bm25 import BM25Okapi

	# ─────────────────────────────────────────────────────────────
	# CONFIG — must match 02_chunk_and_embed.py exactly
	# Changing any of these will cause zero results
	# ─────────────────────────────────────────────────────────────

	# Same model used in chunk_and_embed
	EMBEDDING_MODEL = "all-MiniLM-L6-v2"

	# Same absolute path logic used in 02_chunk_and_embed.py
	BASE_DIR = os.path.dirname(os.path.abspath(__file__))
	REPO_DIR = os.path.dirname(BASE_DIR)
	VECTOR_DB_PATH = os.path.join(REPO_DIR, "data", "vector_db")

	# Same collection name used in chunk_and_embed
	COLLECTION_NAME = "saulgpt_indian_laws"

	# CrossEncoder reranker — scores results by true relevance
	RERANKER_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2"

	# Number of results to return after reranking
	TOP_K = 3

	# ─────────────────────────────────────────────────────────────
	# LAZY LOADING — heavy imports deferred until first query
	# SentenceTransformer (~18s) + ChromaDB (~2s) loaded on demand
	# ─────────────────────────────────────────────────────────────

	_models_loaded = False
	embedder = None
	chroma_client = None
	collection = None
	reranker = None
	_cached_bm25 = None
	_cached_docs = None
	_cached_metas = None


	def _ensure_models():
	"""Load SentenceTransformer + ChromaDB on first call only, with graceful fallbacks."""
	global _models_loaded, embedder, chroma_client, collection, reranker
	if _models_loaded:
	return

	print("Loading retrieval models (first query)...")
	try:
	import chromadb
	from chromadb.config import Settings
	chroma_client = chromadb.PersistentClient(
	path=VECTOR_DB_PATH,
	settings=Settings(anonymized_telemetry=False)
	)
	collection = chroma_client.get_or_create_collection(COLLECTION_NAME)
	print(f"ChromaDB loaded. Collection: {COLLECTION_NAME}")
	print(f"Total chunks in DB: {collection.count()}\n")
	except Exception as e:
	print(f"[Layer 2 Warning] ChromaDB load failed or unavailable: {e}")
	collection = None

	try:
	from sentence_transformers import SentenceTransformer
	embedder = SentenceTransformer(EMBEDDING_MODEL, device="cpu")
	except Exception as e:
	print(f"[Layer 2 Warning] SentenceTransformer load failed: {e}")
	embedder = None

	try:
	from sentence_transformers import CrossEncoder
	reranker = CrossEncoder(RERANKER_MODEL, device="cpu")
	except Exception as e:
	print(f"[Layer 2 Warning] CrossEncoder load skipped: {e}")
	reranker = None

	_models_loaded = True


	# ─────────────────────────────────────────────────────────────
	# BM25 INDEX BUILDER
	# Loads all documents from ChromaDB and builds BM25 index (cached)
	# ─────────────────────────────────────────────────────────────

	def build_bm25_index(act_name_filter: str = None):
	"""
	Builds BM25 keyword index from ChromaDB documents.
	Caches global index to avoid expensive re-indexing and RAM churn.
	"""
	global _cached_bm25, _cached_docs, _cached_metas

	if not act_name_filter and _cached_bm25 is not None:
	return _cached_bm25, _cached_docs, _cached_metas

	if collection is None:
	return None, [], []

	try:
	if act_name_filter:
	result = collection.get(where={"act_name": act_name_filter})
	else:
	result = collection.get()
	except Exception as e:
	print(f"[BM25 Warning] collection.get failed: {e}")
	return None, [], []

	documents = result.get("documents", [])
	metadatas = result.get("metadatas", [])

	if not documents:
	return None, [], []

	tokenized_corpus = [doc.lower().split() for doc in documents]
	bm25 = BM25Okapi(tokenized_corpus)

	if not act_name_filter:
	_cached_bm25 = bm25
	_cached_docs = documents
	_cached_metas = metadatas

	return bm25, documents, metadatas


	# ─────────────────────────────────────────────────────────────
	# MAIN RETRIEVAL FUNCTION
	# ─────────────────────────────────────────────────────────────

	def retrieve_with_hybrid_logic(payload: dict, k: int = TOP_K) -> list:
	"""
	Hybrid retrieval: semantic + BM25 + reranking.

	Takes Layer 1 output payload and returns
	top k most relevant law sections.

	Args:
	payload: dict from layer1_understanding.analyze_query()
	must contain:
	- search_optimized_query (str)
	- explicit_citations (list)
	- domain (str) [optional]

	k: number of results to return (default 8)

	Returns:
	list of result dicts each containing:
	content, act_name, section_number,
	is_repealed, law_type, relevance_score

	Example:
	results = retrieve_with_hybrid_logic({
	"search_optimized_query": "unpaid wages employer",
	"explicit_citations": ["SECTION 15"],
	"domain": "labour"
	})
	"""

	_ensure_models()

	def _safe_top_k(value, default=TOP_K):
	"""Guarantee a positive integer for ChromaDB n_results."""
	try:
	value = int(value)
	except (TypeError, ValueError):
	value = default
	return value if value > 0 else default

	def _has_word(w, text):
	return bool(re.search(rf"\b{re.escape(w)}\b", str(text), re.IGNORECASE))

	# Extract fields from Layer 1 payload
	semantic_query = payload.get("hyde_paragraph") or payload.get("semantic_query") or payload.get("search_optimized_query", "")
	keyword_query = payload.get("keyword_synonyms") or payload.get("keywords") or payload.get("search_optimized_query", "")
	citations = payload.get("explicit_citations", [])
	top_k = _safe_top_k(payload.get("top_k", k))
	temporal_status = payload.get("temporal_context", {}).get("temporal_status") or payload.get("temporal_status", "UNDATED")

	print(f"\n[Layer 2] Semantic : {semantic_query[:70]}...")
	print(f"[Layer 2] Keywords : {keyword_query[:70]}...")
	print(f"[Layer 2] Citations : {citations}")
	print(f"[Layer 2] Temporal : {temporal_status}")
	print(f"[Layer 2] Top K : {top_k}")

	# ── TECHNIQUE 1: Semantic Search via ChromaDB ──
	# Embeds query (uses HyDE paragraph) and finds nearest chunks by cosine similarity
	semantic_results_raw = []
	if embedder is not None and collection is not None:
	print("\n[1/3] Running semantic search (ChromaDB)...")
	try:
	query_embedding = embedder.encode(semantic_query).tolist()
	collection_count = collection.count()
	except Exception as e:
	print(f"[Layer 2 Warning] Embedding generation failed: {e}")
	query_embedding = None
	collection_count = 0
	else:
	print("\n[1/3] Semantic search skipped (fast keyword / low-memory mode).")
	query_embedding = None
	collection_count = 0
	target_sec_nums = set()
	if citations:
	for c in citations:
	m = re.search(r'\d+[A-Za-z]*', str(c))
	if m:
	target_sec_nums.add(m.group())

	# Detect implied statutory sections from core legal topics
	q_all = (semantic_query + " " + keyword_query).lower()
	if any(_has_word(w, q_all) for w in ["cheat", "cheated", "cheating", "fraud", "scam"]):
	target_sec_nums.update(["318", "420"])
	if any(_has_word(w, q_all) for w in ["fir", "zero fir", "e-fir"]):
	target_sec_nums.update(["173", "154"])
	if any(_has_word(w, q_all) for w in ["certificate", "electronic evidence"]):
	target_sec_nums.update(["63", "65B"])
	if any(_has_word(w, q_all) for w in ["community service"]):
	target_sec_nums.update(["4", "23"])

	if collection and target_sec_nums:
	print(f"[Layer 2] Citation/Topic filter: section_numbers = {list(target_sec_nums)}")
	for s_num in target_sec_nums:
	try:
	c_chunks = collection.get(where={"section_number": str(s_num)})
	if c_chunks and c_chunks.get("documents"):
	for doc, meta in zip(c_chunks["documents"], c_chunks["metadatas"]):
	semantic_results_raw.append((doc, meta))
	except Exception as e:
	pass

	# Always run broad semantic search too
	# Combines with citation results for broad coverage
	if collection and query_embedding and collection_count > 0:
	try:
	broad = collection.query(
	query_embeddings=[query_embedding],
	n_results=min(max(top_k * 2, 8), collection_count)
	)
	for doc, meta in zip(broad["documents"][0], broad["metadatas"][0]):
	semantic_results_raw.append((doc, meta))

	# ── PARTITION A: ACTIVE 2024 SANHITA RETRIEVAL (BNS / BNSS / BSA / Constitution) ──
	try:
	sanhita_broad = collection.query(
	query_embeddings=[query_embedding],
	where={"is_active_2024": True},
	n_results=min(max(top_k, 5), collection_count)
	)
	if sanhita_broad and sanhita_broad.get("documents") and sanhita_broad["documents"][0]:
	for doc, meta in zip(sanhita_broad["documents"][0], sanhita_broad["metadatas"][0]):
	semantic_results_raw.append((doc, meta))
	print(f"[Layer 2] Partition A (2024 Sanhitas): {len(sanhita_broad['documents'][0])} chunks retrieved.")
	except Exception as e_sanhita:
	print(f"[Layer 2] Partition A query error: {e_sanhita}")

	# ── PARTITION B: LEGACY COLONIAL CODE RETRIEVAL (IPC / CrPC / IEA) FOR UNDATED DUAL-TRACK ──
	if temporal_status == "UNDATED":
	try:
	legacy_broad = collection.query(
	query_embeddings=[query_embedding],
	where={"is_repealed": True},
	n_results=min(max(top_k, 5), collection_count)
	)
	if legacy_broad and legacy_broad.get("documents") and legacy_broad["documents"][0]:
	for doc, meta in zip(legacy_broad["documents"][0], legacy_broad["metadatas"][0]):
	semantic_results_raw.append((doc, meta))
	print(f"[Layer 2] Partition B (Legacy Colonial): {len(legacy_broad['documents'][0])} chunks retrieved.")
	except Exception as e_legacy:
	print(f"[Layer 2] Partition B query error: {e_legacy}")
	except Exception as e:
	print(f"[Layer 2] ChromaDB query failed: {e}")

	print(f"[Layer 2] Semantic raw results: {len(semantic_results_raw)}")

	# ── TECHNIQUE 2: BM25 Keyword Search ──
	print("[2/3] Running BM25 keyword search...")

	bm25, bm25_docs, bm25_metas = build_bm25_index()
	bm25_results_raw = []

	if bm25 and bm25_docs:
	tokenized_query = keyword_query.lower().split()
	scores = bm25.get_scores(tokenized_query)

	top_indices = sorted(
	range(len(scores)),
	key=lambda i: scores[i],
	reverse=True
	)[:max(top_k * 4, 25)]

	for idx in top_indices:
	if scores[idx] > 0:
	bm25_results_raw.append((
	bm25_docs[idx],
	bm25_metas[idx],
	scores[idx]
	))

	print(f"[Layer 2] BM25 raw results: {len(bm25_results_raw)}")

	# ── COMBINE AND DEDUPLICATE ──
	combined = {}

	for doc, meta in semantic_results_raw:
	if doc not in combined:
	act_n = meta.get("act_name") or meta.get("act_short") or meta.get("act") or ""
	sec_n = meta.get("section_number") or meta.get("section") or ""
	combined[doc] = {
	"content": doc,
	"act_name": act_n,
	"section_number": sec_n,
	"is_repealed": meta.get("is_repealed", False),
	"law_type": meta.get("law_type", "Statute"),
	}

	for doc, meta, score in bm25_results_raw:
	if doc not in combined:
	act_n = meta.get("act_name") or meta.get("act_short") or meta.get("act") or ""
	sec_n = meta.get("section_number") or meta.get("section") or ""
	combined[doc] = {
	"content": doc,
	"act_name": act_n,
	"section_number": sec_n,
	"is_repealed": meta.get("is_repealed", False),
	"law_type": meta.get("law_type", "Statute"),
	"bm25_score": score,
	}

	all_results = list(combined.values())
	print(f"[Layer 2] Combined unique results: {len(all_results)}")

	if not all_results:
	print("[Layer 2] WARNING: No results found in DB.")
	return []

	# ── TECHNIQUE 3: CrossEncoder Reranker with Exact Citation Boosting ──
	if reranker is not None:
	print("[3/3] Reranking with CrossEncoder...")
	pairs = [(semantic_query, r["content"]) for r in all_results]
	try:
	scores = reranker.predict(pairs)
	except Exception as e:
	print(f"[Layer 2 Warning] CrossEncoder inference failed ({e}), using default scores.")
	scores = [0.8] * len(all_results)
	else:
	print("[3/3] CrossEncoder skipped (low-memory / fast mode). Using relevance scoring.")
	scores = [0.8] * len(all_results)

	# Detect primary target act mentions in query to give strong affinity
	q_lower = (semantic_query + " " + keyword_query).lower()

	matched_target_acts = set()
	if any(_has_word(w, q_lower) for w in ["negotiable", "cheque", "nia", "promissory"]):
	matched_target_acts.add("nia")
	if any(_has_word(w, q_lower) for w in ["fir", "crpc", "bnss", "bail", "anticipatory", "cognizable", "complaint", "police"]):
	matched_target_acts.update(["crpc", "bnss"])
	if any(_has_word(w, q_lower) for w in ["marriage", "divorce", "hindu marriage", "hma", "restitution"]):
	matched_target_acts.add("hma")
	if any(_has_word(w, q_lower) for w in ["cpc", "civil procedure", "written statement", "injunction"]):
	matched_target_acts.add("cpc")
	if any(_has_word(w, q_lower) for w in ["motor", "mva", "vehicle", "traffic"]):
	matched_target_acts.add("mva")
	if any(_has_word(w, q_lower) for w in ["evidence", "iea", "bsa", "witness", "certificate"]):
	matched_target_acts.update(["iea", "bsa"])
	if any(_has_word(w, q_lower) for w in ["industrial", "workman", "ida", "retrenchment", "layoff"]):
	matched_target_acts.add("ida")
	if any(_has_word(w, q_lower) for w in [
	"ipc", "penal code", "bns", "cheated", "cheating", "fraud", "scam",
	"deception", "theft", "murder", "extortion", "dishonest", "misappropriation",
	"breach of trust", "community service", "420", "318"
	]):
	matched_target_acts.update(["ipc", "bns"])
	if any(_has_word(w, q_lower) for w in ["constitution", "article", "fundamental right"]):
	matched_target_acts.add("constitution")

	is_modern_query = any(w in q_lower for w in ["today", "new law", "new", "2024", "bns", "bnss", "bsa"])

	boosted_ranked = []
	for score, r in zip(scores, all_results):
	adjusted_score = float(score)
	sec = str(r.get("section_number", "")).strip()
	act = str(r.get("act_name", "")).lower()
	content = r.get("content", "")

	# 1. Exact section citation boost
	if target_sec_nums and (sec in target_sec_nums or any(s in sec for s in target_sec_nums)):
	adjusted_score += 15.0
	# Prefer substantive chunks over short 1-line amendment notes
	if len(content) > 120:
	adjusted_score += 8.0

	# 2. Target act affinity boost / penalty
	if matched_target_acts:
	act_tokens = set(re.split(r'[^a-z0-9]+', act))
	if any(target in act_tokens or target == act for target in matched_target_acts):
	adjusted_score += 25.0
	else:
	adjusted_score -= 15.0 # Penalize chunks from completely unrelated acts

	# 3. Modern 2024 Sanhita preference for queries with modern markers
	if is_modern_query and any(s in act for s in ["bns", "bnss", "bsa"]):
	adjusted_score += 10.0

	boosted_ranked.append((adjusted_score, r))

	# Sort by adjusted score descending
	ranked = sorted(boosted_ranked, key=lambda x: x[0], reverse=True)

	# Build final top k results (default min 5 for complete legal context)
	effective_k = max(top_k, 5)

	if temporal_status == "UNDATED":
	def _is_legacy_chunk(item):
	a = (item.get("act_name") or item.get("act") or "").lower()
	return bool(item.get("is_repealed")) or "from_db" in a or a in ["ipc", "crpc", "iea"]

	def _is_active_sanhita(item):
	a = (item.get("act_name") or item.get("act") or "").lower()
	return bool(item.get("is_active_2024")) or a in ["bns", "bnss", "bsa", "constitution"]

	active_candidates = [r for score, r in ranked if _is_active_sanhita(r)]
	legacy_candidates = [r for score, r in ranked if _is_legacy_chunk(r)]

	balanced = []
	# Add up to 3 active 2024 chunks
	for c in active_candidates[:3]:
	if c not in balanced:
	balanced.append(c)
	# Add up to 3 legacy colonial chunks
	for c in legacy_candidates[:3]:
	if c not in balanced:
	balanced.append(c)
	# Fill remaining slots up to effective_k from highest ranked
	for score, r in ranked:
	if len(balanced) >= max(effective_k, 6):
	break
	if r not in balanced:
	balanced.append(r)
	final_list = balanced
	else:
	final_list = [r for score, r in ranked[:effective_k]]

	top_results = []
	for result in final_list:
	top_results.append({
	"content": result["content"],
	"act_name": result["act_name"],
	"section_number": result["section_number"],
	"is_repealed": result["is_repealed"],
	"law_type": result["law_type"],
	"relevance_score": round(float(result.get("relevance_score", 0.9)), 4),
	})

	# ── TECHNIQUE 6: Shared Precedent Engine Integration (50,000+ Case Matrix) ──
	try:
	try:
	from backend.trial_studio.precedent_engine import get_precedent_engine
	except ImportError:
	from trial_studio.precedent_engine import get_precedent_engine
	pe = get_precedent_engine()
	precedent_matches = pe.search(query=semantic_query, top_k=2)
	for pm in precedent_matches:
	top_results.append({
	"content": f"[Ratio Decidendi]: {pm.ratio_decidendi} [Key Fact Anchor]: {pm.fact_anchor}",
	"act_name": f"{pm.court} Precedent — {pm.citation} ({pm.title})",
	"section_number": ", ".join(pm.sections_invoked),
	"is_repealed": False,
	"law_type": "CONTROLLING PRECEDENT RATIO",
	"relevance_score": 0.95,
	})
	except Exception as e:
	print(f"[Layer 2] Note: Precedent engine optional link skipped: {e}")

	print(f"[Layer 2] Final top {len(top_results)} sections & precedents ready for Layer 3 (Dual-Track balanced: {temporal_status == 'UNDATED'}).\n")
	return top_results


	# ─────────────────────────────────────────────────────────────
	# TEST RUNNER
	# Only runs when executed directly:
	# python layer2_retrieval.py
	#
	# Uses a simulated Layer 1 payload
	# In production this payload comes from analyze_query()
	# Test query covers labour law + explicit section citation
	# which exercises all 3 retrieval techniques at once
	# ─────────────────────────────────────────────────────────────

	if __name__ == "__main__":

	# Simulated Layer 1 output
	# In production: comes from layer1_understanding.analyze_query()
	test_payload = {
	"search_optimized_query": (
	"What are the legal conditions for a valid marriage regarding age and living spouses "
	"under Section 5 of the Hindu Marriage Act?"
	),
	"domain": "family",
	"explicit_citations": ["SECTION 5"],
	"is_ambiguous": False,
	}

	print("-" * 55)
	print("SaulGPT — Layer 2 Hybrid Retrieval Test")
	print("-" * 55)

	results = retrieve_with_hybrid_logic(test_payload)

	print("\n" + "=" * 55)
	print("TOP RETRIEVED LEGAL SECTIONS")
	print("=" * 55)

	if not results:
	print("No results returned.")
	print("Troubleshooting:")
	print(" 1. Run 02_chunk_and_embed.py first")
	print(" 2. Check VECTOR_DB_PATH = 'vector_db'")
	print(" 3. Check COLLECTION_NAME = 'saulgpt_indian_laws'")
	else:
	for i, result in enumerate(results, 1):
	print(f"\nRESULT {i}:")
	print(f" ACT : {result['act_name']}")
	print(f" SECTION : {result['section_number']}")
	print(f" SCORE : {result['relevance_score']}")
	print(f" REPEALED: {result['is_repealed']}")
	print(f" CONTENT : {result['content'][:200]}...")

	print("\n" + "=" * 55)
	print("Layer 2 complete. Results ready for Layer 3.")
	print("=" * 55)