From c0eae3f802a5ab0d0ea5e33088cdbc54d54df0fa Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Wed, 24 Jul 2024 20:53:50 +0800 Subject: [PATCH] fix es update emb problem --- .../storage/llama_index_es_memory_store.py | 44 +++++++++---------- memoryscope/utils/memory_handler.py | 4 +- 2 files changed, 24 insertions(+), 24 deletions(-) diff --git a/memoryscope/storage/llama_index_es_memory_store.py b/memoryscope/storage/llama_index_es_memory_store.py index 4bd7730e..8a6bd771 100644 --- a/memoryscope/storage/llama_index_es_memory_store.py +++ b/memoryscope/storage/llama_index_es_memory_store.py @@ -28,9 +28,9 @@ class LlamaIndexEsMemoryStore(BaseMemoryStore): es_url=es_url, retrieval_strategy=_AsyncDenseVectorStrategy(hybrid=use_hybrid), **kwargs) + # TODO The llamaIndex utilizes some deprecated functions, hence langchain logs warning messages. By # adding the following lines of code, the display of deprecated information is suppressed. - self.index = VectorStoreIndex.from_vector_store(vector_store=self.es_store, embed_model=self.embedding_model.model) @@ -41,7 +41,7 @@ class LlamaIndexEsMemoryStore(BaseMemoryStore): top_k: int = 3, filter_dict: Dict[str, List[str]] | Dict[str, str] = None) -> List[MemoryNode]: # if index is not created, return [] - exists = self.es_store._store.client.indices.exists(index=self.index_name) + exists = self.es_store.client.indices.exists(index=self.index_name) if not exists: return [] @@ -51,30 +51,32 @@ class LlamaIndexEsMemoryStore(BaseMemoryStore): es_filter = _to_elasticsearch_filter(filter_dict) retriever = self.index.as_retriever(vector_store_kwargs={"es_filter": es_filter, "fields": ['embedding']}, similarity_top_k=top_k, - sparse_top_k=top_k, ) - if query is None: - query = QueryBundle(query_str='**--**', - embedding=self.dummy_query_vector()) + sparse_top_k=top_k) + if not query: + query = QueryBundle(query_str='**--**', embedding=self.dummy_query_vector()) text_nodes = retriever.retrieve(query) return [self._text_node_2_memory_node(n) for n in text_nodes] async def a_retrieve_memories(self, - query: str, - top_k: int, + query: Optional[str] = None, + top_k: int = 3, filter_dict: Dict[str, List[str]] | Dict[str, str] = None) -> List[MemoryNode]: - self.logger.info(f"query={query} top_k={top_k} filter_dict={filter_dict}") + # if index is not created, return [] + exists = self.es_store.client.indices.exists(index=self.index_name) + if not exists: + return [] if filter_dict is None: filter_dict = {} - es_filter = _to_elasticsearch_filter(filter_dict) - retriever = self.index.as_retriever( - vector_store_kwargs={"es_filter": es_filter}, - similarity_top_k=top_k) - if query is None: - query = QueryBundle(query_str='**--**', - embedding=self.dummy_query_vector()) + es_filter = _to_elasticsearch_filter(filter_dict) + retriever = self.index.as_retriever(vector_store_kwargs={"es_filter": es_filter, "fields": ['embedding']}, + similarity_top_k=top_k, + sparse_top_k=top_k) + + if not query: + query = QueryBundle(query_str='**--**', embedding=self.dummy_query_vector()) text_nodes: List[NodeWithScore] = await retriever.aretrieve(query) return [self._text_node_2_memory_node(n) for n in text_nodes] @@ -85,7 +87,7 @@ class LlamaIndexEsMemoryStore(BaseMemoryStore): self.insert(node) def batch_update(self, nodes: List[MemoryNode], update_embedding: bool = True): - # TODO batch_update & update_embedding + # TODO batch_update for node in nodes: self.update(node, update_embedding=update_embedding) @@ -100,12 +102,10 @@ class LlamaIndexEsMemoryStore(BaseMemoryStore): def delete(self, node: MemoryNode): return self.es_store.delete(node.memory_id) - def delete_conditional(self, filter_dict: Dict = {}): - nodes = self.retrieve_memories(filter_dict=filter_dict, top_k=10000) - self.batch_delete(nodes) - def update(self, node: MemoryNode, update_embedding: bool = True): - # TODO update without embedding? + if update_embedding: + node.vector = [] + self.delete(node) self.insert(node) diff --git a/memoryscope/utils/memory_handler.py b/memoryscope/utils/memory_handler.py index b554401c..36ec9b9d 100644 --- a/memoryscope/utils/memory_handler.py +++ b/memoryscope/utils/memory_handler.py @@ -78,6 +78,7 @@ class MemoryHandler(object): Args: key (str): The key mapping to memory nodes. nodes (List[MemoryNode]): A single memory node or a list of memory nodes to be updated. + log_repeat: if log_repeat=True, print log info """ if nodes is None: nodes = [] @@ -102,7 +103,7 @@ class MemoryHandler(object): Fetch the memories by keys. Args: - key (str): The key mapping to memory nodes. + keys (str | List[str]): The key mapping to memory nodes. Returns: List[MemoryNode]: Memories mapped to the key. @@ -157,7 +158,6 @@ class MemoryHandler(object): Args: keys (str): The memories. nodes (List[MemoryNode]): A single memory node or a list of memory nodes to be updated. - : """ update_memories: Dict[str, MemoryNode] = {n.memory_id: n for n in self.get_memories(keys=keys)}