Spaces:
Running
Running
File size: 2,840 Bytes
47b5f0c 39f6b9b 47b5f0c 39f6b9b 47b5f0c 97cbab9 47b5f0c |
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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
import os
from app.service.api import baseURL
from qdrant_client import QdrantClient, models
api_key = os.environ.get("QDRANT_API_KEY")
class QdrantConnectionDb:
client = None
_instance = None
_collection_name = "docuRAG"
_vector_size = 384
dense_model = "sentence-transformers/all-MiniLM-L6-v2"
sparse_model = "prithivida/Splade_PP_en_v1"
def __new__(cls, *args, **kwargs):
"""
Create a new instance of QdrantConnectionDb if it does not exist and initialize the collection and models.
"""
if cls._instance is None:
cls._instance = super(QdrantConnectionDb, cls).__new__(cls)
cls.client = QdrantClient(url=baseURL, api_key=api_key)
cls._initialize_collection(
cls.client,
cls._collection_name,
cls._vector_size,
)
cls._set_models(cls.dense_model, cls.sparse_model)
return cls._instance
@classmethod
def _initialize_collection(
cls, client: QdrantClient, collection_name: str, _vector_size: int
):
"""
Initialize collection if it does not exist
:param client: QdrantClient
:param collection_name: str
:param _vector_size: int
:return: None
"""
try:
collections = client.get_collections().collections
if collection_name not in [c.name for c in collections]:
client.create_collection(
collection_name=collection_name,
vectors_config={
"text-dense": models.VectorParams(
size=_vector_size,
distance=models.Distance.COSINE,
)
},
sparse_vectors_config={
"text-sparse": models.SparseVectorParams(
index=models.SparseIndexParams(
on_disk=False,
)
)
},
)
print(f"Collection {collection_name} initialized successfully")
except Exception as e:
print(f"Error while initializing collection: {e}")
def get_client(self) -> QdrantClient:
"""
Get the QdrantClient instance
"""
return self.client
@classmethod
def _set_models(self, model_name: str, sparse_model_name: str):
"""
Set the model and sparse model for the client
"""
self.client.set_model(model_name)
self.client.set_sparse_model(sparse_model_name)
@classmethod
def get_collection_name(cls) -> str:
"""
Get the current collection name
"""
return cls._collection_name
|