app.py

import os
import gradio as gr
import chromadb
from openai import OpenAI
import json
from sentence_transformers import SentenceTransformer
from loguru import logger
from test_embeddings import test_chromadb_content, initialize_chromadb

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):
        try:
            # Initialize and verify ChromaDB content
            logger.info("Initializing LegalAssistant...")
            
            # Try to verify content, if fails, try to initialize
            if not test_chromadb_content():
                logger.warning("ChromaDB verification failed, attempting to initialize...")
                if not initialize_chromadb():
                    raise ValueError("Failed to initialize ChromaDB")
            
            # Initialize ChromaDB client
            base_path = os.path.dirname(os.path.abspath(__file__))
            chroma_path = os.path.join(base_path, 'chroma_db')
            
            self.chroma_client = chromadb.PersistentClient(path=chroma_path)
            self.embedding_function = SentenceTransformerEmbeddings()
            
            # Get existing collection
            self.collection = self.chroma_client.get_collection(
                name="legal_documents",
                embedding_function=self.embedding_function
            )
            
            logger.info(f"Collection loaded with {self.collection.count()} 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"
            )
            
            logger.info("LegalAssistant initialized successfully")
            
        except Exception as e:
            logger.error(f"Error initializing LegalAssistant: {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)."
        return True, ""

    def get_response(self, query: str) -> dict:
        """Process query and get response from Mistral AI"""
        try:
            # Validate query
            is_valid, error_message = self.validate_query(query)
            if not is_valid:
                return {
                    "answer": error_message,
                    "references": [],
                    "summary": "Invalid query",
                    "confidence": "LOW"
                }

            # Search ChromaDB for relevant content
            results = self.collection.query(
                query_texts=[query],
                n_results=3
            )
            
            if not results['documents'][0]:
                return {
                    "answer": "No relevant information found in the document.",
                    "references": [],
                    "summary": "No matching content",
                    "confidence": "LOW"
                }
            
            # Format context with section titles
            context_parts = []
            references = []
            
            for doc, meta in zip(results['documents'][0], results['metadatas'][0]):
                context_parts.append(f"{meta['title']}:\n{doc}")
                references.append(meta['title'])
            
            context = "\n\n".join(context_parts)
            
            # Prepare system prompt with explicit JSON format
            system_prompt = '''You are a specialized legal assistant that MUST follow these STRICT rules:

1. You MUST ONLY use information from the provided context.
2. DO NOT use any external knowledge about laws, IPC, Constitution, or legal matters.
3. Your response MUST be in this EXACT JSON format:
{
    "answer": "Your detailed answer using ONLY information from the context",
    "reference_sections": ["List of section titles used from context"],
    "summary": "Brief 2-3 line summary",
    "confidence": "HIGH/MEDIUM/LOW"
}

Confidence Level Rules:
- HIGH: When exact information is found in context
- MEDIUM: When partial or indirect information is found
- LOW: When information is unclear or not found

If information is not in context, respond with:
{
    "answer": "This information is not present in the provided document.",
    "reference_sections": [],
    "summary": "Information not found in document",
    "confidence": "LOW"
}'''

            # Prepare user content
            content = f'''Context Sections:
{context}

Question: {query}

IMPORTANT:
1. Use ONLY the information from the above context
2. Format your response as a valid JSON object with the exact structure shown above
3. Include ONLY section titles that exist in the context
4. DO NOT add any text outside the JSON structure
5. Ensure the JSON is properly formatted with double quotes'''

            # Get response from Mistral AI
            response = self.mistral_client.chat.completions.create(
                model="mistral-medium",
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": content}
                ],
                temperature=0.1,
                max_tokens=1000,
                response_format={ "type": "json_object" }
            )
            
            # Parse and validate response
            if response.choices and response.choices[0].message.content:
                try:
                    result = json.loads(response.choices[0].message.content)
                    
                    # Validate response structure
                    required_fields = ["answer", "reference_sections", "summary", "confidence"]
                    if not all(field in result for field in required_fields):
                        raise ValueError("Missing required fields in response")
                    
                    # Validate confidence level
                    if result["confidence"] not in ["HIGH", "MEDIUM", "LOW"]:
                        result["confidence"] = "LOW"
                    
                    # Validate references against context
                    valid_references = [ref for ref in result["reference_sections"]
                                     if ref in references]
                    
                    # If references don't match, adjust confidence
                    if len(valid_references) != len(result["reference_sections"]):
                        result["reference_sections"] = valid_references
                        result["confidence"] = "LOW"
                    
                    # Ensure answer and summary are strings
                    result["answer"] = str(result["answer"])
                    result["summary"] = str(result["summary"])
                    
                    return {
                        "answer": result["answer"],
                        "references": valid_references,
                        "summary": result["summary"],
                        "confidence": result["confidence"]
                    }
                    
                except json.JSONDecodeError as e:
                    logger.error(f"JSON parsing error: {str(e)}")
                    return {
                        "answer": "Error: Failed to parse response format",
                        "references": [],
                        "summary": "Response format error",
                        "confidence": "LOW"
                    }
                except ValueError as e:
                    logger.error(f"Validation error: {str(e)}")
                    return {
                        "answer": "Error: Invalid response structure",
                        "references": [],
                        "summary": "Response validation error",
                        "confidence": "LOW"
                    }
            
            return {
                "answer": "Error: No valid response received",
                "references": [],
                "summary": "No response generated",
                "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
try:
    assistant = LegalAssistant()
except Exception as e:
    logger.error(f"Failed to initialize LegalAssistant: {str(e)}")
    raise

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 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=2)
        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]
    )

if __name__ == "__main__":
    demo.launch()