From b32d41dd49707b6b65c0d4cf0e0e92ddb1164d91 Mon Sep 17 00:00:00 2001
From: moon <shawmakesmagic@gmail.com>
Date: Thu, 10 Aug 2023 20:31:57 -0700
Subject: [PATCH] 0.4.4 - fixes clustering tests

---
 agentmemory/clustering.py     | 1 +
 agentmemory/main.py           | 3 ++-
 agentmemory/postgres.py       | 5 +++++
 agentmemory/tests/__init__.py | 2 +-
 setup.py                      | 2 +-
 5 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/agentmemory/clustering.py b/agentmemory/clustering.py
index afe86ba..77cdd00 100644
--- a/agentmemory/clustering.py
+++ b/agentmemory/clustering.py
@@ -11,6 +11,7 @@ def cluster(epsilon, min_samples, category, filter_metadata=None, novel=False):
     cluster_id = 0
     for memory in memories:
         memory_id = memory["id"]
+        print("Memory ID: ", memory_id)
         if visited[memory_id]:
             continue
         visited[memory_id] = True
diff --git a/agentmemory/main.py b/agentmemory/main.py
index 0d6a13a..49a325c 100644
--- a/agentmemory/main.py
+++ b/agentmemory/main.py
@@ -344,7 +344,8 @@ def update_memory(category, id, text=None, metadata=None, embedding=None):
     documents = [text] if text is not None else None
     metadatas = [metadata] if metadata is not None else None
     embeddings = [embedding] if embedding is not None else None
-
+    print('********************** UPDATE')
+    print(id, documents, metadatas, embeddings)
     # Update the memory with the new text and/or metadata
     memories.update(
         ids=[str(id)], documents=documents, metadatas=metadatas, embeddings=embeddings
diff --git a/agentmemory/postgres.py b/agentmemory/postgres.py
index 29b50fc..184c368 100644
--- a/agentmemory/postgres.py
+++ b/agentmemory/postgres.py
@@ -101,9 +101,14 @@ def query(
         return self.client.query(self.category, query_texts, n_results)
 
     def update(self, ids, documents=None, metadatas=None, embeddings=None):
+        self.client.ensure_table_exists(self.category)
         # if embeddings is not None
         if embeddings is None:
+            if documents is None:
+                documents = [None] * len(ids)
             for id_, document, metadata in zip(ids, documents, metadatas):
+                print("updating")
+                print(id_, document, metadata)
                 self.client.update(self.category, id_, document, metadata)
         else:
             for id_, document, metadata, emb in zip(
diff --git a/agentmemory/tests/__init__.py b/agentmemory/tests/__init__.py
index d26b90f..1c4155e 100644
--- a/agentmemory/tests/__init__.py
+++ b/agentmemory/tests/__init__.py
@@ -2,5 +2,5 @@
 from .main import *
 from .persistence import *
 from .events import *
-# from .clustering import *
+from .clustering import *
 from .check_model import *
\ No newline at end of file
diff --git a/setup.py b/setup.py
index f0c69b7..f321699 100644
--- a/setup.py
+++ b/setup.py
@@ -16,7 +16,7 @@
 
 setup(
     name='agentmemory',
-    version='0.4.3',
+    version='0.4.4',
     description='Easy-to-use memory for agents, document search, knowledge graphing and more.',
     long_description=long_description,  # added this line
     long_description_content_type="text/markdown",  # and this line