klas/t6_mem0
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

added on_disk param to qdrant configs (#1677)

Browse Source
This commit is contained in:
dbcontributions
2024-08-12 15:08:44 +05:30
committed by GitHub
parent 442f6d72d0
commit 4aae2b5cca
2 changed files with 9 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
mem0/configs/vector_stores/qdrant.py
View File

@@ -14,6 +14,7 @@ class QdrantConfig(BaseModel):
path: Optional[str] = Field("/tmp/qdrant", description="Path for local Qdrant database") path: Optional[str] = Field("/tmp/qdrant", description="Path for local Qdrant database")
url: Optional[str] = Field(None, description="Full URL for Qdrant server") url: Optional[str] = Field(None, description="Full URL for Qdrant server")
api_key: Optional[str] = Field(None, description="API key for Qdrant server") api_key: Optional[str] = Field(None, description="API key for Qdrant server")
on_disk: Optional[bool]= Field(False, description="Enables persistant storage")
@model_validator(mode="before") @model_validator(mode="before")
def check_host_port_or_path(cls, values): def check_host_port_or_path(cls, values):

10
mem0/vector_stores/qdrant.py
View File

@@ -28,6 +28,7 @@ class Qdrant(VectorStoreBase):
path, path,
url, url,
api_key, api_key,
on_disk
): ):
""" """
Initialize the Qdrant vector store. Initialize the Qdrant vector store.
@@ -39,6 +40,7 @@ class Qdrant(VectorStoreBase):
path (str, optional): Path for local Qdrant database. path (str, optional): Path for local Qdrant database.
url (str, optional): Full URL for Qdrant server. url (str, optional): Full URL for Qdrant server.
api_key (str, optional): API key for Qdrant server. api_key (str, optional): API key for Qdrant server.
on_disk (bool, optional): Enables persistant storage.
""" """
if client: if client:
self.client = client self.client = client
@@ -51,17 +53,17 @@ class Qdrant(VectorStoreBase):
if host and port: if host and port:
params["host"] = host params["host"] = host
params["port"] = port params["port"] = port
if not params: if not params:
params["path"] = path params["path"] = path
if not on_disk:
if os.path.exists(path) and os.path.isdir(path): if os.path.exists(path) and os.path.isdir(path):
shutil.rmtree(path) shutil.rmtree(path)
self.client = QdrantClient(**params) self.client = QdrantClient(**params)
self.create_col(collection_name, embedding_model_dims) self.create_col(collection_name, embedding_model_dims, on_disk)
def create_col(self, name, vector_size, distance=Distance.COSINE): def create_col(self, name, vector_size, on_disk, distance=Distance.COSINE):
""" """
Create a new collection. Create a new collection.
@@ -79,7 +81,7 @@ class Qdrant(VectorStoreBase):
self.client.create_collection( self.client.create_collection(
collection_name=name, collection_name=name,
vectors_config=VectorParams(size=vector_size, distance=distance), vectors_config=VectorParams(size=vector_size, distance=distance, on_disk=on_disk),
) )
def insert(self, name, vectors, payloads=None, ids=None): def insert(self, name, vectors, payloads=None, ids=None):