"""RAG chat service with Gemini 2.5 Flash and streaming support.""" import os import time from typing import List, Dict, Any, Optional, Generator, Tuple from langchain_google_genai import ChatGoogleGenerativeAI from langchain_core.prompts import ChatPromptTemplate from langchain_core.runnables import RunnablePassthrough from langchain_core.output_parsers import StrOutputParser from langchain_core.documents import Document from src.rag.vector_store import vector_store_manager from src.rag.memory import chat_memory_manager from src.core.config import config from src.core.logging_config import get_logger logger = get_logger(__name__) class ChatUsageLimiter: """Manages chat usage limits to prevent abuse.""" def __init__(self, max_messages_per_session: int = 50, max_messages_per_hour: int = 100): """ Initialize usage limiter. Args: max_messages_per_session: Maximum messages per chat session max_messages_per_hour: Maximum messages per hour across all sessions """ self.max_messages_per_session = max_messages_per_session self.max_messages_per_hour = max_messages_per_hour self.hourly_usage = {} # Track usage by hour logger.info(f"Chat usage limiter initialized: {max_messages_per_session}/session, {max_messages_per_hour}/hour") def check_session_limit(self, session_message_count: int) -> Tuple[bool, str]: """ Check if session has exceeded message limit. Args: session_message_count: Number of messages in current session Returns: Tuple of (allowed, reason_if_not_allowed) """ if session_message_count >= self.max_messages_per_session: return False, f"Session limit reached ({self.max_messages_per_session} messages per session). Please start a new chat." return True, "" def check_hourly_limit(self) -> Tuple[bool, str]: """ Check if hourly limit has been exceeded. Returns: Tuple of (allowed, reason_if_not_allowed) """ current_hour = int(time.time()) // 3600 # Clean old entries (keep only last 2 hours) hours_to_keep = [current_hour - 1, current_hour] self.hourly_usage = {h: count for h, count in self.hourly_usage.items() if h in hours_to_keep} current_usage = self.hourly_usage.get(current_hour, 0) if current_usage >= self.max_messages_per_hour: return False, f"Hourly limit reached ({self.max_messages_per_hour} messages per hour). Please try again later." return True, "" def record_usage(self) -> None: """Record a message usage.""" current_hour = int(time.time()) // 3600 self.hourly_usage[current_hour] = self.hourly_usage.get(current_hour, 0) + 1 def can_send_message(self, session_message_count: int) -> Tuple[bool, str]: """ Check if user can send a message. Args: session_message_count: Number of messages in current session Returns: Tuple of (allowed, reason_if_not_allowed) """ # Check session limit session_ok, session_reason = self.check_session_limit(session_message_count) if not session_ok: return False, session_reason # Check hourly limit hourly_ok, hourly_reason = self.check_hourly_limit() if not hourly_ok: return False, hourly_reason return True, "" class RAGChatService: """RAG-powered chat service with document context.""" def __init__(self): """Initialize the RAG chat service.""" self.usage_limiter = ChatUsageLimiter( max_messages_per_session=config.rag.max_messages_per_session, max_messages_per_hour=config.rag.max_messages_per_hour ) self._llm = None self._rag_chain = None self._current_retrieval_method = "similarity" self._default_retrieval_method = "similarity" self._default_retrieval_config = {"k": 4} logger.info("RAG chat service initialized") def get_llm(self) -> ChatGoogleGenerativeAI: """Get or create the Gemini LLM instance.""" if self._llm is None: try: google_api_key = config.api.google_api_key or os.getenv("GOOGLE_API_KEY") if not google_api_key: raise ValueError("Google API key not found. Please set GOOGLE_API_KEY in environment variables.") self._llm = ChatGoogleGenerativeAI( model="gemini-2.5-flash", # Latest Gemini model google_api_key=google_api_key, temperature=config.rag.rag_temperature, max_tokens=config.rag.rag_max_tokens, disable_streaming=False # Enable streaming (new parameter name) ) logger.info("Gemini 2.5 Flash LLM initialized successfully") except Exception as e: logger.error(f"Failed to initialize Gemini LLM: {e}") raise return self._llm def create_rag_chain(self, retrieval_method: str = "similarity", retrieval_config: Optional[Dict[str, Any]] = None): """ Create the RAG chain for document-aware conversations. Args: retrieval_method: Method to use ("similarity", "mmr", "hybrid") retrieval_config: Configuration for the retrieval method """ if self._rag_chain is None or hasattr(self, '_current_retrieval_method') and self._current_retrieval_method != retrieval_method: try: llm = self.get_llm() # Set default retrieval config if retrieval_config is None: retrieval_config = {"k": 4} # Get retriever based on method if retrieval_method == "hybrid": # Use hybrid retriever (semantic + keyword) semantic_weight = retrieval_config.get("semantic_weight", 0.7) keyword_weight = retrieval_config.get("keyword_weight", 0.3) search_type = retrieval_config.get("search_type", "similarity") search_kwargs = {k: v for k, v in retrieval_config.items() if k not in ["semantic_weight", "keyword_weight", "search_type"]} retriever = vector_store_manager.get_hybrid_retriever( k=retrieval_config.get("k", 4), semantic_weight=semantic_weight, keyword_weight=keyword_weight, search_type=search_type, search_kwargs=search_kwargs if search_kwargs else None ) logger.info(f"Using hybrid retriever with weights: semantic={semantic_weight}, keyword={keyword_weight}") elif retrieval_method == "mmr": # Use MMR for diversity search_kwargs = retrieval_config.copy() if "fetch_k" not in search_kwargs: search_kwargs["fetch_k"] = retrieval_config.get("k", 4) * 5 # Default fetch 5x more for MMR if "lambda_mult" not in search_kwargs: search_kwargs["lambda_mult"] = 0.5 # Balance relevance vs diversity retriever = vector_store_manager.get_retriever( search_type="mmr", search_kwargs=search_kwargs ) logger.info(f"Using MMR retriever with config: {search_kwargs}") else: # Default similarity search retriever = vector_store_manager.get_retriever( search_type="similarity", search_kwargs=retrieval_config ) logger.info(f"Using similarity retriever with config: {retrieval_config}") # Store current method for comparison self._current_retrieval_method = retrieval_method # Create a prompt template for RAG prompt_template = ChatPromptTemplate.from_template(""" You are a helpful assistant that can chat naturally while specializing in answering questions about uploaded documents. Instructions: 1. For document-related questions: Use the provided context to give comprehensive answers and always cite your sources 2. For conversational interactions (greetings, introductions, clarifications, follow-ups): Respond naturally and helpfully 3. For questions about topics not covered in the documents: Politely explain that you specialize in the uploaded documents but can still have a conversation 4. When using document information, always cite which parts of the documents you referenced 5. Include relevant tables and code blocks when they help answer the question 6. Be conversational, friendly, and helpful 7. Remember information shared in our conversation (like names, preferences, etc.) Context from documents: {context} Chat History: {chat_history} User Message: {question} """) def format_docs(docs: List[Document]) -> str: """Format retrieved documents for context.""" if not docs: return "No relevant documents found." formatted = [] for i, doc in enumerate(docs, 1): source = doc.metadata.get('source', 'Unknown') chunk_id = doc.metadata.get('chunk_id', f'chunk_{i}') formatted.append(f"Document {i} (Source: {source}, ID: {chunk_id}):\n{doc.page_content}") return "\n\n".join(formatted) def format_chat_history() -> str: """Format chat history for context.""" history = chat_memory_manager.get_conversation_history(max_messages=10) if not history: return "No previous conversation." formatted = [] for user_msg, assistant_msg in history[-5:]: # Last 5 exchanges formatted.append(f"User: {user_msg}") formatted.append(f"Assistant: {assistant_msg}") return "\n".join(formatted) # Create the RAG chain self._rag_chain = ( { "context": retriever | format_docs, "chat_history": lambda _: format_chat_history(), "question": RunnablePassthrough() } | prompt_template | llm | StrOutputParser() ) logger.info("RAG chain created successfully") except Exception as e: logger.error(f"Failed to create RAG chain: {e}") raise def get_rag_chain(self, retrieval_method: str = "similarity", retrieval_config: Optional[Dict[str, Any]] = None): """ Get the RAG chain, creating it if necessary. Args: retrieval_method: Method to use ("similarity", "mmr", "hybrid") retrieval_config: Configuration for the retrieval method """ if self._rag_chain is None or (hasattr(self, '_current_retrieval_method') and self._current_retrieval_method != retrieval_method): self.create_rag_chain(retrieval_method, retrieval_config) return self._rag_chain def chat_stream_with_retrieval(self, user_message: str, retrieval_method: str = "similarity", retrieval_config: Optional[Dict[str, Any]] = None) -> Generator[str, None, None]: """ Stream chat response using RAG with specified retrieval method. Args: user_message: User's message retrieval_method: Method to use ("similarity", "mmr", "hybrid") retrieval_config: Configuration for the retrieval method Yields: Chunks of the response as they're generated """ try: # Check usage limits current_session = chat_memory_manager.current_session session_message_count = len(current_session.messages) if current_session else 0 can_send, reason = self.usage_limiter.can_send_message(session_message_count) if not can_send: yield f"❌ {reason}" return # Record usage self.usage_limiter.record_usage() # Add user message to memory chat_memory_manager.add_message("user", user_message) # Get RAG chain with specified retrieval method rag_chain = self.get_rag_chain(retrieval_method, retrieval_config) # Stream the response response_chunks = [] for chunk in rag_chain.stream(user_message): if chunk: response_chunks.append(chunk) yield chunk # Save complete response to memory complete_response = "".join(response_chunks) if complete_response.strip(): chat_memory_manager.add_message("assistant", complete_response) # Save session periodically chat_memory_manager.save_session() except Exception as e: error_msg = f"Error generating response: {str(e)}" logger.error(error_msg) yield f"❌ {error_msg}" def chat_stream(self, user_message: str) -> Generator[str, None, None]: """ Stream chat response using RAG. Args: user_message: User's message Yields: Chunks of the response as they're generated """ try: # Check usage limits current_session = chat_memory_manager.current_session session_message_count = len(current_session.messages) if current_session else 0 can_send, reason = self.usage_limiter.can_send_message(session_message_count) if not can_send: yield f"❌ {reason}" return # Record usage self.usage_limiter.record_usage() # Add user message to memory chat_memory_manager.add_message("user", user_message) # Get RAG chain rag_chain = self.get_rag_chain() # Stream the response response_chunks = [] for chunk in rag_chain.stream(user_message): if chunk: response_chunks.append(chunk) yield chunk # Save complete response to memory complete_response = "".join(response_chunks) if complete_response.strip(): chat_memory_manager.add_message("assistant", complete_response) # Save session periodically chat_memory_manager.save_session() except Exception as e: error_msg = f"Error generating response: {str(e)}" logger.error(error_msg) yield f"❌ {error_msg}" def chat(self, user_message: str) -> str: """ Get a complete chat response (non-streaming). Args: user_message: User's message Returns: Complete response string """ try: # Check usage limits current_session = chat_memory_manager.current_session session_message_count = len(current_session.messages) if current_session else 0 can_send, reason = self.usage_limiter.can_send_message(session_message_count) if not can_send: return f"❌ {reason}" # Record usage self.usage_limiter.record_usage() # Add user message to memory chat_memory_manager.add_message("user", user_message) # Get RAG chain rag_chain = self.get_rag_chain() # Get response response = rag_chain.invoke(user_message) # Save response to memory if response.strip(): chat_memory_manager.add_message("assistant", response) chat_memory_manager.save_session() return response except Exception as e: error_msg = f"Error generating response: {str(e)}" logger.error(error_msg) return f"❌ {error_msg}" def chat_with_retrieval(self, user_message: str, retrieval_method: str = "similarity", retrieval_config: Optional[Dict[str, Any]] = None) -> str: """ Get a complete chat response with specified retrieval method (non-streaming). Args: user_message: User's message retrieval_method: Method to use ("similarity", "mmr", "hybrid") retrieval_config: Configuration for the retrieval method Returns: Complete response string """ try: # Check usage limits current_session = chat_memory_manager.current_session session_message_count = len(current_session.messages) if current_session else 0 can_send, reason = self.usage_limiter.can_send_message(session_message_count) if not can_send: return f"❌ {reason}" # Record usage self.usage_limiter.record_usage() # Add user message to memory chat_memory_manager.add_message("user", user_message) # Get RAG chain with specified retrieval method rag_chain = self.get_rag_chain(retrieval_method, retrieval_config) # Get response response = rag_chain.invoke(user_message) # Save response to memory if response.strip(): chat_memory_manager.add_message("assistant", response) chat_memory_manager.save_session() return response except Exception as e: error_msg = f"Error generating response: {str(e)}" logger.error(error_msg) return f"❌ {error_msg}" def set_default_retrieval_method(self, method: str, config: Optional[Dict[str, Any]] = None): """ Set the default retrieval method for this service. Args: method: Retrieval method ("similarity", "mmr", "hybrid") config: Configuration for the method """ self._default_retrieval_method = method self._default_retrieval_config = config or {} # Reset the chain to use new method self._rag_chain = None logger.info(f"Default retrieval method set to: {method} with config: {config}") def get_usage_stats(self) -> Dict[str, Any]: """Get current usage statistics.""" current_session = chat_memory_manager.current_session session_message_count = len(current_session.messages) if current_session else 0 current_hour = int(time.time()) // 3600 hourly_count = self.usage_limiter.hourly_usage.get(current_hour, 0) return { "session_messages": session_message_count, "session_limit": self.usage_limiter.max_messages_per_session, "hourly_messages": hourly_count, "hourly_limit": self.usage_limiter.max_messages_per_hour, "session_remaining": max(0, self.usage_limiter.max_messages_per_session - session_message_count), "hourly_remaining": max(0, self.usage_limiter.max_messages_per_hour - hourly_count) } def start_new_session(self, document_sources: Optional[List[str]] = None) -> str: """Start a new chat session.""" session_id = chat_memory_manager.create_session(document_sources) logger.info(f"Started new chat session: {session_id}") return session_id def test_service(self) -> Dict[str, Any]: """Test the RAG service components.""" results = { "llm_available": False, "vector_store_available": False, "embeddings_available": False, "errors": [] } try: # Test LLM llm = self.get_llm() test_response = llm.invoke("Test message") results["llm_available"] = True except Exception as e: results["errors"].append(f"LLM test failed: {str(e)}") try: # Test vector store vector_info = vector_store_manager.get_collection_info() results["vector_store_available"] = "error" not in vector_info results["document_count"] = vector_info.get("document_count", 0) except Exception as e: results["errors"].append(f"Vector store test failed: {str(e)}") try: # Test embeddings from src.rag.embeddings import embedding_manager results["embeddings_available"] = embedding_manager.test_embedding_model() except Exception as e: results["errors"].append(f"Embeddings test failed: {str(e)}") return results # Global RAG chat service instance rag_chat_service = RAGChatService()