diff --git a/agentmemory/main.py b/agentmemory/main.py index 0217c43..815c6fa 100644 --- a/agentmemory/main.py +++ b/agentmemory/main.py @@ -10,7 +10,6 @@ get_include_types, ) - from agentmemory.client import get_client def create_memory(category, text, metadata={}, embedding=None, id=None): @@ -18,48 +17,49 @@ 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") \ No newline at end of file diff --git a/agentmemory/postgres.py b/agentmemory/postgres.py index 8e3f1ee..580a304 100644 --- a/agentmemory/postgres.py +++ b/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,14 +342,9 @@ 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)}) @@ -358,7 +352,7 @@ def insert_memory(self, category, document, metadata={}, embedding=None, id=None """ 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() \ No newline at end of file