"""RAG Engine - Core retrieval-augmented generation service.""" from typing import Optional import hashlib from src.clients.embeddings import get_embedding from src.clients.chat_provider import chat_completion from src.clients.qdrant_client import search_similar, ensure_collection_exists, get_collection_info from src.config.settings import settings from .conversation_context import ConversationContext from .citation_system import CitationSystem from .response_formatter import ResponseFormatter class RAGEngine: """Core RAG engine for question answering with document retrieval.""" def __init__(self, collection_name: Optional[str] = None): """Initialize RAG engine. Args: collection_name: Qdrant collection name for document storage. """ self.collection_name = collection_name or settings.QDRANT_COLLECTION self.context_manager = ConversationContext() self.citation_system = CitationSystem() self.response_formatter = ResponseFormatter() self._initialized = False self._embedding_cache = {} # Cache embeddings to avoid repeated API calls self._collection_has_data = False async def initialize(self) -> bool: """Initialize the RAG engine and ensure collection exists. Returns: True if initialization successful. """ if self._initialized: return True self._initialized = ensure_collection_exists( self.collection_name, vector_size=settings.EMBEDDING_DIM, # OpenRouter embedding dimensions (default 3072) ) # Check if collection already has data if self._initialized: collection_info = get_collection_info(self.collection_name) if collection_info and collection_info.get('points_count', 0) > 0: self._collection_has_data = True print(f"✓ Collection '{self.collection_name}' has {collection_info['points_count']} documents - using existing data") else: print(f"⚠ Collection '{self.collection_name}' is empty - embeddings will be generated for new documents") return self._initialized def _get_cached_embedding(self, text: str) -> Optional[list]: """Get cached embedding for text if available. Args: text: Text to get embedding for. Returns: Cached embedding or None. """ cache_key = hashlib.md5(text.encode()).hexdigest() return self._embedding_cache.get(cache_key) def _cache_embedding(self, text: str, embedding: list) -> None: """Cache embedding for text. Args: text: Text that was embedded. embedding: Embedding vector to cache. """ cache_key = hashlib.md5(text.encode()).hexdigest() self._embedding_cache[cache_key] = embedding # Keep cache size limited if len(self._embedding_cache) > 100: # Remove oldest entry (first inserted) oldest_key = next(iter(self._embedding_cache)) del self._embedding_cache[oldest_key] async def query( self, question: str, conversation_history: Optional[list] = None, selected_text: Optional[str] = None, top_k: int = 5, include_citations: bool = True, language: Optional[str] = "en", ) -> dict: """Process a question using RAG. Args: question: User question to answer. conversation_history: Previous conversation messages. selected_text: Optional selected text for context filtering. top_k: Number of documents to retrieve. include_citations: Whether to include source citations. language: Language code for the response (en, ur, ur-PK, ar, es, ...). Returns: Dictionary with answer, sources, and metadata. """ # Build context-aware query enhanced_query = self.context_manager.build_query( question=question, conversation_history=conversation_history, selected_text=selected_text, ) # Check cache first to avoid unnecessary API calls query_embedding = self._get_cached_embedding(enhanced_query) if query_embedding is None: # Only call embedding API if not in cache and needed if self._collection_has_data: # Collection has data, generate embedding for search query_embedding = get_embedding(enhanced_query) self._cache_embedding(enhanced_query, query_embedding) print("✓ Generated embedding for query (cached for future use)") else: # Collection is empty, use fallback print("⚠ Collection empty - using fallback embedding") from src.clients.embeddings import simple_embedding query_embedding = simple_embedding(enhanced_query) else: print("✓ Using cached embedding for query") # Search for relevant documents search_results = search_similar( collection_name=self.collection_name, query_vector=query_embedding, top_k=top_k, score_threshold=settings.RAG_SIMILARITY_THRESHOLD, ) print(f"✓ Search returned {len(search_results)} results") for i, r in enumerate(search_results[:3]): score = r.get('score', 0) url = r.get('payload', {}).get('url', 'N/A') print(f" Result {i+1}: score={score:.4f}, url={url}") if not search_results: print("⚠ No results above threshold, trying without threshold...") search_results = search_similar( collection_name=self.collection_name, query_vector=query_embedding, top_k=top_k, score_threshold=0.0, ) print(f"✓ Search (no threshold) returned {len(search_results)} results") # Build context from retrieved documents context_text = self._build_context(search_results) # Generate answer with context system_prompt = self._get_system_prompt(context_text, include_citations, language) messages = [] if conversation_history: messages.extend(conversation_history[-6:]) # Last 6 messages for context messages.append({"role": "user", "content": question}) answer = chat_completion( messages=messages, system_prompt=system_prompt, max_tokens=settings.RAG_MAX_RESPONSE_TOKENS, temperature=0.7, ) # Format response with citations citations = [] if include_citations: citations = self.citation_system.extract_citations(search_results) return self.response_formatter.format_response( answer=answer, sources=search_results, citations=citations, query=question, ) def _build_context(self, search_results: list) -> str: """Build context string from search results. Args: search_results: List of search results from Qdrant. Returns: Formatted context string. """ if not search_results: return "No relevant documentation found." context_parts = [] for i, result in enumerate(search_results, 1): payload = result.get("payload", {}) # Support both 'text' (from main.py ingestion) and 'content' (from src/ indexing) content = payload.get("text", payload.get("content", "")) title = payload.get("title", "Document") source = payload.get("url", payload.get("source_url", payload.get("file_path", ""))) context_parts.append( f"[Source {i}] {title}\n" f"Content: {content}\n" f"Reference: {source}\n" ) return "\n---\n".join(context_parts) # Language names for response translation instructions LANGUAGE_NAMES = { "en": "English", "ur": "Urdu (اردو script)", "ur-PK": "Roman Urdu (Urdu written in Latin script)", "ar": "Arabic (العربية)", "es": "Spanish", "fr": "French", "de": "German", "zh": "Chinese (Simplified)", "hi": "Hindi", "pt": "Portuguese", "ru": "Russian", "ja": "Japanese", } def _get_system_prompt(self, context: str, include_citations: bool, language: Optional[str] = "en") -> str: """Generate system prompt for RAG responses. Args: context: Retrieved document context. include_citations: Whether to include citation instructions. language: Response language code. Returns: System prompt string. """ citation_instruction = "" if include_citations: citation_instruction = ( "When answering, cite your sources using [Source N] notation " "where N corresponds to the source number in the context. " ) language_instruction = "" lang = (language or "en").strip() if lang != "en": lang_name = self.LANGUAGE_NAMES.get(lang, lang) language_instruction = ( f"IMPORTANT: Respond ENTIRELY in {lang_name}. " "Translate your answer naturally; keep technical terms and code in English where standard. " ) return f"""You are a helpful AI assistant for Physical AI & Humanoid Robotics in Education. {language_instruction}{citation_instruction} RULES: 1. Answer concisely in 2-4 sentences using the context below 2. Cite sources using [Source N] when referencing specific information 3. Only say "I cannot find information" if context is completely empty 4. Be direct - no unnecessary introductions, tables, or lengthy explanations 5. Use simple, clear language CONTEXT: {context} Keep answers short and to the point. Use ONLY information from the context above."""