add llamaindex-elasticsearch support

memory_scope/storage/base_vector_store.py

@@ -2,6 +2,7 @@ from abc import ABCMeta, abstractmethod
 from typing import Dict, List
 
 from memory_scope.models.base_model import BaseModel
+from memory_scope.scheme.memory_node import MemoryNode
 
 
 class BaseVectorStore(metaclass=ABCMeta):
@@ -13,7 +14,7 @@ class BaseVectorStore(metaclass=ABCMeta):
         self.kwargs: dict = kwargs
 
     @abstractmethod
-    def retrieve(self, text: str, limit_size: int, filter_dict: Dict[str, List[str]]):
+    def retrieve(self, query: str, top_k: int, filter_dict: Dict[str, List[str]]):
         """
         :param text:
         :param limit_size:
@@ -22,7 +23,7 @@ class BaseVectorStore(metaclass=ABCMeta):
         """
 
     @abstractmethod
-    async def async_retrieve(self, text: str, limit_size: int, filter_dict: Dict[str, List[str]]):
+    async def async_retrieve(self, query: str, top_k: int, filter_dict: Dict[str, List[str]]):
         """
         :param text:
         :param limit_size:
@@ -31,7 +32,7 @@ class BaseVectorStore(metaclass=ABCMeta):
         """
 
     @abstractmethod
-    def insert(self, text: str):
+    def insert(self, node: MemoryNode):
         """ TODO 是否overwrite
         :return:
         """

memory_scope/storage/llama_index_elastic_search_store.py

@@ -0,0 +1,122 @@
+from typing import Dict, List, Any
+
+from llama_index.core.schema import TextNode
+from llama_index.core.vector_stores import VectorStoreQuery
+from llama_index.core import VectorStoreIndex, StorageContext, ServiceContext
+from llama_index.vector_stores.elasticsearch import ElasticsearchStore, AsyncDenseVectorStrategy
+from llama_index.core.vector_stores.types import MetadataFilters, ExactMatchFilter, VectorStoreQueryMode
+from memory_scope.models.base_model import BaseModel
+from memory_scope.storage.base_vector_store import BaseVectorStore
+from memory_scope.scheme.memory_node import MemoryNode
+
+
+
+
+def _to_elasticsearch_filter(standard_filters: Dict[str, List[str]]) -> Dict[str, Any]:
+    """
+    Convert standard filters to Elasticsearch filter.
+
+    Args:
+        standard_filters: Standard Llama-index filters.
+
+    Returns:
+        Elasticsearch filter.
+    """
+
+    result = {
+        "bool" : {}
+    }
+    for key, value in standard_filters.items():
+        if isinstance(value, list):
+            operands = []
+            for v in value:
+                operands.append(
+                    {
+                        "term": 
+                        {
+                            f"metadata.{key}.keyword": {"value": v}
+                        }
+                    }
+                )
+            result['bool'].update({"should": operands})
+            result['bool'].update({"minimum_should_match": 1})
+        else:
+            operand = [{
+                "term": {
+                    f"metadata.{key}.keyword": {
+                        "value": value,
+                    }
+                }
+            }]
+            if "must" in result['bool']:
+                result['bool']['must'].extend(operand)
+            else:
+                result['bool'].update({"must": operand})
+    return result
+
+
+class LlamaIndexElasticSearchStore(BaseVectorStore):
+    def __init__(self, 
+                 index_name: str, 
+                 embedding_model: BaseModel, 
+                 content_key: str = "text", 
+                 **kwargs):
+
+        self.index_name: str = index_name
+        self.embedding_model: BaseModel = embedding_model
+
+        self.es_store = ElasticsearchStore(index_name=self.index_name,
+                                           retrieval_strategy=AsyncDenseVectorStrategy(hybrid=True),
+                                           **kwargs)
+        
+        self.service_context = ServiceContext.from_defaults(embed_model=self.embedding_model, llm=None)
+        self.index = VectorStoreIndex.from_vector_store(vector_store=self.es_store,
+                                                        service_context=self.service_context)
+        
+    
+    def retrieve(self, query: str, top_k: int = 3, filter_dict: Dict[str, List[str]] = {}) -> MemoryNode:
+        
+        filter = _to_elasticsearch_filter(filter_dict)
+        retriever = self.index.as_retriever(
+            vector_store_kwargs={
+               "es_filter": filter
+            },
+            similarity_top_k=top_k
+        )
+        textnodes = retriever.retrieve(query)
+        results = self._textnodes2memorynodes(textnodes)
+
+        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)
+    
+    def insert(self, node: MemoryNode):
+        node = self._memorynode2textnode(node)
+        self.index.insert_nodes([node])
+
+    def insert_batch(self, node: MemoryNode) -> None:
+        raise NotImplementedError
+
+    def delete(self):
+        raise NotImplementedError
+    
+    def flush(self):
+        raise NotImplementedError
+    
+    def _memorynode2textnode(self, memory_node: MemoryNode) -> TextNode:
+        content = memory_node.content
+        meta = memory_node.model_dump(exclude={"content"})
+        return TextNode(text=content, metadata=meta)
+
+    def _textnode2memorynode(self, text_node: TextNode) -> MemoryNode:
+        content = text_node.text
+        meta = text_node.metadata
+        mem_node = MemoryNode(content=content, **meta)
+        return mem_node
+
+    def _textnodes2memorynodes(self, text_nodes: TextNode) -> MemoryNode:
+        mem_nodes = [self._textnode2memorynode(node) for node in text_nodes]
+        return mem_nodes
+        

