Change auto-incrementing integer ID's with auto generated UUID's.

agentmemory/main.py

@@ -10,56 +10,56 @@
  get_include_types,
 )
 
-
 from agentmemory.client import get_client
 
 def create_memory(category, text, metadata={}, embedding=None, id=None):
  """
  Create a new memory in a collection.
 
  Arguments:
- category (str): Category of the collection.
- text (str): Document text.
- id (str): Unique id.
- metadata (dict): Metadata.
+ - category (str): Category of the collection.
+ - text (str): Document text.
+ - metadata (dict): Metadata.
+ - embedding: Embedding vector, if available.
+ - id (str or None): Unique ID for the memory. If None, a UUID will be generated.
 
  Returns:
- None
+ - The generated UUID or the provided ID.
 
  Example:
  >>> create_memory('sample_category', 'sample_text', id='sample_id', metadata={'sample_key': 'sample_value'})
  """
-
- # get or create the collection
+ 
+ # Get or create the collection
  memories = get_client().get_or_create_collection(category)
 
- # add timestamps to metadata
+ # Add timestamps to metadata
  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 each field in metadata, if the field is a boolean, convert it to a string
  for key, value in metadata.items():
  if isinstance(value, bool) or isinstance(value, dict) or isinstance(value, list):
- debug_log(f"WARNING: Boolean metadata field {key} converted to string")
+ debug_log(f"WARNING: Non-string metadata field {key} converted to string")
  metadata[key] = str(value)
 
- # insert the document into the collection
- memories.upsert(
- ids=[str(id)],
+ # Prepare a list for the IDs
+ ids = [id] if id is not None else []
+
+ # Insert the document into the collection
+ memories.add(
+ ids=ids,
  documents=[text],
  metadatas=[metadata],
- embeddings=[embedding] if embedding is not None else None,
+ embeddings=[embedding] if embedding is not None else None
  )
+
+ # Here, we assume that `add` method appends the newly generated ID to the `ids` list.
+ memory_id = ids[0]
 
- debug_log(f"Created memory {id}: {text}", metadata)
- return id
+ debug_log(f"Created memory {memory_id}: {text}", metadata)
+
+ return memory_id # This will now be a UUID or the ID you provided
 
 
 def create_unique_memory(category, content, metadata={}, similarity=0.95):
@@ -203,11 +203,11 @@ def get_memory(category, id, include_embeddings=True):
 
  Returns:
  dict: The retrieved memory.
-
- Example:
- >>> get_memory("books", "1")
  """
 
+ # No need to check UUID here as PostgreSQL will handle the type casting
+ debug_log(f"Getting memory with ID: {id}, Type: {type(id)}")
+
  # Get or create the collection for the given category
  memories = get_client().get_or_create_collection(category)
 
@@ -555,4 +555,4 @@ def wipe_all_memories():
  for collection in collections:
  client.delete_collection(collection.name)
 
- debug_log("Wiped all memories", type="system")
+ debug_log("Wiped all memories", type="system")

agentmemory/postgres.py

@@ -1,5 +1,6 @@
 from pathlib import Path
 import psycopg2
+import uuid
 
 from agentmemory.check_model import check_model, infer_embeddings
 
@@ -69,13 +70,13 @@ def count(self):
 
  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)
+ for document, metadata in zip(documents, metadatas):
+ generated_id = self.client.insert_memory(self.category, document, metadata)
+ ids.append(generated_id) # appending the returned id to ids list
  else:
- for id_, document, metadata, emb in zip(
- ids, documents, metadatas, embeddings
- ):
- self.client.insert_memory(self.category, document, metadata, emb)
+ for document, metadata, emb in zip(documents, metadatas, embeddings):
+ generated_id = self.client.insert_memory(self.category, document, metadata, emb)
+ ids.append(generated_id) # appending the returned id to ids list
 
  def get(
  self,
@@ -118,13 +119,10 @@ def get(
  self.client._ensure_metadata_columns_exist(category, parse_metadata(where))
 
  if ids:
- if not all(isinstance(i, str) or isinstance(i, int) for i in ids):
- raise Exception(
- "ids must be a list of integers or strings representing integers"
- )
- ids = [int(i) for i in ids]
- conditions.append("id=ANY(%s)")
- params.append(ids)
+ # Type casting to uuid
+ conditions.append("id=ANY(%s::uuid[])")
+ params.append([str(id_) for id_ in ids])
+
 
  if limit is None:
  limit = 100 # or another default value
@@ -221,13 +219,14 @@ def delete(self, ids=None, where=None, where_document=None):
  params.append(f"%{where_document}%")
 
  if ids:
- if not all(isinstance(i, str) or isinstance(i, int) for i in ids):
- raise Exception(
- "ids must be a list of integers or strings representing integers"
- )
- ids = [int(i) for i in ids]
- conditions.append("id=ANY(%s::int[])") # Added explicit type casting
- params.append(ids)
+ # Validate UUIDs
+ try:
+ ids = [uuid.UUID(str(i)) for i in ids]
+ except ValueError:
+ raise Exception("ids must be a list of valid UUIDs or strings that can be converted to UUIDs")
+
+ conditions.append("id=ANY(%s::uuid[])") # Use uuid[] for PostgreSQL UUID array type
+ params.append([str(id_) for id_ in ids])
 
  if where:
  for key, value in where.items():
@@ -262,7 +261,6 @@ def __init__(self, name):
 
 default_model_path = str(Path.home() / ".cache" / "onnx_models")
 
-
 class PostgresClient:
  def __init__(
  self,
@@ -286,11 +284,12 @@ def ensure_table_exists(self, category):
  self.cur.execute(
  f"""
  CREATE TABLE IF NOT EXISTS {table_name} (
- id SERIAL PRIMARY KEY,
+ id uuid DEFAULT uuid_generate_v4(),
  document TEXT NOT NULL,
- embedding VECTOR(384)
+ embedding VECTOR(384),
+ PRIMARY KEY (id)
  )
- """
+  """
  )
  self.connection.commit()
 
@@ -343,22 +342,17 @@ 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 = self.get_or_create_collection(category).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)})
  RETURNING id;
  """
  self.cur.execute(query, tuple(values))
  self.connection.commit()
- return self.cur.fetchone()[0]
+ return self.cur.fetchone()[0] # This will fetch the generated UUID
 
  def create_embedding(self, document):
  embeddings = infer_embeddings([document], model_path=self.model_path)
@@ -495,4 +489,4 @@ def update(self, category, id_, document=None, metadata=None, embedding=None):
 
  def close(self):
  self.cur.close()
- self.connection.close()
+ self.connection.close()