diff --git a/memory_scope/models/base_model.py b/memory_scope/models/base_model.py index 6bd2e6c4..28b12fe5 100644 --- a/memory_scope/models/base_model.py +++ b/memory_scope/models/base_model.py @@ -2,11 +2,11 @@ import inspect import time from abc import abstractmethod, ABCMeta -from enumeration.model_enum import ModelEnum +from memory_scope.enumeration.model_enum import ModelEnum from memory_scope.models import MODEL_REGISTRY -from .response import ModelResponse, ModelResponseGen -from utils.logger import Logger -from utils.timer import Timer +from memory_scope.models.response import ModelResponse, ModelResponseGen +from memory_scope.utils.logger import Logger +from memory_scope.utils.timer import Timer class BaseModel(metaclass=ABCMeta): diff --git a/memory_scope/models/llama_index_embedding_model.py b/memory_scope/models/llama_index_embedding_model.py index 2e2689c9..a9397116 100644 --- a/memory_scope/models/llama_index_embedding_model.py +++ b/memory_scope/models/llama_index_embedding_model.py @@ -2,10 +2,10 @@ from typing import List from llama_index.embeddings.dashscope import DashScopeEmbedding -from models import MODEL_REGISTRY -from models.base_model import BaseModel -from models.response import ModelResponse, ModelResponseGen -from enumeration.model_enum import ModelEnum +from memory_scope.models import MODEL_REGISTRY +from memory_scope.models.base_model import BaseModel +from memory_scope.models.response import ModelResponse, ModelResponseGen +from memory_scope.enumeration.model_enum import ModelEnum class LlamaIndexEmbeddingModel(BaseModel): diff --git a/memory_scope/models/response.py b/memory_scope/models/response.py index a44e1fab..841024fd 100644 --- a/memory_scope/models/response.py +++ b/memory_scope/models/response.py @@ -3,7 +3,7 @@ from typing import Generator, List, Dict, Any from pydantic import BaseModel, Field -from enumeration.model_enum import ModelEnum +from memory_scope.enumeration.model_enum import ModelEnum class ModelResponse(BaseModel): diff --git a/memory_scope/storage/base_vector_store.py b/memory_scope/storage/base_vector_store.py index 9695084f..36bba692 100644 --- a/memory_scope/storage/base_vector_store.py +++ b/memory_scope/storage/base_vector_store.py @@ -1,7 +1,8 @@ from abc import ABCMeta, abstractmethod from typing import Dict, List -from models.base_model import BaseModel +from memory_scope.models.base_model import BaseModel +from memory_scope.scheme.memory_node import MemoryNode class BaseVectorStore(metaclass=ABCMeta): diff --git a/memory_scope/storage/llama_index_elastic_search_store.py b/memory_scope/storage/llama_index_elastic_search_store.py index 3e1f2cac..f1e6ecd9 100644 --- a/memory_scope/storage/llama_index_elastic_search_store.py +++ b/memory_scope/storage/llama_index_elastic_search_store.py @@ -10,6 +10,22 @@ from memory_scope.storage.base_vector_store import BaseVectorStore from memory_scope.scheme.memory_node import MemoryNode +class _ElasticsearchStore(ElasticsearchStore): + async def adelete(self, ref_doc_id: str, **delete_kwargs: Any) -> None: + """ + Async delete node from Elasticsearch index. + + Args: + ref_doc_id: ID of the node to delete. + delete_kwargs: Optional. Additional arguments to + pass to AsyncElasticsearch delete_by_query. + + Raises: + Exception: If AsyncElasticsearch delete_by_query fails. + """ + return await self._store.delete( + query={"term": {"_id": ref_doc_id}}, **delete_kwargs + ) def _to_elasticsearch_filter(standard_filters: Dict[str, List[str]]) -> Dict[str, Any]: @@ -65,7 +81,7 @@ class LlamaIndexElasticSearchStore(BaseVectorStore): self.index_name: str = index_name self.embedding_model: BaseModel = embedding_model - self.es_store = ElasticsearchStore(index_name=self.index_name, + self.es_store = _ElasticsearchStore(index_name=self.index_name, retrieval_strategy=AsyncDenseVectorStrategy(hybrid=True), **kwargs) @@ -89,8 +105,17 @@ class LlamaIndexElasticSearchStore(BaseVectorStore): return results async def async_retrieve(self, query: str, top_k: int = 3, filter_dict: Dict[str, List[str]] = {}) -> MemoryNode: - raise NotImplementedError - ## return await super().async_retrieve(text, limit_size, filter_dict) + filter = _to_elasticsearch_filter(filter_dict) + retriever = self.index.as_retriever( + vector_store_kwargs={ + "es_filter": filter + }, + similarity_top_k=top_k + ) + textnodes = await retriever.aretrieve(query) + results = self._textnodes2memorynodes(textnodes) + + return results def insert(self, node: MemoryNode): node = self._memorynode2textnode(node) @@ -99,20 +124,26 @@ class LlamaIndexElasticSearchStore(BaseVectorStore): def insert_batch(self, node: MemoryNode) -> None: raise NotImplementedError - def delete(self): - raise NotImplementedError - + def delete(self, node: MemoryNode) -> None: + memory_id = node.memory_id + self.es_store.delete(memory_id) + + def update(self, node: MemoryNode) -> None: + self.delete(node) + self.insert(node) + def flush(self): raise NotImplementedError def _memorynode2textnode(self, memory_node: MemoryNode) -> TextNode: content = memory_node.content + memory_id = memory_node.memory_id meta = memory_node.model_dump(exclude={"content"}) - return TextNode(text=content, metadata=meta) + return TextNode(id_=memory_id, text=content, metadata=meta) def _textnode2memorynode(self, text_node: TextNode) -> MemoryNode: content = text_node.text - meta = text_node.metadata + meta = text_node.metadata mem_node = MemoryNode(content=content, **meta) return mem_node diff --git a/tests/storages/test_storages_lli_es.py b/tests/storages/test_storages_lli_es.py index 4ed04075..e73cd1ca 100644 --- a/tests/storages/test_storages_lli_es.py +++ b/tests/storages/test_storages_lli_es.py @@ -18,7 +18,7 @@ class TestLlamaIndexElasticSearchStore(unittest.TestCase): emb = LlamaIndexEmbeddingModel(**config).model config = { - "index_name" : "0625_3", + "index_name" : "0626_1", "es_url" : "http://localhost:9200", "embedding_model" : emb, @@ -28,63 +28,103 @@ class TestLlamaIndexElasticSearchStore(unittest.TestCase): MemoryNode( content="The lives of two mob hitmen, a boxer, a gangster and his wife, and a pair of diner bandits intertwine in four tales of violence and redemption.", memory_type="observation", - id="0" + user_id="0", + status="valid", + memory_id="aaa123", ), MemoryNode( content="When the menace known as the Joker wreaks havoc and chaos on the people of Gotham, Batman must accept one of the greatest psychological and physical tests of his ability to fight injustice.", memory_type="observation", - id="1" + user_id="1", + status="valid", + memory_id="bbb456", ), MemoryNode( content="An insomniac office worker and a devil-may-care soapmaker form an underground fight club that evolves into something much, much more.", memory_type="insights", - id="2" + user_id="2", + status="valid", + memory_id="ccc789", + ), MemoryNode( content="A thief who steals corporate secrets through the use of dream-sharing technology is given the inverse task of planting an idea into thed of a C.E.O.", memory_type="insights", - id="3" - + user_id="3", + status="valid", + memory_id="ddd012", ), MemoryNode( content="A computer hacker learns from mysterious rebels about the true nature of his reality and his role in the war against its controllers.", memory_type="profile", - id="4" + user_id="4", + status="valid", + memory_id="eee345", ), MemoryNode( content="Two detectives, a rookie and a veteran, hunt a serial killer who uses the seven deadly sins as his motives.", memory_type="profile", - id="5" + user_id="5", + status="valid", + memory_id="fff678" ), MemoryNode( content="An organized crime dynasty's aging patriarch transfers control of his clandestine empire to his reluctant son.", memory_type="insights", - id="6"), + user_id="6", + status="valid", + memory_id="ggg901", + ), MemoryNode( content="ggggggggg", memory_type="profile", - id="6"), + user_id="6", + status="valid", + memory_id="ggg234", + ), ] - # @unittest.skip("tmp") - def test_insert(self, ): - for node in self.data: - self.es_store.insert(node) - - # @unittest.skip("tmp") def test_retrieve(self, ): - filter = { - "id": ["1", "2", "3"], - "memory_type": "insights", + "user_id": "6", } + for node in self.data: + self.es_store.insert(node) + self.es_store.insert(MemoryNode( + content="xxxxxx", + memory_type="profile", + user_id="6", + status="valid", + memory_id="ggg567" + )) res = self.es_store.retrieve(query="hacker", filter_dict=filter, top_k=10) print(len(res)) print(res) - \ No newline at end of file + self.es_store.update(MemoryNode( + content="test update", + memory_type="profile", + user_id="6", + status="invalid", + memory_id="ggg567" + )) + res = self.es_store.retrieve(query="hacker", filter_dict=filter, top_k=10) + print(len(res)) + print(res) + + + self.es_store.delete(MemoryNode( + content="test update", + memory_type="profile", + user_id="6", + status="invalid", + memory_id="ggg567" + )) + res = self.es_store.retrieve(query="hacker", filter_dict=filter, top_k=10) + print(len(res)) + print(res)