Update app.py

app.py

@@ -1,105 +1,144 @@
+#!/usr/bin/env python3
 import os
+from pathlib import Path
 import gradio as gr
-import torch
-from retriever import get_retriever
+
 from langchain.chains import ConversationalRetrievalChain
+from langchain.memory import ConversationBufferMemory
+from langchain.prompts import PromptTemplate
 from langchain_community.llms import HuggingFacePipeline
+from langchain_chroma import Chroma
+from retriever import get_retriever
+
 from transformers import pipeline
-from dotenv import load_dotenv
 
-load_dotenv()
 
-# Load retriever (your ingestion pipeline already created vector_db)
+# =====================================================
+# Auto-build DB if missing
+# =====================================================
+if not Path("vector_db").exists() or not any(Path("vector_db").iterdir()):
+    print("⚠️ Vector DB not found. Running ingestion...")
+    os.system("python src/ingest_documents.py")
+
+
+# =====================================================
+# Load retriever
+# =====================================================
 retriever = get_retriever()
 
-# Use a CPU-friendly model (Flan-T5)
+
+# =====================================================
+# LLM Setup (lighter model for CPU Spaces)
+# =====================================================
 pipe = pipeline(
-    "text2text-generation",
-    model="google/flan-t5-base",   # or flan-t5-large if you want stronger reasoning
+    "text-generation",
+    model="google/flan-t5-base",   # ✅ smaller + CPU friendly
     max_new_tokens=512,
-    temperature=0.2,
-    device=-1  # CPU
+    temperature=0.3
 )
-
 llm = HuggingFacePipeline(pipeline=pipe)
 
-# System prompts
-ENGLISH_SYSTEM_PROMPT = """You are a Nigerian Legal AI Assistant specialized in Nigerian law. 
-You have deep knowledge of:
+
+# =====================================================
+# Prompts
+# =====================================================
+english_system_prompt = """
+You are a Nigerian Legal AI Assistant specialized in Nigerian law. You have deep knowledge of:
 - Nigerian Constitution 1999
-- Labour Act and Employment Laws  
-- Nigeria Data Protection Act / NDPR
+- Labour Act and Employment Laws
+- Nigeria Data Protection Act
 - Federal Competition and Consumer Protection Act (FCCPA)
 
 PERSONALITY: Professional but approachable, uses Nigerian legal terminology, understands local context.
 
 RESPONSE STYLE:
-- Start with a direct answer
+- Start with direct answer to the question
 - Quote specific sections/articles when available
 - Explain in simple terms what the law means
-- Always include a disclaimer about consulting qualified lawyers
+- Always include disclaimer: "⚠️ This is not legal advice. Please consult a qualified lawyer for specific issues."
 - Use Nigerian English expressions naturally (but not forced)
 
-CONVERSATION MEMORY: Remember previous user questions in this chat for context.
+CONVERSATION MEMORY: Remember previous questions in this chat to provide contextual follow-ups.
 """
 
-PIDGIN_SYSTEM_PROMPT = """You be Nigerian Legal Assistant wey sabi Constitution, Labour Law, Data Protection, and FCCPA. 
-You sabi answer questions for Nigerian Pidgin.
+pidgin_system_prompt = """
+You be Nigerian Legal AI Assistant wey sabi Nigerian law well well. You get knowledge of:
+- Nigerian Constitution 1999
+- Labour Act and Employment Laws
+- Nigeria Data Protection Act
+- Federal Competition and Consumer Protection Act (FCCPA)
 
-PERSONALITY: Friendly but serious, explain law like you dey gist but keep am correct.
+PERSONALITY: Friendly, approachable, dey use Naija way of talk but still correct for legal matter.
 
 RESPONSE STYLE:
-- Start with straight answer
-- Always talk the section/article wey back the matter
-- Break am down for simple Pidgin make even person wey no study law fit understand
-- Put small Nigerian English/Pidgin expressions but no change the legal meaning
-- Always add disclaimer say make dem consult lawyer for final advice
-
-CONVERSATION MEMORY: Remember wetin user don ask before so your answer go make sense with the gist.
+- Start with direct answer
+- Mention the exact section/article if available
+- Explain am for clear Pidgin wey anybody fit understand
+- Always add disclaimer: "⚠️ No be legal advice o, abeg meet lawyer if matter serious."
+- Remember wetin dem don ask before, make conversation flow well.
 """
 
