add sync version elasticsearch for llamaindex

memory_scope/storage/llama_index_es_memory_store_sync.py

@@ -0,0 +1,241 @@
+from typing import Dict, List, Any, Optional, cast
+
+from llama_index.core import VectorStoreIndex
+from llama_index.core.schema import TextNode, NodeWithScore
+from llama_index.vector_stores.elasticsearch import ElasticsearchStore, AsyncDenseVectorStrategy
+
+from memory_scope.enumeration.memory_status_enum import MemoryNodeStatus
+from memory_scope.models.base_model import BaseModel
+from memory_scope.scheme.memory_node import MemoryNode
+from memory_scope.storage.base_memory_store import BaseMemoryStore
+from memory_scope.utils.logger import Logger
+
+from memory_scope.storage.llama_index_sync_elasticsearch import SyncElasticsearchStore
+
+class _AsyncDenseVectorStrategy(AsyncDenseVectorStrategy):
+    def _hybrid(self, query: str, knn: Dict[str, Any], filter: List[Dict[str, Any]], top_k: int) -> Dict[str, Any]:
+        # Add a query to the knn query.
+        # RRF is used to even the score from the knn query and text query
+        # RRF has two optional parameters: {'rank_constant':int, 'window_size':int}
+        # https://www.elastic.co/guide/en/elasticsearch/reference/current/rrf.html
+        query_body = {
+            "knn": knn,
+            "query": {
+                "bool": {
+                    "must": [
+                        {
+                            "match": {
+                                self.text_field: {
+                                    "query": query,
+                                }
+                            }
+                        }
+                    ],
+                    "filter": filter,
+                }
+            },
+        }
+
+        if isinstance(self.rrf, Dict):
+            query_body["rank"] = {"rrf": self.rrf}
+        elif isinstance(self.rrf, bool) and self.rrf is True:
+            query_body["rank"] = {"rrf": {"window_size": top_k}}
+        return query_body
+
+    def es_query(
+            self,
+            *,
+            query: Optional[str],
+            query_vector: Optional[List[float]],
+            text_field: str,
+            vector_field: str,
+            k: int,
+            num_candidates: int,
+            filter: List[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        if filter is None:
+            filter = []
+
+        knn = {
+            "filter": filter,
+            "field": vector_field,
+            "k": k,
+            "num_candidates": num_candidates,
+        }
+
+        if query_vector is not None:
+            knn["query_vector"] = query_vector
+        else:
+            # Inference in Elasticsearch. When initializing we make sure to always have
+            # a model_id if we don't have an embedding_service.
+            knn["query_vector_builder"] = {
+                "text_embedding": {
+                    "model_id": self.model_id,
+                    "model_text": query,
+                }
+            }
+
+        if self.hybrid:
+            return self._hybrid(query=cast(str, query), knn=knn, filter=filter, top_k=k)
+
+        return {"knn": knn}
+
+
+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 LlamaIndexEsMemoryStore(BaseMemoryStore):
+    def __init__(self,
+                 embedding_model: BaseModel,
+                 index_name: str,
+                 es_url: str,
+                 use_hybrid: bool = True,
+                 **kwargs):
+
+        self.embedding_model: BaseModel = embedding_model
+        self.es_store = SyncElasticsearchStore(index_name=index_name,
+                                            es_url=es_url,
+                                            retrieval_strategy=_AsyncDenseVectorStrategy(hybrid=use_hybrid),
+                                            **kwargs)
+        self.index = VectorStoreIndex.from_vector_store(vector_store=self.es_store,
+                                                        embed_model=self.embedding_model.model)
+        self.index.build_index_from_nodes([TextNode(text="text")])
+        self.logger = Logger.get_logger()
+
+    def retrieve_memories(self,
+                          query: str,
+                          top_k: int,
+                          filter_dict: Dict[str, List[str]] | Dict[str, str] = None) -> List[MemoryNode]:
+        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,
+                                            sparse_top_k=top_k)
+        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,
+                                  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 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)
+        text_nodes: List[NodeWithScore] = await retriever.aretrieve(query)
+        return [self._text_node_2_memory_node(n) for n in text_nodes]
+
+    def insert(self, node: MemoryNode):
+        self.index.insert_nodes([self._memory_node_2_text_node(node)])
+
+    def delete(self, node: MemoryNode):
+        memory_id = node.memory_id
+        return self.es_store.delete(memory_id)
+
+    def update(self, node: MemoryNode):
+        self.delete(node)
+        self.insert(node)
+
+    def update_batch(self, nodes: List[MemoryNode]):
+        for node in nodes:
+            self.update(node)
+
+    def close(self):
+        self.es_store.close()
+
+    def update_memories(self, nodes: MemoryNode | List[MemoryNode]):
+        if not nodes:
+            self.logger.warning("empty nodes!")
+            return
+
+        if isinstance(nodes, MemoryNode):
+            nodes = [nodes]
+
+        # emb & insert new memories
+        # TODO batch insert
+        new_memories = [n for n in nodes if n.status == MemoryNodeStatus.NEW]
+        if new_memories:
+            for n in new_memories:
+                n.status = MemoryNodeStatus.ACTIVE.value
+                self.insert(n)
+
+        # emb & update new memories
+        # TODO insert overwrite
+        c_modified_memories = [n for n in nodes if n.status == MemoryNodeStatus.CONTENT_MODIFIED]
+        if c_modified_memories:
+            for n in c_modified_memories:
+                n.status = MemoryNodeStatus.ACTIVE.value
+                self.delete(n)
+                self.insert(n)
+
+        # update new memories
+        # TODO no emb
+        modified_memories = [n for n in nodes if n.status == MemoryNodeStatus.MODIFIED]
+        if modified_memories:
+            for n in modified_memories:
+                n.status = MemoryNodeStatus.ACTIVE.value
+                self.delete(n)
+                self.insert(n)
+
+        # set memories expired
+        expired_memories = [n for n in nodes if n.status == MemoryNodeStatus.EXPIRED]
+        if expired_memories:
+            for n in expired_memories:
+                n.status = MemoryNodeStatus.ACTIVE.value
+                self.delete(n)
+                self.insert(n)
+
+    @staticmethod
+    def _memory_node_2_text_node(memory_node: MemoryNode) -> TextNode:
+        return TextNode(id_=memory_node.memory_id,
+                        text=memory_node.content,
+                        metadata=memory_node.model_dump(exclude={"content"}))
+
+    @staticmethod
+    def _text_node_2_memory_node(text_node: NodeWithScore) -> MemoryNode:
+        return MemoryNode(content=text_node.text, **text_node.metadata)

