From c0ab3b0fb729e4de04d992274feeb38bfa3d4101 Mon Sep 17 00:00:00 2001 From: "xianzhe.xxz" Date: Thu, 18 Jul 2024 20:15:45 +0800 Subject: [PATCH 1/5] support es query wo emb --- .../storage/llama_index_es_memory_store.py | 145 +++------------ .../storage/llama_index_sync_elasticsearch.py | 168 +++++++++++++----- 2 files changed, 146 insertions(+), 167 deletions(-) diff --git a/memory_scope/storage/llama_index_es_memory_store.py b/memory_scope/storage/llama_index_es_memory_store.py index f49ba532..41e3a3af 100644 --- a/memory_scope/storage/llama_index_es_memory_store.py +++ b/memory_scope/storage/llama_index_es_memory_store.py @@ -1,137 +1,18 @@ 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.core.schema import TextNode, NodeWithScore, QueryBundle from llama_index.vector_stores.elasticsearch import AsyncDenseVectorStrategy 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,6 +20,7 @@ class LlamaIndexEsMemoryStore(BaseMemoryStore): index_name: str, es_url: str, use_hybrid: bool = True, + emb_dims: int = 1536, **kwargs): self.embedding_model: BaseModel = embedding_model @@ -157,8 +39,8 @@ class LlamaIndexEsMemoryStore(BaseMemoryStore): 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 filter_dict is None: filter_dict = {} @@ -166,6 +48,10 @@ class LlamaIndexEsMemoryStore(BaseMemoryStore): 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) + 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 +67,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] @@ -215,7 +106,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: """ diff --git a/memory_scope/storage/llama_index_sync_elasticsearch.py b/memory_scope/storage/llama_index_sync_elasticsearch.py index bc0c3d36..41c7a3e6 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 @@ -40,6 +40,7 @@ DISTANCE_STRATEGIES = Literal[ ] + def get_elasticsearch_client( url: Optional[str] = None, cloud_id: Optional[str] = None, @@ -80,47 +81,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 +133,129 @@ 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 + 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. @@ -545,6 +628,7 @@ class SyncElasticsearchStore(BasePydanticVectorStore): top_k_nodes = [] top_k_ids = [] top_k_scores = [] + print("hits:", len(hits)) for hit in hits: source = hit["_source"] metadata = source.get("metadata", None) From 2c14b958a902ec7b0a391dec7ca2ad52ee3b427c Mon Sep 17 00:00:00 2001 From: "xianzhe.xxz" Date: Tue, 23 Jul 2024 14:38:52 +0800 Subject: [PATCH 2/5] support query without embedding --- memory_scope/scheme/memory_node.py | 6 +- .../storage/llama_index_es_memory_store.py | 11 ++-- .../storage/llama_index_sync_elasticsearch.py | 55 +++++++++++-------- 3 files changed, 39 insertions(+), 33 deletions(-) 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 41e3a3af..cb16b682 100644 --- a/memory_scope/storage/llama_index_es_memory_store.py +++ b/memory_scope/storage/llama_index_es_memory_store.py @@ -4,7 +4,6 @@ from typing import Dict, List, Any, Optional, cast from llama_index.core import VectorStoreIndex from llama_index.core.schema import TextNode, NodeWithScore, QueryBundle -from llama_index.vector_stores.elasticsearch import AsyncDenseVectorStrategy from memory_scope.models.base_model import BaseModel from memory_scope.scheme.memory_node import MemoryNode @@ -23,6 +22,7 @@ class LlamaIndexEsMemoryStore(BaseMemoryStore): emb_dims: int = 1536, **kwargs): + self.emb_dims = emb_dims self.embedding_model: BaseModel = embedding_model self.es_store = SyncElasticsearchStore(index_name=index_name, es_url=es_url, @@ -30,9 +30,8 @@ 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, + + 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")]) @@ -49,7 +48,7 @@ class LlamaIndexEsMemoryStore(BaseMemoryStore): retriever = self.index.as_retriever(vector_store_kwargs={"es_filter": es_filter}, similarity_top_k=top_k, sparse_top_k=top_k) if query is None: - query = QueryBundle(query_str='-', + query = QueryBundle(query_str='**--**', embedding=self.dummy_query_vector()) text_nodes = retriever.retrieve(query) @@ -69,7 +68,7 @@ class LlamaIndexEsMemoryStore(BaseMemoryStore): similarity_top_k=top_k) if query is None: - query = QueryBundle(query_str='-', + query = QueryBundle(query_str='**--**', embedding=self.dummy_query_vector()) text_nodes: List[NodeWithScore] = await retriever.aretrieve(query) diff --git a/memory_scope/storage/llama_index_sync_elasticsearch.py b/memory_scope/storage/llama_index_sync_elasticsearch.py index 41c7a3e6..9187d6b6 100644 --- a/memory_scope/storage/llama_index_sync_elasticsearch.py +++ b/memory_scope/storage/llama_index_sync_elasticsearch.py @@ -40,7 +40,6 @@ DISTANCE_STRATEGIES = Literal[ ] - def get_elasticsearch_client( url: Optional[str] = None, cloud_id: Optional[str] = None, @@ -133,36 +132,43 @@ 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 - query_body = { - "knn": knn, - "query": { - "bool": { - "must": [ - { - "match": { - self.text_field: { - "query": query, + if query == "**--**": + query_body = { + "query": { + "bool": { + "filter": filter, + } + }, + } + else: + query_body = { + "knn": knn, + "query": { + "bool": { + "must": [ + { + "match": { + self.text_field: { + "query": query, + } } } - } - ], - "filter": filter, - } - }, - } + ], + "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}} + 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( @@ -668,8 +674,9 @@ 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( From 46ce7d8cb713af31be241ee30e6c29d0a9da099c Mon Sep 17 00:00:00 2001 From: "xianzhe.xxz" Date: Tue, 23 Jul 2024 15:08:05 +0800 Subject: [PATCH 3/5] add conditional delete --- memory_scope/storage/llama_index_es_memory_store.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/memory_scope/storage/llama_index_es_memory_store.py b/memory_scope/storage/llama_index_es_memory_store.py index cb16b682..a75d2f3c 100644 --- a/memory_scope/storage/llama_index_es_memory_store.py +++ b/memory_scope/storage/llama_index_es_memory_store.py @@ -95,6 +95,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) From 7d11369f5ad7cdf376d0d80acb3e5ae084ac6bcc Mon Sep 17 00:00:00 2001 From: "xianzhe.xxz" Date: Tue, 23 Jul 2024 15:50:43 +0800 Subject: [PATCH 4/5] es update without reduntant embedding --- memory_scope/storage/llama_index_es_memory_store.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/memory_scope/storage/llama_index_es_memory_store.py b/memory_scope/storage/llama_index_es_memory_store.py index a75d2f3c..d42287e4 100644 --- a/memory_scope/storage/llama_index_es_memory_store.py +++ b/memory_scope/storage/llama_index_es_memory_store.py @@ -124,9 +124,10 @@ class LlamaIndexEsMemoryStore(BaseMemoryStore): Returns: TextNode: The converted TextNode with content and metadata from the MemoryNode. - """ + """ return TextNode(id_=memory_node.memory_id, text=memory_node.content, + embedding=memory_node.vector, metadata=memory_node.model_dump(exclude={"content"})) @staticmethod @@ -140,4 +141,4 @@ class LlamaIndexEsMemoryStore(BaseMemoryStore): Returns: MemoryNode: The converted MemoryNode with text and metadata from the NodeWithScore. """ - return MemoryNode(content=text_node.text, **text_node.metadata) + return MemoryNode(content=text_node.text, vector=text_node.embedding, **text_node.metadata) From 5c28465fc39c2763c007ccae4a754328b02aa594 Mon Sep 17 00:00:00 2001 From: "xianzhe.xxz" Date: Wed, 24 Jul 2024 17:25:08 +0800 Subject: [PATCH 5/5] minor fix --- .../models/llama_index_embedding_model.py | 1 - .../storage/llama_index_es_memory_store.py | 30 +++++++++++++------ .../storage/llama_index_sync_elasticsearch.py | 11 +++++-- 3 files changed, 29 insertions(+), 13 deletions(-) diff --git a/memory_scope/models/llama_index_embedding_model.py b/memory_scope/models/llama_index_embedding_model.py index ddd1e029..bacb5db7 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/storage/llama_index_es_memory_store.py b/memory_scope/storage/llama_index_es_memory_store.py index d42287e4..6fb6e510 100644 --- a/memory_scope/storage/llama_index_es_memory_store.py +++ b/memory_scope/storage/llama_index_es_memory_store.py @@ -21,7 +21,7 @@ class LlamaIndexEsMemoryStore(BaseMemoryStore): 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, @@ -32,21 +32,26 @@ class LlamaIndexEsMemoryStore(BaseMemoryStore): # adding the following lines of code, the display of deprecated information is suppressed. self.index = VectorStoreIndex.from_vector_store(vector_store=self.es_store, - embed_model=self.embedding_model.model) + 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: 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()) @@ -124,11 +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, - embedding=memory_node.vector, - 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: @@ -141,4 +149,8 @@ class LlamaIndexEsMemoryStore(BaseMemoryStore): Returns: MemoryNode: The converted MemoryNode with text and metadata from the NodeWithScore. """ - return MemoryNode(content=text_node.text, vector=text_node.embedding, **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 9187d6b6..de73aa08 100644 --- a/memory_scope/storage/llama_index_sync_elasticsearch.py +++ b/memory_scope/storage/llama_index_sync_elasticsearch.py @@ -169,6 +169,7 @@ class _AsyncDenseVectorStrategy(AsyncDenseVectorStrategy): 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( @@ -590,6 +591,7 @@ class SyncElasticsearchStore(BasePydanticVectorStore): Callable[[Dict, Union[VectorStoreQuery, None]], Dict] ] = None, es_filter: Optional[List[Dict]] = None, + fields: List[str] = [], **kwargs: Any, ) -> VectorStoreQueryResult: """ @@ -629,15 +631,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 = [] - print("hits:", len(hits)) 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"] @@ -645,6 +647,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( @@ -662,6 +665,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, @@ -682,5 +686,6 @@ class SyncElasticsearchStore(BasePydanticVectorStore): 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 )