feat(utils): add cosine similarity functions and integrate into vector operations

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,
         )

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",

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

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}")

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

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,

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)

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"