Update app.py

app.py

@@ -7,28 +7,25 @@ import os
 from astrapy.db import AstraDB
 from dotenv import load_dotenv
 from huggingface_hub import login
-from sentence_transformers import SentenceTransformer
-import json
+import time
+import threading
+from queue import Queue
+import asyncio
 
 # Load environment variables
 load_dotenv()
-
-# Login to Hugging Face Hub
 login(token=os.getenv("HUGGINGFACE_API_TOKEN"))
 
+class SearchCancelled(Exception):
+    pass
+
 class LegalTextSearchBot:
     def __init__(self):
-        # Initialize AstraDB connection
         self.astra_db = AstraDB(
             token=os.getenv("ASTRA_DB_APPLICATION_TOKEN"),
             api_endpoint=os.getenv("ASTRA_DB_API_ENDPOINT")
         )
-        
-        # Set your collection
-        self.collection = self.astra_db.collection(os.getenv("ASTRA_DB_COLLECTION"))
-
-        # Initialize the models
-        self.embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+        self.collection = self.astra_db.collection("legal_content")
         
         pipe = pipeline(
             "text-generation",
@@ -40,7 +37,6 @@ class LegalTextSearchBot:
         )
         self.llm = HuggingFacePipeline(pipeline=pipe)
         
-        # Create prompt template
         self.template = """
         IMPORTANT: You are a legal assistant that provides accurate information based on the Indian legal sections provided in the context.
         
@@ -62,22 +58,21 @@ class LegalTextSearchBot:
         
         self.prompt = ChatPromptTemplate.from_template(self.template)
         self.chat_history = ""
+        self.cancel_search = False
 
     def _search_astra(self, query: str) -> List[Dict]:
-        """Search AstraDB for relevant legal sections"""
-        try:
-            # Generate embedding for the query
-            query_embedding = self.embedding_model.encode(query).tolist()
+        if self.cancel_search:
+            raise SearchCancelled("Search was cancelled by user")
             
-            # First try searching in searchable_text
+        try:
             results = list(self.collection.vector_find(
-                query_embedding,
+                query,
                 limit=5,
-                fields=["section_number", "title", "chapter_info", "content", "searchable_text"]
+                fields=["section_number", "title", "chapter_number", "chapter_title", 
+                       "content", "type", "metadata"]
             ))
             
-            if not results:
-                # If no results, try a more general search
+            if not results and not self.cancel_search:
                 results = list(self.collection.find(
                     {},
                     limit=5
@@ -86,55 +81,59 @@ class LegalTextSearchBot:
             return results
             
         except Exception as e:
-            print(f"Error searching AstraDB: {str(e)}")
-            return []
+            if not isinstance(e, SearchCancelled):
+                print(f"Error searching AstraDB: {str(e)}")
+            raise
 
     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'
+        if self.cancel_search:
+            raise SearchCancelled("Search was cancelled by user")
             
+        try:
             return f"""
+{'='*80}
+Chapter {section.get('chapter_number', 'N/A')}: {section.get('chapter_title', 'N/A')}
 Section {section.get('section_number', 'N/A')}: {section.get('title', 'N/A')}
-Chapter: {chapter_title}
+Type: {section.get('type', 'section')}
 
 Content:
 {section.get('content', 'N/A')}
 
