let postgres set the ID

agentmemory/chroma_client.py

@@ -43,6 +43,12 @@ def update(self, ids, documents=None, metadatas=None, embeddings=None):
         return self.collection.update(ids, embeddings, metadatas, documents)
 
     def upsert(self, ids, documents=None, metadatas=None, embeddings=None):
+        # if no id is provided, generate one based on count of documents in collection
+        if any(id is None for id in ids):
+            origin = self.count()
+            # pad the id with zeros to make it 16 digits long
+            ids = [str(id_).zfill(16) for id_ in range(origin, origin+len(documents))]
+
         return self.collection.upsert(ids, embeddings, metadatas, documents)
 
     def delete(self, ids=None, where=None, where_document=None):

agentmemory/main.py

@@ -37,12 +37,6 @@ def create_memory(category, text, metadata={}, embedding=None, id=None):
     metadata["created_at"] = datetime.datetime.now().timestamp()
     metadata["updated_at"] = datetime.datetime.now().timestamp()
 
-    # if no id is provided, generate one based on count of documents in collection
-    if id is None:
-        id = str(memories.count())
-        # pad the id with zeros to make it 16 digits long
-        id = id.zfill(16)
-
     # for each field in metadata...
     # if the field is a boolean, convert it to a string
     for key, value in metadata.items():
@@ -52,7 +46,7 @@ def create_memory(category, text, metadata={}, embedding=None, id=None):
 
     # insert the document into the collection
     memories.upsert(
-        ids=[str(id)],
+        ids=[id],
         documents=[text],
         metadatas=[metadata],
         embeddings=[embedding] if embedding is not None else None,

agentmemory/postgres.py

@@ -2,6 +2,7 @@
 from pathlib import Path
 import os
 from functools import reduce
+from itertools import repeat
 
 import psycopg2
 
@@ -118,15 +119,14 @@ def count(self):
 
         return self.client.cur.fetchone()[0]
 
-    def add(self, ids, documents=None, metadatas=None, embeddings=None):
-        if embeddings is None:
-            for id_, document, metadata in zip(ids, documents, metadatas):
-                self.client.insert_memory(self.category, document, metadata)
-        else:
-            for id_, document, metadata, emb in zip(
-                ids, documents, metadatas, embeddings
-            ):
-                self.client.insert_memory(self.category, document, metadata, emb)
+    def add(self, ids=None, documents=None, metadatas=None, embeddings=None):
+        # dropping ids, using database serial
+        embeddings = embeddings or repeat(None)
+        metadatas = metadatas or repeat({})
+        for document, metadata, emb in zip(
+            documents, metadatas, embeddings
+        ):
+            self.client.insert_memory(self.category, document, metadata, emb)
 
     def get(
         self,
@@ -335,14 +335,10 @@ def insert_memory(self, category, document, metadata={}, embedding=None, id=None
         if embedding is None:
             embedding = self.create_embedding(document)
 
-        # if the id is None, get the length of the table by counting the number of rows in the category
-        if id is None:
-            id = collection.count()
-
         # Extracting the keys and values from metadata to insert them into respective columns
-        columns = ["id", "document", "embedding"] + list(metadata.keys())
+        columns = ["document", "embedding"] + list(metadata.keys())
         placeholders = ["%s"] * len(columns)
-        values = [id, document, embedding] + list(metadata.values())
+        values = [document, embedding] + list(metadata.values())
 
         query = f"""
         INSERT INTO {table_name} ({', '.join(columns)}) VALUES ({', '.join(placeholders)})