adding param and return types (#1689)

mem0/vector_stores/qdrant.py

@@ -20,27 +20,29 @@ from mem0.vector_stores.base import VectorStoreBase
 class Qdrant(VectorStoreBase):
     def __init__(
         self,
-        collection_name,
-        embedding_model_dims,
-        client,
-        host,
-        port,
-        path,
-        url,
-        api_key,
-        on_disk
+        collection_name: str,
+        embedding_model_dims: int,
+        client: QdrantClient = None,
+        host: str = None,
+        port: int = None,
+        path: str = None,
+        url: str = None,
+        api_key: str = None,
+        on_disk: bool = False
     ):
         """
         Initialize the Qdrant vector store.
 
         Args:
-            client (QdrantClient, optional): Existing Qdrant client instance.
-            host (str, optional): Host address for Qdrant server.
-            port (int, optional): Port for Qdrant server.
-            path (str, optional): Path for local Qdrant database.
-            url (str, optional): Full URL for Qdrant server.
-            api_key (str, optional): API key for Qdrant server.
-            on_disk (bool, optional): Enables persistant storage.
+            collection_name (str): Name of the collection.
+            embedding_model_dims (int): Dimensions of the embedding model.
+            client (QdrantClient, optional): Existing Qdrant client instance. Defaults to None.
+            host (str, optional): Host address for Qdrant server. Defaults to None.
+            port (int, optional): Port for Qdrant server. Defaults to None.
+            path (str, optional): Path for local Qdrant database. Defaults to None.
+            url (str, optional): Full URL for Qdrant server. Defaults to None.
+            api_key (str, optional): API key for Qdrant server. Defaults to None.
+            on_disk (bool, optional): Enables persistent storage. Defaults to False.
         """
         if client:
             self.client = client
@@ -64,13 +66,13 @@ class Qdrant(VectorStoreBase):
         self.collection_name = collection_name
         self.create_col(embedding_model_dims, on_disk)
 
-    def create_col(self, vector_size, on_disk, distance=Distance.COSINE):
+    def create_col(self, vector_size: int, on_disk: bool, distance: Distance = Distance.COSINE):
         """
         Create a new collection.
 
         Args:
-            name (str): Name of the collection.
             vector_size (int): Size of the vectors to be stored.
+            on_disk (bool): Enables persistent storage.
             distance (Distance, optional): Distance metric for vector similarity. Defaults to Distance.COSINE.
         """
         # Skip creating collection if already exists
@@ -85,7 +87,7 @@ class Qdrant(VectorStoreBase):
             vectors_config=VectorParams(size=vector_size, distance=distance, on_disk=on_disk),
         )
 
-    def insert(self, vectors, payloads=None, ids=None):
+    def insert(self, vectors: list, payloads: list = None, ids: list = None):
         """
         Insert vectors into a collection.
 
@@ -104,7 +106,7 @@ class Qdrant(VectorStoreBase):
         ]
         self.client.upsert(collection_name=self.collection_name, points=points)
 
-    def _create_filter(self, filters):
+    def _create_filter(self, filters: dict) -> Filter:
         """
         Create a Filter object from the provided filters.
 
@@ -128,7 +130,7 @@ class Qdrant(VectorStoreBase):
                 )
         return Filter(must=conditions) if conditions else None
 
-    def search(self, query, limit=5, filters=None):
+    def search(self, query: list, limit: int = 5, filters: dict = None) -> list:
         """
         Search for similar vectors.
 
@@ -149,7 +151,7 @@ class Qdrant(VectorStoreBase):
         )
         return hits
 
-    def delete(self, vector_id):
+    def delete(self, vector_id: int):
         """
         Delete a vector by ID.
 
@@ -163,7 +165,7 @@ class Qdrant(VectorStoreBase):
             ),
         )
 
-    def update(self, vector_id, vector=None, payload=None):
+    def update(self, vector_id: int, vector: list = None, payload: dict = None):
         """
         Update a vector and its payload.
 
@@ -175,7 +177,7 @@ class Qdrant(VectorStoreBase):
         point = PointStruct(id=vector_id, vector=vector, payload=payload)
         self.client.upsert(collection_name=self.collection_name, points=[point])
 
-    def get(self, vector_id):
+    def get(self, vector_id: int) -> dict:
         """
         Retrieve a vector by ID.
 
@@ -190,7 +192,7 @@ class Qdrant(VectorStoreBase):
         )
         return result[0] if result else None
 
-    def list_cols(self):
+    def list_cols(self) -> list:
         """
         List all collections.
 
@@ -203,7 +205,7 @@ class Qdrant(VectorStoreBase):
         """ Delete a collection. """
         self.client.delete_collection(collection_name=self.collection_name)
 
-    def col_info(self):
+    def col_info(self) -> dict:
         """
         Get information about a collection.
 
@@ -212,11 +214,12 @@ class Qdrant(VectorStoreBase):
         """
         return self.client.get_collection(collection_name=self.collection_name)
 
-    def list(self, filters=None, limit=100):
+    def list(self, filters: dict = None, limit: int = 100) -> list:
         """
         List all vectors in a collection.
 
         Args:
+            filters (dict, optional): Filters to apply to the list. Defaults to None.
             limit (int, optional): Number of vectors to return. Defaults to 100.
 
         Returns: