Update app.py

app.py

@@ -1,24 +1,20 @@
-#!/usr/bin/env python3
 import os
 from pathlib import Path
 import gradio as gr
 
 from retriever import get_retriever
-
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from langchain_community.llms import HuggingFacePipeline
 from langchain.prompts import PromptTemplate
-from langchain_huggingface import HuggingFaceEmbeddings
 
-PERSIST_DIR = Path(os.getenv("VECTOR_DB_DIR", "vector_db"))
+# Ensure vector DB exists (from complete_ingestion.py output)
+PERSIST_DIR = Path("data/processed/vector_db")
 if not PERSIST_DIR.exists() or not any(PERSIST_DIR.iterdir()):
-    print("⚠️ Vector DB not found. Running ingestion...")
-    os.system("python src/ingest_documents.py")
-
+    raise RuntimeError("⚠️ Vector DB not found. Please run complete_ingestion.py first.")
 
 retriever = get_retriever()
 
-
+# Load lightweight conversational model
 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)
@@ -28,8 +24,8 @@ gen_pipe = pipeline(
     model=model,
     tokenizer=tokenizer,
     device_map="auto" if os.getenv("SPACE_ACCELERATOR") else None,
-    max_new_tokens=120,
-    temperature=0.2,
+    max_new_tokens=180,
+    temperature=0.3,
     top_p=0.9,
     do_sample=True,
     repetition_penalty=1.1,
@@ -38,16 +34,20 @@ gen_pipe = pipeline(
 
 llm = HuggingFacePipeline(pipeline=gen_pipe)
 
+# Conversational + contextual legal prompt
 RAG_PROMPT = PromptTemplate.from_template(
     "You are a helpful Nigerian Legal Assistant.\n"
-    "Answer briefly from ONLY the provided context. If the answer is not in the context, say you don't know.\n\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 (for context):\n{history}\n\n"
     "Question: {question}\n\n"
-    "Context:\n{context}\n\n"
+    "Context from legal documents:\n{context}\n\n"
     "Answer:"
 )
 
 
-def _format_history(turns, max_turns=2):
+def _format_history(turns, max_turns=4):
     if not turns:
         return ""
     turns = turns[-max_turns:]
@@ -62,11 +62,9 @@ def _retrieve(question, k=3):
 
 
 def _generate(question, history):
-    hist = _format_history(history, max_turns=2)
-    if hist:
-        question = f"{hist}\n\nCurrent question: {question}"
+    hist = _format_history(history, max_turns=4)
     context, docs = _retrieve(question, k=3)
-    prompt = RAG_PROMPT.format(question=question, context=context)
+    prompt = RAG_PROMPT.format(question=question, context=context, history=hist)
     out = llm(prompt)
     if isinstance(out, list):
         text = out[0].get("generated_text", "") or out[0].get("text", "")
@@ -81,6 +79,7 @@ def answer_question(user_input, lang_choice, history=[]):
         if not q:
             return history, history
 
+        # Greeting handling
         if q.lower() in ["hi", "hello", "hey"]:
             if lang_choice == "pidgin":
                 ans = "Hello! I be your Nigerian Legal AI Assistant. How I fit help you? No be legal advice o."
@@ -89,24 +88,29 @@ def answer_question(user_input, lang_choice, history=[]):
             history.append((user_input, ans))
             return history, history
 
+        # Trim long queries
         if len(q) > 300:
             q = q[:300] + "..."
 
+        # Generate answer
         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 == "pidgin":
                 answer = "I no sure from the context wey I get. Abeg rephrase your question."
+            else:
+                answer = "I don't know from the available context. Please try rephrasing your question."
 
+        # Add disclaimer
         if lang_choice == "pidgin":
             answer += "\n\n⚠️ No be legal advice o, abeg meet lawyer."
         else:
             answer += "\n\n⚠️ This is not legal advice. Please consult a qualified lawyer."
 
-        sources = [d.metadata.get("source", "Unknown") for d in docs[:2]]
+        # Add sources/sections
+        sources = [d.metadata.get("section", d.metadata.get("source", "Unknown")) for d in docs[:2]]
         if sources:
-            answer += f"\n\nSources: {', '.join(sources)}"
+            answer += f"\n\nReferenced: {', '.join(sources)}"
 
         history.append((user_input, answer))
         history = history[-8:]
@@ -123,23 +127,33 @@ def _reset():
     return [], []
 
 
-with gr.Blocks() as demo:
-    gr.Markdown("## 📜 KnowYourRight Bot — Nigerian Legal Assistant")
+# Minimal full-screen Gradio UI
+def build_ui():
+    with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), css=".gradio-container {max-width: 100% !important}") as demo:
+        gr.Markdown("""
+        # 📜 KnowYourRight Bot — Nigerian Legal Assistant
+        Conversational, contextual answers from Nigerian legal documents.
+        """)
+
+        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")
+
+        with gr.Row():
+            submit = gr.Button("Send", variant="primary")
+            clear = gr.Button("Clear Chat")
+
+        state = gr.State([])
 
-    chatbot = gr.Chatbot(label="Chat with Legal AI", height=400)
-    msg = gr.Textbox(label="Ask your question here...")
-    lang_choice = gr.Radio(["english", "pidgin"], value="english", label="Language")
+        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])
 
-    with gr.Row():
-        submit = gr.Button("Send")
-        clear = gr.Button("Clear Chat")
+    return demo
 
-    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])
+demo = build_ui()
+demo.launch()
 
-demo.launch()