# Main Multimodal RAG System import os import json import uuid import numpy as np import faiss import torch from datetime import datetime from pathlib import Path from typing import List, Dict import streamlit as st from config import (DATA_DIRECTORY, DATA_FILES, DEFAULT_TOP_K, MAX_CONTEXT_LENGTH, LLM_TEMPERATURE, FAISS_TEMP_MEMORY) from models import load_models from document_processor import DocumentProcessor from device_setup import clear_gpu_memory class MultimodalRAG: """Main RAG system that handles document storage, search, and generation""" def __init__(self): # load all models model_components = load_models() if model_components[0] is None: raise Exception("Failed to load models") (self.text_encoder, self.clip_model, self.clip_processor, self.whisper_model, self.ocr_reader, self.llm, self.device, self.has_gpu) = model_components # initialize document processor self.processor = DocumentProcessor( self.clip_model, self.clip_processor, self.whisper_model, self.ocr_reader, self.text_encoder, self.device ) # initialize storage self.documents = [] self.metadata = [] self.embeddings = None self.index = None # setup data directory self.data_dir = DATA_DIRECTORY self.data_dir.mkdir(exist_ok=True) # load existing data self.load_existing_data() def clear_knowledge_base(self): """clear all stored documents and embeddings""" self.documents = [] self.metadata = [] self.embeddings = None self.index = None try: # remove all files in data directory for file in self.data_dir.glob("*"): file.unlink() except Exception as e: st.error(f"Error clearing saved data files: {e}") def save_data(self): """Save documents, metadata, embeddings, and FAISS index to disk""" try: # save documents and metadata as JSON with open(self.data_dir / DATA_FILES["documents"], "w") as f: json.dump(self.documents, f, default=str) with open(self.data_dir / DATA_FILES["metadata"], "w") as f: json.dump(self.metadata, f, default=str) # save FAISS index if self.index is not None: if self.has_gpu: # convert GPU index back to CPU for saving cpu_index = faiss.index_gpu_to_cpu(self.index) faiss.write_index(cpu_index, str(self.data_dir / DATA_FILES["faiss_index"])) else: faiss.write_index(self.index, str(self.data_dir / DATA_FILES["faiss_index"])) # save embeddings if self.embeddings is not None: np.save(self.data_dir / DATA_FILES["embeddings"], self.embeddings) except Exception as e: st.error(f"Error saving data: {e}") def load_existing_data(self): """Load existing documents and metadata from disk""" try: # load documents doc_file = self.data_dir / DATA_FILES["documents"] if doc_file.exists(): with open(doc_file, "r") as f: self.documents = json.load(f) # load metadata meta_file = self.data_dir / DATA_FILES["metadata"] if meta_file.exists(): with open(meta_file, "r") as f: self.metadata = json.load(f) # load FAISS index index_file = self.data_dir / DATA_FILES["faiss_index"] if index_file.exists(): if self.has_gpu: cpu_index = faiss.read_index(str(index_file)) res = faiss.StandardGpuResources() self.index = faiss.index_cpu_to_gpu(res, 0, cpu_index) else: self.index = faiss.read_index(str(index_file)) # load embeddings emb_file = self.data_dir / DATA_FILES["embeddings"] if emb_file.exists(): self.embeddings = np.load(emb_file) except Exception as e: st.error(f"Error loading existing data: {e}") def process_file(self, uploaded_file, file_type: str) -> bool: """Process uploaded file and add to knowledge base""" try: # process the file processed_data = self.processor.process_uploaded_file(uploaded_file, file_type, self.data_dir) if processed_data is None: return False # create document record file_id = str(uuid.uuid4()) timestamp = datetime.now().isoformat() document = { 'id': file_id, 'type': file_type if file_type not in ['png', 'jpg', 'jpeg'] else 'image', 'filename': uploaded_file.name, 'content': processed_data['content'], 'timestamp': timestamp, 'embedding_type': processed_data['embedding_type'] } # handle file-specific storage (images and audio) if file_type in ['png', 'jpg', 'jpeg', 'mp3', 'wav', 'ogg', 'm4a']: # save the actual file file_save_path = self.data_dir / f"{file_id}.{file_type}" uploaded_file.seek(0) # Reset file pointer with open(file_save_path, 'wb') as f: f.write(uploaded_file.read()) document[f"{document['type']}_path"] = str(file_save_path) # add embedding to collection embedding = processed_data['embedding'] if self.embeddings is None: self.embeddings = embedding.reshape(1, -1) else: # handle dimension mismatch (for mixed embedding types) current_dim = self.embeddings.shape[1] if embedding.shape[0] != current_dim: if embedding.shape[0] > current_dim: embedding = embedding[:current_dim] else: padding = np.zeros(current_dim - embedding.shape[0]) embedding = np.concatenate([embedding, padding]) self.embeddings = np.vstack([self.embeddings, embedding.reshape(1, -1)]) # Add to collections self.documents.append(document) self.metadata.append({ 'id': file_id, 'type': document['type'], 'filename': uploaded_file.name, 'timestamp': timestamp, 'content_preview': processed_data['content'][:200] + "..." if len(processed_data['content']) > 200 else processed_data['content'] }) # Rebuild search index self.build_faiss_index() self.save_data() clear_gpu_memory() return True except Exception as e: st.error(f"Error processing file: {e}") return False def build_faiss_index(self): """Build or rebuild FAISS index for similarity search""" try: if self.embeddings is not None and len(self.embeddings) > 0: dimension = self.embeddings.shape[1] if self.has_gpu: try: # Create GPU resources with limited temp memory res = faiss.StandardGpuResources() res.setTempMemory(FAISS_TEMP_MEMORY) cpu_index = faiss.IndexFlatIP(dimension) self.index = faiss.index_cpu_to_gpu(res, 0, cpu_index) except RuntimeError as e: st.warning(f"GPU FAISS failed: {e}. Using CPU indexing.") self.index = faiss.IndexFlatIP(dimension) else: self.index = faiss.IndexFlatIP(dimension) # Normalize embeddings and add to index embeddings_normalized = self.embeddings.copy().astype('float32') faiss.normalize_L2(embeddings_normalized) self.index.add(embeddings_normalized) except Exception as e: st.error(f"Error building FAISS index: {e}") def search(self, query: str, top_k: int = DEFAULT_TOP_K) -> List[Dict]: """Search for relevant documents""" try: if self.index is None or len(self.documents) == 0: return [] # get query embedding query_embedding = self.processor.get_text_embedding(query) query_embedding = query_embedding.reshape(1, -1).astype('float32') faiss.normalize_L2(query_embedding) # search scores, indices = self.index.search(query_embedding, min(top_k, len(self.documents))) # build results results = [] for score, idx in zip(scores[0], indices[0]): if idx < len(self.documents): result = self.documents[idx].copy() result['score'] = float(score) results.append(result) return results except Exception as e: st.error(f"Error searching: {e}") return [] def generate_answer(self, query: str, retrieved_docs: List[Dict]) -> str: """Generate answer using local LLM""" try: if not retrieved_docs: return "No relevant documents found to answer your query." # build context from retrieved documents context = "Context from retrieved documents:\n\n" citations = [] for i, doc in enumerate(retrieved_docs[:3]): # using top 3 documents context += f"[{i+1}] {doc['filename']} ({doc['type']}): {doc['content'][:MAX_CONTEXT_LENGTH]}...\n\n" citations.append(f"[{i+1}] {doc['filename']}") # Create prompt prompt = f"""Based on the following context, answer the user's question clearly and concisely. Use citations [1], [2], etc. Context: {context} Question: {query} Answer:""" # generate response response = self.llm.generate(prompt, temp=LLM_TEMPERATURE) response += "\n\nSources:\n" + "\n".join(citations) return response except Exception as e: st.error(f"Error generating answer: {e}") return "Error generating response." def get_stats(self) -> Dict: """to get knowledge base stats""" if not self.documents: return {"total": 0} stats = {"total": len(self.documents)} doc_types = {} for doc in self.documents: doc_type = doc['type'] doc_types[doc_type] = doc_types.get(doc_type, 0) + 1 stats.update(doc_types) return stats