Closes [FEATURE REQUEST] Expand to other vector stores beyond Pinecone (#102)

pyproject.toml

@@ -26,6 +26,7 @@ dependencies = [
     "anytree==2.12.1",
     "cohere==5.9.2",
     "configargparse",
+    "faiss-cpu==1.9.0",
     "fastapi==0.112.2",
     "google-ai-generativelanguage==0.6.6",
     "gradio>=4.26.0",
@@ -41,6 +42,9 @@ dependencies = [
     "langchain-openai==0.1.25",
     "langchain-text-splitters==0.2.4",
     "langchain-voyageai==0.1.1",
+    "langchain-milvus==0.1.6",
+    "langchain-chroma==0.1.4",
+    "langchain-qdrant==0.1.4",
     "marqo==3.7.0",
     "nbformat==5.10.4",
     "openai==1.42.0",

sage/config.py

@@ -122,12 +122,16 @@ def add_embedding_args(parser: ArgumentParser) -> Callable:
 
 def add_vector_store_args(parser: ArgumentParser) -> Callable:
     """Adds vector store-related arguments to the parser and returns a validator."""
-    parser.add("--vector-store-provider", default="marqo", choices=["pinecone", "marqo"])
     parser.add(
-        "--pinecone-index-name",
-        default=None,
-        help="Pinecone index name. Required if using Pinecone as the vector store. If the index doesn't exist already, "
-        "we will create it.",
+        "--vector-store-provider", default="marqo", choices=["pinecone", "marqo", "chroma", "faiss", "milvus", "qdrant"]
+    )
+    parser.add(
+        "--index-name", default="sage_index", help="Index name for the Vector Store index. We default it to sage_index"
+    )
+    parser.add(
+        "--milvus-uri",
+        default="milvus_sage.db",
+        help="URI for milvus. We default it to milvus_sage.db",
     )
     parser.add(
         "--index-namespace",
@@ -402,8 +406,8 @@ def validate_vector_store_args(args):
     elif args.vector_store_provider == "pinecone":
         if not os.getenv("PINECONE_API_KEY"):
             raise ValueError("Please set the PINECONE_API_KEY environment variable.")
-        if not args.pinecone_index_name:
-            raise ValueError(f"Please set the vector_store.pinecone_index_name value.")
+        if not args.index_name:
+            raise ValueError(f"Please set the vector_store.index_name value.")
 
 
 def validate_indexing_args(args):

sage/index.py

@@ -11,7 +11,7 @@ from sage.chunker import UniversalFileChunker
 from sage.data_manager import GitHubRepoManager
 from sage.embedder import build_batch_embedder_from_flags
 from sage.github import GitHubIssuesChunker, GitHubIssuesManager
-from sage.vector_store import build_vector_store_from_args
+from sage.vector_store import VectorStoreProvider, build_vector_store_from_args
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger()
@@ -41,8 +41,11 @@ def main():
         return
 
     # Additionally validate embedder and vector store compatibility.
-    if args.embedding_provider == "openai" and args.vector_store_provider != "pinecone":
-        parser.error("When using OpenAI embedder, the vector store type must be Pinecone.")
+    vector_store_providers = [member.value for member in VectorStoreProvider]
+    if args.embedding_provider == "openai" and args.vector_store_provider not in vector_store_providers:
+        parser.error(
+            f"When using OpenAI embedder, the vector store type must be from the list {vector_store_providers}."
+        )
     if args.embedding_provider == "marqo" and args.vector_store_provider != "marqo":
         parser.error("When using the marqo embedder, the vector store type must also be marqo.")
 

sage/vector_store.py

@@ -3,20 +3,33 @@
 import logging
 import os
 from abc import ABC, abstractmethod
+from enum import Enum
 from functools import cached_property
 from typing import Dict, Generator, List, Optional, Tuple
+from uuid import uuid4
 
+import chromadb
+import faiss
 import marqo
 import nltk
 from langchain.retrievers import EnsembleRetriever
+from langchain_chroma import Chroma as LangChainChroma
+from langchain_community.docstore.in_memory import InMemoryDocstore
 from langchain_community.retrievers import BM25Retriever
-from langchain_community.vectorstores import Marqo
+from langchain_community.vectorstores import FAISS, Marqo
 from langchain_community.vectorstores import Pinecone as LangChainPinecone
 from langchain_core.documents import Document
 from langchain_core.embeddings import Embeddings
+from langchain_google_genai import GoogleGenerativeAIEmbeddings
+from langchain_milvus import Milvus
+from langchain_openai import OpenAIEmbeddings
+from langchain_qdrant import QdrantVectorStore as LangChainQdrant
+from langchain_voyageai import VoyageAIEmbeddings
 from nltk.data import find
 from pinecone import Pinecone, ServerlessSpec
 from pinecone_text.sparse import BM25Encoder
+from qdrant_client import QdrantClient
+from qdrant_client.http.models import Distance, VectorParams
 
 from sage.constants import TEXT_FIELD
 from sage.data_manager import DataManager
@@ -24,6 +37,15 @@ from sage.data_manager import DataManager
 Vector = Tuple[Dict, List[float]]  # (metadata, embedding)
 
 
+class VectorStoreProvider(Enum):
+    PINECONE = "pinecone"
+    MARQO = "marqo"
+    CHROMA = "chroma"
+    FAISS = "faiss"
+    MILVUS = "milvus"
+    QDRANT = "qdrant"
+
+
 def is_punkt_downloaded():
     try:
         find("tokenizers/punkt_tab")
@@ -156,6 +178,207 @@ class PineconeVectorStore(VectorStore):
             return dense_retriever
 
 
+class ChromaVectorStore(VectorStore):
+    """Vector store implementation using ChromaDB"""
+
+    def __init__(self, index_name: str, alpha: float = None, bm25_cache: Optional[str] = None):
+        """
+        Args:
+            index_name: The name of the Chroma collection/index to use. If it doesn't exist already, we'll create it.
+            alpha: The alpha parameter for hybrid search: alpha == 1.0 means pure dense search, alpha == 0.0 means pure
+                BM25, and 0.0 < alpha < 1.0 means a hybrid of the two.
+        """
+        self.index_name = index_name
+        self.alpha = alpha
+        self.client = chromadb.PersistentClient()
+
+    @cached_property
+    def index(self):
+        index = self.client.get_or_create_collection(self.index_name)
+        return index
+
+    def ensure_exists(self):
+        pass
+
+    def upsert_batch(self, vectors: List[Vector], namespace: str):
+        del namespace
+
+        ids = []
+        embeddings = []
+        metadatas = []
+        documents = []
+
+        for i, (metadata, embedding) in enumerate(vectors):
+            ids.append(metadata.get("id", str(i)))
+            embeddings.append(embedding)
+            metadatas.append(metadata)
+            documents.append(metadata[TEXT_FIELD])
+
+        self.index.upsert(ids=ids, embeddings=embeddings, metadatas=metadatas, documents=documents)
+
+    def as_retriever(self, top_k: int, embeddings: Embeddings = None, namespace: str = None):
+        vector_store = LangChainChroma(
+            collection_name=self.index_name, embedding_function=embeddings, client=self.client
+        )
+
+        return vector_store.as_retriever(search_kwargs={"k": top_k})
+
+
+class FAISSVectorStore(VectorStore):
+    """Vector store implementation using FAISS"""
+
+    def __init__(self, index_name: str, dimension: int, embeddings: Embeddings = None):
+        """
+        Args:
+            index_name: The name of the FAISS index to use. If it doesn't exist already, we'll create it.
+            dimension: The dimension of the vectors.
+            embeddings: The embedding function used to generate embeddings
+        """
+        self.index_name = index_name
+        self.dimension = dimension
+        self.embeddings = embeddings
+
+        # check if the index exists
+        if os.path.exists(self.index_name):
+            # load the existing index
+            self.vector_store = FAISS.load_local(
+                folder_path=self.index_name, embeddings=self.embeddings, allow_dangerous_deserialization=True
+            )
+        # else create a new index
+        else:
+            self.vector_store = FAISS(
+                embedding_function=self.embeddings,
+                index=self.index,
+                docstore=InMemoryDocstore(),
+                index_to_docstore_id={},
+            )
+
+    @cached_property
+    def index(self):
+        index = faiss.IndexFlatL2(self.dimension)
+        return index
+
+    def ensure_exists(self):
+        pass
+
+    def upsert_batch(self, vectors: List[Vector], namespace: str):
+        del namespace
+
+        ids = []
+        documents = []
+
+        for i, (meta_data, embedding) in enumerate(vectors):
+            ids.append(meta_data.get("id", str(i)))
+            document = Document(page_content=meta_data[TEXT_FIELD], metadata=meta_data)
+            documents.append(document)
+
+        self.vector_store.add_documents(documents=documents, ids=ids)
+
+        # saving the index after every batch upsert
+        self.vector_store.save_local(self.index_name)
+        print("Save Local Executed")
+        logging.error("Save Local Got Executed")
+
+    def as_retriever(self, top_k, embeddings, namespace):
+        del embeddings
+        del namespace
+
+        return self.vector_store.as_retriever(search_kwards={"k": top_k})
+
+
+class MilvusVectorStore(VectorStore):
+    """Vector store implementation using Milvus"""
+
+    def __init__(self, uri: str, index_name: str, embeddings: Embeddings = None):
+        """
+        Args:
+            index_name: The name of the Milvus collection to use. If it doesn't exist already, we'll create it.
+            embeddings: The embedding function used to generate embeddings
+        """
+        self.uri = uri
+        self.index_name = index_name
+        self.embeddings = embeddings
+
+        self.vector_store = Milvus(
+            embedding_function=embeddings, connection_args={"uri": self.uri}, collection_name=self.index_name
+        )
+
+    def ensure_exists(self):
+        pass
+
+    def upsert_batch(self, vectors: List[Vector], namespace: str):
+        del namespace
+
+        ids = []
+        documents = []
+
+        for i, (meta_data, embedding) in enumerate(vectors):
+            ids.append(meta_data.get("id", str(i)))
+            # "text" is a reserved keyword. So removing it
+            page_content = meta_data[TEXT_FIELD]
+            meta_data["content"] = meta_data[TEXT_FIELD]
+            del meta_data[TEXT_FIELD]
+
+            document = Document(page_content=page_content, metadata=meta_data)
+            documents.append(document)
+
+        self.vector_store.add_documents(documents=documents, ids=ids)
+
+    def as_retriever(self, top_k, embeddings, namespace):
+        del embeddings
+        del namespace
+
+        return self.vector_store.as_retriever(search_kwards={"k": top_k})
+
+
+class QdrantVectorStore(VectorStore):
+    """Vector store implementation using Qdrant"""
+
+    def __init__(self, index_name: str, dimension: int, embeddings: Embeddings = None):
+        """
+        Args:
+            index_name: The name of the Qdrant collection to use. If it doesn't exist already, we'll create it.
+            embeddings: The embedding function used to generate embeddings
+        """
+        self.index_name = index_name
+        self.dimension = dimension
+        self.embeddings = embeddings
+        self.client = QdrantClient(path="qdrantdb")
+        self.vector_store = self.index
+
+    @cached_property
+    def index(self):
+        self.ensure_exists()
+        vector_store = LangChainQdrant(client=self.client, collection_name=self.index_name, embedding=self.embeddings)
+        return vector_store
+
+    def ensure_exists(self):
+        if not self.client.collection_exists(self.index_name):
+            self.client.create_collection(
+                collection_name=self.index_name,
+                vectors_config=VectorParams(size=self.dimension, distance=Distance.COSINE),
+            )
+
+    def upsert_batch(self, vectors: List[Vector], namespace: str):
+        del namespace
+
+        ids = []
+        documents = []
+
+        for i, (meta_data, embedding) in enumerate(vectors):
+            ids.append(str(uuid4()))
+            document = Document(page_content=meta_data[TEXT_FIELD], metadata=meta_data)
+            documents.append(document)
+
+        self.vector_store.add_documents(documents=documents, ids=ids)
+
+    def as_retriever(self, top_k, embeddings, namespace):
+        del embeddings
+        del namespace
+
+        return self.vector_store.as_retriever(search_kwards={"k": top_k})
+
+
 class MarqoVectorStore(VectorStore):
     """Vector store implementation using Marqo."""
 
@@ -191,12 +414,22 @@ class MarqoVectorStore(VectorStore):
         return vectorstore.as_retriever(search_kwargs={"k": top_k})
 
 
-def build_vector_store_from_args(args: dict, data_manager: Optional[DataManager] = None) -> VectorStore:
+def build_vector_store_from_args(
+    args: dict,
+    data_manager: Optional[DataManager] = None,
+) -> VectorStore:
     """Builds a vector store from the given command-line arguments.
 
     When `data_manager` is specified and hybrid retrieval is requested, we'll use it to fit a BM25 encoder on the corpus
     of documents.
     """
+    if args.embedding_provider == "openai":
+        embeddings = OpenAIEmbeddings(model=args.embedding_model)
+    elif args.embedding_provider == "voyage":
+        embeddings = VoyageAIEmbeddings(model=args.embedding_model)
+    elif args.embedding_provider == "gemini":
+        embeddings = GoogleGenerativeAIEmbeddings(model=args.embedding_model)
+
     if args.vector_store_provider == "pinecone":
         bm25_cache = os.path.join(".bm25_cache", args.index_namespace, "bm25_encoder.json")
         if args.retrieval_alpha < 1.0 and not os.path.exists(bm25_cache) and data_manager:
@@ -217,11 +450,21 @@ def build_vector_store_from_args(args: dict, data_manager: Optional[DataManager]
             bm25_encoder.dump(bm25_cache)
 
         return PineconeVectorStore(
-            index_name=args.pinecone_index_name,
+            index_name=args.index_name,
             dimension=args.embedding_size if "embedding_size" in args else None,
             alpha=args.retrieval_alpha,
             bm25_cache=bm25_cache,
         )
+    elif args.vector_store_provider == "chroma":
+        return ChromaVectorStore(
+            index_name=args.index_name,
+        )
+    elif args.vector_store_provider == "faiss":
+        return FAISSVectorStore(index_name=args.index_name, dimension=args.embedding_size, embeddings=embeddings)
+    elif args.vector_store_provider == "milvus":
+        return MilvusVectorStore(uri=args.milvus_uri, index_name=args.index_name, embeddings=embeddings)
+    elif args.vector_store_provider == "qdrant":
+        return QdrantVectorStore(index_name=args.index_name, dimension=args.embedding_size, embeddings=embeddings)
     elif args.vector_store_provider == "marqo":
         return MarqoVectorStore(url=args.marqo_url, index_name=args.index_namespace)
     else: