Skip to content
Closed
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 12 additions & 18 deletions evaluation/src/memzero/search.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,21 +37,22 @@ def __init__(self, output_path="results.json", top_k=10, filter_memories=False,
def search_memory(self, user_id, query, max_retries=3, retry_delay=1):
start_time = time.time()
retries = 0
filters={"OR": [{"user_id": user_id}]}
while retries < max_retries:
try:
if self.is_graph:
print("Searching with graph")
memories = self.mem0_client.search(
query,
user_id=user_id,
filters=filters,
top_k=self.top_k,
filter_memories=self.filter_memories,
enable_graph=True,
output_format="v1.1",
)
else:
memories = self.mem0_client.search(
query, user_id=user_id, top_k=self.top_k, filter_memories=self.filter_memories
query, filters=filters, top_k=self.top_k, filter_memories=self.filter_memories
)
break
except Exception as e:
Expand All @@ -62,25 +63,17 @@ def search_memory(self, user_id, query, max_retries=3, retry_delay=1):
time.sleep(retry_delay)

end_time = time.time()
semantic_memories = [
{
"memory": memory["memory"],
"timestamp": memory["metadata"]["timestamp"],
"score": round(memory["score"], 2),
}
for memory in memories["results"]
]
if not self.is_graph:
semantic_memories = [
{
"memory": memory["memory"],
"timestamp": memory["metadata"]["timestamp"],
"score": round(memory["score"], 2),
}
for memory in memories
]
graph_memories = None
else:
semantic_memories = [
{
"memory": memory["memory"],
"timestamp": memory["metadata"]["timestamp"],
"score": round(memory["score"], 2),
}
for memory in memories["results"]
]
graph_memories = [
{"source": relation["source"], "relationship": relation["relationship"], "target": relation["target"]}
for relation in memories["relations"]
Expand Down Expand Up @@ -163,6 +156,7 @@ def process_question(self, val, speaker_a_user_id, speaker_b_user_id):
}

# Save results after each question is processed
os.makedirs(os.path.dirname(self.output_path), exist_ok=True)
with open(self.output_path, "w") as f:
json.dump(self.results, f, indent=4)

Expand Down