rebase master

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):

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):

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):

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):

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
 

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)
 
-        
+        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)