Update app.py

app.py

@@ -1,6 +1,31 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
+import os
+from src.data_processor import LegalDocProcessor
+from src.hybrid_retriever import HybridRetriever
 
+# --- Configuration & Initialization ---
+INDEX_DIR = "index_storage"
+PARENT_DATA = "data/parent_docs.json"
+CHILD_DATA = "data/child_docs.json"
+
+# Initialize the retriever (Logic from your provided files)
+def initialize_retriever():
+    if os.path.exists(INDEX_DIR):
+        print("[*] Loading existing index...")
+        return HybridRetriever(index_dir=INDEX_DIR)
+    else:
+        print("[*] Building new index...")
+        processor = LegalDocProcessor(PARENT_DATA, CHILD_DATA)
+        docs = processor.load_and_clean()
+        if not docs:
+            raise ValueError("No documents found to index. Check your data path.")
+        ret = HybridRetriever(documents=docs, index_dir=INDEX_DIR)
+        ret.save_index()
+        return ret
+
+# Global retriever instance
+retriever = initialize_retriever()
 
 def respond(
     message,
@@ -11,19 +36,43 @@ def respond(
     top_p,
     hf_token: gr.OAuthToken,
 ):
-    """
-    For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-    """
-    client = InferenceClient(token=hf_token.token, model="openai/gpt-oss-20b")
+    # 1. RETRIEVAL STEP: Use your HybridRetriever to find relevant law snippets
+    search_results = retriever.hybrid_search(message, top_k=3)
+    
+    # 2. CONTEXT BUILDING: Format the search results into a string
+    context = "\n\nRELEVANT NEPALESE LAW CONTEXT:\n"
+    if not search_results:
+        context += "No specific legal clauses found for this query."
+    for res in search_results:
+        context += f"--- Source: {res['legal_document_source']} ---\n"
+        context += f"Clause/Section: {res['parent_clause_id']}\n"
+        context += f"Text: {res['parent_clause_text']}\n\n"
+
+    # 3. PROMPT ENGINEERING: Inject context into the system message or user message
+    augmented_system_message = (
+        f"{system_message}\n\n"
+        "You are a legal assistant specializing in Nepalese Law. "
+        "Use the following legal context to answer the user's question accurately. "
+        "If the context doesn't contain the answer, state that you are answering based on general knowledge but couldn't find the specific clause."
+        f"{context}"
+    )
+
+    client = InferenceClient(token=hf_token.token, model="meta-llama/Llama-3.1-70B-Instruct")
 
-    messages = [{"role": "system", "content": system_message}]
+    messages = [{"role": "system", "content": augmented_system_message}]
 
-    messages.extend(history)
+    # Maintain history
+    for val in history:
+        if val['role'] == 'user':
+            messages.append({"role": "user", "content": val['content']})
+        else:
+            messages.append({"role": "assistant", "content": val['content']})
 
     messages.append({"role": "user", "content": message})
 
     response = ""
 
+    # 4. GENERATION STEP: Stream the response from the LLM
     for message in client.chat_completion(
         messages,
         max_tokens=max_tokens,
@@ -31,23 +80,21 @@ def respond(
         temperature=temperature,
         top_p=top_p,
     ):
-        choices = message.choices
-        token = ""
-        if len(choices) and choices[0].delta.content:
-            token = choices[0].delta.content
+        token = message.choices[0].delta.content
+        if token:
+            response += token
+            yield response
 
-        response += token
-        yield response
-
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
+# --- Gradio UI Setup ---
 chatbot = gr.ChatInterface(
     respond,
+    type="messages", # Updated for newer Gradio versions
     additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
+        gr.Textbox(
+            value="You are a helpful Nepalese Legal Advisor. Always cite the Source and Clause ID provided in the context.", 
+            label="System message"
+        ),
+        gr.Slider(minimum=1, maximum=2048, value=1024, step=1, label="Max new tokens"),
         gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
         gr.Slider(
             minimum=0.1,
@@ -57,13 +104,23 @@ chatbot = gr.ChatInterface(
             label="Top-p (nucleus sampling)",
         ),
     ],
+    title="Nepal Law Search AI",
+    description="Ask questions about Nepalese Acts, Codes, and the Constitution. The AI will search the official legal database before answering.",
+    examples=[
+        "What are the punishments for cybercrime?",
+        "What does the constitution say about the right to equality?",
+        "What are the rules for divorce in the Civil Code?",
+        "Is witchcraft accusation a crime in Nepal?"
+    ]
 )
 
 with gr.Blocks() as demo:
     with gr.Sidebar():
+        gr.Markdown("### Authentication")
         gr.LoginButton()
+        gr.Markdown("---")
+        gr.Markdown("**Note:** This system uses Hybrid Search (BM25 + Vector) to find relevant Nepalese Law.")
     chatbot.render()
 
-
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()