api/db/vector_store.py

from abc import abstractmethod
import os
from qdrant_client import QdrantClient
from langchain.embeddings import OpenAIEmbeddings, ElasticsearchEmbeddings
from langchain.vectorstores import Qdrant, ElasticVectorSearch, VectorStore
from qdrant_client.models import VectorParams, Distance


class ToyVectorStore:

    @staticmethod
    def get_instance():
        vector_store = os.getenv("STORE")
        if vector_store == "ELASTIC":
            return ElasticVectorStore()
        elif vector_store == "QDRANT":
            return QdrantVectorStore()
        else:
            raise ValueError(f"Invalid vector store {vector_store}")
    
    def __init__(self):
        self.embeddings = OpenAIEmbeddings()

    @abstractmethod
    def get_collection(self, collection: str="test") -> VectorStore:
        """
        get an instance of vector store
        of collection
        """
        pass
    
    @abstractmethod
    def create_collection(self, collection: str) -> None:
        """
        create an instance of vector store
        with collection name
        """
        pass

class ElasticVectorStore(ToyVectorStore):

    def get_collection(self, collection:str) -> ElasticVectorSearch:
        return ElasticVectorSearch(elasticsearch_url= os.getenv("ES_URL"),
                               index_name= collection, embedding=self.embeddings)

    def create_collection(self, collection: str) -> None:
        store = self.get_collection(collection)
        store.create_index(store.client,collection, dict())


class QdrantVectorStore(ToyVectorStore):

    def __init__(self):
        self.client = QdrantClient(url=os.getenv("QDRANT_URL"),
                                        api_key=os.getenv("QDRANT_API_KEY"))

    def get_collection(self, collection: str) -> Qdrant:  
        return Qdrant(client=self.client,collection_name=collection,
                      embeddings=self.embeddings)

    def create_collection(self, collection: str) -> None:
        self.client.create_collection(collection_name=collection, 
                        vectors_config=VectorParams(size=1536, distance=Distance.COSINE))