Spaces:

kn29
/

rag-chat

Sleeping

Update rag.py

fef6ed9 verified about 1 year ago

14.7 kB

# 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

	# 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