app.py

import gradio as gr
from typing import List, Dict, Tuple
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
from transformers import pipeline
import os
from astrapy.db import AstraDB
from dotenv import load_dotenv
from huggingface_hub import login
import time
import threading
from queue import Queue
import asyncio

# Load environment variables
load_dotenv()
login(token=os.getenv("HUGGINGFACE_API_TOKEN"))

class SearchCancelled(Exception):
    pass

class LegalTextSearchBot:
    def __init__(self):
        self.astra_db = AstraDB(
            token=os.getenv("ASTRA_DB_APPLICATION_TOKEN"),
            api_endpoint=os.getenv("ASTRA_DB_API_ENDPOINT")
        )
        self.collection = self.astra_db.collection("legal_content")
        
        pipe = pipeline(
            "text-generation",
            model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
            max_new_tokens=512,
            temperature=0.7,
            top_p=0.95,
            repetition_penalty=1.15
        )
        self.llm = HuggingFacePipeline(pipeline=pipe)
        
        self.template = """
        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 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. 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 (Legal Sections): {context}
        
        Chat History: {chat_history}
        
        Question: {question}
        
        Answer:"""
        
        self.prompt = ChatPromptTemplate.from_template(self.template)
        self.chat_history = ""
        self.cancel_search = False

    def _search_astra(self, query: str) -> List[Dict]:
        if self.cancel_search:
            raise SearchCancelled("Search was cancelled by user")
            
        try:
            results = list(self.collection.vector_find(
                query,
                limit=5,
                fields=["section_number", "title", "chapter_number", "chapter_title", 
                       "content", "type", "metadata"]
            ))
            
            if not results and not self.cancel_search:
                results = list(self.collection.find(
                    {},
                    limit=5
                ))
            
            return results
            
        except Exception as e:
            if not isinstance(e, SearchCancelled):
                print(f"Error searching AstraDB: {str(e)}")
            raise

    def format_section(self, section: Dict) -> str:
        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')}
Type: {section.get('type', 'section')}

Content:
{section.get('content', 'N/A')}

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, progress=gr.Progress()) -> Tuple[str, str]:
        self.cancel_search = False
        
        try:
            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."
            
            progress(0.3, desc="Processing results...")
            raw_results = []
            context_parts = []
            
            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))
                context_parts.append(f"""
Section {result.get('section_number')}: {result.get('title')}
{result.get('content', '')}
""")
                progress((0.3 + (idx * 0.1)), desc="Processing results...")
            
            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
            })
            
            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():
    with gr.Blocks(title="Bharatiya Nyaya Sanhita Search", theme=gr.themes.Soft()) as iface:
        gr.Markdown("""
        # 📚 Bharatiya Nyaya Sanhita Legal Search System
        
        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:
        """)
        
        search_bot = LegalTextSearchBot()
        
        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():
            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.Markdown(
                    label="📜 Relevant Legal Sections"
                )
            with gr.Column():
                ai_output = gr.Markdown(
                    label="🤖 AI Interpretation"
                )
        
        gr.Examples(
            examples=[
                "What are the penalties for public servants who conceal information?",
                "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"
        )
        
        def search(query):
            return search_bot.search_sections(query)
        
        def stop_search():
            search_bot.cancel()
            return "Search cancelled.", "Search stopped by user."
        
        search_button.click(
            fn=search,
            inputs=query_input,
            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
        )
        
        query_input.submit(
            fn=search,
            inputs=query_input,
            outputs=[raw_output, ai_output],
            cancels=[stop_button]
        )
    
    return iface

if __name__ == "__main__":
    demo = create_interface()
    demo.launch()
else:
    demo = create_interface()
    app = demo.launch(share=False)