making classes and test for vectorstore handling

pytest.ini

@@ -0,0 +1,2 @@
+[pytest]
+pythonpath = .

rag_app/vector_store_handler/vectorstores.py

@@ -0,0 +1,203 @@
+from abc import ABC, abstractmethod
+from langchain.vectorstores import Chroma, FAISS
+from langchain.embeddings import OpenAIEmbeddings
+from langchain.text_splitter import CharacterTextSplitter
+from langchain.document_loaders import TextLoader
+
+class BaseVectorStore(ABC):
+    """
+    Abstract base class for vector stores.
+    
+    This class defines the interface for vector stores and implements
+    common functionality.
+    """
+
+    def __init__(self, embedding_model, persist_directory=None):
+        """
+        Initialize the BaseVectorStore.
+
+        Args:
+            embedding_model: The embedding model to use for vectorizing text.
+            persist_directory (str, optional): Directory to persist the vector store.
+        """
+        self.persist_directory = persist_directory
+        self.embeddings = embedding_model
+        self.vectorstore = None
+
+    def load_and_process_documents(self, file_path, chunk_size=1000, chunk_overlap=0):
+        """
+        Load and process documents from a file.
+
+        Args:
+            file_path (str): Path to the file to load.
+            chunk_size (int): Size of text chunks for processing.
+            chunk_overlap (int): Overlap between chunks.
+
+        Returns:
+            list: Processed documents.
+        """
+        loader = TextLoader(file_path)
+        documents = loader.load()
+        text_splitter = CharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
+        return text_splitter.split_documents(documents)
+
+    @abstractmethod
+    def create_vectorstore(self, texts):
+        """
+        Create a new vector store from the given texts.
+
+        Args:
+            texts (list): List of texts to vectorize and store.
+        """
+        pass
+
+    @abstractmethod
+    def load_existing_vectorstore(self):
+        """
+        Load an existing vector store from the persist directory.
+        """
+        pass
+
+    def similarity_search(self, query):
+        """
+        Perform a similarity search on the vector store.
+
+        Args:
+            query (str): The query text to search for.
+
+        Returns:
+            list: Search results.
+
+        Raises:
+            ValueError: If the vector store is not initialized.
+        """
+        if not self.vectorstore:
+            raise ValueError("Vector store not initialized. Call create_vectorstore or load_existing_vectorstore first.")
+        return self.vectorstore.similarity_search(query)
+
+    @abstractmethod
+    def save(self):
+        """
+        Save the current state of the vector store.
+        """
+        pass
+
+class ChromaVectorStore(BaseVectorStore):
+    """
+    Implementation of BaseVectorStore using Chroma as the backend.
+    """
+
+    def create_vectorstore(self, texts):
+        """
+        Create a new Chroma vector store from the given texts.
+
+        Args:
+            texts (list): List of texts to vectorize and store.
+        """
+        self.vectorstore = Chroma.from_documents(
+            texts, 
+            self.embeddings, 
+            persist_directory=self.persist_directory
+        )
+
+    def load_existing_vectorstore(self):
+        """
+        Load an existing Chroma vector store from the persist directory.
+
+        Raises:
+            ValueError: If persist_directory is not set.
+        """
+        if self.persist_directory:
+            self.vectorstore = Chroma(
+                persist_directory=self.persist_directory,
+                embedding_function=self.embeddings
+            )
+        else:
+            raise ValueError("Persist directory is required for loading Chroma.")
+
+    def save(self):
+        """
+        Save the current state of the Chroma vector store.
+
+        Raises:
+            ValueError: If the vector store is not initialized.
+        """
+        if not self.vectorstore:
+            raise ValueError("Vector store not initialized. Nothing to save.")
+        self.vectorstore.persist()
+
+class FAISSVectorStore(BaseVectorStore):
+    """
+    Implementation of BaseVectorStore using FAISS as the backend.
+    """
+
+    def create_vectorstore(self, texts):
+        """
+        Create a new FAISS vector store from the given texts.
+
+        Args:
+            texts (list): List of texts to vectorize and store.
+        """
+        self.vectorstore = FAISS.from_documents(texts, self.embeddings)
+
+    def load_existing_vectorstore(self):
+        """
+        Load an existing FAISS vector store from the persist directory.
+
+        Raises:
+            ValueError: If persist_directory is not set.
+        """
+        if self.persist_directory:
+            self.vectorstore = FAISS.load_local(self.persist_directory, self.embeddings)
+        else:
+            raise ValueError("Persist directory is required for loading FAISS.")
+
+    def save(self):
+        """
+        Save the current state of the FAISS vector store.
+
+        Raises:
+            ValueError: If the vector store is not initialized.
+        """
+        if not self.vectorstore:
+            raise ValueError("Vector store not initialized. Nothing to save.")
+        self.vectorstore.save_local(self.persist_directory)
+
+# Usage example:
+def main():
+    """
+    Example usage of the vector store classes.
+    """
+    # Create an embedding model
+    embedding_model = OpenAIEmbeddings()
+
+    # Using Chroma
+    chroma_store = ChromaVectorStore(embedding_model, persist_directory="./chroma_store")
+    texts = chroma_store.load_and_process_documents("path/to/your/file.txt")
+    chroma_store.create_vectorstore(texts)
+    results = chroma_store.similarity_search("Your query here")
+    print("Chroma results:", results[0].page_content)
+    chroma_store.save()
+
+    # Load existing Chroma store
+    existing_chroma = ChromaVectorStore(embedding_model, persist_directory="./chroma_store")
+    existing_chroma.load_existing_vectorstore()
+    results = existing_chroma.similarity_search("Another query")
+    print("Existing Chroma results:", results[0].page_content)
+
+    # Using FAISS
+    faiss_store = FAISSVectorStore(embedding_model, persist_directory="./faiss_store")
+    texts = faiss_store.load_and_process_documents("path/to/your/file.txt")
+    faiss_store.create_vectorstore(texts)
+    results = faiss_store.similarity_search("Your query here")
+    print("FAISS results:", results[0].page_content)
+    faiss_store.save()
+
+    # Load existing FAISS store
+    existing_faiss = FAISSVectorStore(embedding_model, persist_directory="./faiss_store")
+    existing_faiss.load_existing_vectorstore()
+    results = existing_faiss.similarity_search("Another query")
+    print("Existing FAISS results:", results[0].page_content)
+
+if __name__ == "__main__":
+    main()

