""" rag_pipeline.py - Enhanced RAG pipeline with state-of-the-art features for LegalMind AI """ from langchain_groq import ChatGroq from langchain.prompts import ChatPromptTemplate from langchain.schema import SystemMessage, HumanMessage from utils import clean_response import os from typing import List, Tuple, Dict, Any, Optional import time import json # Import project configuration from config import GROQ_API_KEY, MAX_RETRIES, LLM_MODELS, LOGS_DIR # Enhanced LLM setup with multiple model options def get_llm(model: str = "deepseek-r1-distill-llama-70b", temperature: float = 0.2): """ Get LLM with proper error handling, multiple model options, and temperature control Args: model: The model to use temperature: The temperature setting (0.0 to 1.0) Returns: Configured LLM or None if error """ groq_api_key = GROQ_API_KEY if not groq_api_key: print("Error: GROQ_API_KEY not found in environment variables. Please add it to your .env file.") return None # Validate model name valid_models = list(LLM_MODELS.keys()) if model not in valid_models: print(f"Warning: Model {model} not in known list. Defaulting to deepseek-r1-distill-llama-70b.") model = "deepseek-r1-distill-llama-70b" # Clamp temperature temperature = max(0.0, min(1.0, temperature)) try: return ChatGroq( model=model, api_key=groq_api_key, temperature=temperature ) except Exception as e: print(f"Error initializing Groq LLM: {e}") return None # Enhanced legal prompt template with improved context handling and source attribution LEGAL_PROMPT_TEMPLATE = """ You are LegalMind AI, an advanced AI Legal Assistant specializing in Indian law and legal analysis. Use the legal information provided in the context to answer the user's question about Indian law. Focus only on the legal aspects described in the provided context. Your primary objective is to provide accurate, well-structured legal analysis based on the information you have been given. Important guidelines: 1. Stay focused on the legal information in the provided context - do not invent or assume legal principles not present. 2. For each important point or claim you make, refer to the source document by indicating the document number. 3. Structure your response clearly with appropriate headers and bullet points for complex answers. 4. Use legal terminology correctly and precisely. 5. If the question has multiple parts, address each part systematically. 6. If you don't know the answer, clearly state that you don't have enough information. 7. Always end with a disclaimer that this is not legal advice. Context information is below: --------------------- {context} --------------------- Given the context information and not prior knowledge, answer the question: Question: {question} Format your answer professionally and ensure you cite the document numbers for key information. Remember, you're a legal assistant helping with research and analysis. """ # Advanced retrieval function with hybrid search def retrieve_docs(vector_db, query, k=5, retrieval_method="hybrid"): """ Retrieve relevant documents from vector database with improved retrieval options Args: vector_db: The vector database query: User query k: Number of documents to retrieve retrieval_method: "similarity", "mmr", or "hybrid" Returns: List of retrieved documents """ if vector_db is None: return [] try: if retrieval_method == "similarity": # Standard similarity search return vector_db.similarity_search(query, k=k) elif retrieval_method == "mmr": # Maximum Marginal Relevance - better diversity return vector_db.max_marginal_relevance_search(query, k=k, fetch_k=k*2) elif retrieval_method == "hybrid": # Hybrid approach: combine similarity and MMR similarity_docs = vector_db.similarity_search(query, k=int(k/2)) mmr_docs = vector_db.max_marginal_relevance_search(query, k=k-len(similarity_docs)) # Combine and deduplicate all_docs = similarity_docs + mmr_docs unique_docs = [] content_set = set() for doc in all_docs: content_hash = hash(doc.page_content) if content_hash not in content_set: content_set.add(content_hash) unique_docs.append(doc) # Break if we have enough documents if len(unique_docs) >= k: break return unique_docs else: print(f"Unknown retrieval method: {retrieval_method}. Using similarity search.") return vector_db.similarity_search(query, k=k) except Exception as e: print(f"Error retrieving documents: {e}") return [] def get_context(documents): """ Extract content from documents to create context with document identifiers Args: documents: List of documents Returns: Formatted context string """ if not documents: return "No relevant information found." context_parts = [] for i, doc in enumerate(documents): # Add document identifier and metadata if available doc_id = f"Document {i+1}" source = f"(Source: {doc.metadata.get('source', 'Unknown')})" if hasattr(doc, 'metadata') and doc.metadata.get('source') else "" page = f", Page {doc.metadata.get('page', 'Unknown')}" if hasattr(doc, 'metadata') and doc.metadata.get('page') else "" context_parts.append(f"{doc_id} {source}{page}:\n{doc.page_content}\n") return "\n".join(context_parts) def answer_query( vector_db, query, max_retries=2, retrieval_method="hybrid", show_sources=False, k=5, temperature=0.2, model="deepseek-r1-distill-llama-70b" ) -> Tuple[str, Optional[List[str]]]: """ Generate answer using enhanced RAG pipeline with improved features Args: vector_db: Vector database query: User query max_retries: Number of retries on failure retrieval_method: Method for document retrieval show_sources: Whether to return source information k: Number of documents to retrieve temperature: LLM temperature model: LLM model to use Returns: Tuple of (response text, source documents if requested) """ if not query or not vector_db: return "Either your question or the document is missing. Please check and try again.", None # Initialize LLM if needed llm = get_llm(model=model, temperature=temperature) if not llm: return "Unable to initialize the language model. Please check your API key and try again.", None # Track performance metrics start_time = time.time() # Try answering with retries for attempt in range(max_retries): try: # Retrieve relevant documents with timing retrieval_start = time.time() documents = retrieve_docs(vector_db, query, k=k, retrieval_method=retrieval_method) retrieval_time = time.time() - retrieval_start if not documents: return "I couldn't find relevant information in the document to answer your question. Please try a different question or upload a document with the relevant information.", None # Get context from documents context_start = time.time() context = get_context(documents) context_time = time.time() - context_start # Create prompt and chain llm_start = time.time() prompt = ChatPromptTemplate.from_template(LEGAL_PROMPT_TEMPLATE) chain = prompt | llm # Get and clean response response = chain.invoke({"question": query, "context": context}) llm_time = time.time() - llm_start total_time = time.time() - start_time # Log performance metrics performance_metrics = { "retrieval_time": round(retrieval_time, 2), "context_time": round(context_time, 2), "llm_time": round(llm_time, 2), "total_time": round(total_time, 2) } # Log to console print(f"Performance: Retrieval: {retrieval_time:.2f}s, Context: {context_time:.2f}s, LLM: {llm_time:.2f}s, Total: {total_time:.2f}s") # Prepare source information if requested sources = None if show_sources: sources = [] for i, doc in enumerate(documents): source_info = f"Document {i+1}" if hasattr(doc, 'metadata'): if 'source' in doc.metadata: source_info += f" | Source: {doc.metadata['source']}" if 'page' in doc.metadata: source_info += f" | Page: {doc.metadata['page']}" # Add a snippet of content content_preview = doc.page_content[:200] + "..." if len(doc.page_content) > 200 else doc.page_content source_info += f"\nExcerpt: {content_preview}" sources.append(source_info) # Log the query for improvement save_query_log(query, clean_response(response), sources or [], performance_metrics) # Return the response and sources if requested return clean_response(response), sources except Exception as e: if attempt < max_retries - 1: print(f"Attempt {attempt+1} failed: {e}. Retrying...") continue else: return f"I encountered an error while processing your question: {str(e)}. Please try again with a simpler query.", None def save_query_log(query, response, sources, performance_metrics): """ Save query logs for analysis and improvement Args: query: User query response: Generated response sources: Source documents used performance_metrics: Performance metrics """ log_entry = { "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"), "query": query, "response_length": len(response), "num_sources": len(sources) if sources else 0, "performance_metrics": performance_metrics } log_file = os.path.join(LOGS_DIR, f"query_log_{time.strftime('%Y%m%d')}.jsonl") with open(log_file, "a") as f: f.write(json.dumps(log_entry) + "\n") def get_document_summary(vector_db, max_tokens=500): """ Generate a summary of the document for quick overview Args: vector_db: Vector database containing the document max_tokens: Maximum tokens for summary Returns: Document summary """ if not vector_db: return "No document loaded." # Get a sample of documents from the database try: # Get documents that represent the key sections docs = vector_db.similarity_search("summarize the main topics and key points of this document", k=5) if not docs: return "Unable to generate summary from this document." # Create a summary prompt summary_prompt = """ You are a legal document summarization expert. Based on the following excerpts from a legal document, provide a concise summary (maximum 3 paragraphs) of what the document appears to be about, its key topics, and main legal points. Focus on the factual content only: {context} Brief Summary: """ # Initialize LLM llm = get_llm(temperature=0.1) # Low temperature for factual summary if not llm: return "Unable to initialize the language model for summary generation." # Extract context context = get_context(docs) # Generate summary prompt = ChatPromptTemplate.from_template(summary_prompt) chain = prompt | llm response = chain.invoke({"context": context}) return clean_response(response) except Exception as e: print(f"Error generating document summary: {e}") return f"Unable to generate summary: {str(e)}"