diff --git a/memory_scope/storage/llama_index_es_memory_store_sync.py b/memory_scope/storage/llama_index_es_memory_store_sync.py new file mode 100644 index 00000000..04597b21 --- /dev/null +++ b/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) diff --git a/memory_scope/storage/llama_index_sync_elasticsearch.py b/memory_scope/storage/llama_index_sync_elasticsearch.py new file mode 100644 index 00000000..3dad68bc --- /dev/null +++ b/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 = "" # Found within the deployment page + es_user = "elastic" + es_password = "" # Provided when creating deployment or can be reset + es_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), + ) diff --git a/tests/storages/test_storages_lli_synces.py b/tests/storages/test_storages_lli_synces.py new file mode 100644 index 00000000..0982c818 --- /dev/null +++ b/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() diff --git a/tests/thread_test2.py b/tests/thread_test2.py index d3420291..f58de946 100644 --- a/tests/thread_test2.py +++ b/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):