From 1dd81f9c250393b58b49463ebbc5aa30cc8f9d19 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Fri, 6 Feb 2026 15:02:41 +0800 Subject: [PATCH 1/7] refactor(core): update config parsing and memory management system --- .gitignore | 3 +- benchmark/halumem/eval_reme.py | 2 +- .../personal/personal_v1_retriever.yaml | 8 +-- .../personal/personal_v1_summarizer.yaml | 60 ++++++------------- reme/config/default.yaml | 2 +- reme/core/context/service_context.py | 4 +- reme/core/utils/pydantic_config_parser.py | 18 ++---- reme/reme.py | 4 +- reme/reme_fs.py | 41 +++++++++++++ 9 files changed, 73 insertions(+), 69 deletions(-) create mode 100644 reme/reme_fs.py diff --git a/.gitignore b/.gitignore index 07fbdaca..b20570b5 100644 --- a/.gitignore +++ b/.gitignore @@ -41,4 +41,5 @@ meta_memory/* *.sqlite3 **/data/*.json *.db -memories/* \ No newline at end of file +memories/* +.reme/* \ No newline at end of file diff --git a/benchmark/halumem/eval_reme.py b/benchmark/halumem/eval_reme.py index b4fc42de..89ee32ad 100644 --- a/benchmark/halumem/eval_reme.py +++ b/benchmark/halumem/eval_reme.py @@ -220,7 +220,7 @@ async def answer_question_with_memories( result = await reme.llm.simple_request_for_json( prompt=prompt, - model_name="qwen-flash" + model_name=model_name, ) return result diff --git a/reme/agent/memory/personal/personal_v1_retriever.yaml b/reme/agent/memory/personal/personal_v1_retriever.yaml index 54ece8d0..99df1374 100644 --- a/reme/agent/memory/personal/personal_v1_retriever.yaml +++ b/reme/agent/memory/personal/personal_v1_retriever.yaml @@ -20,7 +20,6 @@ user_message: | * Entity-focused queries (extract and search specific names, places, events) * Keyword-based searches (core concepts, topics) * Related context queries (broader themes) - - Review all results before proceeding to next phase ### Phase 2(Optional): Temporal Search **Tool**: `retrieve_memory` (with time filter) @@ -32,8 +31,7 @@ user_message: | - After date: `20200101,99999999` (from 20200101 onwards) **Approach**: - Identify temporal constraints from the user question - - Refine Phase 1 queries with appropriate time filters - - Try multiple time ranges if initial searches yield no results + - Refine Phase 1 queries with 3-5 diverse appropriate different time filters ### Phase 3: Deep Dive into History **Tool**: `read_history` @@ -57,6 +55,4 @@ user_message: | - If you find sufficient information to answer the user's question, you may output directly without exhausting all search phases - Exhaust all search strategies before concluding information doesn't exist - ### Output any tangentially related findings, Format: - [timestamp] [memory/profile/history] [relevant content1] - [timestamp] [memory/profile/history] [relevant content2] + Output a summary of all retrieved memories, user profile, and history data. diff --git a/reme/agent/memory/personal/personal_v1_summarizer.yaml b/reme/agent/memory/personal/personal_v1_summarizer.yaml index 5721e64f..83befa96 100644 --- a/reme/agent/memory/personal/personal_v1_summarizer.yaml +++ b/reme/agent/memory/personal/personal_v1_summarizer.yaml @@ -15,65 +15,41 @@ user_message_s1: | - Extract all important information comprehensively—do not miss critical details, but avoid any fabrications or unfounded assumptions - The tool will retrieve similar historical memories via vector search to help you in Step 2 - ### Step 2: Update and Add Memories - Review each memory draft from Step 1 and compare it with the retrieved historical memories, then use `update_memory` to manage all memories in one call: + ### Step 2: Add New Memories + Review each memory draft from Step 1 and compare it with the retrieved historical memories, then use `add_memory` to add new memories: - **For memories_to_update** (updating existing memories): - - For each memory to update, fill in the required parameters: - * `memory_id`: ID of the historical memory to update (from retrieved memories in Step 1) - * `message_time`: timestamp from the conversation (e.g., '2020-01-01 00:00:00') - * `memory_content`: updated or consolidated memory content - - Update memories when: - * The draft contains additional information that should be merged with existing memories - * Historical memories need to be corrected or refined based on new information - - **For memories_to_add** (adding new memories): - - For each new memory, fill in the required parameters: - * `message_time`: timestamp from the conversation (e.g., '2020-01-01 00:00:00') - * `memory_content`: memory content - - Add memories when: - * The draft contains completely new information not present in historical memories - * The information cannot be merged into any existing memory - - **General Guidelines:** + **Parameters for each memory:** + - `message_time`: timestamp from the conversation (e.g., '2020-01-01 00:00:00') + - `memory_content`: memory content + **When to skip:** - **Skip** drafts if their content is already fully covered by historical memories (avoid redundancy) - - You can update and add memories in a single `update_memory` tool call user_message_s2: | - You are a Profile Agent responsible for managing profiles about {memory_target}. + You are a User Profile Agent responsible for managing user profiles about {memory_target}. ## Latest Conversation Format: round [] : {context} - ## Current Profiles + ## Current User Profiles {profiles} ## Task - Analyze the Latest Conversation and use `update_profiles` to manage profiles (both updates and additions in one call): + Analyze the Latest Conversation and use `update_profiles` to manage user profiles (both updates and additions in one call): - **For profiles_to_update** (updating existing profiles): + **For profiles_to_update** (updating existing user profiles): - For each profile to update, fill in the required parameters: - * `profile_id`: ID of the profile to update (from Current Profiles) + * `profile_id`: ID of the profile to update (from Current User Profiles) * `message_time`: timestamp from the conversation (e.g., '2020-01-01 00:00:00') - * `profile_key`: profile key or category (e.g., 'name', 'age', 'occupation') - * `profile_value`: updated profile value (e.g., 'John Smith') - - Update profiles when: - * Information in the conversation conflicts with or supersedes existing profiles - * Profiles need to be consolidated or merged with new information - * Existing profile values need to be corrected or refined + * `profile_key`: key (e.g., 'name', 'age', 'occupation') + * `profile_value`: value (e.g., 'John Smith') - **For profiles_to_add** (adding new profiles): + **For profiles_to_add** (adding new user profiles): - For each new profile, fill in the required parameters: * `message_time`: timestamp from the conversation (e.g., '2020-01-01 00:00:00') - * `profile_key`: profile key or category (e.g., 'name', 'age', 'occupation') - * `profile_value`: profile value (e.g., 'John Smith') - - Add profiles when: - * The information represents a new distinct profile not present in Current Profiles - * The profile key doesn't exist in Current Profiles - * The information cannot be merged into existing profiles + * `profile_key`: key (e.g., 'name', 'age', 'occupation') + * `profile_value`: value (e.g., 'John Smith') **General Guidelines:** - - Use actual names from the conversation (e.g., "Bob") instead of generic references (e.g., "user") - - Extract all important information comprehensively—do not miss critical details, but avoid any fabrications or unfounded assumptions - - You can update and add profiles in a single tool call + - Avoid any fabrications or unfounded assumptions + - You can update and add user profiles in a single tool call diff --git a/reme/config/default.yaml b/reme/config/default.yaml index b7ed53d4..561f931a 100644 --- a/reme/config/default.yaml +++ b/reme/config/default.yaml @@ -22,7 +22,7 @@ llm: backend: openai model_name: qwen3-30b-a3b-instruct-2507 request_interval: 1 - temperature: 0.0001 +# temperature: 0.0001 qwen3_max_instruct: backend: openai diff --git a/reme/core/context/service_context.py b/reme/core/context/service_context.py index 05a45355..045858f7 100644 --- a/reme/core/context/service_context.py +++ b/reme/core/context/service_context.py @@ -118,9 +118,7 @@ class ServiceContext(BaseContext): input_args.append(f"config={config_path}") if args: input_args.extend(args) - if kwargs: - input_args.extend([f"{k}={v}" for k, v in kwargs.items()]) - service_config = parser.parse_args(*input_args) + service_config = parser.parse_args(*input_args, **kwargs) service_config.enable_logo = enable_logo if llm: diff --git a/reme/core/utils/pydantic_config_parser.py b/reme/core/utils/pydantic_config_parser.py index 2f4f2ae0..1fd0ab56 100644 --- a/reme/core/utils/pydantic_config_parser.py +++ b/reme/core/utils/pydantic_config_parser.py @@ -145,19 +145,8 @@ class PydanticConfigParser: raise FileNotFoundError(f"config={config_path} not found") return config_path - def parse_args(self, *args: str) -> T: - """Parse CLI arguments and load configs from YAML files. - - Args: - *args: CLI arguments in format "key=value" or "config=file.yaml". - - Returns: - Validated Pydantic config instance. - - Raises: - ValueError: If no config file is specified. - FileNotFoundError: If specified config file does not exist. - """ + def parse_args(self, *args: str, **kwargs) -> T: + """Parse CLI arguments and load configs from YAML files.""" configs_to_merge = [self.config_class().model_dump()] # Separate config file path from other arguments @@ -184,6 +173,9 @@ class PydanticConfigParser: if filter_args: configs_to_merge.append(self.parse_dot_notation(filter_args)) + if kwargs: + configs_to_merge.append(kwargs) + # Merge all configs and validate self.config_dict = self.merge_configs(*configs_to_merge) return self.config_class.model_validate(self.config_dict) diff --git a/reme/reme.py b/reme/reme.py index bd3e2c64..94a9b196 100644 --- a/reme/reme.py +++ b/reme/reme.py @@ -36,7 +36,7 @@ from .tool.memory import ( AddHistory, ReadAllProfiles, UpdateProfilesV1, - UpdateMemoryV1, + AddMemory, ) @@ -222,7 +222,7 @@ class ReMe(Application): enable_when_to_use=False, enable_multiple=True, ), - UpdateMemoryV1( + AddMemory( enable_thinking_params=enable_thinking_params, enable_memory_target=False, enable_when_to_use=False, diff --git a/reme/reme_fs.py b/reme/reme_fs.py new file mode 100644 index 00000000..f5ef42d2 --- /dev/null +++ b/reme/reme_fs.py @@ -0,0 +1,41 @@ +"""ReMe File System""" + +from .config import ReMeConfigParser +from .core import Application + + +class ReMeFs(Application): + """ReMe File System""" + + def __init__( + self, + *args, + llm_api_key: str | None = None, + llm_api_base: str | None = None, + embedding_api_key: str | None = None, + embedding_api_base: str | None = None, + enable_logo: bool = True, + llm: dict | None = None, + embedding_model: dict | None = None, + vector_store: dict | None = None, + token_counter: dict | None = None, + working_dir: str = "./agent", + **kwargs, + ): + """Initialize ReMe with config.""" + super().__init__( + *args, + llm_api_key=llm_api_key, + llm_api_base=llm_api_base, + embedding_api_key=embedding_api_key, + embedding_api_base=embedding_api_base, + enable_logo=enable_logo, + parser=ReMeConfigParser, + llm=llm, + embedding_model=embedding_model, + vector_store=vector_store, + token_counter=token_counter, + **kwargs, + ) + + self.working_dir: str = working_dir From 9e50e71b52225001c011aca76cda5a09d40124d1 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Fri, 6 Feb 2026 15:35:09 +0800 Subject: [PATCH 2/7] feat(core): add memory store and file watcher support to application --- reme/core/application.py | 20 ++++++++++++++ reme/core/context/service_context.py | 29 +++++++++++++++++++- reme/core/file_watcher/base_file_watcher.py | 16 ++++++++--- reme/core/file_watcher/delta_file_watcher.py | 5 +--- reme/core/file_watcher/full_file_watcher.py | 4 +-- reme/core/schema/__init__.py | 2 -- reme/core/schema/memory_index_meta.py | 14 ---------- reme/core/schema/service_config.py | 17 ++++++++++++ 8 files changed, 79 insertions(+), 28 deletions(-) delete mode 100644 reme/core/schema/memory_index_meta.py diff --git a/reme/core/application.py b/reme/core/application.py index b6f1da43..ff093692 100644 --- a/reme/core/application.py +++ b/reme/core/application.py @@ -4,8 +4,10 @@ import asyncio from .context import PromptHandler, ServiceContext from .embedding import BaseEmbeddingModel +from .file_watcher import BaseFileWatcher from .flow import BaseFlow from .llm import BaseLLM +from .memory_storage import BaseMemoryStore from .schema import Response from .token_counter import BaseTokenCounter from .utils import execute_stream_task, PydanticConfigParser @@ -27,7 +29,9 @@ class Application: llm: dict | None = None, embedding_model: dict | None = None, vector_store: dict | None = None, + memory_store: dict | None = None, token_counter: dict | None = None, + file_watcher: dict | None = None, **kwargs, ): self.service_context = ServiceContext( @@ -43,7 +47,9 @@ class Application: llm=llm, embedding_model=embedding_model, vector_store=vector_store, + memory_store=memory_store, token_counter=token_counter, + file_watcher=file_watcher, **kwargs, ) self.prompt_handler = PromptHandler(language=self.service_context.language) @@ -62,7 +68,9 @@ class Application: llm: dict | None = None, embedding_model: dict | None = None, vector_store: dict | None = None, + memory_store: dict | None = None, token_counter: dict | None = None, + file_watcher: dict | None = None, **kwargs, ) -> "Application": """Create and start an Application instance asynchronously.""" @@ -77,7 +85,9 @@ class Application: llm=llm, embedding_model=embedding_model, vector_store=vector_store, + memory_store=memory_store, token_counter=token_counter, + file_watcher=file_watcher, **kwargs, ) await instance.start() @@ -145,6 +155,16 @@ class Application: """Get the default vector store instance.""" return self.service_context.vector_stores.get("default") + @property + def memory_store(self) -> BaseMemoryStore: + """Get the default memory store instance.""" + return self.service_context.memory_stores.get("default") + + @property + def file_watcher(self) -> BaseFileWatcher: + """Get the default file watcher instance.""" + return self.service_context.file_watchers.get("default") + @property def token_counter(self) -> BaseTokenCounter: """Get the default token counter instance.""" diff --git a/reme/core/context/service_context.py b/reme/core/context/service_context.py index 045858f7..a74eac40 100644 --- a/reme/core/context/service_context.py +++ b/reme/core/context/service_context.py @@ -15,10 +15,11 @@ if TYPE_CHECKING: from ..llm import BaseLLM from ..embedding import BaseEmbeddingModel from ..vector_store import BaseVectorStore + from ..memory_storage import BaseMemoryStore from ..token_counter import BaseTokenCounter from ..flow import BaseFlow from ..service import BaseService - from ..memory_storage import BaseMemoryStore + from ..file_watcher import BaseFileWatcher class ServiceContext(BaseContext): @@ -38,7 +39,9 @@ class ServiceContext(BaseContext): llm: dict | None = None, embedding_model: dict | None = None, vector_store: dict | None = None, + memory_store: dict | None = None, token_counter: dict | None = None, + file_watcher: dict | None = None, **kwargs, ): super().__init__() @@ -55,7 +58,9 @@ class ServiceContext(BaseContext): llm=llm, embedding_model=embedding_model, vector_store=vector_store, + memory_store=memory_store, token_counter=token_counter, + file_watcher=file_watcher, **kwargs, ) @@ -79,6 +84,7 @@ class ServiceContext(BaseContext): self.token_counters: dict[str, "BaseTokenCounter"] = {} self.vector_stores: dict[str, "BaseVectorStore"] = {} self.memory_stores: dict[str, "BaseMemoryStore"] = {} + self.file_watchers: dict[str, "BaseFileWatcher"] = {} self.flows: dict[str, "BaseFlow"] = {} self.mcp_server_mapping: dict[str, dict] = {} @@ -100,7 +106,9 @@ class ServiceContext(BaseContext): llm: dict | None = None, embedding_model: dict | None = None, vector_store: dict | None = None, + memory_store: dict | None = None, token_counter: dict | None = None, + file_watcher: dict | None = None, **kwargs, ) -> ServiceConfig: @@ -129,6 +137,10 @@ class ServiceContext(BaseContext): self._update_section_config(service_config, "token_counter", **token_counter) if vector_store: self._update_section_config(service_config, "vector_store", **vector_store) + if memory_store: + self._update_section_config(service_config, "memory_store", **memory_store) + if file_watcher: + self._update_section_config(service_config, "file_watcher", **file_watcher) return service_config @staticmethod @@ -198,10 +210,25 @@ class ServiceContext(BaseContext): self.memory_stores[name] = R.memory_store[config.backend]( store_name=config.store_name, embedding_model=self.embedding_models[config.embedding_model], + fts_enabled=config.fts_enabled, + snippet_max_chars=config.snippet_max_chars, **config.model_extra, ) await self.memory_stores[name].start() + for name, config in self.service_config.file_watcher.items(): + self.file_watchers[name] = R.file_watcher[config.backend]( + watch_paths=config.watch_paths, + suffix_filters=config.suffix_filters, + recursive=config.recursive, + debounce=config.debounce, + chunk_tokens=config.chunk_tokens, + chunk_overlap=config.chunk_overlap, + memory_store=self.memory_stores[config.memory_store], + **config.model_extra, + ) + await self.file_watchers[name].start() + if self.service_config.mcp_servers: await self.prepare_mcp_servers() diff --git a/reme/core/file_watcher/base_file_watcher.py b/reme/core/file_watcher/base_file_watcher.py index 8df6b9b4..620d2e17 100644 --- a/reme/core/file_watcher/base_file_watcher.py +++ b/reme/core/file_watcher/base_file_watcher.py @@ -25,21 +25,25 @@ class BaseFileWatcher: def __init__( self, watch_paths: list[str] | str, + suffix_filters: list[str] | None = None, recursive: bool = False, debounce: int = 500, # Millisecond debounce - suffix_filters: list[str] | None = None, - callback: Callable[[set[tuple[Change, str]]], None | Coroutine[Any, Any, None]] | None = None, + chunk_tokens: int = 400, + chunk_overlap: int = 80, memory_store: BaseMemoryStore | None = None, + callback: Callable[[set[tuple[Change, str]]], None | Coroutine[Any, Any, None]] | None = None, **kwargs, ): """ Initialize the file watcher""" self.watch_paths: list[str] = [watch_paths] if isinstance(watch_paths, str) else watch_paths + self.suffix_filters: list[str] = suffix_filters or [] self.recursive: bool = recursive self.debounce: int = debounce - self.suffix_filters: list[str] = suffix_filters or [] - self.callback = callback + self.chunk_tokens: int = chunk_tokens + self.chunk_overlap: int = chunk_overlap self.memory_store: BaseMemoryStore = memory_store + self.callback = callback self.kwargs: dict = kwargs self._stop_event = asyncio.Event() @@ -81,6 +85,10 @@ class BaseFileWatcher: async def _watch_loop(self): """Core monitoring loop""" + if not self.watch_paths: + logger.warning("No watch paths specified") + return + async for changes in awatch( *self.watch_paths, watch_filter=self.watch_filter, diff --git a/reme/core/file_watcher/delta_file_watcher.py b/reme/core/file_watcher/delta_file_watcher.py index 60807133..0c9cae7c 100644 --- a/reme/core/file_watcher/delta_file_watcher.py +++ b/reme/core/file_watcher/delta_file_watcher.py @@ -29,7 +29,7 @@ class DeltaFileWatcher(BaseFileWatcher): - Delete affected old chunks and insert new chunks """ - def __init__(self, chunk_tokens: int = 400, chunk_overlap: int = 80, overlap_lines: int = 2, **kwargs): + def __init__(self, overlap_lines: int = 2, **kwargs): """ Initialize delta file watcher. @@ -38,10 +38,7 @@ class DeltaFileWatcher(BaseFileWatcher): chunk_overlap: Overlap tokens between chunks """ super().__init__(**kwargs) - self.chunk_tokens = chunk_tokens - self.chunk_overlap = chunk_overlap self.overlap_lines = overlap_lines - self.dirty = False @staticmethod diff --git a/reme/core/file_watcher/full_file_watcher.py b/reme/core/file_watcher/full_file_watcher.py index 69870c2c..c3c08406 100644 --- a/reme/core/file_watcher/full_file_watcher.py +++ b/reme/core/file_watcher/full_file_watcher.py @@ -19,12 +19,10 @@ from ..utils import chunk_markdown, hash_text class FullFileWatcher(BaseFileWatcher): """Full file watcher implementation for full synchronization""" - def __init__(self, chunk_tokens: int = 400, chunk_overlap: int = 80, **kwargs): + def __init__(self, **kwargs): """ Initialize full file watcher""" super().__init__(**kwargs) - self.chunk_tokens = chunk_tokens - self.chunk_overlap = chunk_overlap self.dirty = False @staticmethod diff --git a/reme/core/schema/__init__.py b/reme/core/schema/__init__.py index a4a4214e..d4c5c051 100644 --- a/reme/core/schema/__init__.py +++ b/reme/core/schema/__init__.py @@ -2,7 +2,6 @@ from .file_metadata import FileMetadata from .memory_chunk import MemoryChunk -from .memory_index_meta import MemoryIndexMeta from .memory_node import MemoryNode from .memory_search_result import MemorySearchResult from .message import ContentBlock, Message, Trajectory @@ -34,7 +33,6 @@ __all__ = [ "LLMConfig", "MCPConfig", "MemoryChunk", - "MemoryIndexMeta", "MemoryNode", "MemorySearchResult", "Message", diff --git a/reme/core/schema/memory_index_meta.py b/reme/core/schema/memory_index_meta.py deleted file mode 100644 index a4538b77..00000000 --- a/reme/core/schema/memory_index_meta.py +++ /dev/null @@ -1,14 +0,0 @@ -"""Memory index metadata schema.""" - -from typing import Optional - -from pydantic import BaseModel, Field - - -class MemoryIndexMeta(BaseModel): - """Metadata for memory index configuration.""" - - model: str = Field(..., description="Name of the embedding model") - chunk_tokens: int = Field(..., description="Maximum tokens per chunk") - chunk_overlap: int = Field(..., description="Number of overlapping tokens between chunks") - vector_dims: Optional[int] = Field(default=None, description="Vector embedding dimensions") diff --git a/reme/core/schema/service_config.py b/reme/core/schema/service_config.py index 16a70b3e..cf7da9d6 100644 --- a/reme/core/schema/service_config.py +++ b/reme/core/schema/service_config.py @@ -85,6 +85,8 @@ class MemoryStoreConfig(BaseModel): backend: str = Field(default="sqlite") store_name: str = Field(default="reme") embedding_model: str = Field(default="default") + fts_enabled: bool = Field(default=True) + snippet_max_chars: int = Field(default=700) class TokenCounterConfig(BaseModel): @@ -96,6 +98,20 @@ class TokenCounterConfig(BaseModel): model_name: str = Field(default="") +class FileWatchConfig(BaseModel): + """Configuration for file watch service.""" + + model_config = ConfigDict(extra="allow") + + watch_paths: list[str] = Field(default_factory=list) + suffix_filters: list[str] = Field(default_factory=list) + recursive: bool = Field(default=False) + debounce: int = Field(default=500) + chunk_tokens: int = Field(default=400) + chunk_overlap: int = Field(default=80) + memory_store: str = Field(default="default") + + class ServiceConfig(BaseModel): """Root configuration schema aggregating all service-level settings and components.""" @@ -121,3 +137,4 @@ class ServiceConfig(BaseModel): vector_store: dict[str, VectorStoreConfig] = Field(default_factory=dict) memory_store: dict[str, MemoryStoreConfig] = Field(default_factory=dict) token_counter: dict[str, TokenCounterConfig] = Field(default_factory=dict) + file_watch: dict[str, FileWatchConfig] = Field(default_factory=dict) From b0b800fbd8a8415e64438a663cbf3ad3d748e66f Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Fri, 6 Feb 2026 15:52:09 +0800 Subject: [PATCH 3/7] feat(vector-store): add threshold-based filtering to vector search --- reme/core/vector_store/base_vector_store.py | 17 +- reme/core/vector_store/chroma_vector_store.py | 26 ++- reme/core/vector_store/es_vector_store.py | 34 +++- reme/core/vector_store/local_vector_store.py | 23 ++- reme/core/vector_store/pgvector_store.py | 31 +++- reme/core/vector_store/qdrant_vector_store.py | 25 ++- tests/test_vector_store.py | 175 ++++++++++++++++++ 7 files changed, 304 insertions(+), 27 deletions(-) diff --git a/reme/core/vector_store/base_vector_store.py b/reme/core/vector_store/base_vector_store.py index 62a73a3e..addaa2ea 100644 --- a/reme/core/vector_store/base_vector_store.py +++ b/reme/core/vector_store/base_vector_store.py @@ -73,8 +73,21 @@ class BaseVectorStore(ABC): """Add one or more vector nodes into the current collection.""" @abstractmethod - async def search(self, query: str, limit: int = 5, filters: dict | None = None, **kwargs) -> list[VectorNode]: - """Find the most similar vector nodes based on a text query.""" + async def search( + self, + query: str, + limit: int = 5, + candidates: int | None = None, + filters: dict | None = None, + threshold: float | None = None, + **kwargs, + ) -> list[VectorNode]: + """Find the most similar vector nodes based on a text query. + + When threshold is None, uses default behavior. + When threshold is set, searches max(candidates, limit) nodes, + filters by threshold, then returns top limit results. + """ @abstractmethod async def delete(self, vector_ids: str | list[str], **kwargs) -> None: diff --git a/reme/core/vector_store/chroma_vector_store.py b/reme/core/vector_store/chroma_vector_store.py index 5b0748f5..76b4b81a 100644 --- a/reme/core/vector_store/chroma_vector_store.py +++ b/reme/core/vector_store/chroma_vector_store.py @@ -307,21 +307,35 @@ class ChromaVectorStore(BaseVectorStore): self, query: str, limit: int = 5, + candidates: int | None = None, filters: dict | None = None, + threshold: float | None = None, **kwargs, ) -> list[VectorNode]: - """Search for the most similar vector nodes based on a text query.""" + """Search for the most similar vector nodes based on a text query. + + When threshold is None, uses default behavior. + When threshold is set, searches max(candidates, limit) nodes, + filters by threshold, then returns top limit results. + """ query_vector = await self.get_embedding(query) where_clause = self._generate_where_clause(filters) include_embeddings = kwargs.get("include_embeddings", False) + # When threshold is set, search more candidates + if threshold is not None: + effective_candidates = candidates if candidates is not None else limit * 2 + search_limit = max(effective_candidates, limit) + else: + search_limit = limit + def _search(): include: list = ["documents", "metadatas", "distances"] if include_embeddings: include.append("embeddings") return self.collection.query( query_embeddings=[query_vector], - n_results=limit, + n_results=search_limit, where=where_clause, include=include, ) @@ -329,9 +343,11 @@ class ChromaVectorStore(BaseVectorStore): results = await self._run_sync_in_executor(_search) nodes = self._parse_results(results, include_score=True) - score_threshold = kwargs.get("score_threshold") - if score_threshold is not None: - nodes = [n for n in nodes if n.metadata.get("score", 0) >= score_threshold] + # Apply threshold filtering if specified + if threshold is not None: + nodes = [n for n in nodes if n.metadata.get("score", 0) >= threshold] + nodes = nodes[:limit] + return nodes async def delete(self, vector_ids: str | list[str], **kwargs): diff --git a/reme/core/vector_store/es_vector_store.py b/reme/core/vector_store/es_vector_store.py index 62b28bc1..f9cfad7d 100644 --- a/reme/core/vector_store/es_vector_store.py +++ b/reme/core/vector_store/es_vector_store.py @@ -238,31 +238,48 @@ class ESVectorStore(BaseVectorStore): self, query: str, limit: int = 5, + candidates: int | None = None, filters: dict | None = None, + threshold: float | None = None, **kwargs, ) -> list[VectorNode]: """Perform a kNN similarity search based on a text query. + When threshold is None, uses default behavior. + When threshold is set, searches max(candidates, limit) nodes, + filters by threshold, then returns top limit results. + Args: query: The text to search for. limit: Maximum number of nearest neighbors to return. - filters: Metadata filters for exact match or 'IN' operations. - **kwargs: Search parameters like num_candidates or score_threshold. + candidates: Number of candidates to consider during search. + filters: Metadata filters for exact match or range operations. + threshold: Minimum score threshold for results. + **kwargs: Additional search parameters. Returns: List of VectorNode objects ordered by similarity. """ query_vector = await self.get_embedding(query) - num_candidates = kwargs.get("num_candidates", limit * 2) + + # When threshold is set, search more candidates + if threshold is not None: + effective_candidates = candidates if candidates is not None else limit * 2 + search_candidates = max(effective_candidates, limit) + search_limit = search_candidates + else: + effective_candidates = candidates if candidates is not None else limit * 2 + search_candidates = effective_candidates + search_limit = limit search_query: dict = { "knn": { "field": "vector", "query_vector": query_vector, - "k": limit, - "num_candidates": num_candidates, + "k": search_limit, + "num_candidates": search_candidates, }, - "size": limit, + "size": search_limit, } if filters: @@ -300,6 +317,11 @@ class ESVectorStore(BaseVectorStore): node.metadata["score"] = hit["_score"] results.append(node) + # Apply threshold filtering if specified + if threshold is not None: + results = [n for n in results if n.metadata.get("score", 0) >= threshold] + results = results[:limit] + return results async def delete(self, vector_ids: str | list[str], refresh: bool = True, **kwargs): diff --git a/reme/core/vector_store/local_vector_store.py b/reme/core/vector_store/local_vector_store.py index c117ddf0..22db669e 100644 --- a/reme/core/vector_store/local_vector_store.py +++ b/reme/core/vector_store/local_vector_store.py @@ -193,10 +193,17 @@ class LocalVectorStore(BaseVectorStore): self, query: str, limit: int = 5, + candidates: int | None = None, filters: dict | None = None, + threshold: float | None = None, **kwargs, ) -> list[VectorNode]: - """Search for nodes similar to the query using brute-force cosine similarity.""" + """Search for nodes similar to the query using brute-force cosine similarity. + + When threshold is None, uses default behavior. + When threshold is set, searches max(candidates, limit) nodes, + filters by threshold, then returns top limit results. + """ query_vector = await self.get_embedding(query) all_nodes = self._load_all_nodes() filtered_nodes = [node for node in all_nodes if self._match_filters(node, filters)] @@ -215,11 +222,17 @@ class LocalVectorStore(BaseVectorStore): scored_nodes.sort(key=lambda x: x[1], reverse=True) - score_threshold = kwargs.get("score_threshold") - if score_threshold is not None: - scored_nodes = [(node, score) for node, score in scored_nodes if score >= score_threshold] + # Apply threshold filtering if specified + if threshold is not None: + # When threshold is set, consider more candidates + effective_candidates = candidates if candidates is not None else limit * 2 + search_limit = max(effective_candidates, limit) + scored_nodes = scored_nodes[:search_limit] + scored_nodes = [(node, score) for node, score in scored_nodes if score >= threshold] + scored_nodes = scored_nodes[:limit] + else: + scored_nodes = scored_nodes[:limit] - scored_nodes = scored_nodes[:limit] results = [] for node, score in scored_nodes: node.metadata["score"] = score diff --git a/reme/core/vector_store/pgvector_store.py b/reme/core/vector_store/pgvector_store.py index 0ec7d6a8..dcb49201 100644 --- a/reme/core/vector_store/pgvector_store.py +++ b/reme/core/vector_store/pgvector_store.py @@ -333,10 +333,17 @@ class PGVectorStore(BaseVectorStore): self, query: str, limit: int = 5, + candidates: int | None = None, filters: dict | None = None, + threshold: float | None = None, **kwargs, ) -> list[VectorNode]: - """Perform vector similarity search with optional metadata filtering.""" + """Perform vector similarity search with optional metadata filtering. + + When threshold is None, uses default behavior. + When threshold is set, searches max(candidates, limit) nodes, + filters by threshold, then returns top limit results. + """ await self._ensure_collection_exists() query_vector = await self.get_embedding(query) @@ -351,6 +358,13 @@ class PGVectorStore(BaseVectorStore): new_placeholder = f"${i + 1}" filter_clause = re.sub(rf"\${i}\b", new_placeholder, filter_clause) + # When threshold is set, search more candidates + if threshold is not None: + effective_candidates = candidates if candidates is not None else limit * 2 + search_limit = max(effective_candidates, limit) + else: + search_limit = limit + async with pool.acquire() as conn: sql = f""" SELECT id, content, vector, metadata, vector <=> $1::vector AS distance @@ -359,14 +373,17 @@ class PGVectorStore(BaseVectorStore): ORDER BY distance LIMIT ${len(filter_params) + 2} """ - rows = await conn.fetch(sql, vector_str, *filter_params, limit) + rows = await conn.fetch(sql, vector_str, *filter_params, search_limit) results = [] - score_threshold = kwargs.get("score_threshold") for row in rows: distance = row["distance"] - if score_threshold is not None and distance > score_threshold: + # Convert distance to score (1 - distance for cosine) + score = 1 - distance + + # Apply threshold filtering if specified + if threshold is not None and score < threshold: continue vector_data = None @@ -379,7 +396,7 @@ class PGVectorStore(BaseVectorStore): if isinstance(metadata, str): metadata = json.loads(metadata) - metadata["score"] = 1 - distance + metadata["score"] = score metadata["_distance"] = distance node = VectorNode( @@ -390,6 +407,10 @@ class PGVectorStore(BaseVectorStore): ) results.append(node) + # Apply limit after threshold filtering + if threshold is not None: + results = results[:limit] + return results async def delete(self, vector_ids: str | list[str], **kwargs): diff --git a/reme/core/vector_store/qdrant_vector_store.py b/reme/core/vector_store/qdrant_vector_store.py index 93ccee70..f14dd958 100644 --- a/reme/core/vector_store/qdrant_vector_store.py +++ b/reme/core/vector_store/qdrant_vector_store.py @@ -325,20 +325,33 @@ class QdrantVectorStore(BaseVectorStore): self, query: str, limit: int = 5, + candidates: int | None = None, filters: dict | None = None, + threshold: float | None = None, **kwargs: Any, ) -> list[VectorNode]: - """Search for the most similar vectors based on a text query.""" + """Search for the most similar vectors based on a text query. + + When threshold is None, uses default behavior. + When threshold is set, searches max(candidates, limit) nodes, + filters by threshold, then returns top limit results. + """ query_vector = await self.get_embedding(query) query_filter = self._create_filter(filters) if filters else None - score_threshold = kwargs.get("score_threshold", None) + + # When threshold is set, search more candidates + if threshold is not None: + effective_candidates = candidates if candidates is not None else limit * 2 + search_limit = max(effective_candidates, limit) + else: + search_limit = limit results = await self.client.query_points( collection_name=self.collection_name, query=query_vector, query_filter=query_filter, - limit=limit, - score_threshold=score_threshold, + limit=search_limit, + score_threshold=threshold, ) nodes = [] @@ -353,6 +366,10 @@ class QdrantVectorStore(BaseVectorStore): node.metadata["score"] = point.score nodes.append(node) + # Apply limit after threshold filtering (Qdrant already filters by threshold) + if threshold is not None: + nodes = nodes[:limit] + return nodes async def delete(self, vector_ids: str | list[str], **kwargs: Any): diff --git a/tests/test_vector_store.py b/tests/test_vector_store.py index 14ed1c1d..d004c43f 100644 --- a/tests/test_vector_store.py +++ b/tests/test_vector_store.py @@ -414,6 +414,180 @@ async def test_search_with_multiple_filters(store: BaseVectorStore, _store_name: logger.info("✓ Multiple filters search test passed") +async def test_search_with_threshold(store: BaseVectorStore, _store_name: str): + """Test vector search with threshold parameter. + + When threshold is None, uses default behavior. + When threshold is set, searches max(candidates, limit) nodes, + filters by threshold, then returns top limit results. + """ + logger.info("=" * 20 + " SEARCH WITH THRESHOLD TEST " + "=" * 20) + + # Clean up any existing test data first + try: + existing_nodes = await store.list(filters={"test_type": "threshold_test"}) + if existing_nodes: + await store.delete([node.vector_id for node in existing_nodes]) + logger.info(f"Cleaned up {len(existing_nodes)} existing test nodes") + except Exception as e: + logger.warning(f"Failed to clean up existing nodes: {e}") + + # Insert test nodes with varied content for different similarity scores + threshold_nodes = [ + VectorNode( + vector_id="threshold_node_1", + content="Machine learning and artificial intelligence are the future of technology.", + metadata={"test_type": "threshold_test", "relevance": "high"}, + ), + VectorNode( + vector_id="threshold_node_2", + content="Deep learning uses neural networks for artificial intelligence applications.", + metadata={"test_type": "threshold_test", "relevance": "high"}, + ), + VectorNode( + vector_id="threshold_node_3", + content="Natural language processing is a branch of AI technology.", + metadata={"test_type": "threshold_test", "relevance": "medium"}, + ), + VectorNode( + vector_id="threshold_node_4", + content="Computer vision enables machines to interpret visual information.", + metadata={"test_type": "threshold_test", "relevance": "medium"}, + ), + VectorNode( + vector_id="threshold_node_5", + content="Software development requires good coding practices.", + metadata={"test_type": "threshold_test", "relevance": "low"}, + ), + VectorNode( + vector_id="threshold_node_6", + content="Database management systems store and retrieve data efficiently.", + metadata={"test_type": "threshold_test", "relevance": "low"}, + ), + VectorNode( + vector_id="threshold_node_7", + content="Cloud computing provides scalable infrastructure for applications.", + metadata={"test_type": "threshold_test", "relevance": "low"}, + ), + VectorNode( + vector_id="threshold_node_8", + content="Cooking pasta requires boiling water and adding salt.", + metadata={"test_type": "threshold_test", "relevance": "none"}, + ), + ] + + await store.insert(threshold_nodes) + logger.info(f"✓ Inserted {len(threshold_nodes)} test nodes") + + # Test 1: Search without threshold (baseline) + results_baseline = await store.search( + query="artificial intelligence and machine learning", + limit=3, + candidates=10, + filters={"test_type": "threshold_test"}, + ) + logger.info(f"Test 1 - Search without threshold: {len(results_baseline)} results") + for i, r in enumerate(results_baseline, 1): + score = r.metadata.get("score", 0) + logger.info(f" Result {i}: {r.vector_id}, score={score:.4f}") + + assert len(results_baseline) <= 3, "Should return at most limit results" + logger.info("✓ Baseline search validated") + + # Test 2: Search with threshold (should filter low-score results) + threshold_value = 0.5 + results_with_threshold = await store.search( + query="artificial intelligence and machine learning", + limit=3, + candidates=10, + filters={"test_type": "threshold_test"}, + threshold=threshold_value, + ) + logger.info( + f"Test 2 - Search with threshold={threshold_value}: {len(results_with_threshold)} results", + ) + for i, r in enumerate(results_with_threshold, 1): + score = r.metadata.get("score", 0) + logger.info(f" Result {i}: {r.vector_id}, score={score:.4f}") + # Verify all results meet threshold + assert score >= threshold_value, f"Score {score} should be >= threshold {threshold_value}" + + # Should return at most limit results + assert len(results_with_threshold) <= 3, "Should return at most limit results" + logger.info("✓ Threshold filtering validated") + + # Test 3: Search with high threshold (should filter more strictly) + high_threshold = 0.7 + results_high_threshold = await store.search( + query="artificial intelligence and machine learning", + limit=10, # Use same limit as candidates to get all results above threshold + candidates=10, + filters={"test_type": "threshold_test"}, + threshold=high_threshold, + ) + logger.info( + f"Test 3 - Search with high threshold={high_threshold}: {len(results_high_threshold)} results", + ) + for i, r in enumerate(results_high_threshold, 1): + score = r.metadata.get("score", 0) + logger.info(f" Result {i}: {r.vector_id}, score={score:.4f}") + assert score >= high_threshold, f"Score {score} should be >= threshold {high_threshold}" + + # Compare with same limit and lower threshold to verify stricter filtering + results_low_threshold = await store.search( + query="artificial intelligence and machine learning", + limit=10, + candidates=10, + filters={"test_type": "threshold_test"}, + threshold=0.3, # Lower threshold + ) + # High threshold should return same or fewer results than low threshold + assert len(results_high_threshold) <= len( + results_low_threshold, + ), ( + f"High threshold ({len(results_high_threshold)}) should return fewer results than low threshold " + f"({len(results_low_threshold)})" + ) + logger.info("✓ High threshold filtering validated") + + # Test 4: Search with threshold and large candidates + results_large_candidates = await store.search( + query="artificial intelligence and machine learning", + limit=2, + candidates=20, + filters={"test_type": "threshold_test"}, + threshold=0.5, + ) + logger.info( + f"Test 4 - Search with threshold and large candidates: {len(results_large_candidates)} results", + ) + # Should search max(candidates, limit) = 20 nodes, filter by threshold, return top 2 + assert len(results_large_candidates) <= 2, "Should return at most limit results" + for r in results_large_candidates: + score = r.metadata.get("score", 0) + assert score >= 0.5, f"Score {score} should be >= 0.5" + logger.info("✓ Large candidates with threshold validated") + + # Test 5: Verify that threshold=None behaves same as no threshold + results_none_threshold = await store.search( + query="artificial intelligence and machine learning", + limit=3, + candidates=10, + filters={"test_type": "threshold_test"}, + threshold=None, + ) + logger.info(f"Test 5 - Search with threshold=None: {len(results_none_threshold)} results") + # Should behave same as baseline (Test 1) + assert len(results_none_threshold) <= 3, "Should return at most limit results" + logger.info("✓ threshold=None validated") + + # Clean up test data + await store.delete([node.vector_id for node in threshold_nodes]) + logger.info("Cleaned up test nodes") + + logger.info("✓ Search with threshold test passed") + + async def test_get_by_id(store: BaseVectorStore, _store_name: str): """Test retrieving nodes by vector_id (single and batch).""" logger.info("=" * 20 + " GET BY ID TEST " + "=" * 20) @@ -1671,6 +1845,7 @@ async def run_all_tests_for_store(store_type: str, store_name: str): await test_search_with_single_filter(store, store_name) await test_search_with_exact_match_filter(store, store_name) await test_search_with_multiple_filters(store, store_name) + await test_search_with_threshold(store, store_name) await test_get_by_id(store, store_name) await test_list_all(store, store_name) await test_list_with_filters(store, store_name) From 6e880c5d7ed69c10c51f86597442223ae89307a0 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Fri, 6 Feb 2026 15:56:00 +0800 Subject: [PATCH 4/7] refactor(config): rename FileWatchConfig to FileWatcherConfig --- reme/core/schema/service_config.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/reme/core/schema/service_config.py b/reme/core/schema/service_config.py index cf7da9d6..11c8fca6 100644 --- a/reme/core/schema/service_config.py +++ b/reme/core/schema/service_config.py @@ -98,8 +98,8 @@ class TokenCounterConfig(BaseModel): model_name: str = Field(default="") -class FileWatchConfig(BaseModel): - """Configuration for file watch service.""" +class FileWatcherConfig(BaseModel): + """Configuration for file watcher service.""" model_config = ConfigDict(extra="allow") @@ -137,4 +137,4 @@ class ServiceConfig(BaseModel): vector_store: dict[str, VectorStoreConfig] = Field(default_factory=dict) memory_store: dict[str, MemoryStoreConfig] = Field(default_factory=dict) token_counter: dict[str, TokenCounterConfig] = Field(default_factory=dict) - file_watch: dict[str, FileWatchConfig] = Field(default_factory=dict) + file_watcher: dict[str, FileWatcherConfig] = Field(default_factory=dict) From 02d5154e1c2a4259b51740ae7ee689a638c4f8fb Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Fri, 6 Feb 2026 16:00:37 +0800 Subject: [PATCH 5/7] style(logo_utils): center align logo display in console output --- reme/core/utils/logo_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/reme/core/utils/logo_utils.py b/reme/core/utils/logo_utils.py index 91c46a02..5b4acd8e 100644 --- a/reme/core/utils/logo_utils.py +++ b/reme/core/utils/logo_utils.py @@ -81,4 +81,4 @@ def print_logo(service_config: "ServiceConfig"): expand=False, ) - Console().print(Group("\n", panel, "\n")) + Console().print(Group("\n", panel, "\n"), justify="center") From bc5e87db5f619d4b3270e3b10c0924197f40767e Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Fri, 6 Feb 2026 16:53:57 +0800 Subject: [PATCH 6/7] refactor(vector_store): simplify search method signature and remove threshold parameter --- reme/core/vector_store/base_vector_store.py | 17 +- reme/core/vector_store/chroma_vector_store.py | 26 +-- reme/core/vector_store/es_vector_store.py | 34 +--- reme/core/vector_store/local_vector_store.py | 23 +-- reme/core/vector_store/pgvector_store.py | 31 +--- reme/core/vector_store/qdrant_vector_store.py | 25 +-- tests/test_vector_store.py | 175 ------------------ 7 files changed, 27 insertions(+), 304 deletions(-) diff --git a/reme/core/vector_store/base_vector_store.py b/reme/core/vector_store/base_vector_store.py index addaa2ea..62a73a3e 100644 --- a/reme/core/vector_store/base_vector_store.py +++ b/reme/core/vector_store/base_vector_store.py @@ -73,21 +73,8 @@ class BaseVectorStore(ABC): """Add one or more vector nodes into the current collection.""" @abstractmethod - async def search( - self, - query: str, - limit: int = 5, - candidates: int | None = None, - filters: dict | None = None, - threshold: float | None = None, - **kwargs, - ) -> list[VectorNode]: - """Find the most similar vector nodes based on a text query. - - When threshold is None, uses default behavior. - When threshold is set, searches max(candidates, limit) nodes, - filters by threshold, then returns top limit results. - """ + async def search(self, query: str, limit: int = 5, filters: dict | None = None, **kwargs) -> list[VectorNode]: + """Find the most similar vector nodes based on a text query.""" @abstractmethod async def delete(self, vector_ids: str | list[str], **kwargs) -> None: diff --git a/reme/core/vector_store/chroma_vector_store.py b/reme/core/vector_store/chroma_vector_store.py index 76b4b81a..5b0748f5 100644 --- a/reme/core/vector_store/chroma_vector_store.py +++ b/reme/core/vector_store/chroma_vector_store.py @@ -307,35 +307,21 @@ class ChromaVectorStore(BaseVectorStore): self, query: str, limit: int = 5, - candidates: int | None = None, filters: dict | None = None, - threshold: float | None = None, **kwargs, ) -> list[VectorNode]: - """Search for the most similar vector nodes based on a text query. - - When threshold is None, uses default behavior. - When threshold is set, searches max(candidates, limit) nodes, - filters by threshold, then returns top limit results. - """ + """Search for the most similar vector nodes based on a text query.""" query_vector = await self.get_embedding(query) where_clause = self._generate_where_clause(filters) include_embeddings = kwargs.get("include_embeddings", False) - # When threshold is set, search more candidates - if threshold is not None: - effective_candidates = candidates if candidates is not None else limit * 2 - search_limit = max(effective_candidates, limit) - else: - search_limit = limit - def _search(): include: list = ["documents", "metadatas", "distances"] if include_embeddings: include.append("embeddings") return self.collection.query( query_embeddings=[query_vector], - n_results=search_limit, + n_results=limit, where=where_clause, include=include, ) @@ -343,11 +329,9 @@ class ChromaVectorStore(BaseVectorStore): results = await self._run_sync_in_executor(_search) nodes = self._parse_results(results, include_score=True) - # Apply threshold filtering if specified - if threshold is not None: - nodes = [n for n in nodes if n.metadata.get("score", 0) >= threshold] - nodes = nodes[:limit] - + score_threshold = kwargs.get("score_threshold") + if score_threshold is not None: + nodes = [n for n in nodes if n.metadata.get("score", 0) >= score_threshold] return nodes async def delete(self, vector_ids: str | list[str], **kwargs): diff --git a/reme/core/vector_store/es_vector_store.py b/reme/core/vector_store/es_vector_store.py index f9cfad7d..62b28bc1 100644 --- a/reme/core/vector_store/es_vector_store.py +++ b/reme/core/vector_store/es_vector_store.py @@ -238,48 +238,31 @@ class ESVectorStore(BaseVectorStore): self, query: str, limit: int = 5, - candidates: int | None = None, filters: dict | None = None, - threshold: float | None = None, **kwargs, ) -> list[VectorNode]: """Perform a kNN similarity search based on a text query. - When threshold is None, uses default behavior. - When threshold is set, searches max(candidates, limit) nodes, - filters by threshold, then returns top limit results. - Args: query: The text to search for. limit: Maximum number of nearest neighbors to return. - candidates: Number of candidates to consider during search. - filters: Metadata filters for exact match or range operations. - threshold: Minimum score threshold for results. - **kwargs: Additional search parameters. + filters: Metadata filters for exact match or 'IN' operations. + **kwargs: Search parameters like num_candidates or score_threshold. Returns: List of VectorNode objects ordered by similarity. """ query_vector = await self.get_embedding(query) - - # When threshold is set, search more candidates - if threshold is not None: - effective_candidates = candidates if candidates is not None else limit * 2 - search_candidates = max(effective_candidates, limit) - search_limit = search_candidates - else: - effective_candidates = candidates if candidates is not None else limit * 2 - search_candidates = effective_candidates - search_limit = limit + num_candidates = kwargs.get("num_candidates", limit * 2) search_query: dict = { "knn": { "field": "vector", "query_vector": query_vector, - "k": search_limit, - "num_candidates": search_candidates, + "k": limit, + "num_candidates": num_candidates, }, - "size": search_limit, + "size": limit, } if filters: @@ -317,11 +300,6 @@ class ESVectorStore(BaseVectorStore): node.metadata["score"] = hit["_score"] results.append(node) - # Apply threshold filtering if specified - if threshold is not None: - results = [n for n in results if n.metadata.get("score", 0) >= threshold] - results = results[:limit] - return results async def delete(self, vector_ids: str | list[str], refresh: bool = True, **kwargs): diff --git a/reme/core/vector_store/local_vector_store.py b/reme/core/vector_store/local_vector_store.py index 22db669e..c117ddf0 100644 --- a/reme/core/vector_store/local_vector_store.py +++ b/reme/core/vector_store/local_vector_store.py @@ -193,17 +193,10 @@ class LocalVectorStore(BaseVectorStore): self, query: str, limit: int = 5, - candidates: int | None = None, filters: dict | None = None, - threshold: float | None = None, **kwargs, ) -> list[VectorNode]: - """Search for nodes similar to the query using brute-force cosine similarity. - - When threshold is None, uses default behavior. - When threshold is set, searches max(candidates, limit) nodes, - filters by threshold, then returns top limit results. - """ + """Search for nodes similar to the query using brute-force cosine similarity.""" query_vector = await self.get_embedding(query) all_nodes = self._load_all_nodes() filtered_nodes = [node for node in all_nodes if self._match_filters(node, filters)] @@ -222,17 +215,11 @@ class LocalVectorStore(BaseVectorStore): scored_nodes.sort(key=lambda x: x[1], reverse=True) - # Apply threshold filtering if specified - if threshold is not None: - # When threshold is set, consider more candidates - effective_candidates = candidates if candidates is not None else limit * 2 - search_limit = max(effective_candidates, limit) - scored_nodes = scored_nodes[:search_limit] - scored_nodes = [(node, score) for node, score in scored_nodes if score >= threshold] - scored_nodes = scored_nodes[:limit] - else: - scored_nodes = scored_nodes[:limit] + score_threshold = kwargs.get("score_threshold") + if score_threshold is not None: + scored_nodes = [(node, score) for node, score in scored_nodes if score >= score_threshold] + scored_nodes = scored_nodes[:limit] results = [] for node, score in scored_nodes: node.metadata["score"] = score diff --git a/reme/core/vector_store/pgvector_store.py b/reme/core/vector_store/pgvector_store.py index dcb49201..0ec7d6a8 100644 --- a/reme/core/vector_store/pgvector_store.py +++ b/reme/core/vector_store/pgvector_store.py @@ -333,17 +333,10 @@ class PGVectorStore(BaseVectorStore): self, query: str, limit: int = 5, - candidates: int | None = None, filters: dict | None = None, - threshold: float | None = None, **kwargs, ) -> list[VectorNode]: - """Perform vector similarity search with optional metadata filtering. - - When threshold is None, uses default behavior. - When threshold is set, searches max(candidates, limit) nodes, - filters by threshold, then returns top limit results. - """ + """Perform vector similarity search with optional metadata filtering.""" await self._ensure_collection_exists() query_vector = await self.get_embedding(query) @@ -358,13 +351,6 @@ class PGVectorStore(BaseVectorStore): new_placeholder = f"${i + 1}" filter_clause = re.sub(rf"\${i}\b", new_placeholder, filter_clause) - # When threshold is set, search more candidates - if threshold is not None: - effective_candidates = candidates if candidates is not None else limit * 2 - search_limit = max(effective_candidates, limit) - else: - search_limit = limit - async with pool.acquire() as conn: sql = f""" SELECT id, content, vector, metadata, vector <=> $1::vector AS distance @@ -373,17 +359,14 @@ class PGVectorStore(BaseVectorStore): ORDER BY distance LIMIT ${len(filter_params) + 2} """ - rows = await conn.fetch(sql, vector_str, *filter_params, search_limit) + rows = await conn.fetch(sql, vector_str, *filter_params, limit) results = [] + score_threshold = kwargs.get("score_threshold") for row in rows: distance = row["distance"] - # Convert distance to score (1 - distance for cosine) - score = 1 - distance - - # Apply threshold filtering if specified - if threshold is not None and score < threshold: + if score_threshold is not None and distance > score_threshold: continue vector_data = None @@ -396,7 +379,7 @@ class PGVectorStore(BaseVectorStore): if isinstance(metadata, str): metadata = json.loads(metadata) - metadata["score"] = score + metadata["score"] = 1 - distance metadata["_distance"] = distance node = VectorNode( @@ -407,10 +390,6 @@ class PGVectorStore(BaseVectorStore): ) results.append(node) - # Apply limit after threshold filtering - if threshold is not None: - results = results[:limit] - return results async def delete(self, vector_ids: str | list[str], **kwargs): diff --git a/reme/core/vector_store/qdrant_vector_store.py b/reme/core/vector_store/qdrant_vector_store.py index f14dd958..93ccee70 100644 --- a/reme/core/vector_store/qdrant_vector_store.py +++ b/reme/core/vector_store/qdrant_vector_store.py @@ -325,33 +325,20 @@ class QdrantVectorStore(BaseVectorStore): self, query: str, limit: int = 5, - candidates: int | None = None, filters: dict | None = None, - threshold: float | None = None, **kwargs: Any, ) -> list[VectorNode]: - """Search for the most similar vectors based on a text query. - - When threshold is None, uses default behavior. - When threshold is set, searches max(candidates, limit) nodes, - filters by threshold, then returns top limit results. - """ + """Search for the most similar vectors based on a text query.""" query_vector = await self.get_embedding(query) query_filter = self._create_filter(filters) if filters else None - - # When threshold is set, search more candidates - if threshold is not None: - effective_candidates = candidates if candidates is not None else limit * 2 - search_limit = max(effective_candidates, limit) - else: - search_limit = limit + score_threshold = kwargs.get("score_threshold", None) results = await self.client.query_points( collection_name=self.collection_name, query=query_vector, query_filter=query_filter, - limit=search_limit, - score_threshold=threshold, + limit=limit, + score_threshold=score_threshold, ) nodes = [] @@ -366,10 +353,6 @@ class QdrantVectorStore(BaseVectorStore): node.metadata["score"] = point.score nodes.append(node) - # Apply limit after threshold filtering (Qdrant already filters by threshold) - if threshold is not None: - nodes = nodes[:limit] - return nodes async def delete(self, vector_ids: str | list[str], **kwargs: Any): diff --git a/tests/test_vector_store.py b/tests/test_vector_store.py index d004c43f..14ed1c1d 100644 --- a/tests/test_vector_store.py +++ b/tests/test_vector_store.py @@ -414,180 +414,6 @@ async def test_search_with_multiple_filters(store: BaseVectorStore, _store_name: logger.info("✓ Multiple filters search test passed") -async def test_search_with_threshold(store: BaseVectorStore, _store_name: str): - """Test vector search with threshold parameter. - - When threshold is None, uses default behavior. - When threshold is set, searches max(candidates, limit) nodes, - filters by threshold, then returns top limit results. - """ - logger.info("=" * 20 + " SEARCH WITH THRESHOLD TEST " + "=" * 20) - - # Clean up any existing test data first - try: - existing_nodes = await store.list(filters={"test_type": "threshold_test"}) - if existing_nodes: - await store.delete([node.vector_id for node in existing_nodes]) - logger.info(f"Cleaned up {len(existing_nodes)} existing test nodes") - except Exception as e: - logger.warning(f"Failed to clean up existing nodes: {e}") - - # Insert test nodes with varied content for different similarity scores - threshold_nodes = [ - VectorNode( - vector_id="threshold_node_1", - content="Machine learning and artificial intelligence are the future of technology.", - metadata={"test_type": "threshold_test", "relevance": "high"}, - ), - VectorNode( - vector_id="threshold_node_2", - content="Deep learning uses neural networks for artificial intelligence applications.", - metadata={"test_type": "threshold_test", "relevance": "high"}, - ), - VectorNode( - vector_id="threshold_node_3", - content="Natural language processing is a branch of AI technology.", - metadata={"test_type": "threshold_test", "relevance": "medium"}, - ), - VectorNode( - vector_id="threshold_node_4", - content="Computer vision enables machines to interpret visual information.", - metadata={"test_type": "threshold_test", "relevance": "medium"}, - ), - VectorNode( - vector_id="threshold_node_5", - content="Software development requires good coding practices.", - metadata={"test_type": "threshold_test", "relevance": "low"}, - ), - VectorNode( - vector_id="threshold_node_6", - content="Database management systems store and retrieve data efficiently.", - metadata={"test_type": "threshold_test", "relevance": "low"}, - ), - VectorNode( - vector_id="threshold_node_7", - content="Cloud computing provides scalable infrastructure for applications.", - metadata={"test_type": "threshold_test", "relevance": "low"}, - ), - VectorNode( - vector_id="threshold_node_8", - content="Cooking pasta requires boiling water and adding salt.", - metadata={"test_type": "threshold_test", "relevance": "none"}, - ), - ] - - await store.insert(threshold_nodes) - logger.info(f"✓ Inserted {len(threshold_nodes)} test nodes") - - # Test 1: Search without threshold (baseline) - results_baseline = await store.search( - query="artificial intelligence and machine learning", - limit=3, - candidates=10, - filters={"test_type": "threshold_test"}, - ) - logger.info(f"Test 1 - Search without threshold: {len(results_baseline)} results") - for i, r in enumerate(results_baseline, 1): - score = r.metadata.get("score", 0) - logger.info(f" Result {i}: {r.vector_id}, score={score:.4f}") - - assert len(results_baseline) <= 3, "Should return at most limit results" - logger.info("✓ Baseline search validated") - - # Test 2: Search with threshold (should filter low-score results) - threshold_value = 0.5 - results_with_threshold = await store.search( - query="artificial intelligence and machine learning", - limit=3, - candidates=10, - filters={"test_type": "threshold_test"}, - threshold=threshold_value, - ) - logger.info( - f"Test 2 - Search with threshold={threshold_value}: {len(results_with_threshold)} results", - ) - for i, r in enumerate(results_with_threshold, 1): - score = r.metadata.get("score", 0) - logger.info(f" Result {i}: {r.vector_id}, score={score:.4f}") - # Verify all results meet threshold - assert score >= threshold_value, f"Score {score} should be >= threshold {threshold_value}" - - # Should return at most limit results - assert len(results_with_threshold) <= 3, "Should return at most limit results" - logger.info("✓ Threshold filtering validated") - - # Test 3: Search with high threshold (should filter more strictly) - high_threshold = 0.7 - results_high_threshold = await store.search( - query="artificial intelligence and machine learning", - limit=10, # Use same limit as candidates to get all results above threshold - candidates=10, - filters={"test_type": "threshold_test"}, - threshold=high_threshold, - ) - logger.info( - f"Test 3 - Search with high threshold={high_threshold}: {len(results_high_threshold)} results", - ) - for i, r in enumerate(results_high_threshold, 1): - score = r.metadata.get("score", 0) - logger.info(f" Result {i}: {r.vector_id}, score={score:.4f}") - assert score >= high_threshold, f"Score {score} should be >= threshold {high_threshold}" - - # Compare with same limit and lower threshold to verify stricter filtering - results_low_threshold = await store.search( - query="artificial intelligence and machine learning", - limit=10, - candidates=10, - filters={"test_type": "threshold_test"}, - threshold=0.3, # Lower threshold - ) - # High threshold should return same or fewer results than low threshold - assert len(results_high_threshold) <= len( - results_low_threshold, - ), ( - f"High threshold ({len(results_high_threshold)}) should return fewer results than low threshold " - f"({len(results_low_threshold)})" - ) - logger.info("✓ High threshold filtering validated") - - # Test 4: Search with threshold and large candidates - results_large_candidates = await store.search( - query="artificial intelligence and machine learning", - limit=2, - candidates=20, - filters={"test_type": "threshold_test"}, - threshold=0.5, - ) - logger.info( - f"Test 4 - Search with threshold and large candidates: {len(results_large_candidates)} results", - ) - # Should search max(candidates, limit) = 20 nodes, filter by threshold, return top 2 - assert len(results_large_candidates) <= 2, "Should return at most limit results" - for r in results_large_candidates: - score = r.metadata.get("score", 0) - assert score >= 0.5, f"Score {score} should be >= 0.5" - logger.info("✓ Large candidates with threshold validated") - - # Test 5: Verify that threshold=None behaves same as no threshold - results_none_threshold = await store.search( - query="artificial intelligence and machine learning", - limit=3, - candidates=10, - filters={"test_type": "threshold_test"}, - threshold=None, - ) - logger.info(f"Test 5 - Search with threshold=None: {len(results_none_threshold)} results") - # Should behave same as baseline (Test 1) - assert len(results_none_threshold) <= 3, "Should return at most limit results" - logger.info("✓ threshold=None validated") - - # Clean up test data - await store.delete([node.vector_id for node in threshold_nodes]) - logger.info("Cleaned up test nodes") - - logger.info("✓ Search with threshold test passed") - - async def test_get_by_id(store: BaseVectorStore, _store_name: str): """Test retrieving nodes by vector_id (single and batch).""" logger.info("=" * 20 + " GET BY ID TEST " + "=" * 20) @@ -1845,7 +1671,6 @@ async def run_all_tests_for_store(store_type: str, store_name: str): await test_search_with_single_filter(store, store_name) await test_search_with_exact_match_filter(store, store_name) await test_search_with_multiple_filters(store, store_name) - await test_search_with_threshold(store, store_name) await test_get_by_id(store, store_name) await test_list_all(store, store_name) await test_list_with_filters(store, store_name) From c8bbb54de19aef90fe231a9a395b61cf400e20f8 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Fri, 6 Feb 2026 17:45:59 +0800 Subject: [PATCH 7/7] 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"