diff --git a/memory_scope/models/llama_index_embedding_model.py b/memory_scope/models/llama_index_embedding_model.py index 0eb6dc51..6ed455bd 100644 --- a/memory_scope/models/llama_index_embedding_model.py +++ b/memory_scope/models/llama_index_embedding_model.py @@ -39,7 +39,6 @@ class LlamaIndexEmbeddingModel(BaseModel): model_response.details = "empty embeddings" model_response.status = False return model_response - if len(embeddings) == 1: # return list[float] embeddings = embeddings[0] diff --git a/memory_scope/scheme/memory_node.py b/memory_scope/scheme/memory_node.py index 0edd0248..de884ae7 100644 --- a/memory_scope/scheme/memory_node.py +++ b/memory_scope/scheme/memory_node.py @@ -25,11 +25,11 @@ class MemoryNode(BaseModel): value: str = Field("", description="memory value") - score_similar: float = Field(0, description="es similar score") + score_recall: float = Field(0, description="embedding similarity score used in recall stage") - score_rank: float = Field(0, description="rank model score") + score_rank: float = Field(0, description="rank model score used in rank stage") - score_rerank: float = Field(0, description="rerank score") + score_rerank: float = Field(0, description="rerank score used in rerank stage") memory_type: str = Field("", description="conversation / observation / insight...") diff --git a/memory_scope/storage/llama_index_es_memory_store.py b/memory_scope/storage/llama_index_es_memory_store.py index f49ba532..6fb6e510 100644 --- a/memory_scope/storage/llama_index_es_memory_store.py +++ b/memory_scope/storage/llama_index_es_memory_store.py @@ -1,137 +1,17 @@ import warnings +import random 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 AsyncDenseVectorStrategy +from llama_index.core.schema import TextNode, NodeWithScore, QueryBundle 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.storage.llama_index_sync_elasticsearch import SyncElasticsearchStore +from memory_scope.storage.llama_index_sync_elasticsearch import SyncElasticsearchStore, _AsyncDenseVectorStrategy, _to_elasticsearch_filter from memory_scope.utils.logger import Logger -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]: - """ - Converts standard Llama-index filters into a format compatible with Elasticsearch. - - This function transforms dictionary-based filters, where each key represents a field and - the value is a list of strings, into an Elasticsearch query structure. It supports both - list values (interpreted as 'should' clauses for OR logic) and single values (interpreted - as 'must' clauses for AND logic). - - Args: - standard_filters (Dict[str, List[str]]): A dictionary containing filter criteria, - where keys are field names and values are lists of strings or single string values - representing filter values. - - Returns: - Dict[str, Any]: A dictionary structured as an Elasticsearch filter query. - """ - result = { - "bool": {} - } - for key, value in standard_filters.items(): - if isinstance(value, list): - operands = [] - for v in value: - key_str = f"metadata.{key}.keyword" if isinstance(v, str) else f"metadata.{key}" - operands.append( - { - "term": - { - key_str: {"value": v} - } - } - ) - result['bool'].update({"should": operands}) # ⭐ Add 'should' clause for OR logic - result['bool'].update({"minimum_should_match": 1}) # Ensure at least one 'should' match - else: - key_str = f"metadata.{key}.keyword" if isinstance(value, str) else f"metadata.{key}" - operand = [{ - "term": { - key_str: { - "value": value, - } - } - }] - if "must" in result['bool']: - result['bool']['must'].extend(operand) # Extend existing 'must' clause for AND logic - else: - result['bool'].update({"must": operand}) # Initialize 'must' clause if not present - return result - - class LlamaIndexEsMemoryStore(BaseMemoryStore): def __init__(self, @@ -139,8 +19,10 @@ class LlamaIndexEsMemoryStore(BaseMemoryStore): index_name: str, es_url: str, use_hybrid: bool = True, + emb_dims: int = 1536, **kwargs): - + self.index_name = index_name + self.emb_dims = emb_dims self.embedding_model: BaseModel = embedding_model self.es_store = SyncElasticsearchStore(index_name=index_name, es_url=es_url, @@ -148,24 +30,32 @@ class LlamaIndexEsMemoryStore(BaseMemoryStore): **kwargs) # TODO The llamaIndex utilizes some deprecated functions, hence langchain logs warning messages. By # adding the following lines of code, the display of deprecated information is suppressed. - with warnings.catch_warnings(): - warnings.simplefilter("ignore") - self.index = VectorStoreIndex.from_vector_store(vector_store=self.es_store, - embed_model=self.embedding_model.model) + + 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, + query: Optional[str] = None, + top_k: int = 3, filter_dict: Dict[str, List[str]] | Dict[str, str] = None) -> List[MemoryNode]: + # if index is not created, return [] + exists = self.es_store._store.client.indices.exists(index=self.index_name) + if not exists: + return [] + 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) + retriever = self.index.as_retriever(vector_store_kwargs={"es_filter": es_filter, "fields": ['embedding']}, + similarity_top_k=top_k, + sparse_top_k=top_k, ) + if query is None: + query = QueryBundle(query_str='**--**', + embedding=self.dummy_query_vector()) + text_nodes = retriever.retrieve(query) return [self._text_node_2_memory_node(n) for n in text_nodes] @@ -181,6 +71,11 @@ class LlamaIndexEsMemoryStore(BaseMemoryStore): retriever = self.index.as_retriever( vector_store_kwargs={"es_filter": es_filter}, similarity_top_k=top_k) + + if query is None: + query = QueryBundle(query_str='**--**', + embedding=self.dummy_query_vector()) + text_nodes: List[NodeWithScore] = await retriever.aretrieve(query) return [self._text_node_2_memory_node(n) for n in text_nodes] @@ -205,6 +100,10 @@ class LlamaIndexEsMemoryStore(BaseMemoryStore): def delete(self, node: MemoryNode): return self.es_store.delete(node.memory_id) + def delete_conditional(self, filter_dict: Dict = {}): + nodes = self.retrieve_memories(filter_dict=filter_dict, top_k=10000) + self.batch_delete(nodes) + def update(self, node: MemoryNode, update_embedding: bool = True): # TODO update without embedding? self.delete(node) @@ -215,7 +114,11 @@ class LlamaIndexEsMemoryStore(BaseMemoryStore): Closes the Elasticsearch store, releasing any resources associated with it. """ self.es_store.close() - + + def dummy_query_vector(self): + random_floats = [random.uniform(0, 1) for _ in range(self.emb_dims)] + return random_floats + @staticmethod def _memory_node_2_text_node(memory_node: MemoryNode) -> TextNode: """ @@ -226,10 +129,14 @@ class LlamaIndexEsMemoryStore(BaseMemoryStore): Returns: TextNode: The converted TextNode with content and metadata from the MemoryNode. - """ + """ + embedding = memory_node.vector + if not embedding: + embedding = None return TextNode(id_=memory_node.memory_id, text=memory_node.content, - metadata=memory_node.model_dump(exclude={"content"})) + embedding=embedding, + metadata=memory_node.model_dump(exclude={"content", "vector"})) @staticmethod def _text_node_2_memory_node(text_node: NodeWithScore) -> MemoryNode: @@ -242,4 +149,8 @@ class LlamaIndexEsMemoryStore(BaseMemoryStore): Returns: MemoryNode: The converted MemoryNode with text and metadata from the NodeWithScore. """ - return MemoryNode(content=text_node.text, **text_node.metadata) + embedding = text_node.embedding + print("textnode embedding", embedding) + if not embedding: + embedding = [] + return MemoryNode(content=text_node.text, vector=embedding, **text_node.metadata) diff --git a/memory_scope/storage/llama_index_sync_elasticsearch.py b/memory_scope/storage/llama_index_sync_elasticsearch.py index 059d6970..1532d9f5 100644 --- a/memory_scope/storage/llama_index_sync_elasticsearch.py +++ b/memory_scope/storage/llama_index_sync_elasticsearch.py @@ -1,7 +1,7 @@ """Elasticsearch vector store.""" from logging import getLogger -from typing import Any, Callable, Dict, List, Literal, Optional, Union +from typing import Any, Callable, Dict, List, Literal, Optional, Union, cast import nest_asyncio import numpy as np @@ -80,47 +80,6 @@ def get_elasticsearch_client( return es_client -def _to_elasticsearch_filter(standard_filters: MetadataFilters) -> Dict[str, Any]: - """ - Transforms Llama-index standard filters into an Elasticsearch-compatible filter structure. - - This function supports both single-term filters and multiple operands combined - with a boolean 'should' clause for more complex queries. - - Args: - standard_filters (MetadataFilters): An instance of MetadataFilters containing - the filtering criteria to be applied. - - Returns: - Dict[str, Any]: A dictionary representing the Elasticsearch filter query. - """ - if len(standard_filters.legacy_filters()) == 1: - # For a single filter term, construct a simple term filter. - filter = standard_filters.legacy_filters()[0] - return { - "term": { - f"metadata.{filter.key}.keyword": { - "value": filter.value, - } - } - } - else: - # When multiple filters are present, create a boolean 'should' clause - # with each individual filter as an operand. - 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]: """ Converts a list of similarity scores into a normalized form for LlamaIndex compatibility. @@ -173,6 +132,137 @@ def _mode_must_match_retrieval_strategy( raise ValueError(f"to enable hybrid mode, it must be set in retrieval strategy") +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 + if query == "**--**": + query_body = { + "query": { + "bool": { + "filter": filter, + } + }, + } + else: + 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]: + """ + Converts standard Llama-index filters into a format compatible with Elasticsearch. + + This function transforms dictionary-based filters, where each key represents a field and + the value is a list of strings, into an Elasticsearch query structure. It supports both + list values (interpreted as 'should' clauses for OR logic) and single values (interpreted + as 'must' clauses for AND logic). + + Args: + standard_filters (Dict[str, List[str]]): A dictionary containing filter criteria, + where keys are field names and values are lists of strings or single string values + representing filter values. + + Returns: + Dict[str, Any]: A dictionary structured as an Elasticsearch filter query. + """ + result = { + "bool": {} + } + for key, value in standard_filters.items(): + if isinstance(value, list): + operands = [] + for v in value: + key_str = f"metadata.{key}.keyword" if isinstance(v, str) else f"metadata.{key}" + operands.append( + { + "term": + { + key_str: {"value": v} + } + } + ) + result['bool'].update({"should": operands}) # ⭐ Add 'should' clause for OR logic + result['bool'].update({"minimum_should_match": 1}) # Ensure at least one 'should' match + else: + key_str = f"metadata.{key}.keyword" if isinstance(value, str) else f"metadata.{key}" + operand = [{ + "term": { + key_str: { + "value": value, + } + } + }] + if "must" in result['bool']: + result['bool']['must'].extend(operand) # Extend existing 'must' clause for AND logic + else: + result['bool'].update({"must": operand}) # Initialize 'must' clause if not present + return result + + + class SyncElasticsearchStore(BasePydanticVectorStore): """ Elasticsearch vector store. @@ -499,6 +589,7 @@ class SyncElasticsearchStore(BasePydanticVectorStore): Callable[[Dict, Union[VectorStoreQuery, None]], Dict] ] = None, es_filter: Optional[List[Dict]] = None, + fields: List[str] = [], **kwargs: Any, ) -> VectorStoreQueryResult: """ @@ -538,14 +629,15 @@ class SyncElasticsearchStore(BasePydanticVectorStore): num_candidates=num_candidates, # query.similarity_top_k * 10, filter=filter, custom_query=custom_query, + fields=fields, ) - top_k_nodes = [] top_k_ids = [] top_k_scores = [] for hit in hits: source = hit["_source"] metadata = source.get("metadata", None) + embedding = source.get("embedding", None) text = source.get(self.text_field, None) node_id = hit["_id"] @@ -553,6 +645,7 @@ class SyncElasticsearchStore(BasePydanticVectorStore): # Attempt to parse metadata using the standard method node = metadata_dict_to_node(metadata) node.text = text + node.embedding = embedding except Exception: # Legacy support for old metadata format logger.warning( @@ -570,6 +663,7 @@ class SyncElasticsearchStore(BasePydanticVectorStore): text=text, metadata=metadata, id_=node_id, + embedding=embedding, start_char_idx=start_char_idx, end_char_idx=end_char_idx, relationships=relationships, @@ -582,12 +676,14 @@ class SyncElasticsearchStore(BasePydanticVectorStore): 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] + total_rank = sum(top_k_scores) + top_k_scores = [rank for rank in 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), + # similarities=_to_llama_similarities(top_k_scores), + similarities=top_k_scores )