Upload 9 files

app.py

@@ -0,0 +1,91 @@
+"""
+Main application entry point.
+"""
+
+import logging
+import os
+from fastapi import FastAPI
+import uvicorn
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
+
+# Configure logging
+from config import get_logging_config
+import logging.config
+logging.config.dictConfig(get_logging_config())
+logger = logging.getLogger(__name__)
+
+# Create FastAPI app
+app = FastAPI(
+    title="RAG API System",
+    description="API for RAG-based question answering",
+    version="1.0.0"
+)
+
+# Initialize components
+def initialize_components():
+    """Initialize all system components."""
+    logger.info("Initializing system components")
+    
+    # Create embedding model
+    from embedding.model import create_embedding_model
+    embedding_model = create_embedding_model()
+    logger.info(f"Embedding model initialized with dimension {embedding_model.dimension}")
+    
+    # Create vector database
+    from storage.vector_db import create_vector_database
+    vector_db = create_vector_database(dimension=embedding_model.dimension)
+    logger.info("Vector database initialized")
+    
+    # Create RAG engine
+    from rag.engine import create_rag_engine
+    rag_engine = create_rag_engine(
+        embedder=embedding_model,
+        vector_db=vector_db
+    )
+    logger.info("RAG engine initialized")
+    
+    return rag_engine
+
+# Register API routes
+def register_api_routes(app, rag_engine):
+    """Register API routes."""
+    from api.routes import RAGAPIRouter
+    router = RAGAPIRouter(app, rag_engine)
+    logger.info("API routes registered")
+
+# Add health check route
+@app.get("/", tags=["Root"])
+async def root():
+    """Root endpoint returning basic system information."""
+    return {
+        "name": "RAG API System",
+        "version": "1.0.0",
+        "status": "running"
+    }
+
+# Main entry point
+def main():
+    """Main application entry point."""
+    logger.info("Starting RAG API system")
+    
+    # Initialize components
+    rag_engine = initialize_components()
+    
+    # Register API routes
+    register_api_routes(app, rag_engine)
+    
+    # Run server if executed directly
+    if __name__ == "__main__":
+        host = os.getenv("API_HOST", "0.0.0.0")
+        port = int(os.getenv("API_PORT", "8000"))
+        
+        logger.info(f"Starting server on http://{host}:{port}")
+        uvicorn.run(app, host=host, port=port)
+    
+    return app
+
+# Create and run application
+app = main()

config.py

