Spaces:
Sleeping
Sleeping
File size: 2,135 Bytes
f2932e2 51a7f02 12a040e f2932e2 51a7f02 f2932e2 c7e10e4 f2932e2 12a040e f2932e2 12a040e f2932e2 12a040e f2932e2 12a040e f2932e2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
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))
|