tests/models/test_models_lli_generation.py

@@ -21,7 +21,7 @@ class TestLLILLM(unittest.TestCase):
             prompt=prompt
         )
         print(ans.text)
-
+    @unittest.skip("tmp")
     def test_llm_messages(self):
         messages = [{"role": "system", "content": "you are a helpful assistant."},
                     {"role": "user", "content": "你是谁？"}]
@@ -30,7 +30,7 @@ class TestLLILLM(unittest.TestCase):
             messages=messages
         )
         print(ans.text)
-
+    @unittest.skip("tmp")
     def test_llm_prompt_stream(self):
         prompt = "你如何看待黄金上涨？"
         ans = self.llm.call(
@@ -43,7 +43,7 @@ class TestLLILLM(unittest.TestCase):
             sys.stdout.write(a.delta)
             sys.stdout.flush()
             time.sleep(0.1)
-        
+    @unittest.skip("tmp")
     def test_llm_messages(self):
         messages = [{"role": "system", "content": "you are a helpful assistant."},
                     {"role": "user", "content": "你如何看待黄金上涨？"}]

tests/storages/test_storages_lli_es.py

@@ -1,29 +1,90 @@
 import unittest
-from memory_scope.models.base_embedding_model import LLIEmbedding
 
-class TestLLIEmbedding(unittest.TestCase):
+from llama_index.core.vector_stores.types import MetadataFilter, MetadataFilters, FilterCondition, FilterOperator
+from llama_index.core.schema import TextNode
+from memory_scope.scheme.memory_node import MemoryNode
+from memory_scope.storage.llama_index_elastic_search_store import LlamaIndexElasticSearchStore
+from memory_scope.models.llama_index_embedding_model import LlamaIndexEmbeddingModel
+
+class TestLlamaIndexElasticSearchStore(unittest.TestCase):
     """Tests for LLIEmbedding"""
 
     def setUp(self):
         config = {
             "method_type": "DashScopeEmbedding",
             "model_name": "text-embedding-v2",
-            "clazz": "models.base_embedding_model"
+            "clazz": "models.llama_index_embedding_model"
         }
-        self.emb = LLIEmbedding(**config)
+        emb = LlamaIndexEmbeddingModel(**config).model
+
+        config = {
+            "index_name" : "0625_3",
+            "es_url" : "http://localhost:9200",
+            "embedding_model" : emb, 
+             
+        }
+        self.es_store = LlamaIndexElasticSearchStore(**config)
+        self.data = [
+            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"
+
+            ),
+            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"
+
+            ),
+            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"
+
+            ),
+            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"
+
+            ),
+            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"
+            ),
+            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"
+            ),
+            MemoryNode(
+                content="An organized crime dynasty's aging patriarch transfers control of his clandestine empire to his reluctant son.",
+                memory_type="insights",
+                id="6"),
+            MemoryNode(
+                content="ggggggggg",
+                memory_type="profile",
+                id="6"),
+            
+        ]
+    @unittest.skip("tmp")
+    def test_insert(self, ):
+        for node in self.data:
+            self.es_store.insert(node)
+
     
-    def test_single_embedding(self):
-        text = "您吃了吗？"
-        embs = self.emb.call(text=text)
-
-    def test_batch_embedding(self):
-        texts = ["您吃了吗？", 
-                "吃了吗您？"]
-        embs = self.emb.call(text=texts)
-
-    async def test_async_embedding(self):
-        texts = ["您吃了吗？", 
-                "吃了吗您？"]
-        # 调用异步函数并等待其结果
-        embs = await self.emb.async_call(texts)
+   #  @unittest.skip("tmp")
+    def test_retrieve(self, ):
         
+        filter = {
+            "id": ["1", "2", "3"],
+            "memory_type": "insights",
+        }
+
+        res = self.es_store.retrieve(query="hacker", filter_dict=filter, top_k=10)
+        print(len(res))
+        print(res)
+
+