import gradio as gr import chromadb import os from openai import OpenAI import json from typing import List, Dict import re class LegalAssistant: def __init__(self): # Initialize ChromaDB self.chroma_client = chromadb.Client() self.collection = self.chroma_client.get_or_create_collection("legal_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 trained on Indian law. You MUST follow these strict rules: RESPONSE FORMAT RULES: 1. ALWAYS structure your response in this exact JSON format: { "answer": "Your detailed answer here", "reference_sections": ["Section X of Act Y", ...], "summary": "2-3 line summary", "confidence": "HIGH/MEDIUM/LOW" } CONTENT RULES: 1. NEVER make assumptions or provide information not supported by Indian law 2. ALWAYS cite specific sections, acts, and legal precedents 3. If information is insufficient, explicitly state "Insufficient information" in answer 4. NEVER provide legal advice, only legal information 5. For any constitutional matters, ALWAYS cite relevant Articles ACCURACY RULES: 1. If confidence is less than 80%, mark as LOW confidence 2. If multiple interpretations exist, list ALL with citations 3. If law has been amended, specify the latest amendment date 4. For case law, cite the full case reference PROHIBITED: 1. NO personal opinions 2. NO hypothetical scenarios 3. NO interpretation of ongoing cases 4. NO advice on specific legal situations ERROR HANDLING: 1. If query is unclear: Request clarification 2. If outside Indian law scope: State "Outside scope of Indian law" 3. If conflicting laws exist: List all applicable laws""" 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." if len(query) > 500: return False, "Query too long. Please be more concise." 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.get('metadatas', [[]])[0] sources = [m.get('source', 'Unknown') for m in metadata] return "\n\n".join(documents), sources return "", [] except Exception as e: print(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) # Prepare content content = f"""Context: {context} Sources: {', '.join(sources)} Question: {query}""" if context else query # 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) return { "answer": result.get("answer", "No answer provided"), "references": result.get("reference_sections", []), "summary": result.get("summary", ""), "confidence": result.get("confidence", "LOW") } except json.JSONDecodeError: return { "answer": "Error: Response format invalid", "references": [], "summary": "Response parsing failed", "confidence": "LOW" } return { "answer": "No response received", "references": [], "summary": "Response generation failed", "confidence": "LOW" } except Exception as 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: 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 3. Provide relevant context if available 4. Keep queries between 10-500 characters """) with gr.Row(): query_input = gr.Textbox( label="Enter your legal query", placeholder="e.g., What is the legal age for marriage in India as per current laws?" ) 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="Legal References", lines=3) with gr.Column(): summary_output = gr.Textbox(label="Summary", lines=2) gr.Markdown(""" ### Important Notes: - This assistant provides legal information, not legal advice - Always verify information with a qualified legal professional - Information is based on Indian law only """) submit_btn.click( fn=process_query, inputs=[query_input], outputs=[answer_output, references_output, summary_output, confidence_output] ) # Launch the app demo.launch()