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, AutoModelForCausalLM, AutoTokenizer
import torch
import os
from astrapy.db import AstraDB
from dotenv import load_dotenv
from huggingface_hub import login
import time
import logging
from functools import lru_cache

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

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

# Initialize model with CPU-compatible settings
model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    torch_dtype=torch.float32,  # Use float32 for CPU compatibility
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

class LegalTextSearchBot:
    def __init__(self):
        try:
            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")
            
            # Initialize pipeline with CPU settings
            pipe = pipeline(
                "text-generation",
                model=model,
                tokenizer=tokenizer,
                max_new_tokens=512,
                temperature=0.7,
                top_p=0.95,
                repetition_penalty=1.15,
                device_map="auto"
            )
            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.is_searching = False
            
        except Exception as e:
            logger.error(f"Error initializing LegalTextSearchBot: {str(e)}")
            raise

    @lru_cache(maxsize=100)
    def _cached_search(self, query: str) -> tuple:
        """Cached version of vector search to improve performance"""
        try:
            results = list(self.collection.vector_find(
                query,
                limit=5,
                fields=["section_number", "title", "chapter_number", "chapter_title", 
                       "content", "type", "metadata"]
            ))
            return tuple(results)  # Convert to tuple for caching
        except Exception as e:
            logger.error(f"Error in vector search: {str(e)}")
            return tuple()

    def _search_astra(self, query: str) -> List[Dict]:
        if not self.is_searching:
            return []
            
        try:
            results = list(self._cached_search(query))
            
            if not results and self.is_searching:
                results = list(self.collection.find(
                    {},
                    limit=5
                ))
            
            return results
            
        except Exception as e:
            logger.error(f"Error searching AstraDB: {str(e)}")
            return []

    def format_section(self, section: Dict) -> str:
        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:
            logger.error(f"Error formatting section: {str(e)}")
            return str(section)

    def search_sections(self, query: str, progress=gr.Progress()) -> Tuple[str, str]:
        self.is_searching = True
        start_time = time.time()
        
        try:
            progress(0, desc="Initializing search...")
            if not query.strip():
                return "Please enter a search query.", "Please provide a specific legal question or topic to search for."
            
            progress(0.1, 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."
            
            if not self.is_searching:
                return "Search cancelled.", "Search was stopped by user."
            
            progress(0.3, desc="Processing results...")
            raw_results = []
            context_parts = []
            
            for idx, result in enumerate(search_results):
                if not self.is_searching:
                    return "Search cancelled.", "Search was stopped 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=f"Processing result {idx + 1} of {len(search_results)}...")
            
            if not self.is_searching:
                return "Search cancelled.", "Search was stopped by user."
            
            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"
            
            elapsed_time = time.time() - start_time
            logger.info(f"Search completed in {elapsed_time:.2f} seconds")
            
            progress(1.0, desc="Search complete!")
            return "\n".join(raw_results), ai_response
            
        except Exception as e:
            logger.error(f"Error processing query: {str(e)}")
            return f"Error processing query: {str(e)}", "An error occurred while processing your query."
        finally:
            self.is_searching = False

    def stop_search(self):
        """Stop the current search operation"""
        self.is_searching = False
        return "Search cancelled.", "Search was stopped by user."

def create_interface():
    with gr.Blocks(title="Bharatiya Nyaya Sanhita Search", theme=gr.themes.Soft()) as iface:
        search_bot = LegalTextSearchBot()
        
        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
        
        *Use the Stop button if you want to cancel a long-running search.*
        """)
        
        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", variant="primary", scale=4)
            stop_button = gr.Button("🛑 Stop", variant="stop", scale=1)
        
        with gr.Row():
            raw_output = gr.Markdown(label="📜 Relevant Legal Sections")
            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"
        )
        
        # Handle search
        search_event = search_button.click(
            fn=search_bot.search_sections,
            inputs=query_input,
            outputs=[raw_output, ai_output],
        )
        
        # Handle stop
        stop_button.click(
            fn=search_bot.stop_search,
            outputs=[raw_output, ai_output],
            cancels=[search_event]
        )
        
        # Handle Enter key
        query_input.submit(
            fn=search_bot.search_sections,
            inputs=query_input,
            outputs=[raw_output, ai_output],
        )
    
    return iface

if __name__ == "__main__":
    try:
        demo = create_interface()
        demo.launch()
    except Exception as e:
        logger.error(f"Error launching application: {str(e)}")
else:
    try:
        demo = create_interface()
        app = demo.launch(share=False)
    except Exception as e:
        logger.error(f"Error launching application: {str(e)}")