-# Conversational chain
+
+# =====================================================
+# Conversational QA Chain
+# =====================================================
+memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
+
 qa_chain = ConversationalRetrievalChain.from_llm(
     llm=llm,
     retriever=retriever,
-    return_source_documents=True
+    memory=memory,
+    return_source_documents=True,
 )
 
-def chat_fn(user_message, history, lang_choice):
-    system_prompt = ENGLISH_SYSTEM_PROMPT if lang_choice == "english" else PIDGIN_SYSTEM_PROMPT
-    query_with_prompt = f"{system_prompt}\n\nUser: {user_message}\nAssistant:"
-    
-    result = qa_chain.invoke({"question": query_with_prompt, "chat_history": history})
+
+# =====================================================
+# Chat function
+# =====================================================
+def answer_question(user_input, lang_choice, history=[]):
+    # Pick system prompt
+    if lang_choice == "pidgin":
+        system_prompt = pidgin_system_prompt
+    else:
+        system_prompt = english_system_prompt
+
+    # Run QA
+    result = qa_chain.invoke({"question": f"{system_prompt}\n\nUser: {user_input}"})
     answer = result["answer"]
 
-    # Collect sources with section reference
+    # Collect sources (with sections)
     sources = []
     for doc in result["source_documents"]:
         section = doc.metadata.get("section", "Unknown Section")
-        src = doc.metadata.get("source", "Unknown Document").replace(".pdf", "")
-        sources.append(f"[{section}] {src}")
-    
+        source = doc.metadata.get("source", "Unknown Document").replace(".pdf", "")
+        sources.append(f"[{section}] from {source}")
+
     if sources:
         answer += "\n\n📚 Sources:\n" + "\n".join(sources)
-    
-    return answer
 
-# Gradio UI
-demo = gr.Blocks(theme="default")
+    history.append(("You: " + user_input, "Bot: " + answer))
+    return history, history
+
+
+# =====================================================
+# Gradio Interface (ChatGPT-like UI)
+# =====================================================
+with gr.Blocks(css=".gradio-container {max-width: 800px !important}") as demo:
+    gr.Markdown("## 📜 KnowYourRight Bot — Nigerian Legal Assistant")
+
+    with gr.Row():
+        with gr.Column(scale=4):
+            chatbot = gr.Chatbot(label="Chat with Legal AI", height=500)
+            msg = gr.Textbox(label="Ask your question here...")
+            lang_choice = gr.Radio(["english", "pidgin"], value="english", label="Language")
+            clear = gr.Button("Clear Chat")
 
-with demo:
-    gr.Markdown("# 📜 KnowYourRight Bot\n### Ask Nigerian law questions in English or Pidgin")
-    lang_choice = gr.Radio(["english", "pidgin"], label="Language Mode", value="english")
-    chatbot = gr.Chatbot(label="Conversation", height=500, show_label=False)
-    msg = gr.Textbox(label="Your Question", placeholder="Type your question here...")
-    clear = gr.Button("Clear Chat")
+    state = gr.State([])
 
-    def respond(user_message, history, lang_choice):
-        bot_reply = chat_fn(user_message, history, lang_choice)
-        history = history + [(user_message, bot_reply)]
-        return history, ""
+    def reset():
+        return [], []
 
-    msg.submit(respond, [msg, chatbot, lang_choice], [chatbot, msg])
-    clear.click(lambda: [], None, chatbot)
+    msg.submit(answer_question, [msg, lang_choice, state], [chatbot, state])
+    clear.click(reset, None, [chatbot, state])
 
-# Hugging Face Spaces auto-launches `demo`
+demo.launch()