From 88946c19e253d352e96cba92a9ea2dab5a4c5b3d Mon Sep 17 00:00:00 2001 From: "fuqingxu.fqx" Date: Mon, 19 Aug 2024 16:03:35 +0800 Subject: [PATCH] bug fix: deal with empty pickle dict, reverse _key_id_dict changes --- memoryscope/core/storage/llama_index_es_memory_store.py | 9 +++++++-- memoryscope/core/worker/memory_manager.py | 2 ++ 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/memoryscope/core/storage/llama_index_es_memory_store.py b/memoryscope/core/storage/llama_index_es_memory_store.py index e2dd1196..478e81be 100644 --- a/memoryscope/core/storage/llama_index_es_memory_store.py +++ b/memoryscope/core/storage/llama_index_es_memory_store.py @@ -189,9 +189,14 @@ class LlamaIndexEsMemoryStore(BaseMemoryStore): Returns: MemoryNode: The converted MemoryNode with text and metadata from the NodeWithScore. """ - key_vector = pickle.loads(text_node.metadata["key_vector"].encode('latin1')) + + if text_node.metadata["key_vector"]: + key_vector = pickle.loads(text_node.metadata["key_vector"].encode('latin1')) + else: + key_vector = [] + text_node.metadata["key_vector"] = key_vector + text_node.metadata["vector"] = text_node.embedding if text_node.embedding else [] text_node.metadata["score_recall"] = text_node.score - text_node.metadata["key_vector"] = key_vector return MemoryNode(content=text_node.text, **text_node.metadata) diff --git a/memoryscope/core/worker/memory_manager.py b/memoryscope/core/worker/memory_manager.py index bf053a67..e43714b9 100644 --- a/memoryscope/core/worker/memory_manager.py +++ b/memoryscope/core/worker/memory_manager.py @@ -98,6 +98,8 @@ class MemoryManager(object): self.logger.info(f"add to memory context memory id={node.memory_id} content={node.content} " f"store_status={node.store_status} action_status={node.action_status}") + self._key_id_dict[key] = [n.memory_id for n in nodes] + if nodes: self.logger.info( self.logger.wrap_in_box(