Update main.py

main.py

@@ -82,29 +82,55 @@ async def analyze_document(data: AnalyzeDocumentInput):
         return {"error": str(e)}
 
 async def generate_response_with_context(user_question: str, relevant_context: str, document_id: str):
-    """Send relevant chunks to Gemini for response generation"""
-    try:
-        prompt = f"""You are a legal document assistant. Answer the user's question based ONLY on the provided context from their legal document.
+    """Let Gemini handle intent classification for better flexibility"""
+    
+    user_question_lower = user_question.lower().strip()
+    
+    # Only check for very simple conversational patterns
+    simple_conversational = user_question_lower in ["thank you", "thanks", "hello", "hi", "bye", "ok", "okay"]
+    
+    if simple_conversational:
+        return "You're welcome! I'm here to help you understand your legal document. Feel free to ask about any terms, risks, or clauses you'd like clarified."
+    
+    # Check if context is available
+    has_relevant_context = relevant_context and len(relevant_context.strip()) > 30
+    
+    if has_relevant_context:
+        prompt = f"""You are NyayAI, a legal document analysis assistant.
 
-Context from document {document_id}:
+DOCUMENT CONTEXT:
 {relevant_context}
 
-User Question: {user_question}
+USER QUESTION: "{user_question}"
 
-Instructions:
-- Provide a clear, accurate answer based on the context above
-- If the answer isn't in the context, say "I cannot find information about this in the provided document"
-- Include specific quotes from the document when relevant
-- Keep your answer focused on legal implications and key details
+RESPONSE RULES:
+1. If the question is related to the document content above: Answer using ONLY the provided context
+2. If the question is about legal documents but the context doesn't contain that information: Say "I don't see information about that specific topic in this document"
+3. If the question is completely unrelated to legal documents (like weather, cooking, sports, etc.): Say "Please ask me questions related to your legal document analysis only"
+4. If you're unsure whether it's document-related: Ask for clarification about which aspect of the document they want to know about
 
 Answer:"""
-        
+    
+    else:
+        # No context available - let Gemini decide if it's document-related
+        prompt = f"""You are NyayAI, a legal document analysis assistant.
+
+USER QUESTION: "{user_question}"
+
+The user is asking about a legal document, but I don't have relevant information about their specific question in the document analysis.
+
+RESPONSE RULES:
+1. If the question seems related to legal document analysis: Say "I don't see information about that specific topic in this document. Please ask about the terms, clauses, risks, or other content from the analysis."
+2. If the question is completely unrelated to legal documents: Say "Please ask me questions related to your legal document analysis only"
+
+Answer:"""
+    
+    try:
         model = genai.GenerativeModel('gemini-1.5-flash')
         response = model.generate_content(prompt)
         return response.text
-        
     except Exception as e:
-        return f"Error generating response: {str(e)}"
+        return "I encountered an error processing your question. Please try asking about your document analysis again."
 
 @app.post("/chat")
 async def chat_with_document(data: ChatInput):