Add support for pgvector (#1675)

mem0/vector_stores/chroma.py

@@ -54,6 +54,7 @@ class ChromaDB(VectorStoreBase):
 
             self.client = chromadb.Client(self.settings)
 
+        self.collection_name = collection_name
         self.collection = self.create_col(collection_name)
 
     def _parse_output(self, data):
@@ -109,12 +110,11 @@ class ChromaDB(VectorStoreBase):
         )
         return collection
 
-    def insert(self, name, vectors, payloads=None, ids=None):
+    def insert(self, vectors, payloads=None, ids=None):
         """
         Insert vectors into a collection.
 
         Args:
-            name (str): Name of the collection.
             vectors (list): List of vectors to insert.
             payloads (list, optional): List of payloads corresponding to vectors. Defaults to None.
             ids (list, optional): List of IDs corresponding to vectors. Defaults to None.
@@ -122,12 +122,11 @@ class ChromaDB(VectorStoreBase):
 
         self.collection.add(ids=ids, embeddings=vectors, metadatas=payloads)
 
-    def search(self, name, query, limit=5, filters=None):
+    def search(self, query, limit=5, filters=None):
         """
         Search for similar vectors.
 
         Args:
-            name (str): Name of the collection.
             query (list): Query vector.
             limit (int, optional): Number of results to return. Defaults to 5.
             filters (dict, optional): Filters to apply to the search. Defaults to None.
@@ -139,23 +138,21 @@ class ChromaDB(VectorStoreBase):
         final_results = self._parse_output(results)
         return final_results
 
-    def delete(self, name, vector_id):
+    def delete(self, vector_id):
         """
         Delete a vector by ID.
 
         Args:
-            name (str): Name of the collection.
             vector_id (int): ID of the vector to delete.
         """
 
         self.collection.delete(ids=vector_id)
 
-    def update(self, name, vector_id, vector=None, payload=None):
+    def update(self, vector_id, vector=None, payload=None):
         """
         Update a vector and its payload.
 
         Args:
-            name (str): Name of the collection.
             vector_id (int): ID of the vector to update.
             vector (list, optional): Updated vector. Defaults to None.
             payload (dict, optional): Updated payload. Defaults to None.
@@ -163,12 +160,11 @@ class ChromaDB(VectorStoreBase):
 
         self.collection.update(ids=vector_id, embeddings=vector, metadatas=payload)
 
-    def get(self, name, vector_id):
+    def get(self, vector_id):
         """
         Retrieve a vector by ID.
 
         Args:
-            name (str): Name of the collection.
             vector_id (int): ID of the vector to retrieve.
 
         Returns:
@@ -186,33 +182,24 @@ class ChromaDB(VectorStoreBase):
         """
         return self.client.list_collections()
 
-    def delete_col(self, name):
-        """
-        Delete a collection.
+    def delete_col(self):
+        """ Delete a collection. """
+        self.client.delete_collection(name=self.collection_name)
 
-        Args:
-            name (str): Name of the collection to delete.
-        """
-        self.client.delete_collection(name=name)
-
-    def col_info(self, name):
+    def col_info(self):
         """
         Get information about a collection.
 
-        Args:
-            name (str): Name of the collection.
-
         Returns:
             dict: Collection information.
         """
-        return self.client.get_collection(name=name)
+        return self.client.get_collection(name=self.collection_name)
 
-    def list(self, name, filters=None, limit=100):
+    def list(self, filters=None, limit=100):
         """
         List all vectors in a collection.
 
         Args:
-            name (str): Name of the collection.
             filters (dict, optional): Filters to apply to the list.
             limit (int, optional): Number of vectors to return. Defaults to 100.