tests/vector_store_handler/test_vectorstores.py

@@ -0,0 +1,86 @@
+import unittest
+from unittest.mock import MagicMock, patch
+from langchain.embeddings import OpenAIEmbeddings
+from langchain.schema import Document
+
+# Update the import to reflect your project structure
+from rag_app.vector_store_handler.vectorstores import BaseVectorStore, ChromaVectorStore, FAISSVectorStore
+
+class TestBaseVectorStore(unittest.TestCase):
+    def setUp(self):
+        self.embedding_model = MagicMock(spec=OpenAIEmbeddings)
+        self.base_store = BaseVectorStore(self.embedding_model, "test_dir")
+
+    def test_init(self):
+        self.assertEqual(self.base_store.persist_directory, "test_dir")
+        self.assertEqual(self.base_store.embeddings, self.embedding_model)
+        self.assertIsNone(self.base_store.vectorstore)
+
+    @patch('rag_app.vector_store_handler.vectorstores.TextLoader')
+    @patch('rag_app.vector_store_handler.vectorstores.CharacterTextSplitter')
+    def test_load_and_process_documents(self, mock_splitter, mock_loader):
+        mock_loader.return_value.load.return_value = ["doc1", "doc2"]
+        mock_splitter.return_value.split_documents.return_value = ["split1", "split2"]
+
+        result = self.base_store.load_and_process_documents("test.txt")
+
+        mock_loader.assert_called_once_with("test.txt")
+        mock_splitter.assert_called_once_with(chunk_size=1000, chunk_overlap=0)
+        self.assertEqual(result, ["split1", "split2"])
+
+    def test_similarity_search_not_initialized(self):
+        with self.assertRaises(ValueError):
+            self.base_store.similarity_search("query")
+
+class TestChromaVectorStore(unittest.TestCase):
+    def setUp(self):
+        self.embedding_model = MagicMock(spec=OpenAIEmbeddings)
+        self.chroma_store = ChromaVectorStore(self.embedding_model, "test_dir")
+
+    @patch('rag_app.vector_store_handler.vectorstores.Chroma')
+    def test_create_vectorstore(self, mock_chroma):
+        texts = [Document(page_content="test")]
+        self.chroma_store.create_vectorstore(texts)
+        mock_chroma.from_documents.assert_called_once_with(
+            texts, 
+            self.embedding_model, 
+            persist_directory="test_dir"
+        )
+
+    @patch('rag_app.vector_store_handler.vectorstores.Chroma')
+    def test_load_existing_vectorstore(self, mock_chroma):
+        self.chroma_store.load_existing_vectorstore()
+        mock_chroma.assert_called_once_with(
+            persist_directory="test_dir",
+            embedding_function=self.embedding_model
+        )
+
+    def test_save(self):
+        self.chroma_store.vectorstore = MagicMock()
+        self.chroma_store.save()
+        self.chroma_store.vectorstore.persist.assert_called_once()
+
+class TestFAISSVectorStore(unittest.TestCase):
+    def setUp(self):
+        self.embedding_model = MagicMock(spec=OpenAIEmbeddings)
+        self.faiss_store = FAISSVectorStore(self.embedding_model, "test_dir")
+
+    @patch('rag_app.vector_store_handler.vectorstores.FAISS')
+    def test_create_vectorstore(self, mock_faiss):
+        texts = [Document(page_content="test")]
+        self.faiss_store.create_vectorstore(texts)
+        mock_faiss.from_documents.assert_called_once_with(texts, self.embedding_model)
+
+    @patch('rag_app.vector_store_handler.vectorstores.FAISS')
+    def test_load_existing_vectorstore(self, mock_faiss):
+        self.faiss_store.load_existing_vectorstore()
+        mock_faiss.load_local.assert_called_once_with("test_dir", self.embedding_model)
+
+    @patch('rag_app.vector_store_handler.vectorstores.FAISS')
+    def test_save(self, mock_faiss):
+        self.faiss_store.vectorstore = MagicMock()
+        self.faiss_store.save()
+        self.faiss_store.vectorstore.save_local.assert_called_once_with("test_dir")
+
+if __name__ == '__main__':
+    unittest.main()