# rag_optimized.py - Performance-Optimized RAG System import torch import numpy as np from transformers import AutoTokenizer, AutoModel from typing import List, Dict, Any, Tuple, Optional import faiss import hashlib from tqdm import tqdm from groq import Groq import re import nltk from sklearn.metrics.pairwise import cosine_similarity import networkx as nx from collections import defaultdict import spacy from rank_bm25 import BM25Okapi import asyncio import time from concurrent.futures import ThreadPoolExecutor import logging # Configure logging logger = logging.getLogger(__name__) # Global model instances (shared across sessions) _SHARED_MODEL = None _SHARED_TOKENIZER = None _SHARED_NLP_MODEL = None _DEVICE = None _THREAD_POOL = None # Legal knowledge base (optimized) LEGAL_CONCEPTS = { 'liability': ['negligence', 'strict liability', 'vicarious liability', 'product liability'], 'contract': ['breach', 'consideration', 'offer', 'acceptance', 'damages', 'specific performance'], 'criminal': ['mens rea', 'actus reus', 'intent', 'malice', 'premeditation'], 'procedure': ['jurisdiction', 'standing', 'statute of limitations', 'res judicata'], 'evidence': ['hearsay', 'relevance', 'privilege', 'burden of proof', 'admissibility'], 'constitutional': ['due process', 'equal protection', 'free speech', 'search and seizure'] } QUERY_PATTERNS = { 'precedent': ['case', 'precedent', 'ruling', 'held', 'decision'], 'statute_interpretation': ['statute', 'section', 'interpretation', 'meaning', 'definition'], 'factual': ['what happened', 'facts', 'circumstances', 'events'], 'procedure': ['how to', 'procedure', 'process', 'filing', 'requirements'] } def initialize_models(model_id: str, groq_api_key: str = None): """Initialize shared models (call once at startup)""" global _SHARED_MODEL, _SHARED_TOKENIZER, _SHARED_NLP_MODEL, _DEVICE, _THREAD_POOL try: nltk.download('punkt', quiet=True) nltk.download('stopwords', quiet=True) except: pass _DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') logger.info(f"Using device: {_DEVICE}") logger.info(f"Loading model: {model_id}") _SHARED_TOKENIZER = AutoTokenizer.from_pretrained(model_id) _SHARED_MODEL = AutoModel.from_pretrained(model_id).to(_DEVICE) _SHARED_MODEL.eval() # Initialize thread pool for CPU-bound operations _THREAD_POOL = ThreadPoolExecutor(max_workers=4) try: _SHARED_NLP_MODEL = spacy.load("en_core_web_sm") except: logger.warning("SpaCy model not found, using basic NER") _SHARED_NLP_MODEL = None class OptimizedSessionRAG: """High-performance session-specific RAG instance that loads pre-computed embeddings""" def __init__(self, session_id: str, groq_api_key: str = None): self.session_id = session_id self.groq_client = Groq(api_key=groq_api_key) if groq_api_key else None # Session-specific indices and data self.dense_index = None self.bm25_index = None self.token_to_chunks = None self.chunks_data = [] # Performance tracking self.load_time = None self.index_build_time = None # Verify shared models are initialized if _SHARED_MODEL is None or _SHARED_TOKENIZER is None: raise ValueError("Models not initialized. Call initialize_models() first.") def load_existing_session_data(self, chunks_from_db: List[Dict[str, Any]]): """OPTIMIZED: Load pre-existing chunks with embeddings from database - NO EMBEDDING CREATION""" start_time = time.time() logger.info(f"Loading existing session data for {self.session_id}: {len(chunks_from_db)} chunks...") # Process chunks from MongoDB format - DIRECT LOADING, NO EMBEDDING COMPUTATION self.chunks_data = self._process_db_chunks_fast(chunks_from_db) # Rebuild indices from existing embeddings ONLY self._rebuild_indices_from_precomputed_embeddings() self.load_time = time.time() - start_time logger.info(f"Session {self.session_id} loaded in {self.load_time:.2f}s with PRE-COMPUTED embeddings!") def _process_db_chunks_fast(self, chunks_from_db: List[Dict[str, Any]]) -> List[Dict[str, Any]]: """FAST: Convert MongoDB chunk format to internal format without any computation""" processed_chunks = [] for chunk in chunks_from_db: # Convert embedding from list to numpy array if needed - NO COMPUTATION embedding = chunk.get('embedding') if embedding is None: raise ValueError(f"Missing embedding for chunk {chunk.get('chunk_id', 'unknown')}") if isinstance(embedding, list): embedding = np.array(embedding, dtype=np.float32) processed_chunk = { 'id': chunk.get('chunk_id', chunk.get('id')), 'text': chunk.get('content', chunk.get('text', '')), 'title': chunk.get('title', 'Document'), 'section_type': chunk.get('section_type', 'general'), 'importance_score': chunk.get('importance_score', 1.0), 'entities': chunk.get('entities', []), 'embedding': embedding # PRE-COMPUTED, NO CREATION } processed_chunks.append(processed_chunk) return processed_chunks def _rebuild_indices_from_precomputed_embeddings(self): """OPTIMIZED: Rebuild search indices using ONLY pre-computed embeddings from database""" if not self.chunks_data: raise ValueError("No chunks data available") start_time = time.time() logger.info(f"Rebuilding indices from {len(self.chunks_data)} pre-computed embeddings...") # 1. Build FAISS index from existing embeddings - NO EMBEDDING COMPUTATION embeddings = [] for chunk in self.chunks_data: if chunk['embedding'] is None: raise ValueError(f"Missing embedding for chunk {chunk.get('id', 'unknown')}") embeddings.append(chunk['embedding']) # Stack embeddings efficiently embeddings_matrix = np.vstack(embeddings).astype('float32') logger.info(f"Built embeddings matrix: {embeddings_matrix.shape}") # Build FAISS index self.dense_index = faiss.IndexFlatIP(embeddings_matrix.shape[1]) self.dense_index.add(embeddings_matrix) # 2. Build BM25 index efficiently tokenized_corpus = [chunk['text'].lower().split() for chunk in self.chunks_data] self.bm25_index = BM25Okapi(tokenized_corpus) # 3. Build token-to-chunk mapping efficiently self.token_to_chunks = defaultdict(set) for i, chunk in enumerate(self.chunks_data): tokens = chunk['text'].lower().split() for token in tokens: self.token_to_chunks[token].add(i) self.index_build_time = time.time() - start_time logger.info(f"All indices rebuilt in {self.index_build_time:.2f}s from pre-computed embeddings!") def create_embedding(self, text: str) -> np.ndarray: """Create embedding for query (ONLY used for new queries, not document loading)""" inputs = _SHARED_TOKENIZER(text, padding=True, truncation=True, max_length=512, return_tensors='pt').to(_DEVICE) with torch.no_grad(): outputs = _SHARED_MODEL(**inputs) attention_mask = inputs['attention_mask'] token_embeddings = outputs.last_hidden_state input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float() embeddings = torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9) # Normalize embeddings embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1) return embeddings.cpu().numpy()[0].astype('float32') def analyze_query_fast(self, query: str) -> Dict[str, Any]: """FAST query analysis - minimal processing""" query_lower = query.lower() # Quick query type classification query_type = 'general' for qtype, patterns in QUERY_PATTERNS.items(): if any(pattern in query_lower for pattern in patterns): query_type = qtype break # Extract key concepts quickly key_concepts = [] for concept_category, concepts in LEGAL_CONCEPTS.items(): for concept in concepts: if concept in query_lower: key_concepts.append(concept) # Simple query expansion expanded_queries = [query] if key_concepts: expanded_queries.append(f"{query} {' '.join(key_concepts[:2])}") return { 'original_query': query, 'query_type': query_type, 'key_concepts': key_concepts, 'expanded_queries': expanded_queries[:2] # Limit to 2 for speed } def fast_retrieval(self, query_analysis: Dict[str, Any], top_k: int = 10) -> List[Tuple[Dict[str, Any], float]]: """OPTIMIZED: Fast multi-stage retrieval with minimal overhead""" candidates = {} # Stage 1: Dense retrieval with primary query only query = query_analysis['original_query'] query_emb = self.create_embedding(query) scores, indices = self.dense_index.search( query_emb.reshape(1, -1), min(top_k * 2, len(self.chunks_data)) ) for idx, score in zip(indices[0], scores[0]): if idx < len(self.chunks_data): chunk = self.chunks_data[idx] chunk_id = chunk['id'] candidates[chunk_id] = { 'chunk': chunk, 'score': float(score) * chunk['importance_score'] } # Stage 2: BM25 boost for top candidates if len(candidates) < top_k: query_tokens = query.lower().split() bm25_scores = self.bm25_index.get_scores(query_tokens) top_bm25_indices = np.argsort(bm25_scores)[-top_k:][::-1] for idx in top_bm25_indices: if idx < len(self.chunks_data): chunk = self.chunks_data[idx] chunk_id = chunk['id'] if chunk_id not in candidates: candidates[chunk_id] = { 'chunk': chunk, 'score': float(bm25_scores[idx]) * 0.3 # Lower weight for BM25 } else: candidates[chunk_id]['score'] += float(bm25_scores[idx]) * 0.2 # Convert to list and sort final_scores = [(data['chunk'], data['score']) for data in candidates.values()] final_scores.sort(key=lambda x: x[1], reverse=True) return final_scores[:top_k] def generate_fast_answer(self, query: str, retrieved_chunks: List[Tuple[Dict[str, Any], float]]) -> Dict[str, Any]: """Generate answer with minimal overhead""" if not self.groq_client: return {'error': 'Groq client not initialized'} # Prepare context efficiently context_parts = [] for i, (chunk, score) in enumerate(retrieved_chunks[:3], 1): # Limit to top 3 for speed context_parts.append(f""" Document {i} - Relevance: {score:.2f} {chunk['text'][:600]} """) context = "\n---\n".join(context_parts) system_prompt = """You are a legal AI assistant. Provide concise, accurate answers based ONLY on the provided documents. If information isn't in the documents, state that clearly.""" user_prompt = f"""Query: {query} Documents: {context} Provide a clear, concise answer based on the documents.""" try: response = self.groq_client.chat.completions.create( messages=[ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt} ], model="llama-3.1-8b-instant", temperature=0.1, max_tokens=500 # Limit for speed ) answer = response.choices[0].message.content # Calculate confidence avg_score = sum(score for _, score in retrieved_chunks[:3]) / min(3, len(retrieved_chunks)) confidence = min(avg_score * 100, 100) return { 'answer': answer, 'confidence': confidence, 'sources': [ { 'chunk_id': chunk['id'], 'title': chunk['title'], 'section': chunk['section_type'], 'relevance_score': float(score), 'text_preview': chunk['text'][:200] + '...', 'entities': [e['text'] for e in chunk['entities'][:3]] } for chunk, score in retrieved_chunks[:5] ] } except Exception as e: return {'error': f'Error generating answer: {str(e)}'} def query_documents(self, query: str, top_k: int = 5) -> Dict[str, Any]: """OPTIMIZED: Main query function with minimal processing time""" if not self.chunks_data: return {'error': f'No documents indexed for session {self.session_id}'} start_time = time.time() # Fast query analysis query_analysis = self.analyze_query_fast(query) # Fast retrieval retrieved_chunks = self.fast_retrieval(query_analysis, top_k) if not retrieved_chunks: return { 'error': 'No relevant documents found', 'query_analysis': query_analysis } # Generate answer result = self.generate_fast_answer(query, retrieved_chunks) result['query_analysis'] = query_analysis result['processing_time'] = time.time() - start_time logger.info(f"Query processed in {result['processing_time']:.2f}s") return result # For backward compatibility - replace SessionRAG with OptimizedSessionRAG SessionRAG = OptimizedSessionRAG