Update app.py

app.py

@@ -1,164 +1,206 @@
 import os
 from pathlib import Path
 import gradio as gr
-
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-from langchain_community.llms import HuggingFacePipeline
 from langchain.prompts import PromptTemplate
 from langchain_community.vectorstores import Chroma
-from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_huggingface import HuggingFaceEmbeddings, HuggingFaceHub
+from langchain.schema.runnable import RunnablePassthrough
+from langchain.schema.output_parser import StrOutputParser
+
+# --- 1. CONFIGURATION & INITIALIZATION ---
+
+# Load environment variables (for Hugging Face API token)
+from dotenv import load_dotenv
+load_dotenv()
+
+# Check for the API token
+if not os.getenv("HUGGINGFACEHUB_API_TOKEN"):
+    print(" HUGGINGFACEHUB_API_TOKEN not found in secrets. Please add it.")
+    exit()
 
-# ----------------------------
-# Load vector DB
-# ----------------------------
+# --- 2. LOAD VECTOR DATABASE (Retriever) ---
+
+print("Loading vector database...")
 PERSIST_DIR = Path("data/processed/vector_db")
+
 if not PERSIST_DIR.exists() or not any(PERSIST_DIR.iterdir()):
     print("⚠️ Vector DB not found. Run complete_ingestion.py first.")
     raise SystemExit(1)
 
-embedding_model = HuggingFaceEmbeddings(model_name="BAAI/bge-small-en")
+# Use the same embedding model as in the ingestion script
+embedding_model = HuggingFaceEmbeddings(
+    model_name="BAAI/bge-small-en",
+    model_kwargs={'device': 'cpu'} # Run embeddings on CPU
+)
+
+# Load the Chroma vector store
 vectordb = Chroma(
     persist_directory=str(PERSIST_DIR),
     embedding_function=embedding_model,
     collection_name="legal_documents"
 )
 
-retriever = vectordb.as_retriever(search_kwargs={"k": 3})
-
-# ----------------------------
-# Lightweight LLM
-# ----------------------------
-MODEL_ID = os.getenv("LLM_ID", "TinyLlama/TinyLlama-1.1B-Chat-v1.0")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
-
-gen_pipe = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    max_new_tokens=120,   # reduced for speed
-    temperature=0.2,
-    top_p=0.85,
-    do_sample=True,
-    repetition_penalty=1.05,
-    return_full_text=False,
+# Create a retriever to fetch relevant documents
+# Increasing k to 4 gives the LLM more context to work with
+retriever = vectordb.as_retriever(search_kwargs={"k": 4})
+print("Vector database loaded successfully.")
+
+# --- 3. SETUP THE LIGHTWEIGHT LLM (via Inference API) ---
+
+print("Initializing LLM via Hugging Face Hub...")
+# We use the Inference API to avoid loading the model locally, which is much faster.
+# Mixtral is a powerful model available on the free tier.
+llm = HuggingFaceHub(
+    repo_id="mistralai/Mixtral-8x7B-Instruct-v0.1",
+    model_kwargs={"temperature": 0.1, "max_length": 1024, "max_new_tokens": 512}
 )
+print("LLM initialized.")
+
+# --- 4. CREATE THE IMPROVED PROMPT TEMPLATE ---
+
+# This new prompt is more directive and helps shape the output.
+RAG_PROMPT_TEMPLATE = """
+You are an expert Nigerian Legal Assistant. Your primary goal is to help users understand Nigerian law by providing clear, concise, and helpful explanations.
+
+**TASK:** Analyze the provided legal context below to answer the user's question.
+
+**CONTEXT:**
+{context}
 
-llm = HuggingFacePipeline(pipeline=gen_pipe)
-
-# ----------------------------
-# Prompt
-# ----------------------------
-RAG_PROMPT = PromptTemplate.from_template(
-    "You are a helpful Nigerian Legal Assistant.\n"
-    "Respond conversationally, summarize clearly, and explain in plain English (or Pidgin if chosen).\n"
-    "Always include the referenced section(s) at the end.\n"
-    "If the answer is not in the context, say you don't know.\n\n"
-    "Conversation history:\n{history}\n\n"
-    "Question: {question}\n\n"
-    "Context from legal documents:\n{context}\n\n"
-    "Answer:"
+**RULES:**
+1.  **Explain, Don't Just Quote:** Do not just copy the text from the context. You MUST synthesize, summarize, and explain the relevant laws in simple, easy-to-understand language.
+2.  **Be Conversational:** Respond in a helpful and advisory tone.
+3.  **Use Only Provided Context:** Base your answer SOLELY on the provided context. If the context does not contain the information needed to answer the question, you MUST say "The provided legal documents do not contain specific information on this topic." Do not use outside knowledge.
+4.  **Language:** Respond in the user's chosen language (English or Nigerian Pidgin).
+5.  **Citations:** At the end of your answer, always list the sources you used from the context.
+
+**QUESTION:** {question}
+
+**ANSWER:**
+"""
+
+RAG_PROMPT = PromptTemplate.from_template(RAG_PROMPT_TEMPLATE)
+
+# --- 5. DEFINE THE RAG CHAIN ---
+
+def format_docs(docs):
+    """Helper function to format retrieved documents into a single string."""
+    return "\n\n---\n\n".join(f"Source: {d.metadata.get('source', 'Unknown')}\nSection: {d.metadata.get('section', 'Unknown')}\nContent: {d.page_content}" for d in docs)
+
+# Create the LangChain RAG chain
+rag_chain = (
+    {"context": retriever | format_docs, "question": RunnablePassthrough()}
+    | RAG_PROMPT
+    | llm
+    | StrOutputParser()
 )
 