@@ -0,0 +1,121 @@
+"""
+Centralized configuration for the RAG system.
+"""
+
+import os
+from typing import Optional, Dict, Any
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
+
+# Embedding model settings
+EMBEDDING_MODEL_NAME = os.getenv("EMBEDDING_MODEL", "sentence-transformers/all-MiniLM-L6-v2")
+EMBEDDING_DIMENSION = int(os.getenv("EMBEDDING_DIMENSION", "384"))
+USE_GPU = os.getenv("USE_GPU", "True").lower() in ("true", "1", "t")
+
+# Document processing settings
+CHUNK_SIZE = int(os.getenv("CHUNK_SIZE", "1000"))
+CHUNK_OVERLAP = int(os.getenv("CHUNK_OVERLAP", "200"))
+MAX_LENGTH = int(os.getenv("MAX_LENGTH", "512"))
+
+# Vector database settings
+VECTOR_DB_TYPE = os.getenv("VECTOR_DB_TYPE", "faiss")  # Options: "faiss", "milvus", etc.
+FAISS_INDEX_TYPE = os.getenv("FAISS_INDEX_TYPE", "Flat")  # Options: "Flat", "IVF", "HNSW"
+MONGODB_URI = os.getenv("MONGODB_URI", "mongodb://localhost:27017/")
+DB_NAME = os.getenv("DB_NAME", "rag_db")
+COLLECTION_NAME = os.getenv("COLLECTION_NAME", "documents")
+
+# Retrieval settings
+TOP_K = int(os.getenv("TOP_K", "5"))
+SEARCH_TYPE = os.getenv("SEARCH_TYPE", "hybrid")  # Options: "semantic", "keyword", "hybrid"
+SEMANTIC_SEARCH_WEIGHT = float(os.getenv("SEMANTIC_SEARCH_WEIGHT", "0.7"))
+KEYWORD_SEARCH_WEIGHT = float(os.getenv("KEYWORD_SEARCH_WEIGHT", "0.3"))
+
+# LLM settings
+LLM_MODEL_NAME = os.getenv("LLM_MODEL", "gpt-3.5-turbo")
+LLM_API_KEY = os.getenv("OPENAI_API_KEY")
+LLM_TEMPERATURE = float(os.getenv("LLM_TEMPERATURE", "0.2"))
+LLM_MAX_TOKENS = int(os.getenv("LLM_MAX_TOKENS", "512"))
+
+# Local LLM settings (optional)
+LOCAL_LLM_MODEL_NAME = os.getenv("LOCAL_LLM_MODEL", "google/flan-t5-base")
+USE_LOCAL_LLM = os.getenv("USE_LOCAL_LLM", "False").lower() in ("true", "1", "t")
+
+# API settings
+API_HOST = os.getenv("API_HOST", "0.0.0.0")
+API_PORT = int(os.getenv("API_PORT", "8000"))
+
+# Logging settings
+LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO")
+LOG_FORMAT = os.getenv("LOG_FORMAT", "%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+
+# Default prompt template
+DEFAULT_PROMPT_TEMPLATE = """
+Answer the following question based ONLY on the provided context. 
+If you cannot answer the question based on the context, say "I don't have enough information to answer this question."
+
+Context:
+{context}
+
+Question: {query}
+
+Answer:
+"""
+
+
+def get_logging_config() -> Dict[str, Any]:
+    """Get logging configuration dictionary."""
+    return {
+        "version": 1,
+        "disable_existing_loggers": False,
+        "formatters": {
+            "standard": {
+                "format": LOG_FORMAT
+            },
+        },
+        "handlers": {
+            "console": {
+                "class": "logging.StreamHandler",
+                "level": LOG_LEVEL,
+                "formatter": "standard",
+                "stream": "ext://sys.stdout"
+            },
+        },
+        "loggers": {
+            "": {
+                "handlers": ["console"],
+                "level": LOG_LEVEL,
+                "propagate": True
+            }
+        }
+    }
+
+
+def get_model_config(model_name: Optional[str] = None) -> Dict[str, Any]:
+    """Get model-specific configuration."""
+    # Default to the configured model if none specified
+    if model_name is None:
+        model_name = EMBEDDING_MODEL_NAME
+        
+    # Common configurations for popular models
+    config_map = {
+        "sentence-transformers/all-MiniLM-L6-v2": {
+            "dimension": 384,
+            "max_length": 512,
+            "normalize": True,
+        },
+        "sentence-transformers/all-mpnet-base-v2": {
+            "dimension": 768,
+            "max_length": 512,
+            "normalize": True,
+        },
+        # Add more models as needed
+    }
+    
+    # Return specific config if available, otherwise return default values
+    return config_map.get(model_name, {
+        "dimension": EMBEDDING_DIMENSION,
+        "max_length": MAX_LENGTH,
+        "normalize": True,
+    })

document-processor.py

@@ -0,0 +1,239 @@
+"""
+Document processing utilities for text extraction and chunking.
+"""
+
+import os
+import logging
+from typing import List, Dict, Any, Optional, Tuple, Union
+import uuid
+
+# Configure logging
+logger = logging.getLogger(__name__)
+
+
+class DocumentProcessor:
+    """
+    Class to handle document processing, chunking, and text extraction.
+    """
+    
+    def __init__(
+        self,
+        chunk_size: int = 1000,
+        chunk_overlap: int = 200
+    ):
+        """
+        Initialize the document processor.
+        
+        Args:
+            chunk_size: Maximum size of text chunks in characters
+            chunk_overlap: Overlap between chunks in characters
+        """
+        self.chunk_size = chunk_size
+        self.chunk_overlap = chunk_overlap
+    
+    def process_file(
+        self,
+        file_path: str,
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> Tuple[List[str], List[Dict[str, Any]]]:
+        """
+        Process a document file: extract text and chunk it.
+        
+        Args:
+            file_path: Path to the document file
+            metadata: Optional metadata about the document
+            
+        Returns:
+            Tuple of (list of text chunks, list of metadata dictionaries)
+        """
+        if not os.path.exists(file_path):
+            raise FileNotFoundError(f"Document not found: {file_path}")
+        
+        # Extract text from document
+        logger.info(f"Processing document: {file_path}")
+        text = self._extract_text(file_path)
+        
+        if not text:
+            logger.warning(f"No text could be extracted from {file_path}")
+            return [], []
+        
+        # Create base metadata if not provided
+        base_metadata = {"source": os.path.basename(file_path)}
+        if metadata:
+            base_metadata.update(metadata)
+        
+        # Chunk the document
+        chunks = self._chunk_text(text, self.chunk_size, self.chunk_overlap)
+        logger.info(f"Created {len(chunks)} chunks from document")
+        
+        # Create chunk-specific metadata
+        chunk_metadata = []
+        for i, _ in enumerate(chunks):
+            metadata_item = {
+                **base_metadata,
+                "chunk_id": i,
+                "total_chunks": len(chunks),
+                "document_id": str(uuid.uuid4())  # Unique ID for tracking
+            }
+            chunk_metadata.append(metadata_item)
+        
+        return chunks, chunk_metadata
+    
+    def _extract_text(self, file_path: str) -> str:
+        """
+        Extract text from a document file based on its extension.
+        
+        Args:
+            file_path: Path to the document file
+            
+        Returns:
+            Extracted text
+        """
+        _, ext = os.path.splitext(file_path)
+        ext = ext.lower()
+        
+        if ext == '.pdf':
+            return self._extract_text_from_pdf(file_path)
+        elif ext == '.txt':
+            return self._extract_text_from_txt(file_path)
+        elif ext == '.md':
+            return self._extract_text_from_txt(file_path)
+        elif ext == '.docx':
+            return self._extract_text_from_docx(file_path)
+        else:
+            raise ValueError(f"Unsupported file format: {ext}")
+    
+    def _extract_text_from_pdf(self, file_path: str) -> str:
+        """
+        Extract text from a PDF file.
+        
+        Args:
+            file_path: Path to the PDF file
+            
+        Returns:
+            Extracted text
+        """
+        try:
+            import PyPDF2
+        except ImportError:
+            raise ImportError(
+                "PyPDF2 is not installed. "
+                "Please install it with `pip install PyPDF2`."
+            )
+        
+        text = ""
+        try:
+            with open(file_path, "rb") as f:
+                pdf_reader = PyPDF2.PdfReader(f)
+                num_pages = len(pdf_reader.pages)
+                logger.info(f"PDF has {num_pages} pages")
+                
+                for page in pdf_reader.pages:
+                    page_text = page.extract_text()
+                    if page_text:
+                        text += page_text + "\n\n"
+        except Exception as e:
+            logger.error(f"Error reading PDF file {file_path}: {e}")
+        
+        logger.info(f"Extracted {len(text)} characters from PDF")
+        return text
+    
+    def _extract_text_from_txt(self, file_path: str) -> str:
+        """
+        Extract text from a plain text file.
+        
+        Args:
+            file_path: Path to the text file
+            
+        Returns:
+            Extracted text
+        """
+        try:
+            with open(file_path, 'r', encoding='utf-8') as f:
+                text = f.read()
+            
+            logger.info(f"Extracted {len(text)} characters from text file")
+            return text
+        except Exception as e:
+            logger.error(f"Error reading text file {file_path}: {e}")
+            return ""
+    
+    def _extract_text_from_docx(self, file_path: str) -> str:
+        """
+        Extract text from a DOCX file.
+        
+        Args:
+            file_path: Path to the DOCX file
+            
+        Returns:
+            Extracted text
+        """
+        try:
+            import docx
+        except ImportError:
+            raise ImportError(
+                "python-docx is not installed. "
+                "Please install it with `pip install python-docx`."
+            )
+        
+        try:
+            doc = docx.Document(file_path)
+            text = "\n\n".join([paragraph.text for paragraph in doc.paragraphs if paragraph.text])
+            
+            logger.info(f"Extracted {len(text)} characters from DOCX")
+            return text
+        except Exception as e:
+            logger.error(f"Error reading DOCX file {file_path}: {e}")
+            return ""
+    
+    @staticmethod
+    def _chunk_text(
+        text: str,
+        chunk_size: int = 1000,
+        overlap: int = 200
+    ) -> List[str]:
+        """
+        Split text into overlapping chunks.
+        
+        Args:
+            text: The text to chunk
+            chunk_size: Maximum chunk size in characters
+            overlap: Overlap between chunks in characters
+            
+        Returns:
+            List of text chunks
+        """
+        if not text or not text.strip():
+            return []
+            
+        chunks = []
+        start = 0
+        text_len = len(text)
+        
+        while start < text_len:
+            # Define the initial chunk end
+            end = min(start + chunk_size, text_len)
+            
+            # Try to find a natural break point if not at the end of text
+            if end < text_len:
+                # Look for paragraph break
+                next_para = text.find('\n\n', end - overlap, end + 100)
+                if next_para != -1:
+                    end = next_para + 2
+                else:
+                    # Look for sentence break
+                    for punct in ['. ', '! ', '? ', '.\n', '!\n', '?\n']:
+                        next_sent = text.find(punct, end - overlap, end + 100)
+                        if next_sent != -1:
+                            end = next_sent + len(punct)
+                            break
+            
+            # Extract the chunk
+            chunk = text[start:end].strip()
+            if chunk:  # Only add non-empty chunks
+                chunks.append(chunk)
+            
+            # Move to next chunk with overlap
+            start = max(end - overlap, start + 1)
+            
+        return chunks

embedding-model.py

@@ -0,0 +1,286 @@
+"""
+Unified embedding model implementation supporting multiple backends.
+"""
+
+from typing import List, Union, Optional, Dict, Any
+import logging
+import numpy as np
+import torch
+from abc import ABC, abstractmethod
+
+# Configure logging
+logger = logging.getLogger(__name__)
+
+
+class EmbeddingModel(ABC):
+    """Abstract base class for embedding models."""
+    
+    @abstractmethod
+    def embed(self, texts: Union[str, List[str]], batch_size: int = 32) -> np.ndarray:
+        """
+        Convert text(s) to embedding vector(s).
+        
+        Args:
+            texts: Input text(s) to embed
+            batch_size: Batch size for processing
+            
+        Returns:
+            Embedding vector(s) as numpy array
+        """
+        pass
+    
+    @property
+    @abstractmethod
+    def dimension(self) -> int:
+        """Get the dimension of the embedding vectors."""
+        pass
+
+
+class SentenceTransformerEmbedding(EmbeddingModel):
+    """Embedding model using sentence-transformers library."""
+    
+    def __init__(
+        self, 
+        model_name: str = "all-MiniLM-L6-v2",
+        device: Optional[str] = None,
+        normalize: bool = True,
+        **kwargs
+    ):
+        """
+        Initialize the sentence transformer embedding model.
+        
+        Args:
+            model_name: Sentence transformer model name or path
+            device: Device to run model on ('cpu', 'cuda', 'cuda:0', etc.)
+            normalize: Whether to L2-normalize embeddings
+            **kwargs: Additional arguments for the model
+        """
+        try:
+            from sentence_transformers import SentenceTransformer
+        except ImportError:
+            raise ImportError(
+                "sentence-transformers is not installed. "
+                "Please install it with `pip install sentence-transformers`."
+            )
+        
+        self.model_name = model_name
+        self.normalize = normalize
+        
+        # Determine device
+        if device is None:
+            self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        else:
+            self.device = device
+            
+        logger.info(f"Loading SentenceTransformer model: {model_name} on {self.device}")
+        try:
+            self.model = SentenceTransformer(model_name, device=self.device)
+            self._dimension = self.model.get_sentence_embedding_dimension()
+            logger.info(f"Model loaded successfully. Embedding dimension: {self._dimension}")
+        except Exception as e:
+            logger.error(f"Failed to load model: {e}")
+            raise
+    
+    def embed(self, texts: Union[str, List[str]], batch_size: int = 32) -> np.ndarray:
+        """
+        Convert text(s) to embedding vector(s).
+        
+        Args:
+            texts: Input text(s) to embed
+            batch_size: Batch size for processing
+            
+        Returns:
+            Embedding vector(s) as numpy array
+        """
+        # Handle single text input
+        if isinstance(texts, str):
+            texts = [texts]
+            
+        # Validate input
+        if not texts:
+            logger.warning("Empty texts provided for embedding")
+            return np.array([])
+            
+        try:
+            # Generate embeddings
+            embeddings = self.model.encode(
+                texts, 
+                batch_size=batch_size,
+                show_progress_bar=False,
+                convert_to_numpy=True
+            )
+            
+            # Normalize if requested
+            if self.normalize:
+                embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)
+                
+            return embeddings
+        except Exception as e:
+            logger.error(f"Error during embedding generation: {e}")
+            raise
+    
+    @property
+    def dimension(self) -> int:
+        """Get the dimension of the embedding vectors."""
+        return self._dimension
+
+
+class HuggingFaceEmbedding(EmbeddingModel):
+    """Embedding model using HuggingFace transformers directly."""
+    
+    def __init__(
+        self, 
+        model_name: str = "sentence-transformers/all-MiniLM-L6-v2",
+        device: Optional[str] = None,
+        normalize: bool = True,
+        max_length: int = 512,
+        **kwargs
+    ):
+        """
+        Initialize the HuggingFace embedding model.
+        
+        Args:
+            model_name: HuggingFace model name or path
+            device: Device to run model on ('cpu', 'cuda', 'cuda:0', etc.)
+            normalize: Whether to L2-normalize embeddings
+            max_length: Maximum token length for inputs
+            **kwargs: Additional arguments for the model
+        """
+        try:
+            from transformers import AutoTokenizer, AutoModel
+        except ImportError:
+            raise ImportError(
+                "transformers is not installed. "
+                "Please install it with `pip install transformers`."
+            )
+        
+        self.model_name = model_name
+        self.normalize = normalize
+        self.max_length = max_length
+        
+        # Determine device
+        if device is None:
+            self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        else:
+            self.device = device
+            
+        logger.info(f"Loading HuggingFace model: {model_name} on {self.device}")
+        try:
+            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+            self.model = AutoModel.from_pretrained(model_name)
+            self.model.to(self.device)
+            self.model.eval()
+            
+            # Get embedding dimension from model config
+            self._dimension = self.model.config.hidden_size
+            logger.info(f"Model loaded successfully. Embedding dimension: {self._dimension}")
+        except Exception as e:
+            logger.error(f"Failed to load model: {e}")
+            raise
+    
+    def _mean_pooling(self, model_output, attention_mask):
+        """Perform mean pooling on token embeddings."""
+        token_embeddings = model_output.last_hidden_state
+        input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+        return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
+    
+    def embed(self, texts: Union[str, List[str]], batch_size: int = 32) -> np.ndarray:
+        """
+        Convert text(s) to embedding vector(s).
+        
+        Args:
+            texts: Input text(s) to embed
+            batch_size: Batch size for processing
+            
+        Returns:
+            Embedding vector(s) as numpy array
+        """
+        # Handle single text input
+        if isinstance(texts, str):
+            texts = [texts]
+            
+        # Validate input
+        if not texts:
+            logger.warning("Empty texts provided for embedding")
+            return np.array([])
+            
+        try:
+            all_embeddings = []
+            
+            # Process in batches
+            for i in range(0, len(texts), batch_size):
+                batch_texts = texts[i:i+batch_size]
+                
+                # Tokenize and move to device
+                inputs = self.tokenizer(
+                    batch_texts,
+                    padding=True,
+                    truncation=True,
+                    max_length=self.max_length,
+                    return_tensors="pt"
+                )
+                inputs = {k: v.to(self.device) for k, v in inputs.items()}
+                
+                # Generate embeddings
+                with torch.no_grad():
+                    outputs = self.model(**inputs)
+                    embeddings = self._mean_pooling(outputs, inputs["attention_mask"])
+                    
+                    # Normalize if requested
+                    if self.normalize:
+                        embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
+                        
+                    # Move to CPU and convert to numpy
+                    embeddings = embeddings.cpu().numpy()
+                    all_embeddings.append(embeddings)
+            
+            # Concatenate all batches
+            return np.vstack(all_embeddings) if all_embeddings else np.array([])
+            
+        except Exception as e:
+            logger.error(f"Error during embedding generation: {e}")
+            raise
+    
+    @property
+    def dimension(self) -> int:
+        """Get the dimension of the embedding vectors."""
+        return self._dimension
+
+
+# Factory function to create embedding models
+def create_embedding_model(
+    backend: str = "sentence-transformers",
+    model_name: Optional[str] = None,
+    **kwargs
+) -> EmbeddingModel:
+    """
+    Factory function to create an embedding model.
+    
+    Args:
+        backend: Backend to use ('sentence-transformers' or 'huggingface')
+        model_name: Model name or path
+        **kwargs: Additional arguments for the model
+        
+    Returns:
+        An EmbeddingModel instance
+    """
+    from ..config import EMBEDDING_MODEL_NAME, get_model_config
+    
+    # Use config model if not specified
+    if model_name is None:
+        model_name = EMBEDDING_MODEL_NAME
+    
+    # Get model-specific config
+    model_config = get_model_config(model_name)
+    
+    # Override with provided kwargs
+    for k, v in kwargs.items():
+        model_config[k] = v
+    
+    # Create the model
+    if backend.lower() == "sentence-transformers":
+        return SentenceTransformerEmbedding(model_name=model_name, **model_config)
+    elif backend.lower() in ["huggingface", "hf", "transformers"]:
+        return HuggingFaceEmbedding(model_name=model_name, **model_config)
+    else:
+        raise ValueError(f"Unsupported backend: {backend}")

package-init.py

@@ -0,0 +1,23 @@
+# embedding/__init__.py
+from .model import create_embedding_model, EmbeddingModel, SentenceTransformerEmbedding, HuggingFaceEmbedding
+
+# storage/__init__.py
+from .vector_db import Document, VectorDatabase, FaissVectorDatabase, create_vector_database
+
+# document/__init__.py
+from .processor import DocumentProcessor
+
+# retrieval/__init__.py
+# Import relevant classes if needed
+
+# rag/__init__.py
+from .engine import RAGEngine, create_rag_engine
+
+# api/__init__.py
+from .routes import RAGAPIRouter
+
+# ui/__init__.py
+# No exports needed
+
+# utils/__init__.py
+# Import utility functions if needed

rag-engine.py

@@ -0,0 +1,357 @@
+"""
+Main RAG (Retrieval-Augmented Generation) engine implementation.
+"""
+
+import os
+import logging
+from typing import List, Dict, Any, Optional, Tuple, Union
+import numpy as np
+
+# Configure logging
+logger = logging.getLogger(__name__)
+
+
+class RAGEngine:
+    """Retrieval-Augmented Generation (RAG) engine for question answering."""
+    
+    def __init__(
+        self,
+        embedder,
+        vector_db,
+        llm=None,
+        top_k: int = 5,
+        search_type: str = "hybrid",
+        prompt_template: Optional[str] = None
+    ):
+        """
+        Initialize the RAG engine.
+        
+        Args:
+            embedder: Embedding model
+            vector_db: Vector database for document storage and retrieval
+            llm: Language model for text generation (optional)
+            top_k: Number of documents to retrieve
+            search_type: Type of search ('semantic', 'keyword', 'hybrid')
+            prompt_template: Optional custom prompt template
+        """
+        self.embedder = embedder
+        self.vector_db = vector_db
+        self.llm = llm
+        self.top_k = top_k
+        self.search_type = search_type
+        
+        # Set default prompt template if none provided
+        if prompt_template is None:
+            from ..config import DEFAULT_PROMPT_TEMPLATE
+            self.prompt_template = DEFAULT_PROMPT_TEMPLATE
+        else:
+            self.prompt_template = prompt_template
+    
+    def add_documents(
+        self,
+        texts: List[str],
+        metadata: Optional[List[Dict[str, Any]]] = None,
+        batch_size: int = 32
+    ) -> List[str]:
+        """
+        Add documents to the database.
+        
+        Args:
+            texts: List of text chunks
+            metadata: Optional list of metadata dictionaries for each text
+            batch_size: Batch size for embedding generation
+            
+        Returns:
+            List of document IDs
+        """
+        from ..storage.vector_db import Document
+        
+        # Handle metadata
+        if metadata is None:
+            metadata = [{} for _ in texts]
+        elif len(metadata) != len(texts):
+            raise ValueError(f"Length mismatch: got {len(texts)} texts but {len(metadata)} metadata entries")
+        
+        # Generate embeddings in batches
+        doc_ids = []
+        
+        for i in range(0, len(texts), batch_size):
+            batch_texts = texts[i:i+batch_size]
+            batch_metadata = metadata[i:i+batch_size]
+            
+            # Generate embeddings
+            logger.info(f"Generating embeddings for batch {i//batch_size + 1}/{(len(texts)-1)//batch_size + 1}")
+            batch_embeddings = self.embedder.embed(batch_texts)
+            
+            # Create document objects
+            documents = []
+            for text, meta, embedding in zip(batch_texts, batch_metadata, batch_embeddings):
+                doc = Document(text=text, metadata=meta, embedding=embedding)
+                documents.append(doc)
+            
+            # Add to database
+            batch_ids = self.vector_db.add_documents(documents)
+            doc_ids.extend(batch_ids)
+        
+        logger.info(f"Added {len(doc_ids)} documents to database")
+        return doc_ids
+    
+    def search(
+        self,
+        query: str,
+        top_k: Optional[int] = None,
+        search_type: Optional[str] = None,
+        filter_dict: Optional[Dict[str, Any]] = None
+    ) -> List[Dict[str, Any]]:
+        """
+        Search for relevant documents.
+        
+        Args:
+            query: Query string
+            top_k: Number of results to return (defaults to self.top_k)
+            search_type: Type of search (defaults to self.search_type)
+            filter_dict: Dictionary of metadata filters
+            
+        Returns:
+            List of document dictionaries
+        """
+        if top_k is None:
+            top_k = self.top_k
+            
+        if search_type is None:
+            search_type = self.search_type
+        
+        # Create filter function if filter_dict is provided
+        filter_func = None
+        if filter_dict:
+            def filter_func(doc):
+                for key, value in filter_dict.items():
+                    # Handle nested keys (e.g., "metadata.source")
+                    if "." in key:
+                        parts = key.split(".")
+                        current = doc.metadata
+                        for part in parts[:-1]:
+                            if part not in current:
+                                return False
+                            current = current[part]
+                        if parts[-1] not in current or current[parts[-1]] != value:
+                            return False
+                    elif key not in doc.metadata or doc.metadata[key] != value:
+                        return False
+                return True
+        
+        # Generate query embedding
+        query_embedding = self.embedder.embed(query)
+        
+        # Perform search
+        results = self.vector_db.search(query_embedding, top_k, filter_func)
+        
+        # Convert results to dictionaries
+        return [
+            {
+                "id": doc.id,
+                "text": doc.text,
+                "metadata": doc.metadata,
+                "score": score
+            }
+            for doc, score in results
+        ]
+    
+    def generate_response(
+        self,
+        query: str,
+        top_k: Optional[int] = None,
+        search_type: Optional[str] = None,
+        filter_dict: Optional[Dict[str, Any]] = None,
+        max_tokens: int = 512
+    ) -> Dict[str, Any]:
+        """
+        Generate a response to a query using RAG.
+        
+        Args:
+            query: Query string
+            top_k: Number of documents to retrieve
+            search_type: Type of search
+            filter_dict: Optional filter for document retrieval
+            max_tokens: Maximum number of tokens in the response
+            
+        Returns:
+            Dictionary with query, response, and retrieved documents
+        """
+        # Retrieve relevant documents
+        retrieved_docs = self.search(query, top_k, search_type, filter_dict)
+        
+        # If no documents were found, return a default message
+        if not retrieved_docs:
+            return {
+                "query": query,
+                "response": "I couldn't find any relevant information to answer your question.",
+                "retrieved_documents": [],
+                "search_type": search_type or self.search_type
+            }
+        
+        # Format context from retrieved documents
+        context = self._format_context(retrieved_docs)
+        
+        # Format prompt with context and query
+        prompt = self.prompt_template.format(context=context, query=query)
+        
+        # Generate response using LLM
+        if self.llm is None:
+            logger.warning("No LLM provided, returning only retrieved documents")
+            response = "No language model available to generate a response. Here's what I found in the documents."
+        else:
+            response = self._generate_llm_response(prompt, max_tokens)
+        
+        # Return the results
+        return {
+            "query": query,
+            "response": response,
+            "retrieved_documents": retrieved_docs,
+            "search_type": search_type or self.search_type
+        }
+    
+    def _format_context(self, documents: List[Dict[str, Any]]) -> str:
+        """
+        Format retrieved documents into context for the prompt.
+        
+        Args:
+            documents: List of retrieved documents
+            
+        Returns:
+            Formatted context string
+        """
+        context_parts = []
+        
+        for i, doc in enumerate(documents):
+            # Extract relevant fields
+            text = doc["text"]
+            metadata = doc["metadata"]
+            source = metadata.get("source", "Unknown")
+            
+            # Format the document
+            doc_text = f"Document {i+1}: [Source: {source}]\n{text}\n"
+            context_parts.append(doc_text)
+        
+        return "\n".join(context_parts)
+    
+    def _generate_llm_response(self, prompt: str, max_tokens: int) -> str:
+        """
+        Generate a response using the LLM.
+        
+        Args:
+            prompt: The formatted prompt
+            max_tokens: Maximum number of tokens in the response
+            
+        Returns:
+            Generated response
+        """
+        if hasattr(self.llm, "generate_openai_response"):
+            # OpenAI-compatible LLM
+            return self.llm.generate_openai_response(prompt, max_tokens)
+        elif hasattr(self.llm, "generate_huggingface_response"):
+            # HuggingFace-compatible LLM
+            return self.llm.generate_huggingface_response(prompt, max_tokens)
+        else:
+            # Default implementation
+            try:
+                return self.llm.generate_response(prompt, max_tokens)
+            except Exception as e:
+                logger.error(f"Error generating response: {e}")
+                return "I encountered an error while generating a response."
+    
+    def update_prompt_template(self, new_template: str) -> None:
+        """
+        Update the prompt template.
+        
+        Args:
+            new_template: New prompt template
+        """
+        self.prompt_template = new_template
+        logger.info("Updated prompt template")
+    
+    def count_documents(self) -> int:
+        """
+        Get the number of documents in the database.
+        
+        Returns:
+            Number of documents
+        """
+        return self.vector_db.count_documents()
+    
+    def clear_documents(self) -> None:
+        """Clear all documents from the database."""
+        self.vector_db.clear()
+        logger.info("Cleared all documents from database")
+
+
+# Factory function to create the RAG engine
+def create_rag_engine(
+    embedder=None,
+    vector_db=None,
+    llm=None,
+    config=None
+) -> RAGEngine:
+    """
+    Factory function to create a RAG engine.
+    
+    Args:
+        embedder: Embedding model (if None, created based on config)
+        vector_db: Vector database (if None, created based on config)
+        llm: Language model (if None, created based on config)
+        config: Configuration module or dictionary
+        
+    Returns:
+        Configured RAGEngine instance
+    """
+    # Load configuration if provided
+    if config is None:
+        from ..config import (
+            TOP_K,
+            SEARCH_TYPE,
+            DEFAULT_PROMPT_TEMPLATE
+        )
+    else:
+        TOP_K = config.get("TOP_K", 5)
+        SEARCH_TYPE = config.get("SEARCH_TYPE", "hybrid")
+        DEFAULT_PROMPT_TEMPLATE = config.get(
+            "DEFAULT_PROMPT_TEMPLATE",
+            """
+            Answer the following question based ONLY on the provided context.
+            
+            Context:
+            {context}
+            
+            Question: {query}
+            
+            Answer:
+            """
+        )
+    
+    # Create embedding model if not provided
+    if embedder is None:
+        from ..embedding.model import create_embedding_model
+        embedder = create_embedding_model()
+    
+    # Create vector database if not provided
+    if vector_db is None:
+        from ..storage.vector_db import create_vector_database
+        vector_db = create_vector_database(dimension=embedder.dimension)
+    
+    # Create language model if not provided and requested
+    if llm is None:
+        try:
+            from ..llm.model import create_llm
+            llm = create_llm()
+        except (ImportError, ModuleNotFoundError):
+            logger.warning("LLM module not found, proceeding without an LLM")
+    
+    # Create and return the RAG engine
+    return RAGEngine(
+        embedder=embedder,
+        vector_db=vector_db,
+        llm=llm,
+        top_k=TOP_K,
+        search_type=SEARCH_TYPE,
+        prompt_template=DEFAULT_PROMPT_TEMPLATE
+    )

readme.md

@@ -0,0 +1,105 @@
+# RAG System
+
+A modular Retrieval-Augmented Generation (RAG) system for document-based question answering.
+
+## Features
+
+- **Document Processing**: Extract and chunk text from PDF, DOCX, and TXT files
+- **Semantic Search**: Embed and search documents based on meaning, not just keywords
+- **Flexible Architecture**: Support for multiple embedding models and vector databases
+- **REST API**: API for integrating with other applications
+- **Web UI**: User-friendly Streamlit interface for document upload and querying
+
+## Architecture
+
+The system consists of the following components:
+
+- **Embedding Model**: Converts text to vector embeddings
+- **Vector Database**: Stores and searches document embeddings
+- **Document Processor**: Extracts and chunks text from documents
+- **RAG Engine**: Combines retrieval and generation for question answering
+- **API**: Exposes functionality through a RESTful API
+- **UI**: Provides a user interface for interacting with the system
+
+## Installation
+
+### Prerequisites
+
+- Python 3.8+
+- pip
+
+### Setup
+
+1. Clone the repository:
+   ```bash
+   git clone https://github.com/yourusername/rag-system.git
+   cd rag-system
+   ```
+
+2. Install dependencies:
+   ```bash
+   pip install -r requirements.txt
+   ```
+
+3. Set up environment variables (optional):
+   ```bash
+   cp .env.example .env
+   # Edit .env with your settings
+   ```
+
+## Usage
+
+### API Server
+
+Run the API server:
+
+```bash
+python app.py
+```
+
+The API will be available at http://localhost:8000
+
+### Streamlit UI
+
+Run the Streamlit UI:
+
+```bash
+streamlit run ui/app.py
+```
+
+The UI will be available at http://localhost:8501
+
+## API Endpoints
+
+- `POST /documents`: Add documents
+- `POST /upload`: Upload and process document files
+- `POST /query`: Query the RAG system
+- `GET /search`: Search for documents
+- `DELETE /documents`: Clear all documents
+- `GET /health`: Check system health
+
+## Configuration
+
+The system can be configured through environment variables or the `config.py` file:
+
+- `EMBEDDING_MODEL_NAME`: Name of the embedding model
+- `VECTOR_DB_TYPE`: Type of vector database to use
+- `CHUNK_SIZE`: Size of document chunks
+- `CHUNK_OVERLAP`: Overlap between chunks
+- `TOP_K`: Number of documents to retrieve
+- `SEARCH_TYPE`: Type of search (semantic, keyword, hybrid)
+- `LLM_MODEL_NAME`: Name of the language model for generation
+- `LLM_API_KEY`: API key for the language model
+
+## Extending
+
+The modular architecture makes it easy to extend the system:
+
+- Add new embedding models in `embedding/model.py`
+- Add new vector databases in `storage/vector_db.py`
+- Add support for new document types in `document/processor.py`
+- Add new LLM integrations in `llm/model.py`
+
+## License
+
+[MIT License](LICENSE)

streamlit-app.py

@@ -0,0 +1,229 @@
+"""
+Streamlit UI for the RAG system.
+"""
+
+import os
+import streamlit as st
+import tempfile
+import logging
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
+
+# Configure logging
+from config import get_logging_config
+import logging.config
+logging.config.dictConfig(get_logging_config())
+logger = logging.getLogger(__name__)
+
+# Set page config
+st.set_page_config(
+    page_title="RAG Document QA System",
+    page_icon="📚",
+    layout="wide",
+    initial_sidebar_state="expanded"
+)
+
+# Initialize session state
+if "document_count" not in st.session_state:
+    st.session_state.document_count = 0
+if "initialized" not in st.session_state:
+    st.session_state.initialized = False
+
+# Initialize RAG engine
+@st.cache_resource
+def initialize_rag_engine():
+    """Initialize RAG engine."""
+    from embedding.model import create_embedding_model
+    from storage.vector_db import create_vector_database
+    from rag.engine import create_rag_engine
+    
+    # Create components
+    embedding_model = create_embedding_model()
+    vector_db = create_vector_database(dimension=embedding_model.dimension)
+    rag_engine = create_rag_engine(
+        embedder=embedding_model,
+        vector_db=vector_db
+    )
+    
+    st.session_state.initialized = True
+    return rag_engine
+
+# Initialize document processor
+@st.cache_resource
+def initialize_document_processor():
+    """Initialize document processor."""
+    from document.processor import DocumentProcessor
+    return DocumentProcessor()
+
+# Main application
+def main():
+    """Main Streamlit application."""
+    # Initialize components
+    rag_engine = initialize_rag_engine()
+    doc_processor = initialize_document_processor()
+    
+    # Update document count
+    st.session_state.document_count = rag_engine.count_documents()
+    
+    # Sidebar
+    st.sidebar.title("📚 RAG Document QA")
+    
+    # Document upload
+    st.sidebar.header("Upload Documents")
+    uploaded_file = st.sidebar.file_uploader(
+        "Choose a document file (PDF, TXT, DOCX)",
+        type=["pdf", "txt", "md", "docx"]
+    )
+    
+    # Upload settings
+    st.sidebar.subheader("Document Settings")
+    chunk_size = st.sidebar.slider(
+        "Chunk Size",
+        min_value=100,
+        max_value=2000,
+        value=1000,
+        step=100,
+        help="Size of text chunks in characters"
+    )
+    chunk_overlap = st.sidebar.slider(
+        "Chunk Overlap",
+        min_value=0,
+        max_value=500,
+        value=200,
+        step=50,
+        help="Overlap between chunks in characters"
+    )
+    
+    # Search settings
+    st.sidebar.header("Search Settings")
+    top_k = st.sidebar.slider(
+        "Results to Return",
+        min_value=1,
+        max_value=10,
+        value=3,
+        help="Number of document chunks to retrieve"
+    )
+    search_type = st.sidebar.selectbox(
+        "Search Type",
+        options=["hybrid", "semantic", "keyword"],
+        index=0,
+        help="Type of search to perform"
+    )
+    
+    # Document info
+    st.sidebar.header("Document Store")
+    st.sidebar.metric("Documents Stored", st.session_state.document_count)
+    
+    if st.sidebar.button("Clear All Documents"):
+        rag_engine.clear_documents()
+        st.session_state.document_count = 0
+        st.sidebar.success("Document store cleared!")
+        st.experimental_rerun()
+    
+    # Process uploaded file
+    if uploaded_file is not None:
+        with st.sidebar.expander("Upload Status", expanded=True):
+            with st.spinner('Processing document...'):
+                # Save to temporary file
+                with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(uploaded_file.name)[1]) as tmp_file:
+                    tmp_file.write(uploaded_file.getvalue())
+                    tmp_file_path = tmp_file.name
+                
+                try:
+                    # Process document
+                    doc_processor.chunk_size = chunk_size
+                    doc_processor.chunk_overlap = chunk_overlap
+                    
+                    chunks, chunk_metadata = doc_processor.process_file(
+                        tmp_file_path,
+                        metadata={"filename": uploaded_file.name, "source": "UI upload"}
+                    )
+                    
+                    if not chunks:
+                        st.sidebar.error("No text could be extracted from the document.")
+                    else:
+                        # Add chunks to RAG engine
+                        doc_ids = rag_engine.add_documents(chunks, chunk_metadata)
+                        
+                        # Update document count
+                        st.session_state.document_count = rag_engine.count_documents()
+                        
+                        st.sidebar.success(f"Added {len(chunks)} document chunks!")
+                except Exception as e:
+                    st.sidebar.error(f"Error processing document: {str(e)}")
+                finally:
+                    # Clean up temporary file
+                    os.unlink(tmp_file_path)
+    
+    # Main content
+    st.title("📚 Document Query System")
+    
+    if st.session_state.document_count == 0:
+        st.info("👈 Please upload documents using the sidebar to get started.")
+        
+        # Sample documents
+        st.subheader("Sample Text")
+        sample_text = st.text_area(
+            "Or try adding some sample text directly:",
+            height=200
+        )
+        
+        if sample_text and st.button("Add Sample Text"):
+            with st.spinner('Processing text...'):
+                # Chunk the text
+                chunks = doc_processor._chunk_text(sample_text, chunk_size, chunk_overlap)
+                
+                # Create metadata
+                chunk_metadata = [
+                    {"source": "Sample text", "chunk_id": i, "total_chunks": len(chunks)}
+                    for i in range(len(chunks))
+                ]
+                
+                # Add to RAG engine
+                doc_ids = rag_engine.add_documents(chunks, chunk_metadata)
+                
+                # Update document count
+                st.session_state.document_count = rag_engine.count_documents()
+                
+                st.success(f"Added {len(chunks)} text chunks!")
+                st.experimental_rerun()
+    else:
+        # Question answering
+        st.subheader("Ask a Question")
+        question = st.text_input("Enter your question:")
+        
+        if question:
+            with st.spinner('Searching for answer...'):
+                try:
+                    # Generate response
+                    result = rag_engine.generate_response(
+                        query=question,
+                        top_k=top_k,
+                        search_type=search_type
+                    )
+                    
+                    # Display response
+                    st.markdown("### Answer")
+                    st.write(result["response"])
+                    
+                    # Display sources
+                    st.markdown("### Sources")
+                    for i, doc in enumerate(result["retrieved_documents"]):
+                        with st.expander(f"Source {i+1} (Score: {doc['score']:.2f})"):
+                            st.markdown(f"**Source:** {doc['metadata'].get('source', 'Unknown')}")
+                            st.text(doc["text"])
+                except Exception as e:
+                    st.error(f"Error generating response: {str(e)}")
+    
+    # About section
+    st.sidebar.markdown("---")
+    st.sidebar.info(
+        "This application allows you to upload documents and ask questions about their content. "
+        "The system uses embedding models for semantic search and retrieval."
+    )
+
+# Run the application
+if __name__ == "__main__":
+    main()

