From c8bbb54de19aef90fe231a9a395b61cf400e20f8 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Fri, 6 Feb 2026 17:45:59 +0800 Subject: [PATCH] feat(utils): add cosine similarity functions and integrate into vector operations --- reme/core/schema/memory_node.py | 4 + reme/core/utils/__init__.py | 4 +- reme/core/utils/common_utils.py | 52 +++++++++++++ reme/core/vector_store/local_vector_store.py | 18 +---- reme/tool/memory/history/read_history_v2.py | 23 +----- reme/tool/memory/vector/memory_handler.py | 80 ++++++++++++++++++-- reme/tool/memory/vector/retrieve_memory.py | 16 +++- tests/test_vector_store.py | 8 +- 8 files changed, 153 insertions(+), 52 deletions(-) diff --git a/reme/core/schema/memory_node.py b/reme/core/schema/memory_node.py index e6fde012..bc1ef672 100644 --- a/reme/core/schema/memory_node.py +++ b/reme/core/schema/memory_node.py @@ -42,6 +42,7 @@ class MemoryNode(BaseModel): time_modified: Last modification timestamp. author: Author or source of this memory. score: Relevance or importance score. + vector: Vector embedding of the memory content. metadata: Additional metadata for extensibility. """ @@ -58,6 +59,7 @@ class MemoryNode(BaseModel): author: str = Field(default="", description="Author or source of the memory") score: float = Field(default=0, description="Relevance or importance score") + vector: list[float] | None = Field(default=None, description="Vector embedding of the memory content") metadata: dict[str, Any] = Field(default_factory=dict, description="Additional metadata") def _update_modified_time(self) -> "MemoryNode": @@ -145,6 +147,7 @@ class MemoryNode(BaseModel): return VectorNode( vector_id=self.memory_id, content=vector_content, + vector=self.vector, metadata=metadata, ) @@ -226,5 +229,6 @@ class MemoryNode(BaseModel): time_modified=metadata.pop("time_modified", ""), author=metadata.pop("author", ""), score=metadata.pop("score", 0), + vector=node.vector, metadata=metadata, ) diff --git a/reme/core/utils/__init__.py b/reme/core/utils/__init__.py index 642238c0..6803a725 100644 --- a/reme/core/utils/__init__.py +++ b/reme/core/utils/__init__.py @@ -3,7 +3,7 @@ from .cache_handler import CacheHandler from .case_converter import snake_to_camel, camel_to_snake from .chunking_utils import chunk_markdown -from .common_utils import run_coro_safely, execute_stream_task, hash_text +from .common_utils import run_coro_safely, execute_stream_task, hash_text, cosine_similarity, batch_cosine_similarity from .env_utils import load_env from .execute_utils import exec_code, run_shell_command from .http_client import HttpClient @@ -24,6 +24,8 @@ __all__ = [ "run_coro_safely", "execute_stream_task", "hash_text", + "cosine_similarity", + "batch_cosine_similarity", "load_env", "exec_code", "run_shell_command", diff --git a/reme/core/utils/common_utils.py b/reme/core/utils/common_utils.py index bf35d3db..db12b43c 100644 --- a/reme/core/utils/common_utils.py +++ b/reme/core/utils/common_utils.py @@ -5,6 +5,7 @@ import hashlib from collections.abc import AsyncGenerator, Coroutine from typing import Any +import numpy as np from loguru import logger from ..enumeration import ChunkEnum @@ -94,3 +95,54 @@ def hash_text(text: str) -> str: Hexadecimal representation of the SHA-256 hash """ return hashlib.sha256(text.encode("utf-8")).hexdigest() + + +def cosine_similarity(vec1: list[float], vec2: list[float]) -> float: + """Calculate the cosine similarity between two numeric vectors.""" + if len(vec1) != len(vec2): + raise ValueError(f"Vectors must have same length: {len(vec1)} != {len(vec2)}") + + dot_product = sum(a * b for a, b in zip(vec1, vec2)) + magnitude1 = sum(a * a for a in vec1) ** 0.5 + magnitude2 = sum(b * b for b in vec2) ** 0.5 + + if magnitude1 == 0 or magnitude2 == 0: + return 0.0 + + return dot_product / (magnitude1 * magnitude2) + + +def batch_cosine_similarity(nd_array1: np.ndarray, nd_array2: np.ndarray) -> np.ndarray: + """Calculate cosine similarity matrix between two batches of vectors. + + Args: + nd_array1: Matrix of shape (batch_size1, emb_size) + nd_array2: Matrix of shape (batch_size2, emb_size) + + Returns: + Similarity matrix of shape (batch_size1, batch_size2) where + result[i, j] is the cosine similarity between nd_array1[i] and nd_array2[j] + + Raises: + ValueError: If embedding dimensions don't match + """ + if nd_array1.shape[1] != nd_array2.shape[1]: + raise ValueError(f"Embedding dimensions must match: {nd_array1.shape[1]} != {nd_array2.shape[1]}") + + # Compute dot products: (batch_size1, emb_size) @ (emb_size, batch_size2) + # Result shape: (batch_size1, batch_size2) + dot_products = np.dot(nd_array1, nd_array2.T) + + # Compute L2 norms for each vector + norms1 = np.linalg.norm(nd_array1, axis=1) # Shape: (batch_size1,) + norms2 = np.linalg.norm(nd_array2, axis=1) # Shape: (batch_size2,) + + # Compute outer product of norms: (batch_size1, 1) @ (1, batch_size2) + # Result shape: (batch_size1, batch_size2) + norm_products = np.outer(norms1, norms2) + + # Avoid division by zero + norm_products = np.where(norm_products == 0, 1e-10, norm_products) + + # Compute cosine similarities + return dot_products / norm_products diff --git a/reme/core/vector_store/local_vector_store.py b/reme/core/vector_store/local_vector_store.py index c117ddf0..e3af5a49 100644 --- a/reme/core/vector_store/local_vector_store.py +++ b/reme/core/vector_store/local_vector_store.py @@ -9,6 +9,7 @@ from loguru import logger from .base_vector_store import BaseVectorStore from ..embedding import BaseEmbeddingModel from ..schema import VectorNode +from ..utils import cosine_similarity class LocalVectorStore(BaseVectorStore): @@ -79,21 +80,6 @@ class LocalVectorStore(BaseVectorStore): return nodes - @staticmethod - def _cosine_similarity(vec1: list[float], vec2: list[float]) -> float: - """Calculate the cosine similarity between two numeric vectors.""" - if len(vec1) != len(vec2): - raise ValueError(f"Vectors must have same length: {len(vec1)} != {len(vec2)}") - - dot_product = sum(a * b for a, b in zip(vec1, vec2)) - magnitude1 = sum(a * a for a in vec1) ** 0.5 - magnitude2 = sum(b * b for b in vec2) ** 0.5 - - if magnitude1 == 0 or magnitude2 == 0: - return 0.0 - - return dot_product / (magnitude1 * magnitude2) - @staticmethod def _match_filters(node: VectorNode, filters: dict | None) -> bool: """Check if a vector node matches the provided metadata filters. @@ -208,7 +194,7 @@ class LocalVectorStore(BaseVectorStore): continue try: - score = self._cosine_similarity(query_vector, node.vector) + score = cosine_similarity(query_vector, node.vector) scored_nodes.append((node, score)) except ValueError as e: logger.warning(f"Failed to calculate similarity for node {node.vector_id}: {e}") diff --git a/reme/tool/memory/history/read_history_v2.py b/reme/tool/memory/history/read_history_v2.py index d0499727..7ca267e6 100644 --- a/reme/tool/memory/history/read_history_v2.py +++ b/reme/tool/memory/history/read_history_v2.py @@ -6,7 +6,7 @@ from loguru import logger from ..base_memory_tool import BaseMemoryTool from ....core.schema import MemoryNode, ToolCall, Message -from ....core.utils import format_messages +from ....core.utils import format_messages, cosine_similarity class ReadHistoryV2(BaseMemoryTool): @@ -103,7 +103,7 @@ class ReadHistoryV2(BaseMemoryTool): for block in message_blocks: block_text = format_messages(block, add_index=False) block_embedding = await self.embedding_model.get_embedding(block_text) - similarity = self._calculate_cosine_similarity(query_embedding, block_embedding) + similarity = cosine_similarity(query_embedding, block_embedding) block_similarities.append((similarity, block_text)) block_similarities.sort(key=lambda x: x[0], reverse=True) @@ -121,22 +121,3 @@ class ReadHistoryV2(BaseMemoryTool): history_ids = [item["history_id"] for item in history_items] logger.info(f"Successfully read {len(all_results)} history result(s): {history_ids}") return output - - @staticmethod - def _calculate_cosine_similarity(vec1: list[float], vec2: list[float]) -> float: - """Calculate cosine similarity between two vectors""" - if len(vec1) != len(vec2): - raise ValueError(f"Vectors must have same length: {len(vec1)} != {len(vec2)}") - - try: - dot_product = sum(a * b for a, b in zip(vec1, vec2)) - magnitude1 = sum(a * a for a in vec1) ** 0.5 - magnitude2 = sum(b * b for b in vec2) ** 0.5 - - if magnitude1 == 0 or magnitude2 == 0: - return 0.0 - - return dot_product / (magnitude1 * magnitude2) - except Exception as e: - logger.error(f"Error calculating cosine similarity: {e}") - return 0.0 diff --git a/reme/tool/memory/vector/memory_handler.py b/reme/tool/memory/vector/memory_handler.py index 0fd04c21..2b304df7 100644 --- a/reme/tool/memory/vector/memory_handler.py +++ b/reme/tool/memory/vector/memory_handler.py @@ -1,8 +1,12 @@ """Memory handler""" +import numpy as np +from loguru import logger + from ....core.context import ServiceContext from ....core.enumeration import MemoryType from ....core.schema import MemoryNode +from ....core.utils.common_utils import batch_cosine_similarity from ....core.vector_store import BaseVectorStore @@ -195,17 +199,77 @@ class MemoryHandler: return list(seen_ids.values()) - async def batch_search(self, searches: list[dict]) -> list[MemoryNode]: + async def batch_search(self, searches: list[dict], hybrid_threshold: float = None) -> list[MemoryNode]: """Execute multiple search queries in batch and return deduplicated results.""" - seen_ids: dict[str, MemoryNode] = {} + if hybrid_threshold is not None: + # Extract query list from searches + query_list: list[str] = [search["query"] for search in searches] - for search_params in searches: - search_result = await self.search(**search_params) - for memory_node in search_result: - if memory_node.memory_id not in seen_ids: - seen_ids[memory_node.memory_id] = memory_node + # Step 1: Get embeddings for all queries using the embedding model + # Shape: [query_size X emb_size] + embedding_model = self.vector_store.embedding_model + query_embeddings_list: list[list[float]] = await embedding_model.get_embeddings(query_list) + query_embeddings = np.array(query_embeddings_list) # Convert to numpy array - return list(seen_ids.values()) + # Step 2: Use self.search to get search results for each query and deduplicate + seen_ids: dict[str, MemoryNode] = {} + for search_params in searches: + search_result = await self.search(**search_params) + for memory_node in search_result: + if memory_node.memory_id not in seen_ids: + seen_ids[memory_node.memory_id] = memory_node + + # Step 3: Get deduplicated results + deduplicated_results = list(seen_ids.values()) + + # If no results, return empty list + if not deduplicated_results: + return [] + + # Step 4: Extract embeddings from results + # Shape: [result_size X emb_size] + result_embeddings_list = [node.vector for node in deduplicated_results if node.vector] + + # Filter out nodes without embeddings + results_with_embeddings = [node for node in deduplicated_results if node.vector] + + if not result_embeddings_list: + logger.warning("No results with embeddings found") + return deduplicated_results + + result_embeddings = np.array(result_embeddings_list) + + # Step 5: Compute cosine similarity matrix + # Shape: [query_size X result_size] + similarity_matrix = batch_cosine_similarity(query_embeddings, result_embeddings) + + # Step 6: Calculate average score for each result across all queries + # Shape: [result_size] + avg_scores = np.mean(similarity_matrix, axis=0) + + # Step 7: Filter results by hybrid_threshold and sort by average score + filtered_results = [] + for idx, node in enumerate(results_with_embeddings): + if avg_scores[idx] >= hybrid_threshold: + node.score = float(avg_scores[idx]) + filtered_results.append(node) + + # Sort by score in descending order + filtered_results.sort(key=lambda x: x.score, reverse=True) + + return filtered_results + + else: + # Original behavior: simple deduplication without hybrid scoring + seen_ids: dict[str, MemoryNode] = {} + + for search_params in searches: + search_result = await self.search(**search_params) + for memory_node in search_result: + if memory_node.memory_id not in seen_ids: + seen_ids[memory_node.memory_id] = memory_node + + return list(seen_ids.values()) async def list( self, diff --git a/reme/tool/memory/vector/retrieve_memory.py b/reme/tool/memory/vector/retrieve_memory.py index af19bc1e..4fd30c9d 100644 --- a/reme/tool/memory/vector/retrieve_memory.py +++ b/reme/tool/memory/vector/retrieve_memory.py @@ -11,11 +11,19 @@ from ....core.utils import deduplicate_memories class RetrieveMemory(BaseMemoryTool): """Tool to retrieve memories using similarity search""" - def __init__(self, top_k: int = 20, enable_memory_target: bool = False, enable_time_filter: bool = False, **kwargs): + def __init__( + self, + top_k: int = 20, + enable_memory_target: bool = False, + enable_time_filter: bool = False, + hybrid_threshold: float | None = None, + **kwargs, + ): super().__init__(**kwargs) self.top_k: int = top_k self.enable_memory_target: bool = enable_memory_target self.enable_time_filter: bool = enable_time_filter + self.hybrid_threshold: float | None = hybrid_threshold def _build_query_parameters(self) -> dict: """Build the query parameters schema based on enabled features.""" @@ -111,7 +119,11 @@ class RetrieveMemory(BaseMemoryTool): memory_nodes: list[MemoryNode] = [] for target, searches in queries_by_target.items(): handler = MemoryHandler(target, self.service_context) - nodes = await handler.batch_search(searches) + if self.hybrid_threshold is not None: + nodes = await handler.batch_search(searches, self.hybrid_threshold) + nodes = nodes[: self.top_k] + else: + nodes = await handler.batch_search(searches) memory_nodes.extend(nodes) memory_nodes = deduplicate_memories(memory_nodes) diff --git a/tests/test_vector_store.py b/tests/test_vector_store.py index 14ed1c1d..a5c2a849 100644 --- a/tests/test_vector_store.py +++ b/tests/test_vector_store.py @@ -26,7 +26,7 @@ from loguru import logger from reme.core.embedding import OpenAIEmbeddingModel from reme.core.schema import VectorNode -from reme.core.utils import load_env +from reme.core.utils import load_env, cosine_similarity from reme.core.vector_store import ( BaseVectorStore, ChromaVectorStore, @@ -657,19 +657,19 @@ async def test_cosine_similarity(store_name: str): vec3 = [1.0, 0.0, 0.0] # Test perpendicular vectors (similarity = 0) - sim1 = LocalVectorStore._cosine_similarity(vec1, vec2) # pylint: disable=protected-access + sim1 = cosine_similarity(vec1, vec2) # pylint: disable=protected-access logger.info(f"Similarity between perpendicular vectors: {sim1:.4f}") assert abs(sim1) < 0.0001, "Perpendicular vectors should have similarity close to 0" # Test identical vectors (similarity = 1) - sim2 = LocalVectorStore._cosine_similarity(vec1, vec3) # pylint: disable=protected-access + sim2 = cosine_similarity(vec1, vec3) # pylint: disable=protected-access logger.info(f"Similarity between identical vectors: {sim2:.4f}") assert abs(sim2 - 1.0) < 0.0001, "Identical vectors should have similarity close to 1" # Test with real-world like vectors vec4 = [0.5, 0.5, 0.5] vec5 = [0.6, 0.4, 0.5] - sim3 = LocalVectorStore._cosine_similarity(vec4, vec5) # pylint: disable=protected-access + sim3 = cosine_similarity(vec4, vec5) # pylint: disable=protected-access logger.info(f"Similarity between similar vectors: {sim3:.4f}") assert sim3 > 0.9, "Similar vectors should have high similarity"