app.py

import gradio as gr
import chromadb
import os
from openai import OpenAI
import json
from typing import List, Dict
import re
from sentence_transformers import SentenceTransformer
from loguru import logger

class SentenceTransformerEmbeddings:
    def __init__(self, model_name: str = 'all-MiniLM-L6-v2'):
        self.model = SentenceTransformer(model_name)

    def __call__(self, input: List[str]) -> List[List[float]]:
        embeddings = self.model.encode(input)
        return embeddings.tolist()

class LegalAssistant:
    def __init__(self):
        # Initialize ChromaDB
        self.chroma_client = chromadb.Client()
        
        # Initialize embedding function
        self.embedding_function = SentenceTransformerEmbeddings()
        
        # Create or get collection with proper embedding function
        self.collection = self.chroma_client.get_or_create_collection(
            name="legal_documents",
            embedding_function=self.embedding_function
        )
        
        # Load documents if collection is empty
        if self.collection.count() == 0:
            self._load_documents()
        
        # Initialize Mistral AI client
        self.mistral_client = OpenAI(
            api_key=os.environ.get("MISTRAL_API_KEY", "dfb2j1YDsa298GXTgZo3juSjZLGUCfwi"),
            base_url="https://api.mistral.ai/v1"
        )
        
        # Define system prompt with strict rules
        self.system_prompt = """You are a specialized legal assistant that MUST follow these STRICT rules:

CRITICAL RULE:
YOU MUST ONLY USE INFORMATION FROM THE PROVIDED CONTEXT. DO NOT USE ANY EXTERNAL KNOWLEDGE, INCLUDING KNOWLEDGE ABOUT IPC, CONSTITUTION, OR ANY OTHER LEGAL DOCUMENTS.

RESPONSE FORMAT RULES:
1. ALWAYS structure your response in this exact JSON format:
   {
     "answer": "Your detailed answer here using ONLY information from the provided context",
     "reference_sections": ["Exact section titles from the context"],
     "summary": "2-3 line summary using ONLY information from context",
     "confidence": "HIGH/MEDIUM/LOW based on context match"
   }

STRICT CONTENT RULES:
1. NEVER mention or reference IPC, Constitution, or any laws not present in the context
2. If the information is not in the context, respond ONLY with:
   {
     "answer": "This information is not present in the provided document.",
     "reference_sections": [],
     "summary": "Information not found in document",
     "confidence": "LOW"
   }
3. ONLY cite sections that are explicitly present in the provided context
4. DO NOT make assumptions or inferences beyond the context
5. DO NOT combine information from external knowledge

CONTEXT USAGE RULES:
1. HIGH confidence: Only when exact information is found in context
2. MEDIUM confidence: When partial information is found
3. LOW confidence: When information is unclear or not found
4. If multiple sections are relevant, cite ALL relevant sections from context

PROHIBITED ACTIONS:
1. NO references to IPC sections
2. NO references to Constitutional articles
3. NO mentions of case law not in context
4. NO legal interpretations beyond context
5. NO combining document information with external knowledge

ERROR HANDLING:
1. If query is about laws not in context: State "This topic is not covered in the provided document"
2. If query is unclear: Request specific clarification about which part of the document to check
3. If context is insufficient: State "The document does not contain this information"
"""

    def _load_documents(self):
        """Load and index documents from a2023-45.txt and index.txt"""
        try:
            # Read the main document
            with open('a2023-45.txt', 'r', encoding='utf-8') as f:
                document = f.read()
            
            # Read the index
            with open('index.txt', 'r', encoding='utf-8') as f:
                index_content = f.readlines()
            
            # Parse index and split document
            sections = []
            current_section = ""
            current_title = ""
            
            for line in document.split('\n'):
                if any(index_line.strip() in line for index_line in index_content):
                    if current_section:
                        sections.append({
                            "title": current_title,
                            "content": current_section.strip()
                        })
                    current_title = line.strip()
                    current_section = ""
                else:
                    current_section += line + "\n"
            
            # Add the last section
            if current_section:
                sections.append({
                    "title": current_title,
                    "content": current_section.strip()
                })
            
            # Add to ChromaDB
            documents = [section["content"] for section in sections]
            metadatas = [{"title": section["title"], "source": "a2023-45.txt", "section_number": i + 1} 
                        for i, section in enumerate(sections)]
            ids = [f"section_{i+1}" for i in range(len(sections))]
            
            self.collection.add(
                documents=documents,
                metadatas=metadatas,
                ids=ids
            )
            
            logger.info(f"Loaded {len(sections)} sections into ChromaDB")
            
        except Exception as e:
            logger.error(f"Error loading documents: {str(e)}")
            raise

    def validate_query(self, query: str) -> tuple[bool, str]:
        """Validate the input query"""
        if not query or len(query.strip()) < 10:
            return False, "Query too short. Please provide more details (minimum 10 characters)."
        if len(query) > 500:
            return False, "Query too long. Please be more concise (maximum 500 characters)."
        if not re.search(r'[?.]$', query):
            return False, "Query must end with a question mark or period."
        return True, ""

    def _search_documents(self, query: str) -> tuple[str, List[str]]:
        """Search ChromaDB for relevant documents"""
        try:
            results = self.collection.query(
                query_texts=[query],
                n_results=3
            )
            
            if results and results['documents']:
                documents = results['documents'][0]
                metadata = results['metadatas'][0]
                
                # Format the context with section titles
                formatted_docs = []
                references = []
                
                for doc, meta in zip(documents, metadata):
                    formatted_docs.append(f"{meta['title']}:\n{doc}")
                    references.append(f"{meta['title']} (Section {meta['section_number']})")
                
                return "\n\n".join(formatted_docs), references
            return "", []
            
        except Exception as e:
            logger.error(f"Search error: {str(e)}")
            return "", []

    def get_response(self, query: str) -> Dict:
        """Get response from Mistral AI with context from ChromaDB"""
        # Validate query
        is_valid, error_message = self.validate_query(query)
        if not is_valid:
            return {
                "answer": error_message,
                "references": [],
                "summary": "Invalid query",
                "confidence": "LOW"
            }

        try:
            # Get relevant context from ChromaDB
            context, sources = self._search_documents(query)
            
            if not context:
                return {
                    "answer": "This information is not present in the provided document.",
                    "references": [],
                    "summary": "Information not found in document",
                    "confidence": "LOW"
                }
            
            # Prepare content with explicit instructions
            content = f"""IMPORTANT: ONLY use information from the following context to answer the question. DO NOT use any external knowledge.

Context Sections:
{context}

Available Document Sections:
{', '.join(sources)}

Question: {query}

Remember: ONLY use information from the above context. If the information is not in the context, state that it's not in the document."""
            
            # Get response from Mistral AI
            response = self.mistral_client.chat.completions.create(
                model="mistral-medium",
                messages=[
                    {
                        "role": "system",
                        "content": self.system_prompt
                    },
                    {
                        "role": "user",
                        "content": content
                    }
                ],
                temperature=0.1,
                max_tokens=1000
            )
            
            # Parse response
            if response.choices and len(response.choices) > 0:
                try:
                    result = json.loads(response.choices[0].message.content)
                    
                    # Validate that references only contain sections from sources
                    valid_references = [ref for ref in result.get("reference_sections", []) 
                                     if any(source.split(" (Section")[0] in ref for source in sources)]
                    
                    # If references mention unauthorized sources, return error
                    if len(valid_references) != len(result.get("reference_sections", [])):
                        logger.warning("Response contained unauthorized references")
                        return {
                            "answer": "Error: Response contained unauthorized references. Only information from the provided document is allowed.",
                            "references": [],
                            "summary": "Invalid response generated",
                            "confidence": "LOW"
                        }
                    
                    return {
                        "answer": result.get("answer", "No answer provided"),
                        "references": valid_references,
                        "summary": result.get("summary", ""),
                        "confidence": result.get("confidence", "LOW")
                    }
                except json.JSONDecodeError:
                    logger.error("Failed to parse response JSON")
                    return {
                        "answer": "Error: Response format invalid",
                        "references": [],
                        "summary": "Response parsing failed",
                        "confidence": "LOW"
                    }
            
            return {
                "answer": "No valid response received",
                "references": [],
                "summary": "Response generation failed",
                "confidence": "LOW"
            }
            
        except Exception as e:
            logger.error(f"Error in get_response: {str(e)}")
            return {
                "answer": f"Error: {str(e)}",
                "references": [],
                "summary": "System error occurred",
                "confidence": "LOW"
            }

