botInfinity/NEPAL_Constitution_Assistant_AI

1import os2import streamlit as st3from qdrant_client import QdrantClient4from langchain_qdrant import (5    QdrantVectorStore,6    RetrievalMode,7    FastEmbedSparse8)9from langchain_huggingface import HuggingFaceEmbeddings10from sentence_transformers import CrossEncoder11from langchain_groq import ChatGroq12 13# ------------------------------14# Streamlit Config (MUST RUN FAST)15# ------------------------------16st.set_page_config(17    page_title="Nepal Constitution AI",18    page_icon="🧑‍⚖️",19    layout="wide"20)21 22st.title("🧑‍⚖️ Nepal Constitution – AI Legal Assistant")23st.caption("Hybrid RAG (Dense + BM25) + Cross-Encoder Reranking")24 25# 🔥 EARLY VISIBILITY (HF health check helper)26st.write("✅ App booted successfully.")27 28# ------------------------------29# Hard stop if DB missing (NO SILENT FAIL)30# ------------------------------31if not os.path.exists("./qdrant_db"):32    st.error("❌ qdrant_db folder not found. You must commit it to the repo.")33    st.stop()34 35# ------------------------------36# User Input37# ------------------------------38query = st.text_input(39    "Ask a constitutional or legal question:",40    placeholder="e.g. What does Article 275 say about local governance?"41)42 43# ------------------------------44# Cached Heavy Stuff45# ------------------------------46@st.cache_resource47def load_embeddings():48    return HuggingFaceEmbeddings(49        model_name="BAAI/bge-m3",50        model_kwargs={"device": "cpu"},51        encode_kwargs={"normalize_embeddings": True}52    )53 54@st.cache_resource55def load_sparse_embeddings():56    return FastEmbedSparse(model_name="Qdrant/bm25")57 58@st.cache_resource59def load_reranker():60    return CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")61 62@st.cache_resource63def load_vector_store():64    embeddings = load_embeddings()65    sparse_embeddings = load_sparse_embeddings()66    client = QdrantClient(path="./qdrant_db")67 68    return QdrantVectorStore(69        client = client,70        collection_name="nepal_law",71        embedding=embeddings,72        sparse_embedding=sparse_embeddings,73        retrieval_mode=RetrievalMode.HYBRID74    )75 76@st.cache_resource77def load_llm():78    return ChatGroq(79        model="llama-3.1-8b-instant",80        temperature=0.2,81        max_tokens=60082    )83 84# ------------------------------85# Reranking86# ------------------------------87def rerank(query, docs, top_k=8):88    reranker = load_reranker()89    pairs = [(query, d.page_content) for d in docs]90    scores = reranker.predict(pairs)91 92    ranked = sorted(93        zip(docs, scores),94        key=lambda x: x[1],95        reverse=True96    )97 98    return [doc for doc, _ in ranked[:top_k]]99 100 101if query:102    with st.spinner("🔍 Searching constitution..."):103        vector_store = load_vector_store()104        retrieved = vector_store.similarity_search(query, k=20)105        reranked = rerank(query, retrieved)106 107        context = "\n\n".join(108            f"[Source {i+1}]\n{doc.page_content}"109            for i, doc in enumerate(reranked)110        )111 112    prompt = f"""113You are a constitutional law assistant for Nepal.114 115RULES:116- Use ONLY the provided context.117- Do NOT invent articles, clauses, or interpretations.118- If the answer is not found, say so explicitly.119- Use formal, neutral legal language.120- Reference article/section numbers when mentioned.121 122CONTEXT:123{context}124 125QUESTION:126{query}127 128ANSWER:129"""130 131    with st.spinner("🧠 Generating answer..."):132        llm = load_llm()133        response = llm.invoke(prompt)134 135    st.markdown("### ✅ Answer")136    st.write(response.content)137 138    with st.expander("📚 Retrieved Constitutional Sources"):139        for i, doc in enumerate(reranked):140            st.markdown(f"**Source {i+1}**")141            st.write(doc.page_content)142            st.markdown("---")143