import os import json import tempfile from typing import List, Dict, Any, Optional from pathlib import Path # LangChain imports for RAG from langchain.text_splitter import RecursiveCharacterTextSplitter from langchain_community.vectorstores import Chroma from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings from langchain.chains import RetrievalQA from langchain.prompts import PromptTemplate from langchain.schema import Document # Google Gemini imports from google import genai class RAGSystem: """ Complete RAG (Retrieval-Augmented Generation) system using Google Gemini Handles document ingestion, chunking, embedding, and question answering """ def __init__(self, persist_directory: str = "./chroma_db"): """Initialize the RAG system with Google Gemini and ChromaDB""" self.persist_directory = persist_directory self.gemini_api_key = None # Initialize components (lazy loading) self.embeddings = None self.llm = None self.vectorstore = None self.retriever = None self.qa_chain = None # Text splitter for document chunking self.text_splitter = RecursiveCharacterTextSplitter( chunk_size=1000, chunk_overlap=200, length_function=len, separators=["\n\n", "\n", " ", ""] ) # Track ingested documents self.ingested_documents = [] def _initialize_components(self): """Lazy initialization of Gemini components""" if self.llm is None: self.gemini_api_key = os.getenv('GEMINI_API_KEY') if not self.gemini_api_key: raise ValueError("GEMINI_API_KEY environment variable must be set") # Initialize Google Gemini LLM self.llm = ChatGoogleGenerativeAI( model="gemini-2.5-flash", temperature=0.1, max_tokens=2048, google_api_key=self.gemini_api_key ) # Initialize Google embeddings self.embeddings = GoogleGenerativeAIEmbeddings( model="models/text-embedding-004", google_api_key=self.gemini_api_key ) # Initialize or load existing vector store self._initialize_vectorstore() def _initialize_vectorstore(self): """Initialize ChromaDB vector store""" try: # Try to load existing vectorstore if os.path.exists(self.persist_directory): self.vectorstore = Chroma( persist_directory=self.persist_directory, embedding_function=self.embeddings ) else: # Create new empty vectorstore self.vectorstore = Chroma( persist_directory=self.persist_directory, embedding_function=self.embeddings ) # Set up retriever self.retriever = self.vectorstore.as_retriever( search_type="similarity", search_kwargs={"k": 5} # Retrieve top 5 most similar chunks ) except Exception as e: raise Exception(f"Failed to initialize vector store: {str(e)}") def ingest_document(self, text_content: str, metadata: Dict[str, Any]) -> Dict[str, Any]: """ Ingest a document into the RAG system Args: text_content: The full text content of the document metadata: Document metadata (filename, type, etc.) Returns: Dict with ingestion results """ try: # Initialize components if needed self._initialize_components() # Create document object document = Document( page_content=text_content, metadata=metadata ) # Split document into chunks chunks = self.text_splitter.split_documents([document]) # Add chunk numbers to metadata for i, chunk in enumerate(chunks): chunk.metadata.update({ 'chunk_id': i, 'total_chunks': len(chunks) }) # Add chunks to vector store self.vectorstore.add_documents(chunks) # Persist the changes self.vectorstore.persist() # Track ingested document doc_info = { 'filename': metadata.get('filename', 'Unknown'), 'document_type': metadata.get('document_type', 'Unknown'), 'chunks_created': len(chunks), 'ingestion_timestamp': metadata.get('ingestion_timestamp', 'Unknown') } self.ingested_documents.append(doc_info) return { 'status': 'success', 'chunks_created': len(chunks), 'document_info': doc_info } except Exception as e: return { 'status': 'error', 'error': str(e) } def query(self, question: str, return_source_docs: bool = True) -> Dict[str, Any]: """ Query the RAG system with a question Args: question: User's question return_source_docs: Whether to return source documents Returns: Dict with answer and source information """ try: # Initialize components if needed self._initialize_components() if not self.vectorstore: return { 'status': 'error', 'error': 'No documents have been ingested yet. Please upload and process some PDFs first.' } # Create RAG chain if not exists if not self.qa_chain: self._setup_qa_chain() # Execute query result = self.qa_chain.invoke({ "query": question, "return_source_documents": return_source_docs }) # Format response response = { 'status': 'success', 'answer': result.get('result', ''), 'question': question } # Add source documents if requested if return_source_docs and 'source_documents' in result: response['sources'] = [] for doc in result['source_documents']: response['sources'].append({ 'content': doc.page_content[:200] + '...', # Preview 'metadata': doc.metadata }) return response except Exception as e: return { 'status': 'error', 'error': f"Query failed: {str(e)}" } def _setup_qa_chain(self): """Set up the question-answering chain with custom prompt""" # Custom prompt template for better responses prompt_template = """ You are an AI assistant that answers questions based on the provided document context. Use the following context to answer the question accurately and comprehensively. If the answer cannot be found in the context, say "I don't have enough information in the provided documents to answer this question." Context: {context} Question: {question} Answer:""" prompt = PromptTemplate( template=prompt_template, input_variables=["context", "question"] ) # Create RetrievalQA chain self.qa_chain = RetrievalQA.from_llm( llm=self.llm, retriever=self.retriever, prompt=prompt, return_source_documents=True ) def get_document_list(self) -> List[Dict[str, Any]]: """Get list of ingested documents""" return self.ingested_documents.copy() def get_vector_store_stats(self) -> Dict[str, Any]: """Get statistics about the vector store""" try: self._initialize_components() if not self.vectorstore: return {'total_chunks': 0, 'status': 'empty'} # Get collection info collection = self.vectorstore._collection stats = { 'total_chunks': collection.count(), 'total_documents': len(self.ingested_documents), 'status': 'active' } return stats except Exception as e: return { 'status': 'error', 'error': str(e) } def clear_knowledge_base(self) -> Dict[str, Any]: """Clear all documents from the knowledge base""" try: # Delete vector store directory import shutil if os.path.exists(self.persist_directory): shutil.rmtree(self.persist_directory) # Reset components self.vectorstore = None self.qa_chain = None self.ingested_documents = [] return {'status': 'success', 'message': 'Knowledge base cleared successfully'} except Exception as e: return {'status': 'error', 'error': str(e)} def search_similar_chunks(self, query: str, k: int = 5) -> List[Dict[str, Any]]: """Search for similar document chunks""" try: self._initialize_components() if not self.vectorstore: return [] # Perform similarity search docs = self.vectorstore.similarity_search(query, k=k) results = [] for doc in docs: results.append({ 'content': doc.page_content, 'metadata': doc.metadata, 'preview': doc.page_content[:150] + '...' }) return results except Exception as e: print(f"Search error: {e}") return []