Update app.py

app.py

@@ -8,6 +8,7 @@ from astrapy.db import AstraDB
 from dotenv import load_dotenv
 from huggingface_hub import login
 from sentence_transformers import SentenceTransformer
+import json
 
 # Load environment variables
 load_dotenv()
@@ -68,16 +69,45 @@ class LegalTextSearchBot:
             # Generate embedding for the query
             query_embedding = self.embedding_model.encode(query).tolist()
             
-            # Perform vector search in AstraDB
-            results = self.collection.vector_find(
+            # First try searching in searchable_text
+            results = list(self.collection.vector_find(
                 query_embedding,
-                limit=5
-            )
-            return list(results)
+                limit=5,
+                fields=["section_number", "title", "chapter_info", "content", "searchable_text"]
+            ))
+            
+            if not results:
+                # If no results, try a more general search
+                results = list(self.collection.find(
+                    {},
+                    limit=5
+                ))
+            
+            return results
+            
         except Exception as e:
             print(f"Error searching AstraDB: {str(e)}")
             return []
 
+    def format_section(self, section: Dict) -> str:
+        """Format a section for display"""
+        try:
+            chapter_info = section.get('chapter_info', {})
+            chapter_title = chapter_info.get('title', 'N/A') if isinstance(chapter_info, dict) else 'N/A'
+            
+            return f"""
+Section {section.get('section_number', 'N/A')}: {section.get('title', 'N/A')}
+Chapter: {chapter_title}
+
+Content:
+{section.get('content', 'N/A')}
+
+{"="*80}
+"""
+        except Exception as e:
+            print(f"Error formatting section: {str(e)}")
+            return str(section)
+
     def search_sections(self, query: str) -> Tuple[str, str]:
         """Search legal sections and return both raw results and AI interpretation"""
         try:
@@ -85,24 +115,24 @@ class LegalTextSearchBot:
             search_results = self._search_astra(query)
             
             if not search_results:
-                return "No relevant sections found.", ""
+                return "No relevant sections found.", "I apologize, but I cannot find relevant sections in the database."
             
             # Format raw results
             raw_results = []
+            context_parts = []
+            
             for result in search_results:
-                section_info = f"""
-                Section {result.get('section_number')}: {result.get('title')}
-                Chapter: {result.get('chapter_info', {}).get('title', 'N/A')}
-                
-                Content:
-                {result.get('content', 'N/A')}
+                # Format for display
+                raw_results.append(self.format_section(result))
                 
-                {"="*80}
-                """
-                raw_results.append(section_info)
+                # Add to context for AI
+                context_parts.append(f"""
+Section {result.get('section_number')}: {result.get('title')}
+{result.get('content', '')}
+""")
             
-            # Combine relevant content for AI interpretation
-            context = "\n\n".join([result.get('content', '') for result in search_results])
+            # Combine context for AI
+            context = "\n\n".join(context_parts)
             
             # Generate AI interpretation
             chain = self.prompt | self.llm
@@ -117,7 +147,9 @@ class LegalTextSearchBot:
             return "\n".join(raw_results), ai_response
             
         except Exception as e:
-            return f"Error processing query: {str(e)}", ""
+            error_msg = f"Error processing query: {str(e)}"
+            print(error_msg)
+            return error_msg, "An error occurred while processing your query."
 
 def create_interface():
     """Create the Gradio interface"""