Update app.py

app.py

@@ -18,7 +18,8 @@ def initialize_retriever():
         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.")
+            # Create a dummy doc if files are missing to prevent crash
+            return None
         ret = HybridRetriever(documents=docs, index_dir=INDEX_DIR)
         ret.save_index()
         return ret
@@ -28,39 +29,40 @@ retriever = initialize_retriever()
 
 def respond(
     message,
-    history,  # In Gradio 4.x, this is a list of lists: [[user, bot], [user, bot]]
+    history,
     system_message,
     max_tokens,
     temperature,
     top_p,
-    hf_token: gr.OAuthToken,
+    hf_token: gr.OAuthToken, # Gradio automatically injects this from the Login button
 ):
     # 1. RETRIEVAL STEP
-    search_results = retriever.hybrid_search(message, top_k=3)
-    
-    # 2. CONTEXT BUILDING
-    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"
+    context = ""
+    if retriever:
+        search_results = retriever.hybrid_search(message, top_k=3)
+        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
+    # 2. PROMPT ENGINEERING
     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. "
-        "Always cite the 'Source' and 'Clause/Section' if you use the context.\n"
+        "Use the provided legal context to answer accurately. "
+        "Always cite the 'Source' and 'Clause/Section'.\n"
         f"{context}"
     )
 
-    client = InferenceClient(token=hf_token.token, model="meta-llama/Llama-3.1-70B-Instruct")
+    # Use the OAuth token or fall back to environment variable
+    token = hf_token.token if hf_token else os.getenv("HF_TOKEN")
+    client = InferenceClient(token=token, model="meta-llama/Llama-3.1-70B-Instruct")
 
     messages = [{"role": "system", "content": augmented_system_message}]
 
-    # 4. HISTORY CONVERSION (Convert [[u, b]] to [{"role": "user", "content": u}, ...])
     for user_msg, assistant_msg in history:
         if user_msg:
             messages.append({"role": "user", "content": user_msg})
@@ -71,21 +73,23 @@ def respond(
 
     response = ""
 
-    # 5. GENERATION STEP
-    for msg in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = msg.choices[0].delta.content
-        if token:
-            response += token
-            yield response
+    # 3. GENERATION STEP
+    try:
+        for msg in client.chat_completion(
+            messages,
+            max_tokens=max_tokens,
+            stream=True,
+            temperature=temperature,
+            top_p=top_p,
+        ):
+            token_text = msg.choices[0].delta.content
+            if token_text:
+                response += token_text
+                yield response
+    except Exception as e:
+        yield f"Error calling AI: {str(e)}. Please make sure you are logged in or provide a valid HF Token."
 
 # --- Gradio UI Setup ---
-# Removed type="messages" to support Gradio 4.x
 chatbot = gr.ChatInterface(
     respond,
     additional_inputs=[
@@ -95,20 +99,16 @@ chatbot = gr.ChatInterface(
         ),
         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,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p"),
     ],
     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.",
+    description="Ask questions about Nepalese Acts, Codes, and the Constitution.",
+    # FIX: Examples must be a list of lists because we have additional_inputs
     examples=[
-        "What are the punishments for cybercrime?",
-        "What does the constitution say about the right to equality?",
-        "Is witchcraft accusation a crime in Nepal?"
+        ["What are the punishments for cybercrime?"],
+        ["What does the constitution say about the right to equality?"],
+        ["Is witchcraft accusation a crime in Nepal?"],
+        ["What is the legal age for marriage in Nepal?"]
     ]
 )
 
@@ -117,7 +117,7 @@ with gr.Blocks() as demo:
         gr.Markdown("### Authentication")
         gr.LoginButton()
         gr.Markdown("---")
-        gr.Markdown("**Search Status:** Database Loaded ✅")
+        gr.Markdown("**Status:** Database Ready ✅")
     chatbot.render()
 
 if __name__ == "__main__":