Skip to content

Commit

Permalink
0.4.7 - fix postgres bug, get_memories returns embeddings
Browse files Browse the repository at this point in the history
  • Loading branch information
lalalune committed Aug 13, 2023
1 parent 4c3472f commit 70581a4
Show file tree
Hide file tree
Showing 4 changed files with 26 additions and 6 deletions.
3 changes: 3 additions & 0 deletions agentmemory/persistence.py
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,9 @@ def export_memory_to_file(path="./memory.json", include_embeddings=True):
# Export the database to a dictionary
collections_dict = export_memory_to_json(include_embeddings)

print('collections_dict')
print(collections_dict)

# Write the dictionary to a JSON file
with open(path, "w") as outfile:
json.dump(collections_dict, outfile)
Expand Down
20 changes: 19 additions & 1 deletion agentmemory/postgres.py
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,8 @@ def get(
where_document=None,
include=["metadatas", "documents"],
):
# TODO: Mirrors Chroma API, but could be optimized a lot

category = self.category
table_name = self.client._table_name(category)
conditions = []
Expand Down Expand Up @@ -152,12 +154,28 @@ def get(
item["metadata"] = metadata
result.append(item)

return {
output = {
"ids": [row["id"] for row in result],
"documents": [row["document"] for row in result],
"metadatas": [row["metadata"] for row in result],
}

if len(result) == 0 or include is None:
return output

# embeddings is an array, check if include includes "embeddings"
if 'embeddings' in include and result[0].get("embedding", None) is not None:
output["embeddings"] = [row["embedding"] for row in result]
# transform from ndarray to list
output["embeddings"] = [emb.tolist() for emb in output["embeddings"]]

if 'distances' in include and result[0].get("distance", None) is not None:
output["distances"] = [row["distances"] for row in result]
# transform to list
output["distances"] = [dist.tolist() for dist in output["distances"]]

return output

def peek(self, limit=10):
return self.get(limit=limit)

Expand Down
7 changes: 3 additions & 4 deletions agentmemory/tests/clustering.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,8 @@ def test_cluster_no_neighbors():

memories_data = get_memories("numbers")

assert memories_data[0]['embedding'] is not None

# All memories should be marked as noise since they have no neighbors
for memory in memories_data:
assert memory["metadata"].get("cluster") == "noise"
Expand All @@ -62,13 +64,10 @@ def test_cluster_insufficient_neighbors():

memories_data = get_memories("fruits")

# Print metadata for debugging
for memory in memories_data:
print(f"Memory Document: {memory['document']}, Metadata: {memory['metadata']}")

# Only 'banana' should be marked as noise since 'apple' has 2 neighbors but needs 3
for memory in memories_data:
assert memory["metadata"].get("cluster") == "noise"

wipe_category("fruits")


Expand Down
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@

setup(
name='agentmemory',
version='0.4.6',
version='0.4.7',
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
Expand Down

0 comments on commit 70581a4

Please sign in to comment.