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
from sentence_transformers import SentenceTransformer
import torch
import os
from astrapy.db import AstraDB
from dotenv import load_dotenv
from huggingface_hub import login
import time
import logging
import numpy as np
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"))

class LegalTextSearchBot:
    def __init__(self):
        try:
            # Initialize AstraDB connection
            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(os.getenv("ASTRA_DB_COLLECTION"))
            
            # Initialize language model
            model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
            model = AutoModelForCausalLM.from_pretrained(
                model_name,
                device_map="auto",
                torch_dtype=torch.float32,
            )
            tokenizer = AutoTokenizer.from_pretrained(model_name)
            
            # Initialize text generation pipeline
            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)
            
            # Initialize sentence transformer for embeddings
            self.embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
            
            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
            
            logger.info("Successfully initialized LegalTextSearchBot")
            
        except Exception as e:
            logger.error(f"Error initializing LegalTextSearchBot: {str(e)}")
            raise

    def get_embedding(self, text: str) -> List[float]:
        """Generate embedding vector for text"""
        try:
            # Clean and prepare text
            text = text.replace('\n', ' ').strip()
            if not text:
                text = " "  # Ensure non-empty input
            
            # Generate embedding
            embedding = self.embedding_model.encode(text)
            
            # Pad or truncate to 1024 dimensions
            if len(embedding) < 1024:
                embedding = np.pad(embedding, (0, 1024 - len(embedding)))
            elif len(embedding) > 1024:
                embedding = embedding[:1024]
            
            return embedding.tolist()
            
        except Exception as e:
            logger.error(f"Error generating embedding: {str(e)}")
            raise

    @lru_cache(maxsize=100)
    def _cached_search(self, query: str) -> tuple:
        """Cached version of vector search"""
        try:
            # Generate embedding for query
            query_embedding = self.get_embedding(query)
            
            results = list(self.collection.vector_find(
                query_embedding,
                top_k=5,  # Using top_k instead of limit
                fields=["section_number", "title", "chapter_number", "chapter_title", 
                       "content", "type", "metadata"]
            ))
            return tuple(results)
        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:
                # Fallback to regular search
                cursor = self.collection.find({})
                results = []
                for doc in cursor:
                    if len(results) >= 5:
                        break
                    results.append(doc)
            
            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 generate_ai_response(self, context: str, query: str) -> str:
        """Generate AI interpretation with error handling"""
        try:
            chain = self.prompt | self.llm
            response = chain.invoke({
                "context": context,
                "chat_history": self.chat_history,
                "question": query
            })
            
            # Handle different response types
            if isinstance(response, dict):
                return response.get('text', str(response))
            elif isinstance(response, list):
                return response[0] if response else "No response generated"
            else:
                return str(response)
                
        except Exception as e:
            logger.error(f"Error generating AI response: {str(e)}")
            return "I apologize, but I encountered an error while interpreting the legal sections. Please try rephrasing your question."

    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', 'N/A')}: {result.get('title', 'N/A')}
{result.get('content', 'N/A')}
""")
                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)
            
            ai_response = self.generate_ai_response(context, 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)}")