Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from typing import List, Dict
+from typing import List, Dict, Tuple
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
 from transformers import pipeline
@@ -7,6 +7,7 @@ import os
 from astrapy.db import AstraDB
 from dotenv import load_dotenv
 from huggingface_hub import login
+from sentence_transformers import SentenceTransformer
 
 # Load environment variables
 load_dotenv()
@@ -14,7 +15,7 @@ load_dotenv()
 # Login to Hugging Face Hub
 login(token=os.getenv("HUGGINGFACE_API_TOKEN"))
 
-class AstraDBChatbot:
+class LegalTextSearchBot:
     def __init__(self):
         # Initialize AstraDB connection
         self.astra_db = AstraDB(
@@ -25,7 +26,9 @@ class AstraDBChatbot:
         # Set your collection
         self.collection = self.astra_db.collection(os.getenv("ASTRA_DB_COLLECTION"))
 
-        # Initialize the model - using a smaller model suitable for CPU
+        # Initialize the models
+        self.embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+        
         pipe = pipeline(
             "text-generation",
             model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
@@ -38,16 +41,17 @@ class AstraDBChatbot:
         
         # Create prompt template
         self.template = """
-        IMPORTANT: You are a helpful assistant that provides information based on the retrieved context.
+        IMPORTANT: You are a legal assistant that provides accurate information based on the Indian legal sections provided in the context.
         
         STRICT RULES:
-        1. Base your response ONLY on the provided context
-        2. If you cannot find relevant information, respond with: "I apologize, but I cannot find information about that in the database."
+        1. Base your response ONLY on the provided legal sections
+        2. If you cannot find relevant information, respond with: "I apologize, but I cannot find information about that in the legal database."
         3. Do not make assumptions or use external knowledge
-        4. Be concise and accurate in your responses
-        5. If quoting from the context, clearly indicate it
+        4. Always cite the specific section numbers you're referring to
+        5. Be precise and accurate in your legal interpretations
+        6. If quoting from the sections, use quotes and cite the section number
         
-        Context: {context}
+        Context (Legal Sections): {context}
         
         Chat History: {chat_history}
         
@@ -59,55 +63,138 @@ class AstraDBChatbot:
         self.chat_history = ""
 
     def _search_astra(self, query: str) -> List[Dict]:
-        """Search AstraDB for relevant documents"""
+        """Search AstraDB for relevant legal sections"""
         try:
+            # Generate embedding for the query
+            query_embedding = self.embedding_model.encode(query).tolist()
+            
             # Perform vector search in AstraDB
             results = self.collection.vector_find(
-                query,
-                limit=5  # Adjust the limit based on your needs
+                query_embedding,
+                limit=5
             )
-            return results
+            return list(results)
         except Exception as e:
             print(f"Error searching AstraDB: {str(e)}")
             return []
 
-    def chat(self, query: str, history) -> str:
-        """Process a query and return a response"""
+    def search_sections(self, query: str) -> Tuple[str, str]:
+        """Search legal sections and return both raw results and AI interpretation"""
         try:
-            # Search AstraDB for relevant content
+            # Search AstraDB for relevant sections
             search_results = self._search_astra(query)
             
             if not search_results:
-                return "I apologize, but I cannot find information about that in the database."
+                return "No relevant sections found.", ""
+            
+            # Format raw results
+            raw_results = []
+            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')}
+                
+                {"="*80}
+                """
+                raw_results.append(section_info)
             
-            # Extract and combine relevant content from search results
+            # Combine relevant content for AI interpretation
             context = "\n\n".join([result.get('content', '') for result in search_results])
             
-            # Generate response using LLM
+            # Generate AI interpretation
             chain = self.prompt | self.llm
-            result = chain.invoke({
+            ai_response = chain.invoke({
                 "context": context,
                 "chat_history": self.chat_history,
                 "question": query
             })
             
-            self.chat_history += f"\nUser: {query}\nAI: {result}\n"
+            self.chat_history += f"\nUser: {query}\nAI: {ai_response}\n"
+            
+            return "\n".join(raw_results), ai_response
             
-            return result
         except Exception as e:
-            return f"Error processing query: {str(e)}"
-
-# Initialize the chatbot
-chatbot = AstraDBChatbot()
+            return f"Error processing query: {str(e)}", ""
 
-# Create the Gradio interface
-iface = gr.ChatInterface(
-    chatbot.chat,
-    title="AstraDB-powered Q&A Chatbot",
-    description="Ask questions and get answers from your AstraDB database.",
-    examples=["What information do you have about this topic?", "Can you tell me more about specific details?"],
-    theme=gr.themes.Soft()
-)
+def create_interface():
+    """Create the Gradio interface"""
+    with gr.Blocks(title="Legal Text Search System", theme=gr.themes.Soft()) as iface:
+        gr.Markdown("""
+        # 📚 Legal Text 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
+        
+        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",
+                placeholder="e.g., What are the penalties for public servants who conceal information?",
+                lines=2
+            )
+        
+        with gr.Row():
+            search_button = gr.Button("🔍 Search Legal Sections", variant="primary")
+        
+        with gr.Row():
+            with gr.Column():
+                raw_output = gr.Textbox(
+                    label="📜 Relevant Legal Sections",
+                    lines=15,
+                    max_lines=30
+                )
+            with gr.Column():
+                ai_output = gr.Textbox(
+                    label="🤖 AI Interpretation",
+                    lines=15,
+                    max_lines=30
+                )
+        
+        # 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"
+            ],
+            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
+        
+        # Connect the button to the search function
+        search_button.click(
+            fn=search,
+            inputs=query_input,
+            outputs=[raw_output, ai_output]
+        )
+        
+        # Also allow pressing Enter to search
+        query_input.submit(
+            fn=search,
+            inputs=query_input,
+            outputs=[raw_output, ai_output]
+        )
+    
+    return iface
 
-# Launch the interface
-iface.launch()
+# Create and launch the interface
+if __name__ == "__main__":
+    demo = create_interface()
+    demo.launch()