Update app.py

app.py

@@ -34,25 +34,39 @@ pipe = pipeline(
     "text-generation",
     model="google/flan-t5-base",   # ✅ smaller + CPU friendly
     max_new_tokens=256,  # Reduced from 512 to fit within context
-    temperature=0.3,
+    temperature=0.7,
     do_sample=True,
-    pad_token_id=0  # Add padding token
+    pad_token_id=0,  # Add padding token
+    truncation=True,
+    return_full_text=False  # Only return generated text, not the prompt
 )
 llm = HuggingFacePipeline(pipeline=pipe)
 
 
 # =====================================================
-# Prompts (shortened to reduce token usage)
+# Custom prompt template for better responses
 # =====================================================
-english_system_prompt = """You are a Nigerian Legal AI Assistant. Provide direct answers about Nigerian law with relevant sections/articles. Always end with: "⚠️ This is not legal advice. Consult a qualified lawyer."
-"""
+custom_template = """Based on the following Nigerian law documents, answer the user's question clearly and directly.
 
-pidgin_system_prompt = """You be Nigerian Legal AI Assistant. Give direct answer about Nigerian law with correct section/article. Always end with: "⚠️ No be legal advice o, abeg meet lawyer."
-"""
+Context: {context}
+
+Question: {question}
+
+Instructions:
+- Give a direct, helpful answer
+- Quote specific sections when relevant
+- Use simple, clear language
+- For greetings, respond politely and ask how you can help with Nigerian law
+
+Answer:"""
+
+PROMPT = PromptTemplate(
+    template=custom_template, input_variables=["context", "question"]
+)
 
 
 # =====================================================
-# Conversational QA Chain with fixed memory
+# Conversational QA Chain with custom prompt
 # =====================================================
 memory = ConversationBufferMemory(
     memory_key="chat_history", 
@@ -65,32 +79,51 @@ qa_chain = ConversationalRetrievalChain.from_llm(
     retriever=retriever,
     memory=memory,
     return_source_documents=True,
+    combine_docs_chain_kwargs={"prompt": PROMPT}  # Use custom prompt
 )
 
 
 # =====================================================
-# Chat function with better token management
+# Chat function with better response handling
 # =====================================================
 def answer_question(user_input, lang_choice, history=[]):
     try:
-        # Pick system prompt
-        if lang_choice == "pidgin":
-            system_prompt = pidgin_system_prompt
-        else:
-            system_prompt = english_system_prompt
+        # Handle greetings and simple queries
+        user_lower = user_input.lower().strip()
+        if user_lower in ["hello", "hi", "hey", "good morning", "good afternoon", "good evening"]:
+            if lang_choice == "pidgin":
+                response = "Hello! How far? I be your Nigerian Legal AI Assistant. Wetin you wan know about Nigerian law today? ⚠️ No be legal advice o, abeg meet lawyer if matter serious."
+            else:
+                response = "Hello! I'm your Nigerian Legal AI Assistant. How can I help you with Nigerian law today? ⚠️ This is not legal advice. Please consult a qualified lawyer for specific issues."
+            
+            history.append(("You: " + user_input, "Bot: " + response))
+            return history, history
 
         # Truncate user input if too long
         max_input_length = 200  # Limit user input length
         if len(user_input) > max_input_length:
             user_input = user_input[:max_input_length] + "..."
 
-        # Create shorter question format
-        question = f"{system_prompt}\nQ: {user_input}"
-        
-        # Run QA
-        result = qa_chain.invoke({"question": question})
+        # Run QA with simple question
+        result = qa_chain.invoke({"question": user_input})
         answer = result["answer"]
 
+        # Clean up the answer - remove any retrieval artifacts
+        if "Use the following pieces of context" in answer:
+            # If the model returns retrieval instructions, provide a fallback
+            if lang_choice == "pidgin":
+                answer = "I dey try find information about your question for Nigerian law documents. Wetin specifically you wan know? ⚠️ No be legal advice o."
+            else:
+                answer = "I'm searching through Nigerian law documents for your question. Could you be more specific about what you'd like to know? ⚠️ This is not legal advice."
+
+        # Add disclaimer if not present
+        if lang_choice == "pidgin":
+            if "No be legal advice" not in answer:
+                answer += "\n\n⚠️ No be legal advice o, abeg meet lawyer if matter serious."
+        else:
+            if "not legal advice" not in answer.lower():
+                answer += "\n\n⚠️ This is not legal advice. Please consult a qualified lawyer for specific issues."
+
         # Collect sources (with sections) - limit to top 3
         sources = []
         for doc in result["source_documents"][:3]:  # Limit to top 3 sources
@@ -98,11 +131,11 @@ def answer_question(user_input, lang_choice, history=[]):
             source = doc.metadata.get("source", "Unknown Document").replace(".pdf", "")
             sources.append(f"[{section}] from {source}")
 
-        if sources:
+        if sources and len(answer) < 400:  # Only add sources if answer isn't too long
             answer += "\n\n📚 Sources:\n" + "\n".join(sources)
 
         # Truncate answer if too long
-        max_answer_length = 800
+        max_answer_length = 600
         if len(answer) > max_answer_length:
             answer = answer[:max_answer_length] + "...\n\n⚠️ Response truncated due to length limits."
 
@@ -115,7 +148,10 @@ def answer_question(user_input, lang_choice, history=[]):
         return history, history
 
     except Exception as e:
-        error_msg = f"Sorry, I encountered an error: {str(e)[:100]}..."
+        if lang_choice == "pidgin":
+            error_msg = f"Sorry o, I get small wahala: {str(e)[:50]}... Try ask again."
+        else:
+            error_msg = f"Sorry, I encountered an error: {str(e)[:50]}... Please try asking again."
         history.append(("You: " + user_input, "Bot: " + error_msg))
         return history, history