-# ----------------------------
-# Helpers
-# ----------------------------
-def _format_history(turns, max_turns=4):
-    if not turns:
-        return ""
-    turns = turns[-max_turns:]
-    return "\n".join([f"User: {u}\nAssistant: {a}" for u, a in turns])
-
-def _retrieve(question, k=3):
-    docs = retriever.invoke(question)  # ✅ fixed deprecation
-    texts = [d.page_content.strip() for d in docs[:k]]
-    context = "\n\n---\n\n".join(texts)
-    return context, docs
-
-def _generate(question, history):
-    hist = _format_history(history, max_turns=4)
-    context, docs = _retrieve(question, k=3)
-    prompt = RAG_PROMPT.format(question=question, context=context, history=hist)
-    out = llm.invoke(prompt)  # ✅ fixed deprecation
-    if isinstance(out, list) and out and "generated_text" in out[0]:
-        text = out[0]["generated_text"]
-    else:
-        text = str(out)
-    return text.strip(), docs
-
-# ----------------------------
-# Main logic
-# ----------------------------
+# --- 6. MAIN APPLICATION LOGIC ---
+
 def answer_question(user_input, lang_choice, history=[]):
+    """Main function to handle user queries, run the RAG chain, and format the output."""
     try:
-        q = (user_input or "").strip()
-        if not q:
+        query = (user_input or "").strip()
+        if not query:
             return history, history
 
-        if q.lower() in ["hi", "hello", "hey"]:
-            ans = "Hello! I'm your Nigerian Legal AI Assistant. How can I help you?" \
-                  if lang_choice == "english" else \
-                  "Hello! I be your Nigerian Legal AI Assistant. How I fit help you? No be legal advice o."
+        # Simple conversational starters
+        if query.lower() in ["hi", "hello", "hey"]:
+            ans = ("Hello! I'm your Nigerian Legal AI Assistant. How can I help you today?"
+                   if lang_choice == "english" else
+                   "Howfa! I be your Nigerian Legal AI Assistant. How I fit help you today? No be legal advice o.")
             history.append((user_input, ans))
             return history, history
 
