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()