Update app.py

app.py

@@ -170,6 +170,28 @@ References: {', '.join(section.get('metadata', {}).get('references', [])) or 'No
             logger.error(f"Error formatting section: {str(e)}")
             return str(section)
 
+    def generate_ai_response(self, context: str, query: str) -> str:
+        """Generate AI interpretation with error handling"""
+        try:
+            chain = self.prompt | self.llm
+            response = chain.invoke({
+                "context": context,
+                "chat_history": self.chat_history,
+                "question": query
+            })
+            
+            # Handle different response types
+            if isinstance(response, dict):
+                return response.get('text', str(response))
+            elif isinstance(response, list):
+                return response[0] if response else "No response generated"
+            else:
+                return str(response)
+                
+        except Exception as e:
+            logger.error(f"Error generating AI response: {str(e)}")
+            return "I apologize, but I encountered an error while interpreting the legal sections. Please try rephrasing your question."
+
     def search_sections(self, query: str, progress=gr.Progress()) -> Tuple[str, str]:
         self.is_searching = True
         start_time = time.time()
@@ -198,8 +220,8 @@ References: {', '.join(section.get('metadata', {}).get('references', [])) or 'No
                     
                 raw_results.append(self.format_section(result))
                 context_parts.append(f"""
-Section {result.get('section_number')}: {result.get('title')}
-{result.get('content', '')}
+Section {result.get('section_number', 'N/A')}: {result.get('title', 'N/A')}
+{result.get('content', 'N/A')}
 """)
                 progress((0.3 + (idx * 0.1)), desc=f"Processing result {idx + 1} of {len(search_results)}...")
             
@@ -209,13 +231,7 @@ Section {result.get('section_number')}: {result.get('title')}
             progress(0.8, desc="Generating AI interpretation...")
             context = "\n\n".join(context_parts)
             
-            chain = self.prompt | self.llm
-            ai_response = chain.invoke({
-                "context": context,
-                "chat_history": self.chat_history,
-                "question": query
-            })
-            
+            ai_response = self.generate_ai_response(context, query)
             self.chat_history += f"\nUser: {query}\nAI: {ai_response}\n"
             
             elapsed_time = time.time() - start_time