-        if len(q) > 300:
-            q = q[:300] + "..."
-
-        answer, docs = _generate(q, history)
-
-        if not answer or len(answer) < 5:
-            answer = "I don't know from the available context. Please try rephrasing your question." \
-                if lang_choice == "english" else \
-                "I no sure from the context wey I get. Abeg rephrase your question."
-
-        disclaimer = "⚠️ This is not legal advice. Please consult a qualified lawyer." \
-            if lang_choice == "english" else \
-            "⚠️ No be legal advice o, abeg meet lawyer."
-        answer += f"\n\n{disclaimer}"
-
-        # ✅ references improved
-        refs = []
-        for d in docs[:2]:
-            src = d.metadata.get("source", "Unknown Source")
-            sec = d.metadata.get("section", "Unknown Section")
-            refs.append(f"{src} – {sec}")
-        if refs:
-            answer += "\n\nReferenced: " + "; ".join(refs)
-
-        history.append((user_input, answer))
+        print(f"Received query: {query}")
+
+        # Retrieve documents first to build references
+        docs = retriever.invoke(query)
+        if not docs:
+            print("No documents retrieved.")
+            answer = "I could not find any relevant information in the legal documents for your query. Please try rephrasing."
+        else:
+            # Invoke the RAG chain to get the answer
+            print("Invoking RAG chain...")
+            answer = rag_chain.invoke(query)
+            print("RAG chain finished.")
+
+
+        # Add a disclaimer
+        disclaimer = ("\n\n--- \n*⚠️ Disclaimer: This is AI-generated information and not legal advice. Please consult a qualified lawyer for professional guidance.*"
+                      if lang_choice == "english" else
+                      "\n\n--- \n*⚠️ No be legal advice o, abeg find lawyer for proper advice.*")
+        
+        # Build robust references
+        # Use a set to avoid duplicate references
+        references = set()
+        for doc in docs:
+            source = doc.metadata.get("source", "Unknown Source")
+            section = doc.metadata.get("section", "Unknown Section")
+            # Only add if both source and section are known
+            if source != "Unknown Source" and section != "Unknown Section":
+                references.add(f"- {source} ({section})")
+
+        if references:
+            answer += "\n\n**References:**\n" + "\n".join(sorted(list(references)))
+
+        answer += disclaimer
+        
+        history.append((user_input, answer.strip()))
+        
+        # Keep chat history to a reasonable length
         return history[-8:], history[-8:]
 
     except Exception as e:
-        print(f"Error: {e}")
-        err = "Sorry, an error occurred. Please try again."
-        history.append((user_input, err))
+        print(f"An error occurred: {e}")
+        error_message = "Sorry, an unexpected error occurred. Please try again or rephrase your question."
+        history.append((user_input, error_message))
         return history, history
 
 def _reset():
+    """Resets the chat state."""
     return [], []
 
-# ----------------------------
-# UI
-# ----------------------------
+# --- 7. GRADIO UI ---
+
 def build_ui():
-    with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
+    """Builds the Gradio web interface."""
+    with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), title="KnowYourRight Bot") as demo:
         gr.Markdown("# 📜 KnowYourRight Bot — Nigerian Legal Assistant")
-        chatbot = gr.Chatbot(label="Chat with Legal AI", height=600, bubble_full_width=False)
-        msg = gr.Textbox(label="Ask your question...", placeholder="Type your legal question here", lines=2)
-        lang_choice = gr.Radio(["english", "pidgin"], value="english", label="Language")
-
+        gr.Markdown("Ask questions about the Nigerian Constitution, Labour Act, and more. *Powered by AI.*")
+        
+        chatbot = gr.Chatbot(label="Chat History", height=600, bubble_full_width=False, avatar_images=("user.png", "bot.png"))
+        
         with gr.Row():
-            submit = gr.Button("Send", variant="primary")
-            clear = gr.Button("Clear Chat")
-
-        state = gr.State([])
-        submit.click(answer_question, [msg, lang_choice, state], [chatbot, state])
-        submit.click(lambda: "", None, msg)
-        msg.submit(answer_question, [msg, lang_choice, state], [chatbot, state])
-        msg.submit(lambda: "", None, msg)
-        clear.click(_reset, None, [chatbot, state])
+            msg = gr.Textbox(
+                label="Your Question",
+                placeholder="e.g., 'What are my rights if I am arrested?'",
+                lines=2,
+                scale=4,
+            )
+            submit_btn = gr.Button("▶️ Send", variant="primary", scale=1)
+
+        lang_choice = gr.Radio(["english", "pidgin"], value="english", label="Response Language")
+        clear_btn = gr.Button("🗑️ Clear Chat")
+
+        # State to store the conversation history
+        chat_state = gr.State([])
+
+        # Event handlers
+        submit_btn.click(answer_question, [msg, lang_choice, chat_state], [chatbot, chat_state])
+        msg.submit(answer_question, [msg, lang_choice, chat_state], [chatbot, chat_state])
+
+        # Clear the input textbox after submission
+        clear_on_submit = [submit_btn, msg]
+        for component in clear_on_submit:
+            component.click(lambda: "", None, msg)
+
+        clear_btn.click(_reset, None, [chatbot, chat_state])
+    
     return demo
 
-demo = build_ui()
-demo.launch()
+if __name__ == "__main__":
+    print("Building Gradio UI...")
+    demo = build_ui()
+    print("Launching Gradio app...")
+    demo.launch(debug=True) # Set debug=False for production