-{"="*80}
+References: {', '.join(section.get('metadata', {}).get('references', [])) or 'None'}
+{'='*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"""
+    def search_sections(self, query: str, progress=gr.Progress()) -> Tuple[str, str]:
+        self.cancel_search = False
+        
         try:
-            # Search AstraDB for relevant sections
+            progress(0, desc="Searching relevant sections...")
             search_results = self._search_astra(query)
             
             if not search_results:
                 return "No relevant sections found.", "I apologize, but I cannot find relevant sections in the database."
             
-            # Format raw results
+            progress(0.3, desc="Processing results...")
             raw_results = []
             context_parts = []
             
-            for result in search_results:
-                # Format for display
+            for idx, result in enumerate(search_results):
+                if self.cancel_search:
+                    raise SearchCancelled("Search was cancelled by user")
+                    
                 raw_results.append(self.format_section(result))
-                
-                # Add to context for AI
                 context_parts.append(f"""
 Section {result.get('section_number')}: {result.get('title')}
 {result.get('content', '')}
 """)
+                progress((0.3 + (idx * 0.1)), desc="Processing results...")
             
-            # Combine context for AI
+            progress(0.8, desc="Generating AI interpretation...")
             context = "\n\n".join(context_parts)
             
-            # Generate AI interpretation
             chain = self.prompt | self.llm
             ai_response = chain.invoke({
                 "context": context,
@@ -144,30 +143,33 @@ Section {result.get('section_number')}: {result.get('title')}
             
             self.chat_history += f"\nUser: {query}\nAI: {ai_response}\n"
             
+            progress(1.0, desc="Complete!")
             return "\n".join(raw_results), ai_response
             
+        except SearchCancelled:
+            return "Search cancelled by user.", "Search was stopped. Please try again with a new query."
         except Exception as e:
             error_msg = f"Error processing query: {str(e)}"
             print(error_msg)
             return error_msg, "An error occurred while processing your query."
 
+    def cancel(self):
+        self.cancel_search = True
+
 def create_interface():
-    """Create the Gradio interface"""
-    with gr.Blocks(title="Legal Text Search System", theme=gr.themes.Soft()) as iface:
+    with gr.Blocks(title="Bharatiya Nyaya Sanhita Search", theme=gr.themes.Soft()) as iface:
         gr.Markdown("""
-        # 📚 Legal Text Search System
+        # 📚 Bharatiya Nyaya Sanhita Legal Search System
         
-        This system allows you to search through Indian legal sections and get both:
-        1. 📜 Raw section contents that match your query
-        2. 🤖 AI-powered interpretation of the relevant sections
+        Search through the Bharatiya Nyaya Sanhita, 2023 and get:
+        1. 📜 Relevant sections, explanations, and illustrations
+        2. 🤖 AI-powered interpretation of the legal content
         
         Enter your legal query below:
         """)
         
-        # Initialize the search bot
         search_bot = LegalTextSearchBot()
         
-        # Create input and output components
         with gr.Row():
             query_input = gr.Textbox(
                 label="Your Query",
@@ -176,57 +178,65 @@ def create_interface():
             )
         
         with gr.Row():
-            search_button = gr.Button("🔍 Search Legal Sections", variant="primary")
+            with gr.Column(scale=4):
+                search_button = gr.Button("🔍 Search Legal Sections", variant="primary")
+            with gr.Column(scale=1):
+                stop_button = gr.Button("🛑 Stop Search", variant="stop")
         
         with gr.Row():
             with gr.Column():
-                raw_output = gr.Textbox(
-                    label="📜 Relevant Legal Sections",
-                    lines=15,
-                    max_lines=30
+                raw_output = gr.Markdown(
+                    label="📜 Relevant Legal Sections"
                 )
             with gr.Column():
-                ai_output = gr.Textbox(
-                    label="🤖 AI Interpretation",
-                    lines=15,
-                    max_lines=30
+                ai_output = gr.Markdown(
+                    label="🤖 AI Interpretation"
                 )
         
-        # Add example queries
         gr.Examples(
             examples=[
                 "What are the penalties for public servants who conceal information?",
-                "What is the punishment for corruption?",
-                "What happens if a public servant fails to prevent an offense?",
-                "What are the legal consequences for concealing design to commit offence?",
-                "Explain the duties and responsibilities of public servants"
+                "What constitutes criminal conspiracy?",
+                "Explain the provisions related to culpable homicide",
+                "What are the penalties for causing death by negligence?",
+                "What are the punishments for corruption?"
             ],
             inputs=query_input,
             label="Example Queries"
         )
         
-        # Set up the search function
         def search(query):
-            raw_results, ai_response = search_bot.search_sections(query)
-            return raw_results, ai_response
+            return search_bot.search_sections(query)
+        
+        def stop_search():
+            search_bot.cancel()
+            return "Search cancelled.", "Search stopped by user."
         
-        # Connect the button to the search function
         search_button.click(
             fn=search,
             inputs=query_input,
-            outputs=[raw_output, ai_output]
+            outputs=[raw_output, ai_output],
+            cancels=[stop_button]  # Cancel any ongoing search when stop is clicked
+        )
+        
+        stop_button.click(
+            fn=stop_search,
+            outputs=[raw_output, ai_output],
+            cancels=[search_button]  # Cancel the search button when stop is clicked
         )
         
-        # Also allow pressing Enter to search
         query_input.submit(
             fn=search,
             inputs=query_input,
-            outputs=[raw_output, ai_output]
+            outputs=[raw_output, ai_output],
+            cancels=[stop_button]
         )
     
     return iface
 
-# Create and launch the interface
 if __name__ == "__main__":
     demo = create_interface()
-    demo.launch()
+    demo.launch()
+else:
+    demo = create_interface()
+    app = demo.launch(share=False)