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)