# Initialize the assistant
assistant = LegalAssistant()

# Create Gradio interface
def process_query(query: str) -> tuple:
    """Process the query and return formatted response"""
    response = assistant.get_response(query)
    return (
        response["answer"],
        ", ".join(response["references"]) if response["references"] else "No specific references",
        response["summary"] if response["summary"] else "No summary available",
        response["confidence"]
    )

# Create the Gradio interface with a professional theme
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # Indian Legal Assistant
    ## Guidelines for Queries:
    1. Be specific and clear in your questions
    2. End questions with a question mark or period
    3. Keep queries between 10-500 characters
    4. Questions will be answered based ONLY on the provided legal document
    """)
    
    with gr.Row():
        query_input = gr.Textbox(
            label="Enter your legal query",
            placeholder="e.g., What are the main provisions in this document?"
        )
    
    with gr.Row():
        submit_btn = gr.Button("Submit", variant="primary")
    
    with gr.Row():
        confidence_output = gr.Textbox(label="Confidence Level")
    
    with gr.Row():
        answer_output = gr.Textbox(label="Answer", lines=5)
    
    with gr.Row():
        with gr.Column():
            references_output = gr.Textbox(label="Document References", lines=3)
        with gr.Column():
            summary_output = gr.Textbox(label="Summary", lines=2)
    
    gr.Markdown("""
    ### Important Notes:
    - Responses are based ONLY on the provided document
    - No external legal knowledge is used
    - All references are from the document itself
    - Confidence levels indicate how well the answer matches the document content
    """)
    
    submit_btn.click(
        fn=process_query,
        inputs=[query_input],
        outputs=[answer_output, references_output, summary_output, confidence_output]
    )

# Launch the app
if __name__ == "__main__":
    demo.launch()