vector-db.py

@@ -0,0 +1,545 @@
+"""
+Vector database implementation for document storage and retrieval.
+"""
+
+from typing import List, Dict, Any, Optional, Union, Tuple, Callable
+import logging
+import os
+import json
+import uuid
+import numpy as np
+from dataclasses import dataclass, field, asdict
+
+# Configure logging
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Document:
+    """Class to represent a document or text chunk with metadata and embeddings."""
+    text: str
+    metadata: Dict[str, Any] = field(default_factory=dict)
+    embedding: Optional[np.ndarray] = None
+    id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for serialization."""
+        result = asdict(self)
+        # Convert numpy array to list for JSON serialization
+        if self.embedding is not None:
+            result['embedding'] = self.embedding.tolist()
+        return result
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> 'Document':
+        """Create Document from dictionary."""
+        if 'embedding' in data and data['embedding'] is not None:
+            data['embedding'] = np.array(data['embedding'], dtype=np.float32)
+        return cls(**data)
+
+
+class VectorDatabase:
+    """Base class for vector databases."""
+    
+    def __init__(self, dimension: int = 384):
+        """
+        Initialize the vector database.
+        
+        Args:
+            dimension: Dimension of the embedding vectors
+        """
+        self.dimension = dimension
+    
+    def add_document(self, document: Document) -> str:
+        """
+        Add a document to the database.
+        
+        Args:
+            document: Document to add
+            
+        Returns:
+            Document ID
+        """
+        raise NotImplementedError("Subclasses must implement add_document")
+    
+    def add_documents(self, documents: List[Document]) -> List[str]:
+        """
+        Add multiple documents to the database.
+        
+        Args:
+            documents: List of documents to add
+            
+        Returns:
+            List of document IDs
+        """
+        return [self.add_document(doc) for doc in documents]
+    
+    def search(
+        self,
+        query_embedding: np.ndarray,
+        top_k: int = 5,
+        filter_func: Optional[Callable[[Document], bool]] = None
+    ) -> List[Tuple[Document, float]]:
+        """
+        Search for similar documents.
+        
+        Args:
+            query_embedding: Query embedding vector
+            top_k: Number of results to return
+            filter_func: Optional function to filter results
+            
+        Returns:
+            List of (document, score) tuples
+        """
+        raise NotImplementedError("Subclasses must implement search")
+    
+    def delete_document(self, doc_id: str) -> bool:
+        """
+        Delete a document from the database.
+        
+        Note: FAISS doesn't support direct deletion, so we handle this
+        by rebuilding the index when needed.
+        
+        Args:
+            doc_id: Document ID to delete
+            
+        Returns:
+            True if document was found and deleted
+        """
+        if doc_id not in self.documents:
+            return False
+        
+        # Remove from documents dictionary
+        del self.documents[doc_id]
+        
+        # If document was in index, mark for rebuild
+        if doc_id in self.id_to_index:
+            # Remove from mappings
+            del self.id_to_index[doc_id]
+            # We'll rebuild the index on the next query
+            
+        return True
+    
+    def _rebuild_index(self):
+        """Rebuild the FAISS index from scratch."""
+        # Re-initialize the index
+        self._initialize_index()
+        self.id_to_index = {}
+        self.index_to_id = {}
+        
+        # Collect all documents with embeddings
+        docs_with_embeddings = [doc for doc in self.documents.values() if doc.embedding is not None]
+        
+        if not docs_with_embeddings:
+            logger.warning("No documents with embeddings to rebuild index")
+            return
+        
+        # Extract embeddings
+        embeddings = np.array([doc.embedding for doc in docs_with_embeddings], dtype=np.float32)
+        
+        # Train if needed
+        if self.needs_training and len(docs_with_embeddings) >= 100:
+            logger.info("Training FAISS index during rebuild")
+            train_data = embeddings[:min(1000, len(embeddings))]
+            self.index.train(train_data)
+        
+        # Add to index if trained or doesn't need training
+        if not self.needs_training or not self.needs_training or self.index.is_trained:
+            self.index.add(embeddings)
+            
+            # Update mappings
+            for i, doc in enumerate(docs_with_embeddings):
+                self.id_to_index[doc.id] = i
+                self.index_to_id[i] = doc.id
+    
+    def get_document(self, doc_id: str) -> Optional[Document]:
+        """
+        Get a document by ID.
+        
+        Args:
+            doc_id: Document ID to get
+            
+        Returns:
+            Document if found, None otherwise
+        """
+        return self.documents.get(doc_id)
+    
+    def count_documents(self) -> int:
+        """
+        Get the number of documents in the database.
+        
+        Returns:
+            Number of documents
+        """
+        return len(self.documents)
+    
+    def clear(self) -> None:
+        """Clear all documents from the database."""
+        self.documents = {}
+        self.id_to_index = {}
+        self.index_to_id = {}
+        self._initialize_index()
+    
+    def save(self, directory: str) -> None:
+        """
+        Save the database to disk.
+        
+        Args:
+            directory: Directory to save to
+        """
+        import faiss
+        
+        os.makedirs(directory, exist_ok=True)
+        
+        # Save documents
+        documents_data = {doc_id: doc.to_dict() for doc_id, doc in self.documents.items()}
+        with open(os.path.join(directory, "documents.json"), "w") as f:
+            json.dump(documents_data, f)
+        
+        # Save mappings
+        mappings = {
+            "id_to_index": self.id_to_index,
+            "index_to_id": {str(k): v for k, v in self.index_to_id.items()}  # Convert int keys to strings for JSON
+        }
+        with open(os.path.join(directory, "mappings.json"), "w") as f:
+            json.dump(mappings, f)
+        
+        # Save index
+        faiss.write_index(self.index, os.path.join(directory, "faiss_index.bin"))
+        
+        # Save metadata
+        metadata = {
+            "dimension": self.dimension,
+            "index_type": self.index_type,
+            "document_count": len(self.documents)
+        }
+        with open(os.path.join(directory, "metadata.json"), "w") as f:
+            json.dump(metadata, f)
+    
+    @classmethod
+    def load(cls, directory: str) -> 'FaissVectorDatabase':
+        """
+        Load a database from disk.
+        
+        Args:
+            directory: Directory to load from
+            
+        Returns:
+            Loaded FaissVectorDatabase
+        """
+        import faiss
+        
+        # Load metadata
+        with open(os.path.join(directory, "metadata.json"), "r") as f:
+            metadata = json.load(f)
+        
+        # Create instance
+        db = cls(dimension=metadata["dimension"], index_type=metadata["index_type"])
+        
+        # Load documents
+        with open(os.path.join(directory, "documents.json"), "r") as f:
+            documents_data = json.load(f)
+        
+        db.documents = {doc_id: Document.from_dict(doc_data) for doc_id, doc_data in documents_data.items()}
+        
+        # Load mappings
+        with open(os.path.join(directory, "mappings.json"), "r") as f:
+            mappings = json.load(f)
+        
+        db.id_to_index = mappings["id_to_index"]
+        db.index_to_id = {int(k): v for k, v in mappings["index_to_id"].items()}  # Convert string keys back to int
+        
+        # Load index
+        db.index = faiss.read_index(os.path.join(directory, "faiss_index.bin"))
+        
+        return db
+
+
+# Factory function to create vector databases
+def create_vector_database(
+    db_type: str = "faiss",
+    dimension: int = 384,
+    **kwargs
+) -> VectorDatabase:
+    """
+    Factory function to create a vector database.
+    
+    Args:
+        db_type: Database type ('faiss')
+        dimension: Dimension of the embedding vectors
+        **kwargs: Additional arguments for the database
+        
+    Returns:
+        A VectorDatabase instance
+    """
+    if db_type.lower() == "faiss":
+        return FaissVectorDatabase(dimension=dimension, **kwargs)
+    else:
+        raise ValueError(f"Unsupported database type: {db_type}")
+        Args:
+            doc_id: Document ID to delete
+            
+        Returns:
+            True if document was deleted, False otherwise
+        """
+        raise NotImplementedError("Subclasses must implement delete_document")
+    
+    def get_document(self, doc_id: str) -> Optional[Document]:
+        """
+        Get a document by ID.
+        
+        Args:
+            doc_id: Document ID to get
+            
+        Returns:
+            Document if found, None otherwise
+        """
+        raise NotImplementedError("Subclasses must implement get_document")
+    
+    def count_documents(self) -> int:
+        """
+        Get the number of documents in the database.
+        
+        Returns:
+            Number of documents
+        """
+        raise NotImplementedError("Subclasses must implement count_documents")
+    
+    def clear(self) -> None:
+        """Clear all documents from the database."""
+        raise NotImplementedError("Subclasses must implement clear")
+    
+    def save(self, directory: str) -> None:
+        """
+        Save the database to disk.
+        
+        Args:
+            directory: Directory to save to
+        """
+        raise NotImplementedError("Subclasses must implement save")
+    
+    @classmethod
+    def load(cls, directory: str) -> 'VectorDatabase':
+        """
+        Load a database from disk.
+        
+        Args:
+            directory: Directory to load from
+            
+        Returns:
+            Loaded database
+        """
+        raise NotImplementedError("Subclasses must implement load")
+
+
+class FaissVectorDatabase(VectorDatabase):
+    """Vector database implementation using FAISS."""
+    
+    def __init__(self, dimension: int = 384, index_type: str = "Flat"):
+        """
+        Initialize the FAISS vector database.
+        
+        Args:
+            dimension: Dimension of the embedding vectors
+            index_type: FAISS index type (e.g., "Flat", "IVF", "HNSW")
+        """
+        super().__init__(dimension)
+        self.index_type = index_type
+        self.documents: Dict[str, Document] = {}
+        self.id_to_index: Dict[str, int] = {}
+        self.index_to_id: Dict[int, str] = {}
+        
+        # Initialize FAISS index
+        self._initialize_index()
+    
+    def _initialize_index(self):
+        """Initialize FAISS index based on the specified type."""
+        try:
+            import faiss
+        except ImportError:
+            raise ImportError(
+                "faiss-cpu is not installed. "
+                "Please install it with `pip install faiss-cpu` or `pip install faiss-gpu`."
+            )
+        
+        if self.index_type == "Flat":
+            self.index = faiss.IndexFlatL2(self.dimension)
+        elif self.index_type == "IVF":
+            # IVF requires training, so we'll use a placeholder
+            # This would need to be trained on actual data
+            quantizer = faiss.IndexFlatL2(self.dimension)
+            n_cells = 100  # Number of centroids
+            self.index = faiss.IndexIVFFlat(quantizer, self.dimension, n_cells)
+            self.index.nprobe = 10  # Number of cells to probe at search time
+        elif self.index_type == "HNSW":
+            self.index = faiss.IndexHNSWFlat(self.dimension, 32)  # 32 neighbors per node
+        else:
+            logger.warning(f"Unknown index type {self.index_type}, falling back to Flat")
+            self.index = faiss.IndexFlatL2(self.dimension)
+        
+        # Mark if index needs training
+        self.needs_training = self.index_type in ["IVF"]
+    
+    def add_document(self, document: Document) -> str:
+        """
+        Add a document to the database.
+        
+        Args:
+            document: Document to add
+            
+        Returns:
+            Document ID
+        """
+        # If no embedding is provided, log warning
+        if document.embedding is None:
+            logger.warning(f"Document {document.id} has no embedding - skipping indexing")
+            self.documents[document.id] = document
+            return document.id
+        
+        # Ensure embedding is in the right format
+        embedding = np.array([document.embedding], dtype=np.float32)
+        
+        # Train index if needed and we have enough data
+        if self.needs_training and len(self.documents) >= 100 and not self.index.is_trained:
+            logger.info("Training FAISS index")
+            # Collect 1000 embeddings for training
+            train_data = np.vstack([doc.embedding for doc in list(self.documents.values())[:1000]])
+            self.index.train(train_data)
+        
+        # Add to FAISS index if it's trained or doesn't need training
+        if not self.needs_training or self.index.is_trained:
+            idx = len(self.id_to_index)
+            self.index.add(embedding)
+            
+            # Update mapping dictionaries
+            self.id_to_index[document.id] = idx
+            self.index_to_id[idx] = document.id
+        
+        # Store document
+        self.documents[document.id] = document
+        
+        return document.id
+    
+    def add_documents(self, documents: List[Document]) -> List[str]:
+        """
+        Add multiple documents to the database.
+        
+        Args:
+            documents: List of Document objects
+            
+        Returns:
+            List of document IDs
+        """
+        doc_ids = []
+        
+        # First, collect all valid documents with embeddings
+        valid_docs = []
+        valid_embeddings = []
+        
+        for doc in documents:
+            if doc.embedding is not None:
+                valid_docs.append(doc)
+                valid_embeddings.append(doc.embedding)
+        
+        if not valid_docs:
+            logger.warning("No valid documents with embeddings to add")
+            return []
+        
+        # Train index if needed and we have enough data
+        if self.needs_training and not self.index.is_trained:
+            if len(valid_embeddings) >= 100 or (len(self.documents) + len(valid_docs)) >= 100:
+                logger.info("Training FAISS index")
+                # Use available embeddings for training
+                train_data = np.vstack([
+                    *[doc.embedding for doc in list(self.documents.values()) if doc.embedding is not None],
+                    *valid_embeddings
+                ])
+                train_data = train_data[:min(1000, len(train_data))]  # Limit to 1000 samples
+                self.index.train(train_data)
+        
+        # Add embeddings to FAISS index if it's trained or doesn't need training
+        if not self.needs_training or self.index.is_trained:
+            embeddings_array = np.array(valid_embeddings, dtype=np.float32)
+            start_idx = len(self.id_to_index)
+            self.index.add(embeddings_array)
+            
+            # Update mappings
+            for i, doc in enumerate(valid_docs):
+                idx = start_idx + i
+                self.id_to_index[doc.id] = idx
+                self.index_to_id[idx] = doc.id
+        
+        # Store all documents (with or without embeddings)
+        for doc in documents:
+            self.documents[doc.id] = doc
+            doc_ids.append(doc.id)
+        
+        return doc_ids
+    
+    def search(
+        self,
+        query_embedding: np.ndarray,
+        top_k: int = 5,
+        filter_func: Optional[Callable[[Document], bool]] = None
+    ) -> List[Tuple[Document, float]]:
+        """
+        Search for similar documents.
+        
+        Args:
+            query_embedding: Query embedding vector
+            top_k: Number of results to return
+            filter_func: Optional function to filter results
+            
+        Returns:
+            List of (document, score) tuples
+        """
+        if not self.documents or not self.id_to_index:
+            logger.warning("Cannot search: database is empty")
+            return []
+        
+        # Ensure index is trained if needed
+        if self.needs_training and not self.index.is_trained:
+            logger.warning("Cannot search: index not trained")
+            return []
+        
+        # Convert to correct format if needed
+        if len(query_embedding.shape) == 1:
+            query_embedding = np.array([query_embedding], dtype=np.float32)
+        
+        # Check if we need to rebuild the index
+        if len(self.id_to_index) != self.index.ntotal:
+            logger.info("Rebuilding index before search")
+            self._rebuild_index()
+        
+        # Adjust top_k based on available items
+        effective_top_k = min(top_k, self.index.ntotal)
+        if effective_top_k < top_k:
+            logger.warning(f"Requested top_k={top_k} but only {effective_top_k} items in index")
+        
+        # Perform search
+        distances, indices = self.index.search(query_embedding, effective_top_k)
+        
+        # Retrieve documents
+        results = []
+        for i, idx in enumerate(indices[0]):
+            if idx != -1:  # FAISS uses -1 for padding when there aren't enough results
+                doc_id = self.index_to_id.get(idx)
+                if doc_id and doc_id in self.documents:
+                    doc = self.documents[doc_id]
+                    
+                    # Apply filter if provided
+                    if filter_func is None or filter_func(doc):
+                        # Convert L2 distance to similarity score (1 / (1 + distance))
+                        score = 1.0 / (1.0 + distances[0][i])
+                        results.append((doc, score))
+        
+        # Sort by score in descending order
+        results.sort(key=lambda x: x[1], reverse=True)
+        
+        return results
+    
+    def delete_document(self, doc_id: str) -> bool:
+        """
+        Delete a document from the database.
+        
+