memory_scope/storage/llama_index_sync_elasticsearch.py

@@ -0,0 +1,546 @@
+"""Elasticsearch vector store."""
+
+import asyncio
+from logging import getLogger
+from typing import Any, Callable, Dict, List, Literal, Optional, Union
+
+import nest_asyncio
+import numpy as np
+from elasticsearch import AsyncElasticsearch, Elasticsearch
+
+from llama_index.core.bridge.pydantic import PrivateAttr
+from llama_index.core.schema import BaseNode, MetadataMode, TextNode
+from llama_index.core.vector_stores.types import (
+    BasePydanticVectorStore,
+    MetadataFilters,
+    VectorStoreQuery,
+    VectorStoreQueryMode,
+    VectorStoreQueryResult,
+)
+from llama_index.core.vector_stores.utils import (
+    metadata_dict_to_node,
+    node_to_metadata_dict,
+)
+from elasticsearch.helpers.vectorstore import AsyncVectorStore, VectorStore
+from elasticsearch.helpers.vectorstore import (
+    AsyncBM25Strategy,
+    AsyncSparseVectorStrategy,
+    AsyncDenseVectorStrategy,
+    AsyncRetrievalStrategy,
+    DistanceMetric,
+)
+
+from llama_index.vector_stores.elasticsearch.utils import (
+    get_user_agent,
+)
+
+
+logger = getLogger(__name__)
+
+DISTANCE_STRATEGIES = Literal[
+    "COSINE",
+    "DOT_PRODUCT",
+    "EUCLIDEAN_DISTANCE",
+]
+
+def get_elasticsearch_client(
+    url: Optional[str] = None,
+    cloud_id: Optional[str] = None,
+    api_key: Optional[str] = None,
+    username: Optional[str] = None,
+    password: Optional[str] = None,
+    use_async: Optional[bool] = False,
+) -> AsyncElasticsearch:
+    if url and cloud_id:
+        raise ValueError(
+            "Both es_url and cloud_id are defined. Please provide only one."
+        )
+
+    connection_params: Dict[str, Any] = {}
+
+    if url:
+        connection_params["hosts"] = [url]
+    elif cloud_id:
+        connection_params["cloud_id"] = cloud_id
+    else:
+        raise ValueError("Please provide either elasticsearch_url or cloud_id.")
+
+    if api_key:
+        connection_params["api_key"] = api_key
+    elif username and password:
+        connection_params["basic_auth"] = (username, password)
+    if use_async:
+        es_client = AsyncElasticsearch(
+            **connection_params, headers={"user-agent": get_user_agent()}
+        )
+    else:
+        es_client = Elasticsearch(
+            **connection_params, headers={"user-agent": get_user_agent()}
+        )
+
+    es_client.info()  # use sync client so don't have to 'await' to just get info
+
+    return es_client
+
+
+def _to_elasticsearch_filter(standard_filters: MetadataFilters) -> Dict[str, Any]:
+    """
+    Convert standard filters to Elasticsearch filter.
+
+    Args:
+        standard_filters: Standard Llama-index filters.
+
+    Returns:
+        Elasticsearch filter.
+    """
+    if len(standard_filters.legacy_filters()) == 1:
+        filter = standard_filters.legacy_filters()[0]
+        return {
+            "term": {
+                f"metadata.{filter.key}.keyword": {
+                    "value": filter.value,
+                }
+            }
+        }
+    else:
+        operands = []
+        for filter in standard_filters.legacy_filters():
+            operands.append(
+                {
+                    "term": {
+                        f"metadata.{filter.key}.keyword": {
+                            "value": filter.value,
+                        }
+                    }
+                }
+            )
+        return {"bool": {"should": operands}}
+
+
+def _to_llama_similarities(scores: List[float]) -> List[float]:
+    if scores is None or len(scores) == 0:
+        return []
+
+    scores_to_norm: np.ndarray = np.array(scores)
+    return np.exp(scores_to_norm - np.max(scores_to_norm)).tolist()
+
+
+def _mode_must_match_retrieval_strategy(
+    mode: VectorStoreQueryMode, retrieval_strategy: AsyncRetrievalStrategy
+) -> None:
+    """
+    Different retrieval strategies require different ways of indexing that must be known at the
+    time of adding data. The query mode is known at query time. This function checks if the
+    retrieval strategy (and way of indexing) is compatible with the query mode and raises and
+    exception in the case of a mismatch.
+    """
+    if mode == VectorStoreQueryMode.DEFAULT:
+        # it's fine to not specify an explicit other mode
+        return
+
+    mode_retrieval_dict = {
+        VectorStoreQueryMode.SPARSE: AsyncSparseVectorStrategy,
+        VectorStoreQueryMode.TEXT_SEARCH: AsyncBM25Strategy,
+        VectorStoreQueryMode.HYBRID: AsyncDenseVectorStrategy,
+    }
+
+    required_strategy = mode_retrieval_dict.get(mode)
+    if not required_strategy:
+        raise NotImplementedError(f"query mode {mode} currently not supported")
+
+    if not isinstance(retrieval_strategy, required_strategy):
+        raise ValueError(
+            f"query mode {mode} incompatible with retrieval strategy {type(retrieval_strategy)}, "
+            f"expected {required_strategy}"
+        )
+
+    if mode == VectorStoreQueryMode.HYBRID and not retrieval_strategy.hybrid:
+        raise ValueError(f"to enable hybrid mode, it must be set in retrieval strategy")
+
+
+class SyncElasticsearchStore(BasePydanticVectorStore):
+    """
+    Elasticsearch vector store.
+
+    Args:
+        index_name: Name of the Elasticsearch index.
+        es_client: Optional. Pre-existing AsyncElasticsearch client.
+        es_url: Optional. Elasticsearch URL.
+        es_cloud_id: Optional. Elasticsearch cloud ID.
+        es_api_key: Optional. Elasticsearch API key.
+        es_user: Optional. Elasticsearch username.
+        es_password: Optional. Elasticsearch password.
+        text_field: Optional. Name of the Elasticsearch field that stores the text.
+        vector_field: Optional. Name of the Elasticsearch field that stores the
+                    embedding.
+        batch_size: Optional. Batch size for bulk indexing. Defaults to 200.
+        distance_strategy: Optional. Distance strategy to use for similarity search.
+                        Defaults to "COSINE".
+        retrieval_strategy: Retrieval strategy to use. AsyncBM25Strategy /
+            AsyncSparseVectorStrategy / AsyncDenseVectorStrategy / AsyncRetrievalStrategy.
+            Defaults to AsyncDenseVectorStrategy.
+
+    Raises:
+        ConnectionError: If AsyncElasticsearch client cannot connect to Elasticsearch.
+        ValueError: If neither es_client nor es_url nor es_cloud_id is provided.
+
+    Examples:
+        `pip install llama-index-vector-stores-elasticsearch`
+
+        ```python
+        from llama_index.vector_stores import ElasticsearchStore
+
+        # Additional setup for ElasticsearchStore class
+        index_name = "my_index"
+        es_url = "http://localhost:9200"
+        es_cloud_id = "<cloud-id>"  # Found within the deployment page
+        es_user = "elastic"
+        es_password = "<password>"  # Provided when creating deployment or can be reset
+        es_api_key = "<api-key>"  # Create an API key within Kibana (Security -> API Keys)
+
+        # Connecting to ElasticsearchStore locally
+        es_local = ElasticsearchStore(
+            index_name=index_name,
+            es_url=es_url,
+        )
+
+        # Connecting to Elastic Cloud with username and password
+        es_cloud_user_pass = ElasticsearchStore(
+            index_name=index_name,
+            es_cloud_id=es_cloud_id,
+            es_user=es_user,
+            es_password=es_password,
+        )
+
+        # Connecting to Elastic Cloud with API Key
+        es_cloud_api_key = ElasticsearchStore(
+            index_name=index_name,
+            es_cloud_id=es_cloud_id,
+            es_api_key=es_api_key,
+        )
+        ```
+
+    """
+
+    class Config:
+        # allow pydantic to tolarate its inability to validate AsyncRetrievalStrategy
+        arbitrary_types_allowed = True
+
+    stores_text: bool = True
+    index_name: str
+    es_client: Optional[Any]
+    es_url: Optional[str]
+    es_cloud_id: Optional[str]
+    es_api_key: Optional[str]
+    es_user: Optional[str]
+    es_password: Optional[str]
+    text_field: str = "content"
+    vector_field: str = "embedding"
+    batch_size: int = 200
+    distance_strategy: Optional[DISTANCE_STRATEGIES] = "COSINE"
+    retrieval_strategy: AsyncRetrievalStrategy
+
+    _store = PrivateAttr()
+    def __init__(
+        self,
+        index_name: str,
+        es_client: Optional[Any] = None,
+        es_url: Optional[str] = None,
+        es_cloud_id: Optional[str] = None,
+        es_api_key: Optional[str] = None,
+        es_user: Optional[str] = None,
+        es_password: Optional[str] = None,
+        text_field: str = "content",
+        vector_field: str = "embedding",
+        batch_size: int = 200,
+        distance_strategy: Optional[DISTANCE_STRATEGIES] = "COSINE",
+        retrieval_strategy: Optional[AsyncRetrievalStrategy] = None,
+    ) -> None:
+        nest_asyncio.apply()
+
+        if not es_client:
+            es_client = get_elasticsearch_client(
+                url=es_url,
+                cloud_id=es_cloud_id,
+                api_key=es_api_key,
+                username=es_user,
+                password=es_password,
+            )
+
+        if retrieval_strategy is None:
+            retrieval_strategy = AsyncDenseVectorStrategy(
+                distance=DistanceMetric[distance_strategy]
+            )
+
+        metadata_mappings = {
+            "document_id": {"type": "keyword"},
+            "doc_id": {"type": "keyword"},
+            "ref_doc_id": {"type": "keyword"},
+        }
+
+        self._store = VectorStore(
+            user_agent=get_user_agent(),
+            client=es_client,
+            index=index_name,
+            retrieval_strategy=retrieval_strategy,
+            text_field=text_field,
+            vector_field=vector_field,
+            metadata_mappings=metadata_mappings,
+        )
+
+        super().__init__(
+            index_name=index_name,
+            es_client=es_client,
+            es_url=es_url,
+            es_cloud_id=es_cloud_id,
+            es_api_key=es_api_key,
+            es_user=es_user,
+            es_password=es_password,
+            text_field=text_field,
+            vector_field=vector_field,
+            batch_size=batch_size,
+            distance_strategy=distance_strategy,
+            retrieval_strategy=retrieval_strategy,
+        )
+
+    @property
+    def client(self) -> Any:
+        """Get async elasticsearch client."""
+        return self._store.client
+
+    def close(self) -> None:
+        return self._store.close()
+    
+    def add(
+        self,
+        nodes: List[BaseNode],
+        *,
+        create_index_if_not_exists: bool = True,
+        **add_kwargs: Any,
+    ) -> List[str]:
+        """
+        Add nodes to Elasticsearch index.
+
+        Args:
+            nodes: List of nodes with embeddings.
+            create_index_if_not_exists: Optional. Whether to create
+                                        the Elasticsearch index if it
+                                        doesn't already exist.
+                                        Defaults to True.
+
+        Returns:
+            List of node IDs that were added to the index.
+
+        Raises:
+            ImportError: If elasticsearch['async'] python package is not installed.
+            BulkIndexError: If AsyncElasticsearch async_bulk indexing fails.
+        """
+        
+        return self.sync_add(nodes, create_index_if_not_exists=create_index_if_not_exists)
+    
+    def sync_add(
+        self,
+        nodes: List[BaseNode],
+        *,
+        create_index_if_not_exists: bool = True,
+        **add_kwargs: Any,
+    ) -> List[str]:
+        """
+        Asynchronous method to add nodes to Elasticsearch index.
+
+        Args:
+            nodes: List of nodes with embeddings.
+            create_index_if_not_exists: Optional. Whether to create
+                                        the AsyncElasticsearch index if it
+                                        doesn't already exist.
+                                        Defaults to True.
+
+        Returns:
+            List of node IDs that were added to the index.
+
+        Raises:
+            ImportError: If elasticsearch python package is not installed.
+            BulkIndexError: If AsyncElasticsearch async_bulk indexing fails.
+        """
+        if len(nodes) == 0:
+            return []
+
+        embeddings: List[List[float]] = []
+        texts: List[str] = []
+        metadatas: List[dict] = []
+        ids: List[str] = []
+        for node in nodes:
+            ids.append(node.node_id)
+            embeddings.append(node.get_embedding())
+            texts.append(node.get_content(metadata_mode=MetadataMode.NONE))
+            metadatas.append(node_to_metadata_dict(node, remove_text=True))
+
+        if not self._store.num_dimensions:
+            self._store.num_dimensions = len(embeddings[0])
+
+        return self._store.add_texts(
+            texts=texts,
+            metadatas=metadatas,
+            vectors=embeddings,
+            ids=ids,
+            create_index_if_not_exists=create_index_if_not_exists,
+            bulk_kwargs=add_kwargs,
+        )
+
+    def delete(self, ref_doc_id: str, **delete_kwargs: Any) -> None:
+        """
+        Delete node from Elasticsearch index.
+
+        Args:
+            ref_doc_id: ID of the node to delete.
+            delete_kwargs: Optional. Additional arguments to
+                        pass to Elasticsearch delete_by_query.
+
+        Raises:
+            Exception: If Elasticsearch delete_by_query fails.
+        """
+        return self.sync_delete(ref_doc_id, **delete_kwargs)
+
+    def sync_delete(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 self._store.delete(
+        #     query={"term": {"metadata.ref_doc_id": ref_doc_id}}, **delete_kwargs
+        # )
+        return self._store.delete(query={"term": {"_id": ref_doc_id}}, **delete_kwargs)
+
+    def query(
+        self,
+        query: VectorStoreQuery,
+        custom_query: Optional[
+            Callable[[Dict, Union[VectorStoreQuery, None]], Dict]
+        ] = None,
+        es_filter: Optional[List[Dict]] = None,
+        **kwargs: Any,
+    ) -> VectorStoreQueryResult:
+        """
+        Query index for top k most similar nodes.
+
+        Args:
+            query_embedding (List[float]): query embedding
+            custom_query: Optional. custom query function that takes in the es query
+                        body and returns a modified query body.
+                        This can be used to add additional query
+                        parameters to the Elasticsearch query.
+            es_filter: Optional. Elasticsearch filter to apply to the
+                        query. If filter is provided in the query,
+                        this filter will be ignored.
+
+        Returns:
+            VectorStoreQueryResult: Result of the query.
+
+        Raises:
+            Exception: If Elasticsearch query fails.
+
+        """
+        return self.sync_query(query, custom_query, es_filter, **kwargs)    
+    
+    def sync_query(
+        self,
+        query: VectorStoreQuery,
+        custom_query: Optional[
+            Callable[[Dict, Union[VectorStoreQuery, None]], Dict]
+        ] = None,
+        es_filter: Optional[List[Dict]] = None,
+        **kwargs: Any,
+    ) -> VectorStoreQueryResult:
+        """
+        Asynchronous query index for top k most similar nodes.
+
+        Args:
+            query_embedding (VectorStoreQuery): query embedding
+            custom_query: Optional. custom query function that takes in the es query
+                        body and returns a modified query body.
+                        This can be used to add additional query
+                        parameters to the AsyncElasticsearch query.
+            es_filter: Optional. AsyncElasticsearch filter to apply to the
+                        query. If filter is provided in the query,
+                        this filter will be ignored.
+
+        Returns:
+            VectorStoreQueryResult: Result of the query.
+
+        Raises:
+            Exception: If AsyncElasticsearch query fails.
+
+        """
+        _mode_must_match_retrieval_strategy(query.mode, self.retrieval_strategy)
+
+        if query.filters is not None and len(query.filters.legacy_filters()) > 0:
+            filter = [_to_elasticsearch_filter(query.filters)]
+        else:
+            filter = es_filter or []
+        hits = self._store.search(
+            query=query.query_str,
+            query_vector=query.query_embedding,
+            k=query.similarity_top_k,
+            num_candidates=100, # query.similarity_top_k * 10,
+            filter=filter,
+            custom_query=custom_query,
+        )
+
+        top_k_nodes = []
+        top_k_ids = []
+        top_k_scores = []
+        for hit in hits:
+            source = hit["_source"]
+            metadata = source.get("metadata", None)
+            text = source.get(self.text_field, None)
+            node_id = hit["_id"]
+
+            try:
+                node = metadata_dict_to_node(metadata)
+                node.text = text
+            except Exception:
+                # Legacy support for old metadata format
+                logger.warning(
+                    f"Could not parse metadata from hit {hit['_source']['metadata']}"
+                )
+                node_info = source.get("node_info")
+                relationships = source.get("relationships", {})
+                start_char_idx = None
+                end_char_idx = None
+                if isinstance(node_info, dict):
+                    start_char_idx = node_info.get("start", None)
+                    end_char_idx = node_info.get("end", None)
+
+                node = TextNode(
+                    text=text,
+                    metadata=metadata,
+                    id_=node_id,
+                    start_char_idx=start_char_idx,
+                    end_char_idx=end_char_idx,
+                    relationships=relationships,
+                )
+            top_k_nodes.append(node)
+            top_k_ids.append(node_id)
+            top_k_scores.append(hit.get("_rank", hit["_score"]))
+
+        if (
+            isinstance(self.retrieval_strategy, AsyncDenseVectorStrategy)
+            and self.retrieval_strategy.hybrid
+        ):
+            total_rank = sum(top_k_scores)
+            top_k_scores = [(total_rank - rank) / total_rank for rank in top_k_scores]
+            # top_k_scores = [total_rank - rank / total_rank for rank in top_k_scores]
+
+
+        return VectorStoreQueryResult(
+            nodes=top_k_nodes,
+            ids=top_k_ids,
+            similarities=_to_llama_similarities(top_k_scores),
+        )

tests/storages/test_storages_lli_synces.py

@@ -0,0 +1,188 @@
+import unittest
+
+from memory_scope.models.llama_index_embedding_model import LlamaIndexEmbeddingModel
+from memory_scope.scheme.memory_node import MemoryNode
+from memory_scope.storage.llama_index_es_memory_store_sync import LlamaIndexEsMemoryStore
+
+class TestLlamaIndexElasticSearchStore(unittest.TestCase):
+    """Tests for LLIEmbedding"""
+
+    def setUp(self):
+        config = {
+            "module_name": "dashscope_embedding",
+            "model_name": "text-embedding-v2",
+            "clazz": "models.llama_index_embedding_model",
+        }
+        emb = LlamaIndexEmbeddingModel(**config)
+
+        config = {
+            "index_name": "0708_5",
+            "es_url": "http://localhost:9200",
+            "embedding_model": emb,
+            "use_hybrid": True
+
+        }
+        self.es_store = LlamaIndexEsMemoryStore(**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",
+                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",
+                user_id="1",
+                status="valid",
+                memory_id="bbb456",
+                meta_data={"1": "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",
+                user_id="2",
+                status="valid",
+                memory_id="ccc789",
+                meta_data={"2": "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",
+                user_id="3",
+                status="valid",
+                memory_id="ddd012",
+                meta_data={"3": "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",
+                user_id="4",
+                status="valid",
+                memory_id="eee345",
+                meta_data={"4": "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",
+                user_id="5",
+                status="valid",
+                memory_id="fff678",
+                meta_data={"5": "5"},
+
+            ),
+            MemoryNode(
+                content="An organized crime dynasty's aging patriarch transfers control of his clandestine "
+                        "empire to his reluctant son.",
+                memory_type="insights",
+                user_id="6",
+                status="valid",
+                memory_id="ggg901",
+                meta_data={"5": "5"}
+
+            ),
+            MemoryNode(
+                content="ggggggggg",
+                memory_type="profile",
+                user_id="6",
+                status="valid",
+                memory_id="ggg234",
+                meta_data={"5": "5"}
+
+            ),
+            MemoryNode(
+                content="ggggggggg",
+                memory_type="profile",
+                user_id="6",
+                status="valid",
+                memory_id="hhh234",
+                meta_data={"5": "5"}
+
+            ),
+            MemoryNode(
+                content="ggggggggg",
+                memory_type="profile",
+                user_id="6",
+                status="valid",
+                memory_id="iii234",
+                meta_data={"5": "5"}
+
+            ),
+            MemoryNode(
+                content="ggggggggg",
+                memory_type="profile",
+                user_id="6",
+                status="valid",
+                memory_id="jjj234",
+                meta_data={"5": "5"}
+
+            ),
+            MemoryNode(
+                content="ggggggggg",
+                memory_type="profile",
+                user_id="6",
+                status="valid",
+                memory_id="kkk234",
+                meta_data={"5": "5"}
+
+            ),
+        ]
+
+    def test_retrieve(self):
+        # filter_dict = {
+        #     "user_id": "6",
+        # }
+        filter_dict = {}
+
+        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",
+            meta_data={"5": "5"}
+        ))
+        res = self.es_store.retrieve_memories(query="hacker", filter_dict=filter_dict, top_k=15)
+        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_memories(query="hacker", filter_dict=filter_dict, top_k=15)
+        print(len(res))
+        print(res)
+
+        self.es_store.delete(MemoryNode(
+            content="test update",
+            memory_type="profile",
+            user_id="6",
+            status="invalid",
+            memory_id="ggg567"
+        ))
+        import asyncio
+        res = asyncio.run(self.es_store.a_retrieve_memories(query="hacker", filter_dict=filter_dict, top_k=15))
+        # res = self.es_store.async_retrieve(query="hacker", filter_dict=filter_dict, top_k=10)
+        print(len(res))
+        print(res)
+
+    def tearDown(self):
+        self.es_store.close()

tests/thread_test2.py

@@ -6,7 +6,9 @@ import asyncio
 from concurrent.futures import ThreadPoolExecutor
 
 from memory_scope.models.llama_index_embedding_model import LlamaIndexEmbeddingModel
+
 from memory_scope.storage.llama_index_es_memory_store import LlamaIndexEsMemoryStore
+from memory_scope.storage.llama_index_es_memory_store_sync import LlamaIndexEsMemoryStore as SyncLlamaIndexEsMemoryStore
 from memory_scope.utils.logger import Logger
 logger = Logger.get_logger("default")
 
@@ -16,24 +18,26 @@ logger = Logger.get_logger("default")
 class ThreadTest(object):
     def __init__(self):
         self.task_list = []
-        embedding_model_conf = {
+        config = {
             "module_name": "dashscope_embedding",
             "model_name": "text-embedding-v2",
             "clazz": "models.llama_index_embedding_model",
         }
+        emb = LlamaIndexEmbeddingModel(**config)
 
         config = {
             "index_name": "0708_2",
             "es_url": "http://localhost:9200",
-            "embedding_model_conf": embedding_model_conf,
+            "embedding_model": emb,
             "use_hybrid": True
 
         }
-        self.es_store = LlamaIndexEsMemoryStore(**config)   # 不能在async中初始化
+        self.es_store = SyncLlamaIndexEsMemoryStore(**config)   # 不能在async中初始化
         self.logger = logger
 
     def major_func(self, i: int):
-        result = self.es_store.retrieve_memories("_", top_k=10, filter_dict={"memory_id": "ggg567"})
+        result = self.es_store.retrieve_memories("_", top_k=10, filter_dict={})
+        print(result)
         return result
 
     def run(self):