diff --git a/README.md b/README.md index 7774597b..e2f9d04c 100644 --- a/README.md +++ b/README.md @@ -269,7 +269,7 @@ everything under `metadata/` is rebuildable. | ------------------------------------------- | ----------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------- | | [`auto_memory`](docs/en/auto_memory.md) | Agent hook or `reme auto_memory` | Distills useful conversation facts while preserving a filtered conversation source record. | `session/dialog/*.jsonl`, `daily//.md` | | [`auto_resource`](docs/en/auto_resource.md) | Resource watcher or `reme auto_resource` | Turns files under `resource/` into source-linked, content-named daily cards. | `daily//.md` | -| [`auto_index`](docs/en/memory_search.md) | Background watcher or `reme reindex` | Live-indexes Markdown in `daily/` and `digest/`; a full rebuild also scans `resource/` and JSONL. | Searchable chunks, BM25, wikilink graph, and optional vectors | +| [`auto_index`](docs/en/memory_search.md) | Background watcher or `reme reindex` | The watcher ingests Markdown from `daily/` and `digest/`; `reindex` only rebuilds BM25 and embeddings from already-ingested chunks. | Searchable chunks, BM25, wikilink graph, and optional vectors | | [`auto_dream`](docs/en/auto_dream.md) | `dream_cron` or `reme auto_dream` | By default, extracts up to five reusable units from changed files in the latest two-day window, then creates, corroborates, refines, or corrects digest nodes. | `digest/**`, `daily//interests.yaml` | | [`proactive`](docs/en/proactive.md) | `reme proactive` before an agent decides to act | Reads topics generated by `auto_dream`; the host agent decides whether and how to mention them. | Structured topics from `daily//interests.yaml` | @@ -360,7 +360,7 @@ Run `reme help` for the full job list. Common workspace and maintenance commands | `reme read` / `reme write` / `reme edit` | Inspect and maintain Markdown memory files. | | `reme traverse` / `reme graph_snapshot` | Explore wikilink neighborhoods or the category-rooted digest graph. | | `reme chat` | Stream a read-only, workspace-aware agent conversation. Requires LLM credentials. | -| `reme reindex` | Rebuild search and wikilink indexes from existing files. | +| `reme reindex` | Rebuild BM25 and embedding indexes from already-ingested chunks. | ## 🤝 Community and Contributing diff --git a/README_ZH.md b/README_ZH.md index 1530f3ff..65d5551c 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -265,7 +265,7 @@ ReMe 遵循 capture → index → consolidate → recall 的循环。workspace | ------------------------------------------- | ----------------------------------------- | -------------------------------------------------------------------------------------------- | ------------------------------------------------------------ | | [`auto_memory`](docs/zh/auto_memory.md) | Agent hook 或 `reme auto_memory` | 提炼有长期价值的对话事实，同时保留过滤后的对话来源记录。 | `session/dialog/*.jsonl`、`daily//.md` | | [`auto_resource`](docs/zh/auto_resource.md) | 资源监听或 `reme auto_resource` | 将 `resource/` 下的文件转为带来源链接、按内容命名的 daily 卡片。 | `daily//.md` | -| [`auto_index`](docs/zh/memory_search.md) | 后台监听或 `reme reindex` | 实时索引 `daily/` 和 `digest/` 中的 Markdown；全量重建还会扫描 `resource/` 和 JSONL。 | 可检索的 chunks、BM25、wikilink 图谱和可选向量 | +| [`auto_index`](docs/zh/memory_search.md) | 后台监听或 `reme reindex` | watcher 摄取 `daily/` 和 `digest/` 中的 Markdown；`reindex` 只基于已摄取的 chunks 重建 BM25 和 Embedding。 | 可检索的 chunks、BM25、wikilink 图谱和可选向量 | | [`auto_dream`](docs/zh/auto_dream.md) | `dream_cron` 或 `reme auto_dream` | 默认从最近两天内变化的文件中最多提取 5 个可复用 unit，再创建、印证、补充或修正 digest 节点。 | `digest/**`、`daily//interests.yaml` | | [`proactive`](docs/zh/proactive.md) | Agent 决定主动行动前调用 `reme proactive` | 读取 `auto_dream` 生成的 topics；是否以及如何提醒用户由宿主 Agent 决定。 | 来自 `daily//interests.yaml` 的结构化 topics | @@ -350,7 +350,7 @@ ReMe 通过 Agent 多轮搜索与读取的方式，评测多会话和超长上 | `reme read` / `reme write` / `reme edit` | 检查和维护 Markdown 记忆文件。 | | `reme traverse` / `reme graph_snapshot` | 浏览 wikilink 邻域或按类别组织的 digest 图。 | | `reme chat` | 与可感知 workspace 的只读 Agent 进行流式对话；需要 LLM 凭证。 | -| `reme reindex` | 基于已有文件重建检索和 wikilink 索引。 | +| `reme reindex` | 基于已摄取的 chunks 重建 BM25 和 Embedding 索引。 | ## 🤝 社区与贡献 diff --git a/docs/en/auto_resource.md b/docs/en/auto_resource.md index 9d9495b1..cdef773d 100644 --- a/docs/en/auto_resource.md +++ b/docs/en/auto_resource.md @@ -99,7 +99,6 @@ therefore enter the daily memory flow while their source files stay in their ori ## What Happens Next Auto Resource only creates resource interpretations in the daily layer. To distill long-term knowledge from resources -into -`digest/`, use [Auto Dream](./auto_dream.md). The default live index covers daily cards and digest nodes. Run -`reme reindex` -when original resource files must also be directly searchable; see [Memory Search](./memory_search.md). +into `digest/`, use [Auto Dream](./auto_dream.md). The default live index covers daily cards and digest nodes. Manual +`reindex` only rebuilds search indexes from chunks already accepted by an ingestion path; it does not add the original +resource files to search. See [Memory Search](./memory_search.md). diff --git a/docs/en/memory_as_file.md b/docs/en/memory_as_file.md index f519396a..3ca407b9 100644 --- a/docs/en/memory_as_file.md +++ b/docs/en/memory_as_file.md @@ -237,8 +237,9 @@ FileLink Older documents containing wrappers such as `related:: [[path]]`, `- related:: [[path]]`, or `[related:: [[path]]]` remain readable. ReMe ignores the surrounding text and indexes the -inner `[[path]]` as an ordinary link. After upgrading from a version that stored typed links, run `reme reindex` -once to rebuild the derived graph without the removed relationship field. +inner `[[path]]` as an ordinary link. Graph changes are applied when source files pass through the normal ingestion +path. `reme reindex` only rebuilds BM25 and embedding indexes from existing chunks; it does not reparse files or rebuild +the derived graph. ### Sources and Relationships diff --git a/docs/en/memory_search.md b/docs/en/memory_search.md index bdc4692e..2560467d 100644 --- a/docs/en/memory_search.md +++ b/docs/en/memory_search.md @@ -3,8 +3,8 @@ Memory Search is ReMe's memory retrieval entry point. The default background loop continuously builds Markdown under `daily/` and `digest/` into a searchable chunk index and wikilink graph. At query time, it first recalls the most relevant fragments and then expands context along the bidirectional links of the files containing those fragments. -`reme reindex` has a broader rebuild scope that also scans `resource/` and JSONL; it is intentionally different from the -live watcher. +`reme reindex` rebuilds derived BM25 and embedding indexes from the authoritative in-memory `file_chunks`; it does not +rescan workspace files, rechunk content, or rewrite the wikilink graph.

ReMe Auto Index and Memory Search indexing, recall, fusion, and link expansion @@ -29,11 +29,8 @@ The default `index_update_loop` watches two memory directories: - `digest_dir`: long-term distilled digest nodes. The live watcher handles only the `md` suffix. A separate `resource_watch_loop` watches `resource_dir`, and Auto -Resource turns those inputs into daily cards that enter the live index. When `reme reindex` is run manually, its -configuration scans -`daily_dir`, `digest_dir`, and `resource_dir` for `md` and `jsonl`; Markdown uses the `markdown` chunker and JSONL uses -the -`jsonl` chunker. +Resource turns those inputs into daily cards that enter the live index. Manual `reindex` operates on chunks already +accepted by those ingestion paths and therefore does not expand the set of searched files. ## How the Index Is Built @@ -121,9 +118,11 @@ The embedding store accepts `health_check_timeout` for its startup probe. A temp backfill while keeping BM25 available; a later successful provider request resumes the missing-vector backfill automatically. -Embedded integrations that have already verified a provider can call `resume_embedding(verified=True)`. When changing -the embedding vector space, pass `rebuild=True`; persisted vectors are invalidated before a serial background rebuild, -and vector search remains unavailable until the rebuilt vectors are safely persisted. +Embedded integrations that have already verified a provider can call `resume_embedding(verified=True)` to repair +missing vectors in the same vector space. Vector-space changes must use the explicit `reindex` job with +`scope: embedding`; vector search remains unavailable until that job finishes successfully. +Use `scope: bm25` to rebuild only keyword search. `scope: all` runs the BM25 rebuild first and then the embedding +rebuild; all scopes use the current `file_chunks` snapshot. ## How to Search diff --git a/docs/en/quick_start.md b/docs/en/quick_start.md index 0d890c73..a15edef2 100644 --- a/docs/en/quick_start.md +++ b/docs/en/quick_start.md @@ -113,12 +113,15 @@ Related link: [[digest/wiki/search-demo.md]]" and `description` are written to frontmatter. -The background watcher builds the index automatically. You can also rebuild it manually: +The background watcher ingests workspace files automatically. You can manually rebuild the derived BM25 and embedding +indexes from the chunks it has already ingested: ```bash reme reindex ``` +This command does not scan workspace files, rechunk content, or rebuild the wikilink graph. + Search: ```bash diff --git a/docs/en/reme-blog.md b/docs/en/reme-blog.md index 9c3fea5a..bec49207 100644 --- a/docs/en/reme-blog.md +++ b/docs/en/reme-blog.md @@ -178,8 +178,9 @@ Knowledge evolves and links are created in the same workflow. Relationships are Markdown is easy for people to read, but if files are merely piled into directories, agents still struggle to find them quickly. The default live index watches Markdown under `daily/` and `digest/`. A separate resource workflow watches -`resource/` and turns those files into daily cards that enter the same index. For a full rebuild from existing files, -`reme reindex` also scans `resource/` and JSONL. +`resource/` and turns those files into daily cards that enter the same index. Manual `reindex` rebuilds BM25 and +embedding indexes from the chunks those ingestion paths have already accepted; it does not rescan files or rebuild the +Wikilink graph. A Markdown file is parsed into: diff --git a/docs/figure/auto-index-and-memory-search.svg b/docs/figure/auto-index-and-memory-search.svg index 1fdf821f..7843fe66 100644 --- a/docs/figure/auto-index-and-memory-search.svg +++ b/docs/figure/auto-index-and-memory-search.svg @@ -73,7 +73,7 @@ Rebuild scope - reindex adds resource + JSONL. + BM25 + vectors from current chunks. BM25 is enabled by default Recall cap: 200; embeddings opt-in. diff --git a/docs/figure/reme-blog/reme-blog-memory-index.svg b/docs/figure/reme-blog/reme-blog-memory-index.svg index d1625911..fef3203b 100644 --- a/docs/figure/reme-blog/reme-blog-memory-index.svg +++ b/docs/figure/reme-blog/reme-blog-memory-index.svg @@ -1,8 +1,8 @@ Memory Index - Watch daily and digest Markdown, rebuild broader workspace content on demand, split files into - structural semantic chunks, and build BM25, optional vector, and Wikilink graph indexes. + Watch daily and digest Markdown, transform resources through the resource workflow, split files + into structural semantic chunks, and build BM25, optional vector, and Wikilink graph indexes. @@ -27,7 +27,7 @@ digest/ resource/ - via reindex + via resource workflow Create · update · delete diff --git a/docs/zh/auto_resource.md b/docs/zh/auto_resource.md index 5a669c09..6403014d 100644 --- a/docs/zh/auto_resource.md +++ b/docs/zh/auto_resource.md @@ -91,4 +91,5 @@ Auto Resource 不会把原始文件挪走：它仍然留在 `resource/` 下的 ## 后续流向 Auto Resource 只生成 daily 层的资源解读。要把资源中的长期知识沉淀进 `digest/`，使用 [Auto Dream](./auto_dream.md)；默认实时检索会 -索引 daily 卡片和 digest 节点。若还要直接检索原始资源文件，可运行 `reme reindex`，详见 [Memory Search](./memory_search.md)。 +索引 daily 卡片和 digest 节点。手动 `reindex` 只基于摄取流程已经接受的 chunks 重建检索索引，不会把原始资源文件加入检索范围。 +详见 [Memory Search](./memory_search.md)。 diff --git a/docs/zh/memory_as_file.md b/docs/zh/memory_as_file.md index 6d792ed9..0cf36a01 100644 --- a/docs/zh/memory_as_file.md +++ b/docs/zh/memory_as_file.md @@ -219,7 +219,8 @@ FileLink 旧文档中的 `related:: [[path]]`、`- related:: [[path]]` 或 `[related:: [[path]]]` 仍然可以读取。ReMe 会忽略外围文本，把内部 `[[path]]` -作为普通链接建立索引。从曾存储 typed link 的版本升级后，应执行一次 `reme reindex`，用源文件重建不含旧关系字段的派生图索引。 +作为普通链接建立索引。源文件通过正常摄取路径时会应用图谱变更。`reme reindex` 只基于现有 chunks 重建 BM25 和 Embedding +索引，不会重新解析文件或重建派生图谱。 ### 来源和关系 diff --git a/docs/zh/memory_search.md b/docs/zh/memory_search.md index 67fdac8c..702bd18b 100644 --- a/docs/zh/memory_search.md +++ b/docs/zh/memory_search.md @@ -1,8 +1,8 @@ # Memory Search Memory Search 是 ReMe 的记忆检索入口。默认后台持续把 `daily/`、`digest/` 里的 Markdown 构建成可搜索的 chunk 索引和 -wikilink 图谱；查询时先召回最相关的片段，再沿着片段所在文件的双向链接展开上下文。`reme reindex` 的重建范围更宽，会额外扫描 -`resource/` 和 JSONL；这与实时 watcher 的默认范围不同。 +wikilink 图谱；查询时先召回最相关的片段，再沿着片段所在文件的双向链接展开上下文。`reme reindex` 以权威的内存态 +`file_chunks` 为输入重建派生的 BM25 和 Embedding 索引；它不会重新扫描工作区、重新分块或改写 wikilink 图谱。

ReMe Auto Index and Memory Search 索引、召回、融合与链接展开流程 @@ -26,8 +26,7 @@ workspace files - `digest_dir`：长期沉淀后的 digest 节点。 默认实时后缀只有 `md`。`resource_dir` 由独立的 `resource_watch_loop` 监听，并经 Auto Resource 转换成 daily 卡片后进入实时索引。 -如果手动运行 `reme reindex`，其配置会扫描 `daily_dir`、`digest_dir`、`resource_dir` 下的 `md` 和 `jsonl`；Markdown 用 -`markdown` chunker，JSONL 用 `jsonl` chunker。 +手动 `reindex` 只处理这些摄取路径已经接受的 chunk，因此不会扩大搜索文件范围。 ## 索引怎么构建 @@ -109,8 +108,10 @@ file_store: Embedding store 可通过 `health_check_timeout` 配置启动探测。临时失败只会跳过本次向量回填，BM25 仍可使用； 后续真实请求成功后会自动恢复缺失向量的回填。 -已经完成真实服务验证的嵌入式集成可以调用 `resume_embedding(verified=True)`。切换 Embedding 向量空间时应同时传入 -`rebuild=True`；ReMe 会先使旧向量失效，再串行后台重建，并在新向量安全持久化前暂停向量搜索。 +已经完成真实服务验证的嵌入式集成可以调用 `resume_embedding(verified=True)`，修复同一向量空间内缺失的向量。 +切换 Embedding 向量空间必须显式运行 `reindex` Job，并传入 `scope: embedding`；该 Job 成功完成前向量搜索保持不可用。 +`scope: bm25` 只重建关键词索引；`scope: all` 先重建 BM25，再重建 Embedding。所有 scope 都使用当前的 +`file_chunks` 快照。 ## 怎么搜索 diff --git a/docs/zh/quick_start.md b/docs/zh/quick_start.md index 6267e2d9..ee41626a 100644 --- a/docs/zh/quick_start.md +++ b/docs/zh/quick_start.md @@ -109,12 +109,14 @@ reme write \ `path` 是 workspace 内路径；没有后缀时会自动补 `.md`；Markdown 文件会写入 `name` 和 `description` front matter。 -后台 watcher 会自动建索引；也可以手动重建： +后台 watcher 会自动摄取 workspace 文件。也可以基于它已经摄取的 chunks，手动重建派生的 BM25 和 Embedding 索引： ```bash reme reindex ``` +该命令不会扫描 workspace 文件、重新分块或重建 wikilink 图谱。 + 搜索： ```bash diff --git a/docs/zh/reme-blog.md b/docs/zh/reme-blog.md index 5deac4ff..a616b939 100644 --- a/docs/zh/reme-blog.md +++ b/docs/zh/reme-blog.md @@ -184,8 +184,8 @@ Auto Dream 默认查看以目标日期结尾的最近两天，只把相对上次

Markdown 适合人读，但如果只是把文件堆进目录，Agent 仍然很难快速找到它们。默认实时索引持续监听 `daily/` 与 `digest/` 中的 -Markdown；`resource/` 由独立资源流程监听，转成 daily 卡片后进入同一索引。需要从现有文件完整重建时，`reme reindex` 还会扫描 -`resource/` 与 JSONL。 +Markdown；`resource/` 由独立资源流程监听，转成 daily 卡片后进入同一索引。手动 `reindex` 只基于这些摄取路径已经接受的 +chunks 重建 BM25 和 Embedding 索引，不会重新扫描文件或重建 Wikilink 图谱。 一份 Markdown 会被解析为： diff --git a/reme/__init__.py b/reme/__init__.py index 680cc4f1..c4bd9bf8 100644 --- a/reme/__init__.py +++ b/reme/__init__.py @@ -1,6 +1,6 @@ """ReMe CLI package.""" -__version__ = "0.4.1.9" +__version__ = "0.4.1.10" from . import config from . import constants diff --git a/reme/components/file_store/base_file_store.py b/reme/components/file_store/base_file_store.py index b89dd183..91bcaa73 100644 --- a/reme/components/file_store/base_file_store.py +++ b/reme/components/file_store/base_file_store.py @@ -1,6 +1,9 @@ """Abstract base for file store backends.""" +import asyncio from abc import abstractmethod +from contextlib import asynccontextmanager +from functools import wraps from ..base_component import BaseComponent from ...enumeration import ComponentEnum, LinkScopeEnum @@ -18,6 +21,36 @@ class BaseFileStore(BaseComponent): component_type = ComponentEnum.FILE_STORE + def __init__(self, **kwargs): + super().__init__(**kwargs) + self._maintenance_lock = asyncio.Lock() + self._maintenance_lock_owner = None + + @asynccontextmanager + async def _maintenance_guard(self): + """Serialize maintenance and mutations, allowing nested backend overrides.""" + task = asyncio.current_task() + if self._maintenance_lock_owner is task: + yield + return + async with self._maintenance_lock: + self._maintenance_lock_owner = task + try: + yield + finally: + self._maintenance_lock_owner = None + + @staticmethod + def serialized(method): + """Mark a mutation or maintenance method as mutually exclusive.""" + + @wraps(method) + async def wrapped(self, *args, **kwargs): + async with self._maintenance_guard(): # pylint: disable=protected-access + return await method(self, *args, **kwargs) + + return wrapped + # -- CRUD ----------------------------------------------------------------- @abstractmethod @@ -72,3 +105,11 @@ class BaseFileStore(BaseComponent): Meant to be invoked off the request path (cron / idle schedulers). Backends without derived index state keep the default no-op. """ + + async def require_embedding_rebuild(self) -> None: + """Disable vector reads and writes until a full manual rebuild.""" + raise NotImplementedError + + async def reindex(self, scope: str) -> dict: + """Rebuild derived search indexes from current chunks without rescanning files.""" + raise NotImplementedError diff --git a/reme/components/file_store/faiss_local_file_store.py b/reme/components/file_store/faiss_local_file_store.py index 06fe6eda..12684e81 100644 --- a/reme/components/file_store/faiss_local_file_store.py +++ b/reme/components/file_store/faiss_local_file_store.py @@ -9,6 +9,7 @@ from uuid import uuid4 import aiofiles import numpy as np +from .base_file_store import BaseFileStore from .local_file_store import LocalFileStore from ..component_registry import R from ...schema import FileChunk, FileNode @@ -257,6 +258,15 @@ class FaissLocalFileStore(LocalFileStore): await self._stop_reindex_worker() self._rebuild_index() + async def _finalize_embedding_reindex(self) -> None: + """Build and publish the complete FAISS snapshot before job success.""" + await self._stop_reindex_worker() + while True: + self._reindex_event.clear() + await self._reindex_async() + if not self._reindex_event.is_set(): + return + # -- async reindex ---------------------------------------------------- def _submit_reindex(self) -> None: @@ -508,6 +518,7 @@ class FaissLocalFileStore(LocalFileStore): self.logger.info(f"Saved FAISS index: {self._faiss_index.ntotal} vectors to {self.faiss_path}") except Exception as e: self.logger.exception(f"Failed to write FAISS index: {e}") + raise async def _write_sidecar(self) -> None: token = uuid4().hex @@ -537,6 +548,7 @@ class FaissLocalFileStore(LocalFileStore): # -- CRUD overrides --------------------------------------------------- + @BaseFileStore.serialized async def upsert(self, files: list[tuple[FileNode, list[FileChunk]]]) -> None: if not files: return @@ -553,7 +565,7 @@ class FaissLocalFileStore(LocalFileStore): } await super().upsert(files) - if self._faiss_index is None or self.embedding_store is None: + if self._embedding_rebuild_pending or self._faiss_index is None or self.embedding_store is None: return self._sync_index_after_upsert(files, old_ids_by_path, old_text_by_id) @@ -587,13 +599,16 @@ class FaissLocalFileStore(LocalFileStore): self._add_to_index([c.id for c in to_add], vectors) self._compact_if_needed() + @BaseFileStore.serialized async def delete(self, path: str | list[str]) -> None: assert self.file_graph is not None paths = [path] if isinstance(path, str) else path nodes = await self.file_graph.get_nodes(paths) deleted_ids = [cid for n in nodes for cid in n.chunk_ids] await self._delete_nodes(nodes) # reuse resolved nodes; avoids a second get_nodes - if self._faiss_index is None: + if nodes: + self._mutation_generation += 1 + if self._embedding_rebuild_pending or self._faiss_index is None: return for cid in deleted_ids: self._tombstone(cid) @@ -609,6 +624,7 @@ class FaissLocalFileStore(LocalFileStore): await self._stop_reindex_worker() await super()._close() + @BaseFileStore.serialized async def clear(self) -> None: # Serialize with dump so a concurrent _write_sidecar cannot re-create the # sidecar files we are about to unlink, or persist a half-reset index. @@ -640,19 +656,9 @@ class FaissLocalFileStore(LocalFileStore): ): return [] - query_embedding = None - was_healthy = bool(getattr(self.embedding_store, "is_healthy", True)) - try: - query_embedding = await self.embedding_store.get_embedding(query) - except Exception as e: - self._mark_embedding_unhealthy(f"search: {type(e).__name__}: {e}") - if query_embedding is None or not self._embedding_dim_matches(query_embedding): - if query_embedding is not None: - self._mark_embedding_unhealthy( - f"search: query embedding dimension {len(query_embedding)} != {self.embedding_store.dimensions}", - ) + query_embedding = await self._get_query_embedding(query) + if query_embedding is None: return [] - await self._recover_after_real_request(was_healthy) # get_embedding above yielded control; a concurrent clear() drops the # index to None once embedding is disabled, and a reindex may have swapped diff --git a/reme/components/file_store/local_file_store.py b/reme/components/file_store/local_file_store.py index 2b11ea27..cab8024f 100644 --- a/reme/components/file_store/local_file_store.py +++ b/reme/components/file_store/local_file_store.py @@ -19,6 +19,7 @@ from ..keyword_index import BaseKeywordIndex from ...enumeration import LinkScopeEnum from ...schema import FileChunk, FileLink, FileNode from ...utils import batch_cosine_similarity +from ...utils.async_utils import complete_in_thread from ...utils.jsonl_zst import read_jsonl_zst, write_jsonl_zst CachedEmbedding = tuple[str, np.ndarray] @@ -46,6 +47,7 @@ class LocalFileStore(BaseFileStore): file_graph: str = "default", encoding: str = "utf-8", store_version: str = "v1", + embedding_rebuild_required: bool = False, **kwargs, ): super().__init__(**kwargs) @@ -67,8 +69,10 @@ class LocalFileStore(BaseFileStore): self.file_chunks: dict[str, FileChunk] = {} self.chunks_path = self.component_metadata_path / f"file_chunks_{self.name}_{self.store_version}.jsonl.zst" self._embedding_backfill_task: asyncio.Task | None = None - self._embedding_backfill_pending: tuple[bool, bool] | None = None - self._embedding_rebuild_pending = False + self._embedding_backfill_pending = False + self._embedding_rebuild_pending = bool(embedding_rebuild_required) + self._embedding_space_generation = 0 + self._mutation_generation = 0 self._closing = False # -- lifecycle ------------------------------------------------------------ @@ -122,29 +126,25 @@ class LocalFileStore(BaseFileStore): async def _recover_after_real_request(self, was_healthy: bool) -> None: """Schedule repair when a real, non-cache provider request recovers.""" - if self.embedding_store is None or was_healthy or not getattr(self.embedding_store, "is_healthy", True): + if ( + self.embedding_store is None + or self._embedding_rebuild_pending + or was_healthy + or not getattr(self.embedding_store, "is_healthy", True) + ): return self.logger.info(f"{self.name}: embedding provider recovered; scheduling missing-vector backfill") await self.resume_embedding(verified=True) - async def resume_embedding(self, *, verified: bool = False, rebuild: bool = False) -> bool: - """Resume a configured provider and schedule a deduplicated repair. - - Embedded applications may pass ``verified=True`` after they have already - made a successful real provider request, avoiding a redundant ping. Pass - ``rebuild=True`` when the active vector space changed; existing vectors - are derived data and are discarded before a full background rebuild. - """ + async def resume_embedding(self, *, verified: bool = False) -> bool: + """Resume same-vector-space repair after provider recovery.""" if self.embedding_store is None or self._closing: return False if verified: self.embedding_store.is_healthy = True - if rebuild: - await self._prepare_embedding_rebuild() - if not self.file_chunks: - self._embedding_rebuild_pending = False - return True - self._start_embedding_backfill(skip_health_check=verified, rebuild=rebuild) + if self._embedding_rebuild_pending: + return True + self._start_embedding_backfill(skip_health_check=verified) return True def _embedding_dim_matches(self, embedding: np.ndarray | None) -> bool: @@ -171,6 +171,46 @@ class LocalFileStore(BaseFileStore): for chunk in chunks: self._drop_stale_embedding(chunk, context) + def _embedding_request_is_current( + self, + embedding_store: BaseEmbeddingStore, + embedding_generation: int, + ) -> bool: + """Return whether an in-flight provider request still belongs to the active vector space.""" + return ( + not self._embedding_rebuild_pending + and embedding_generation == self._embedding_space_generation + and embedding_store is self.embedding_store + ) + + async def _get_query_embedding(self, query: str) -> np.ndarray | None: + """Embed a query only while its provider and vector space remain current.""" + embedding_store = self.embedding_store + if embedding_store is None or self._embedding_rebuild_pending or not query: + return None + + embedding_generation = self._embedding_space_generation + was_healthy = bool(getattr(embedding_store, "is_healthy", True)) + try: + query_embedding = await embedding_store.get_embedding(query) + except Exception as e: + if self._embedding_request_is_current(embedding_store, embedding_generation): + self._mark_embedding_unhealthy(f"search: {type(e).__name__}: {e}") + return None + + if query_embedding is None or not self._embedding_request_is_current( + embedding_store, + embedding_generation, + ): + return None + if not self._embedding_dim_matches(query_embedding): + self._mark_embedding_unhealthy( + f"search: query embedding dimension {len(query_embedding)} != {embedding_store.dimensions}", + ) + return None + await self._recover_after_real_request(was_healthy) + return query_embedding if self._embedding_request_is_current(embedding_store, embedding_generation) else None + # -- persistence ---------------------------------------------------------- async def load(self) -> None: @@ -269,11 +309,12 @@ class LocalFileStore(BaseFileStore): return self._drop_stale_embeddings(self.file_chunks.values(), "load") - def _start_embedding_backfill(self, *, skip_health_check: bool = False, rebuild: bool = False) -> None: + def _start_embedding_backfill(self, *, skip_health_check: bool = False) -> None: """Schedule startup embedding repair without delaying component readiness.""" started_at = time.monotonic() - if self._closing: - self.logger.info(f"{self.name}: embedding backfill skipped: reason=closing") + if self._closing or self._embedding_rebuild_pending: + reason = "closing" if self._closing else "manual_reindex_required" + self.logger.info(f"{self.name}: embedding backfill skipped: reason={reason}") return if not self.embedding_store: self.logger.info( @@ -282,14 +323,7 @@ class LocalFileStore(BaseFileStore): ) return if self._embedding_backfill_task is not None and not self._embedding_backfill_task.done(): - pending_verified = skip_health_check or bool( - self._embedding_backfill_pending and self._embedding_backfill_pending[0], - ) - pending_rebuild = rebuild or bool( - self._embedding_backfill_pending and self._embedding_backfill_pending[1], - ) - if pending_verified or pending_rebuild: - self._embedding_backfill_pending = (pending_verified, pending_rebuild) + self._embedding_backfill_pending |= skip_health_check self.logger.info( f"{self.name}: embedding backfill scheduling skipped: reason=already_running, " f"elapsed={time.monotonic() - started_at:.3f}s", @@ -302,7 +336,7 @@ class LocalFileStore(BaseFileStore): ) return self._embedding_backfill_task = asyncio.create_task( - self._run_embedding_backfill(skip_health_check=skip_health_check, rebuild=rebuild), + self._run_embedding_backfill(skip_health_check=skip_health_check), name=f"embedding-backfill:{self.name}", ) self.logger.info( @@ -310,37 +344,92 @@ class LocalFileStore(BaseFileStore): f"elapsed={time.monotonic() - started_at:.3f}s", ) - async def _run_embedding_backfill(self, *, skip_health_check: bool, rebuild: bool) -> None: + async def require_embedding_rebuild(self) -> None: + """Make all existing vectors unavailable without rebuilding them.""" + self._embedding_space_generation += 1 + self._embedding_rebuild_pending = True + await self._cancel_embedding_backfill() + + @BaseFileStore.serialized + async def reindex(self, scope: str) -> dict: + """Rebuild derived search indexes from ``file_chunks`` without touching files or the graph.""" + if scope not in {"all", "bm25", "embedding"}: + raise ValueError("reindex scope must be one of: all, bm25, embedding") + if scope == "bm25": + return await self._reindex_bm25() + if scope == "embedding": + return await self._reindex_embedding() + return { + "scope": "all", + "bm25": await self._reindex_bm25(), + "embedding": await self._reindex_embedding(), + } + + async def _reindex_bm25(self) -> dict: + if self.keyword_index is None: + return {"indexed": 0, "scope": "bm25"} + while True: + generation = self._mutation_generation + docs = {chunk_id: chunk.text for chunk_id, chunk in self.file_chunks.items() if chunk.text} + await self._rebuild_keyword_index(docs) + if generation == self._mutation_generation: + return {"indexed": len(docs), "scope": "bm25"} + + async def _reindex_embedding(self) -> dict: + await self.require_embedding_rebuild() + try: + while True: + generation = self._mutation_generation + embedding_generation = self._embedding_space_generation + chunks = tuple(self.file_chunks.values()) + for start in range(0, len(chunks), 512): + for chunk in chunks[start : start + 512]: + chunk.embedding = None + await asyncio.sleep(0) + await self._reset_vector_index() + if self.embedding_store is not None: + await self._backfill_missing_embeddings_inner( + skip_health_check=False, + started_at=time.monotonic(), + ) + missing = [ + chunk.id + for chunk in self.file_chunks.values() + if chunk.text and not self._embedding_dim_matches(chunk.embedding) + ] + if missing: + raise RuntimeError(f"embedding reindex incomplete: {len(missing)} chunks failed") + await self._finalize_embedding_reindex() + await self._dump_owned_state() + if self.embedding_store is not None: + await self.embedding_store.dump() + if generation == self._mutation_generation and embedding_generation == self._embedding_space_generation: + self._embedding_rebuild_pending = False + indexed = len(chunks) if self.embedding_store is not None else 0 + return {"indexed": indexed, "scope": "embedding"} + except BaseException: + self._embedding_rebuild_pending = True + raise + + async def _run_embedding_backfill(self, *, skip_health_check: bool) -> None: """Run one repair and honor a verified request queued behind it.""" current_task = asyncio.current_task() try: - if rebuild: - # A task that was already running when rebuild was requested - # may have written a stale provider result after the first - # invalidation. Clear once more at the queue boundary. - await self._prepare_embedding_rebuild() await self._backfill_missing_embeddings(skip_health_check=skip_health_check) finally: if self._embedding_backfill_task is current_task: self._embedding_backfill_task = None pending = self._embedding_backfill_pending - self._embedding_backfill_pending = None - if pending is not None and not self._closing and self.embedding_store is not None: - pending_verified, pending_rebuild = pending - if pending_verified: - self.embedding_store.is_healthy = True - if pending_rebuild: - self._embedding_rebuild_pending = True - self._start_embedding_backfill( - skip_health_check=pending_verified, - rebuild=pending_rebuild, - ) + self._embedding_backfill_pending = False + if pending and not self._closing and self.embedding_store is not None: + self.embedding_store.is_healthy = True + self._start_embedding_backfill(skip_health_check=True) async def _cancel_embedding_backfill(self) -> None: """Cancel and collect the startup repair task during component shutdown.""" task = self._embedding_backfill_task self._embedding_backfill_task = None - self._embedding_backfill_pending = None + self._embedding_backfill_pending = False if task is None: return if not task.done(): @@ -367,16 +456,10 @@ class LocalFileStore(BaseFileStore): async def _backfill_missing_embeddings(self, *, skip_health_check: bool = False) -> None: """Background-repair persisted chunks that do not have usable vectors.""" started_at = time.monotonic() - try: - await self._backfill_missing_embeddings_inner(skip_health_check=skip_health_check, started_at=started_at) - finally: - if self._embedding_rebuild_pending and not self._closing: - try: - await self._after_embedding_backfill() - await self.dump() - self._embedding_rebuild_pending = False - except Exception: - self.logger.exception(f"{self.name}: failed to finalize embedding rebuild") + await self._backfill_missing_embeddings_inner( + skip_health_check=skip_health_check, + started_at=started_at, + ) async def _backfill_missing_embeddings_inner(self, *, skip_health_check: bool, started_at: float) -> None: """Perform one backfill pass; the caller owns rebuild finalization.""" @@ -462,20 +545,16 @@ class LocalFileStore(BaseFileStore): except Exception: self.logger.exception(f"{self.name}: failed to persist completed embedding backfill") - async def _prepare_embedding_rebuild(self) -> None: - """Invalidate and persist vectors from the previous vector space.""" - self._embedding_rebuild_pending = True - for chunk in self.file_chunks.values(): - chunk.embedding = None - await self._reset_vector_index() - await self.dump() - async def _reset_vector_index(self) -> None: """Drop a derived vector index before rebuilding a changed vector space.""" async def _after_embedding_backfill(self) -> None: """Backend hook for refreshing derived vector indexes after backfill.""" + async def _finalize_embedding_reindex(self) -> None: + """Synchronously publish derived indexes for an explicit reindex job.""" + await self._after_embedding_backfill() + async def _sync_keyword_index_from_chunks(self) -> None: """Repair keyword index when its persisted state does not match chunks.""" if not self.keyword_index: @@ -541,15 +620,24 @@ class LocalFileStore(BaseFileStore): async def _dump_owned_state(self) -> None: """Persist state owned by this store, excluding dependency snapshots.""" try: - write_jsonl_zst( - self.chunks_path, - (self._serialize_chunk(c) for c in self.file_chunks.values()), - self.encoding, - ) + # Keep snapshotting synchronous so concurrent mutation cannot produce a + # mixed-generation checkpoint. Move it off-loop only if profiling shows + # this copy, rather than serialization/compression, is a material stall. + chunks = tuple(chunk.model_copy(deep=True) for chunk in self.file_chunks.values()) + await complete_in_thread(self._dump_chunks_sync, chunks) self.logger.info(f"Saved {len(self.file_chunks)} chunks to {self.chunks_path}") except Exception as e: self.logger.exception(f"Failed to write {self.chunks_path}: {e}") + raise + def _dump_chunks_sync(self, chunks: tuple[FileChunk, ...]) -> None: + write_jsonl_zst( + self.chunks_path, + (self._serialize_chunk(chunk) for chunk in chunks), + self.encoding, + ) + + @BaseFileStore.serialized async def dump(self) -> None: """Persist a complete store/index/graph consistency checkpoint.""" assert self.file_graph is not None @@ -572,6 +660,7 @@ class LocalFileStore(BaseFileStore): # -- CRUD ----------------------------------------------------------------- + @BaseFileStore.serialized async def upsert(self, files: list[tuple[FileNode, list[FileChunk]]]) -> None: if not files: return @@ -587,6 +676,7 @@ class LocalFileStore(BaseFileStore): await self.keyword_index.delete_docs(list(old_chunk_ids)) if self.keyword_index and keyword_docs: await self.keyword_index.add_docs(keyword_docs) + self._mutation_generation += 1 def _stage_upsert( self, @@ -630,6 +720,9 @@ class LocalFileStore(BaseFileStore): cached: dict[str, CachedEmbedding], needs_embed: list[FileChunk], ) -> None: + if self._embedding_rebuild_pending: + chunk.embedding = None + return if not self.embedding_store: return if chunk.embedding is not None: @@ -646,23 +739,35 @@ class LocalFileStore(BaseFileStore): needs_embed.append(chunk) async def _embed_pending(self, chunks: list[FileChunk]) -> None: - if not (chunks and self.embedding_store): + if not (chunks and self.embedding_store) or self._embedding_rebuild_pending: return - was_healthy = bool(getattr(self.embedding_store, "is_healthy", True)) + embedding_store = self.embedding_store + embedding_generation = self._embedding_space_generation + was_healthy = bool(getattr(embedding_store, "is_healthy", True)) try: - await self.embedding_store.get_node_embeddings(chunks) + await embedding_store.get_node_embeddings(chunks) except Exception as e: - self._mark_embedding_unhealthy(f"upsert: {type(e).__name__}: {e}") + if self._embedding_request_is_current(embedding_store, embedding_generation): + self._mark_embedding_unhealthy(f"upsert: {type(e).__name__}: {e}") + return + if not self._embedding_request_is_current(embedding_store, embedding_generation): + for chunk in chunks: + chunk.embedding = None + if not self._embedding_rebuild_pending: + self._start_embedding_backfill(skip_health_check=True) return self._drop_stale_embeddings(chunks, "upsert") if any(chunk.embedding is not None for chunk in chunks): await self._recover_after_real_request(was_healthy) + @BaseFileStore.serialized async def delete(self, path: str | list[str]) -> None: assert self.file_graph is not None paths = [path] if isinstance(path, str) else path nodes: list[FileNode] = await self.file_graph.get_nodes(paths) await self._delete_nodes(nodes) + if nodes: + self._mutation_generation += 1 async def _delete_nodes(self, nodes: list[FileNode]) -> None: """Delete already-resolved nodes and their chunks. @@ -701,6 +806,7 @@ class LocalFileStore(BaseFileStore): assert self.file_graph is not None return await self.file_graph.get_inlinks(path, scope) + @BaseFileStore.serialized async def clear(self) -> None: assert self.file_graph is not None self.file_chunks.clear() @@ -708,27 +814,17 @@ class LocalFileStore(BaseFileStore): if self.keyword_index: await self.keyword_index.clear() await self.file_graph.clear() + self._mutation_generation += 1 # -- search --------------------------------------------------------------- async def vector_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]: - if self.embedding_store is None or self._embedding_rebuild_pending or not query or limit <= 0: + if limit <= 0: return [] - was_healthy = bool(getattr(self.embedding_store, "is_healthy", True)) - try: - query_embedding = await self.embedding_store.get_embedding(query) - except Exception as e: - self._mark_embedding_unhealthy(f"search: {type(e).__name__}: {e}") - return [] + query_embedding = await self._get_query_embedding(query) if query_embedding is None: return [] - if not self._embedding_dim_matches(query_embedding): - self._mark_embedding_unhealthy( - f"search: query embedding dimension {len(query_embedding)} != {self.embedding_store.dimensions}", - ) - return [] - await self._recover_after_real_request(was_healthy) top: list[tuple[float, int, FileChunk]] = [] candidates: list[FileChunk] = [] diff --git a/reme/components/file_store/zvec_local_file_store.py b/reme/components/file_store/zvec_local_file_store.py index 5925fd84..afad4b4d 100644 --- a/reme/components/file_store/zvec_local_file_store.py +++ b/reme/components/file_store/zvec_local_file_store.py @@ -9,9 +9,11 @@ from uuid import uuid4 import aiofiles import numpy as np +from .base_file_store import BaseFileStore from .local_file_store import LocalFileStore from ..component_registry import R from ...schema import FileChunk, FileNode +from ...utils.async_utils import complete_in_thread # Batch size for bulk inserts during a rebuild. _ZVEC_INSERT_BATCH_SIZE = 1024 @@ -105,14 +107,19 @@ class ZvecLocalFileStore(LocalFileStore): def _create_collection(self): """Create a fresh (empty) collection, replacing any directory on disk.""" + self._discard_collection() + return self._zvec.create_and_open(path=str(self.zvec_path), schema=self._collection_schema()) + + def _discard_collection(self) -> None: + """Release and remove the derived collection and its generation sidecar.""" # Release any open handle first: zvec holds an in-process lock on the # collection directory, so the old object must be dropped before the - # directory is wiped and re-created. + # directory is wiped. self._collection = None if self.zvec_path.exists(): shutil.rmtree(self.zvec_path, ignore_errors=True) self._indexed_ids = set() - return self._zvec.create_and_open(path=str(self.zvec_path), schema=self._collection_schema()) + self.zvec_sidecar_path.unlink(missing_ok=True) def _to_doc(self, chunk: FileChunk): """Build a zvec Doc carrying only the id and the float32 vector.""" @@ -178,7 +185,14 @@ class ZvecLocalFileStore(LocalFileStore): async def _reset_vector_index(self) -> None: """Discard all vectors before rebuilding a changed vector space.""" - self._collection = self._create_collection() + if self.embedding_store is None or self._dim == 0: + await complete_in_thread(self._discard_collection) + return + self._collection = await complete_in_thread(self._create_collection) + + async def _finalize_embedding_reindex(self) -> None: + """Publish the complete zvec snapshot before explicit job success.""" + await complete_in_thread(self._rebuild_collection) # -- maintenance ------------------------------------------------------ @@ -330,6 +344,7 @@ class ZvecLocalFileStore(LocalFileStore): self.logger.info(f"Saved zvec collection: {len(self._indexed_ids)} vectors to {self.zvec_path}") except Exception as e: self.logger.exception(f"Failed to persist zvec collection: {e}") + raise async def _write_sidecar(self) -> None: """Atomically write the digest sidecar binding the collection to the chunk generation.""" @@ -349,6 +364,7 @@ class ZvecLocalFileStore(LocalFileStore): # -- CRUD overrides --------------------------------------------------- + @BaseFileStore.serialized async def upsert(self, files: list[tuple[FileNode, list[FileChunk]]]) -> None: if not files: return @@ -359,7 +375,7 @@ class ZvecLocalFileStore(LocalFileStore): old_ids_by_path = {n.path: set(n.chunk_ids) for n in old_nodes} await super().upsert(files) - if self._collection is None or self.embedding_store is None: + if self._embedding_rebuild_pending or self._collection is None or self.embedding_store is None: return self._sync_collection_after_upsert(files, old_ids_by_path) @@ -392,14 +408,20 @@ class ZvecLocalFileStore(LocalFileStore): self._delete_docs(to_delete) self._upsert_docs(to_upsert) + @BaseFileStore.serialized async def delete(self, path: str | list[str]) -> None: assert self.file_graph is not None paths = [path] if isinstance(path, str) else path nodes = await self.file_graph.get_nodes(paths) deleted_ids = [cid for n in nodes for cid in n.chunk_ids] await self._delete_nodes(nodes) # reuse resolved nodes; avoids a second get_nodes + if nodes: + self._mutation_generation += 1 + if self._embedding_rebuild_pending: + return self._delete_docs(deleted_ids) + @BaseFileStore.serialized async def clear(self) -> None: await super().clear() if self._collection is not None: @@ -422,19 +444,9 @@ class ZvecLocalFileStore(LocalFileStore): if index_empty and getattr(self.embedding_store, "is_healthy", True): return [] - query_embedding = None - was_healthy = bool(getattr(self.embedding_store, "is_healthy", True)) - try: - query_embedding = await self.embedding_store.get_embedding(query) - except Exception as e: - self._mark_embedding_unhealthy(f"search: {type(e).__name__}: {e}") - if query_embedding is None or not self._embedding_dim_matches(query_embedding): - if query_embedding is not None: - self._mark_embedding_unhealthy( - f"search: query embedding dimension {len(query_embedding)} != {self.embedding_store.dimensions}", - ) + query_embedding = await self._get_query_embedding(query) + if query_embedding is None: return [] - await self._recover_after_real_request(was_healthy) # get_embedding above yielded control; a concurrent clear() may have # swapped or dropped the collection. Re-read before dereferencing. diff --git a/reme/components/keyword_index/bm25_index.py b/reme/components/keyword_index/bm25_index.py index 0cb3db86..0169726f 100644 --- a/reme/components/keyword_index/bm25_index.py +++ b/reme/components/keyword_index/bm25_index.py @@ -15,6 +15,7 @@ lists keep the stale entries until `optimize_index` rewrites them. Updating an existing doc_id retires the old slot first, then allocates a fresh idx. """ +import asyncio import hashlib import json import math @@ -23,11 +24,13 @@ import re from collections import Counter from collections.abc import KeysView from pathlib import Path +from uuid import uuid4 import numpy as np from .base_keyword_index import BaseKeywordIndex from ..component_registry import R +from ...utils.async_utils import complete_in_thread @R.register("bm25") @@ -52,6 +55,7 @@ class BM25Index(BaseKeywordIndex): # IDF cache; invalidated whenever live-doc count or postings change. self._idf_cache: dict[int, float] = {} + self._dump_lock = asyncio.Lock() # -- Properties ----------------------------------------------------------- @@ -322,14 +326,16 @@ class BM25Index(BaseKeywordIndex): return { "tokenizer_config": self._tokenizer_config(), "tokenizer_fingerprint": self._tokenizer_fingerprint(), - "vocab": self.vocab, - "doc_ids": self._doc_ids, - "doc_id_to_idx": self._doc_id_to_idx, - "doc_lens": self._doc_lens, - "deleted": self._deleted, - "doc_token_ids": self._doc_token_ids, - "posting_doc_idxs": self._posting_doc_idxs, - "posting_tfs": self._posting_tfs, + "vocab": dict(self.vocab), + "doc_ids": list(self._doc_ids), + "doc_id_to_idx": dict(self._doc_id_to_idx), + "doc_lens": self._doc_lens.copy(), + "deleted": self._deleted.copy(), + "doc_token_ids": [token_ids.copy() for token_ids in self._doc_token_ids], + "posting_doc_idxs": {token_id: doc_idxs.copy() for token_id, doc_idxs in self._posting_doc_idxs.items()}, + "posting_tfs": { + token_id: term_frequencies.copy() for token_id, term_frequencies in self._posting_tfs.items() + }, "k1": self.k1, "b": self.b, } @@ -354,27 +360,37 @@ class BM25Index(BaseKeywordIndex): async def dump(self) -> None: """Persist the index via temp file + atomic rename to avoid torn writes.""" - if self.n_docs == 0 and not self.vocab: - self.index_file.unlink(missing_ok=True) - return + async with self._dump_lock: + if self.n_docs == 0 and not self.vocab: + self.index_file.unlink(missing_ok=True) + return + try: + # Keep snapshotting synchronous so the worker receives one coherent + # index generation. Move it off-loop only if profiling identifies this + # copy, rather than pickle/file I/O, as a material event-loop stall. + snapshot = self._snapshot() + await complete_in_thread(self._dump_sync, snapshot) + self.logger.info(f"Saved {self.n_docs} docs to {self.index_file}") + except Exception as e: + self.logger.exception(f"Failed to write {self.index_file}: {e}") + raise + + def _dump_sync(self, snapshot: dict) -> None: + self.index_file.parent.mkdir(parents=True, exist_ok=True) + tmp = self.index_file.with_name(f".{self.index_file.name}.{uuid4().hex}.tmp") try: - self.index_file.parent.mkdir(parents=True, exist_ok=True) - tmp = self.index_file.with_suffix(".tmp") - with open(tmp, "wb") as f: - pickle.dump(self._snapshot(), f) + with open(tmp, "wb") as file: + pickle.dump(snapshot, file) tmp.replace(self.index_file) - self.logger.info(f"Saved {self.n_docs} docs to {self.index_file}") - except Exception as e: - self.logger.exception(f"Failed to write {self.index_file}: {e}") - raise + finally: + tmp.unlink(missing_ok=True) async def load(self) -> None: """Load from disk; missing file is a no-op, corrupt file resets state.""" if not self.index_file.exists(): return try: - with open(self.index_file, "rb") as f: - data = pickle.load(f) + data = await asyncio.to_thread(self._load_sync) self._restore(data) self.logger.info(f"Loaded {self.n_docs} docs from {self.index_file}") except Exception as e: @@ -382,18 +398,23 @@ class BM25Index(BaseKeywordIndex): self.index_file.unlink(missing_ok=True) await self.clear() + def _load_sync(self) -> dict: + with open(self.index_file, "rb") as file: + return pickle.load(file) + async def clear(self) -> None: """Reset in-memory state and remove the persisted file.""" - self.vocab = {} - self._doc_ids = [] - self._doc_id_to_idx = {} - self._doc_lens = np.zeros(0, dtype=np.int32) - self._deleted = np.zeros(0, dtype=bool) - self._doc_token_ids = [] - self._posting_doc_idxs = {} - self._posting_tfs = {} - self._idf_cache = {} - self.index_file.unlink(missing_ok=True) + async with self._dump_lock: + self.vocab = {} + self._doc_ids = [] + self._doc_id_to_idx = {} + self._doc_lens = np.zeros(0, dtype=np.int32) + self._deleted = np.zeros(0, dtype=bool) + self._doc_token_ids = [] + self._posting_doc_idxs = {} + self._posting_tfs = {} + self._idf_cache = {} + self.index_file.unlink(missing_ok=True) # -- Compaction ----------------------------------------------------------- diff --git a/reme/config/beam.yaml b/reme/config/beam.yaml index 22bd9fe4..c3d6bda1 100644 --- a/reme/config/beam.yaml +++ b/reme/config/beam.yaml @@ -192,21 +192,19 @@ jobs: - backend: compressor_step as_llm: compressor - # ── Reindex (full rebuild) ── + # ── Reindex (derived search indexes only) ── reindex: backend: base - description: "wipe the file store and rebuild it from the existing files" - watch_dirs: [daily_dir, digest_dir, resource_dir] - watch_suffixes: [md, jsonl] + description: "rebuild BM25 and/or embedding indexes from current file_chunks" parameters: type: object - properties: {} + properties: + scope: + type: string + enum: [all, bm25, embedding] + default: all steps: - - backend: clear_store_step - - backend: init_changes_step - monitor_type: file_store - monitor_name: default - dispatch_steps: [update_index_step] + - backend: reindex_step # ── Search ── # start_date: diff --git a/reme/config/default.yaml b/reme/config/default.yaml index 4faa9b60..855af208 100644 --- a/reme/config/default.yaml +++ b/reme/config/default.yaml @@ -335,18 +335,16 @@ jobs: reindex: backend: base - description: "wipe the file store and rebuild it from the existing files" - watch_dirs: [daily_dir, digest_dir, resource_dir] - watch_suffixes: [md, jsonl] + description: "rebuild BM25 and/or embedding indexes from current file_chunks" parameters: type: object - properties: { } + properties: + scope: + type: string + enum: [all, bm25, embedding] + default: all steps: - - backend: clear_store_step - - backend: init_changes_step - monitor_type: file_store - monitor_name: default - dispatch_steps: [update_index_step] + - backend: reindex_step search: backend: base diff --git a/reme/config/lme.yaml b/reme/config/lme.yaml index 2c823946..24ec4e40 100644 --- a/reme/config/lme.yaml +++ b/reme/config/lme.yaml @@ -185,21 +185,19 @@ jobs: - backend: compressor_step as_llm: compressor - # ── Reindex (full rebuild) ── + # ── Reindex (derived search indexes only) ── reindex: backend: base - description: "wipe the file store and rebuild it from the existing files" - watch_dirs: [daily_dir, digest_dir, resource_dir] - watch_suffixes: [md, jsonl] + description: "rebuild BM25 and/or embedding indexes from current file_chunks" parameters: type: object - properties: {} + properties: + scope: + type: string + enum: [all, bm25, embedding] + default: all steps: - - backend: clear_store_step - - backend: init_changes_step - monitor_type: file_store - monitor_name: default - dispatch_steps: [update_index_step] + - backend: reindex_step # ── Search ── # start_date: diff --git a/reme/steps/index/__init__.py b/reme/steps/index/__init__.py index 4c66c608..b46dc9c5 100644 --- a/reme/steps/index/__init__.py +++ b/reme/steps/index/__init__.py @@ -10,6 +10,7 @@ from .log_changes import LogChangesStep from .node_search import NodeSearchStep from .init_changes import InitChangesStep from .optimize_index import OptimizeIndexStep +from .reindex import ReindexStep from .search import SearchStep from .search_v2 import SearchV2Step from .traverse import TraverseStep @@ -38,6 +39,7 @@ __all__ = [ "NodeSearchStep", "normalize_posix_path", "ReadAllDraftStep", + "ReindexStep", "OptimizeIndexStep", "SearchStep", "SearchV2Step", diff --git a/reme/steps/index/reindex.py b/reme/steps/index/reindex.py new file mode 100644 index 00000000..1f003a2b --- /dev/null +++ b/reme/steps/index/reindex.py @@ -0,0 +1,19 @@ +"""Explicit scoped rebuild of search indexes from already-ingested chunks.""" + +from ..base_step import BaseStep +from ...components import R + + +@R.register("reindex_step") +class ReindexStep(BaseStep): + """Rebuild BM25 and/or embeddings without scanning files or changing the graph.""" + + async def execute(self): + assert self.context is not None + scope = str(self.context.get("scope", "all")) + details = await self.file_store.reindex(scope) + + self.context.response.answer = details + self.context.response.metadata.update(details) + self.context.response.metadata["scope"] = scope + return self.context.response diff --git a/reme/utils/async_utils.py b/reme/utils/async_utils.py new file mode 100644 index 00000000..290a4456 --- /dev/null +++ b/reme/utils/async_utils.py @@ -0,0 +1,17 @@ +"""Small asyncio helpers.""" + +import asyncio +from collections.abc import Callable +from contextlib import suppress +from typing import Any + + +async def complete_in_thread(func: Callable[..., Any], /, *args) -> Any: + """Finish a side-effecting thread call before propagating cancellation.""" + task = asyncio.create_task(asyncio.to_thread(func, *args)) + try: + return await asyncio.shield(task) + except asyncio.CancelledError: + with suppress(Exception): + await task + raise diff --git a/reme_studio/README_ZH.md b/reme_studio/README_ZH.md index 8195702c..d238f788 100644 --- a/reme_studio/README_ZH.md +++ b/reme_studio/README_ZH.md @@ -102,7 +102,7 @@ NEXT_PUBLIC_REME_WORKSPACE_EXTENSIONS=md,txt,mdx VITE_REME_WORKSPACE_EXTENSIONS=md,txt,mdx ``` -记忆图谱依赖 ReMe 构建的索引。在 Studio 设置中重建索引时，只会根据工作区文件重新生成派生数据，不会修改记忆源文件。 +记忆图谱依赖 ReMe 的摄取流程。在 Studio 设置中重建索引时，只会基于已摄取的 chunks 重建 BM25 和 Embedding 索引；不会扫描工作区文件、重新分块、修改记忆源文件或重建 wikilink 图谱。 ## 检查 diff --git a/reme_studio/app/i18n.ts b/reme_studio/app/i18n.ts index cf7ea14d..adcf704b 100644 --- a/reme_studio/app/i18n.ts +++ b/reme_studio/app/i18n.ts @@ -120,12 +120,12 @@ const messages = { vocabulary: "词汇量", memoryUsage: "内存占用", indexTitle: "工作区索引", - indexDescription: "从现有文件重新构建搜索索引。记忆文件不会被修改。", + indexDescription: "基于已摄取的内容块重建 BM25 和 Embedding 索引。", rebuildIndex: "重建索引", rebuildingIndex: "正在重建…", confirmReindexTitle: "确定重建索引？", confirmReindexDescription: - "现有派生索引会被清空，然后根据工作区文件重新生成。", + "现有 BM25 和 Embedding 索引会基于当前已摄取的内容块重建。不会扫描文件、重新分块或重建图谱。", cancel: "取消", confirmReindex: "确认重建", indexRebuilt: "索引重建完成", @@ -257,12 +257,12 @@ const messages = { memoryUsage: "Memory", indexTitle: "Workspace index", indexDescription: - "Rebuild the search index from existing files. Memory files are not modified.", + "Rebuild BM25 and embedding indexes from already-ingested chunks.", rebuildIndex: "Rebuild index", rebuildingIndex: "Rebuilding…", confirmReindexTitle: "Rebuild the index?", confirmReindexDescription: - "The derived index will be cleared and regenerated from workspace files.", + "BM25 and embedding indexes will be rebuilt from the currently ingested chunks. Files are not scanned or rechunked, and the graph is not rebuilt.", cancel: "Cancel", confirmReindex: "Confirm rebuild", indexRebuilt: "Index rebuilt", diff --git a/tests/integration/test_reme_e2e.py b/tests/integration/test_reme_e2e.py index 77c2f3cd..6afacc16 100644 --- a/tests/integration/test_reme_e2e.py +++ b/tests/integration/test_reme_e2e.py @@ -160,7 +160,7 @@ async def _run_loop(env, reme) -> None: print(f"[3/5] digest nodes: {[str(p.relative_to(env.workspace_dir)) for p in digest_paths]}") assert len(carried) >= 3, f"consolidation lost the topic; only {carried!r} survived" - # ---- 4. recall: rebuild the index, then search it ----------- + # ---- 4. recall: rebuild derived search indexes, then search ----------- reindex = await reme.run_job("reindex") assert reindex.success is True, f"reindex failed: {reindex.answer!r}" diff --git a/tests/unit/test_background_steps.py b/tests/unit/test_background_steps.py index 4a326378..342e42b2 100644 --- a/tests/unit/test_background_steps.py +++ b/tests/unit/test_background_steps.py @@ -225,8 +225,8 @@ def test_collect_existing_filters(): # --------------------------------------------------------------------------- -def test_clear_and_scan_defaults_include_jsonl(): - """Full reindex should include jsonl files when no explicit suffix filter is passed.""" +def test_clear_and_scan_flow_includes_jsonl_when_configured(): + """The legacy clear-and-scan primitives accept configured JSONL inputs.""" async def run(): with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): @@ -1913,7 +1913,7 @@ if __name__ == "__main__": test_match_file_no_suffix_filter() test_collect_existing_filters() # InitChangesStep - test_clear_and_scan_defaults_include_jsonl() + test_clear_and_scan_flow_includes_jsonl_when_configured() test_scan_changes_initial_all_added() test_scan_changes_no_changes() test_scan_changes_detect_modify_delete() diff --git a/tests/unit/test_file_store_consistency.py b/tests/unit/test_file_store_consistency.py index eb693bdb..26f2dce6 100644 --- a/tests/unit/test_file_store_consistency.py +++ b/tests/unit/test_file_store_consistency.py @@ -10,11 +10,16 @@ import os import tempfile import threading import time +from unittest.mock import AsyncMock import numpy as np import pytest -from reme.components.file_store import FaissLocalFileStore, LocalFileStore, ZvecLocalFileStore +from reme.components.file_store import ( + FaissLocalFileStore, + LocalFileStore, + ZvecLocalFileStore, +) from reme.components.file_store import local_file_store as local_file_store_module from reme.components.file_graph import local_file_graph as local_file_graph_module from reme.components.keyword_index import bm25_index as bm25_index_module @@ -64,6 +69,9 @@ class FakeEmbeddingStore: chunk_node.embedding = self._embed(chunk_node.text) return nodes + async def dump(self) -> None: + """Persist no state for the in-memory fake.""" + class CountingFakeEmbeddingStore(FakeEmbeddingStore): """Fake embedding store that records node backfill requests.""" @@ -133,6 +141,19 @@ class BlockingEmbeddingStore(FakeEmbeddingStore): return await super().get_node_embeddings(nodes, **kwargs) +class BlockingQueryEmbeddingStore(FakeEmbeddingStore): + """Fake provider that holds a query request across a vector-space change.""" + + def __init__(self): + self.query_started = asyncio.Event() + self.release_query = asyncio.Event() + + async def get_embedding(self, input_text: str, **kwargs) -> np.ndarray: + self.query_started.set() + await self.release_query.wait() + return await super().get_embedding(input_text, **kwargs) + + class CancellationResistantHealthStore(CountingFakeEmbeddingStore): """Startup probe that completes stale after cancellation is requested.""" @@ -395,7 +416,7 @@ def test_load_rebuilds_keyword_index_from_persisted_chunks_when_missing(): def test_load_clears_graph_when_persisted_chunks_are_missing(): - """A surviving graph must not hide a missing chunk store from reindex.""" + """A surviving graph must not hide a missing chunk store from automatic file ingestion.""" async def go(): with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): @@ -773,7 +794,7 @@ def test_verified_resume_supersedes_inflight_startup_health_check(): recovery = asyncio.create_task(store.resume_embedding(verified=True)) await asyncio.sleep(0) assert await recovery is True - assert store._embedding_backfill_pending == (True, False) + assert store._embedding_backfill_pending is True fake.release_health.set() await startup_task @@ -804,7 +825,7 @@ def test_verified_resume_without_chunks_supersedes_inflight_health_check(): await store.clear() assert await store.resume_embedding(verified=True) is True - assert store._embedding_backfill_pending == (True, False) + assert store._embedding_backfill_pending is True fake.release_health.set() await startup_task @@ -817,116 +838,484 @@ def test_verified_resume_without_chunks_supersedes_inflight_health_check(): @pytest.mark.parametrize("store_factory", [_new_local_store, _new_faiss_store, _new_zvec_store]) -def test_verified_rebuild_discards_same_dimension_vectors_before_backfill(store_factory): - """A changed vector space never searches compatible-shaped stale vectors.""" +def test_embedding_reindex_is_explicit_and_keeps_vectors_disabled_until_success( + store_factory, +): + """A pending vector-space change can only be completed by scoped reindex.""" async def go(): with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): - store = store_factory("t_embedding_verified_rebuild") + store = store_factory("t_explicit_embedding_reindex") await store.start() stale = chunk("a", "a.md", "alpha text") stale.embedding = np.array([0.0, 1.0], dtype=np.float16) await set_chunks_with_graph(store, {"a": stale}) fake = CountingFakeEmbeddingStore() - fake.is_healthy = False store.embedding_store = fake + _ensure_zvec_collection(store) if isinstance(store, FaissLocalFileStore): store._rebuild_index() elif isinstance(store, ZvecLocalFileStore): store._rebuild_collection() - assert await store.resume_embedding(verified=True, rebuild=True) is True - assert store._embedding_rebuild_pending is True - assert store.file_chunks["a"].embedding is None + await store.require_embedding_rebuild() assert await store.vector_search("alpha", 5, {}) == [] + assert await store.resume_embedding(verified=True) is True + assert store._embedding_backfill_task is None - await store._embedding_backfill_task + result = await store.reindex("embedding") + + assert result == {"indexed": 1, "scope": "embedding"} assert store._embedding_rebuild_pending is False assert fake.node_embedding_calls == [["a"]] - assert store.file_chunks["a"].embedding.tolist() == [1.0, 0.0] assert [item.id for item in await store.vector_search("alpha", 5, {})] == ["a"] await store.close() run(go()) -def test_verified_rebuild_discards_late_result_from_previous_vector_space(): - """A queued rebuild clears old-space vectors written by an in-flight batch.""" +@pytest.mark.parametrize("store_factory", [_new_local_store, _new_faiss_store, _new_zvec_store]) +def test_embedding_gate_rejects_query_that_started_in_previous_vector_space( + store_factory, +): + """A query crossing the gate boundary cannot return results from the old index.""" async def go(): with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): - store = _new_local_store("t_embedding_verified_rebuild_race") + store = store_factory("t_query_crosses_embedding_gate") await store.start() - await set_chunks_with_graph(store, {"a": chunk("a", "a.md", "alpha text")}) - fake = DelayedOldVectorStore() - store.embedding_store = fake - store._start_embedding_backfill(skip_health_check=True) - old_task = store._embedding_backfill_task - await fake.first_batch_started.wait() + indexed = chunk("a", "a.md", "alpha") + indexed.embedding = np.array([1.0, 0.0], dtype=np.float16) + await set_chunks_with_graph(store, {"a": indexed}) + blocking = BlockingQueryEmbeddingStore() + store.embedding_store = blocking + _ensure_zvec_collection(store) + if isinstance(store, FaissLocalFileStore): + store._rebuild_index() + elif isinstance(store, ZvecLocalFileStore): + store._rebuild_collection() - assert await store.resume_embedding(verified=True, rebuild=True) is True - assert store._embedding_backfill_pending == (True, True) - fake.release_first_batch.set() + search_task = asyncio.create_task(store.vector_search("alpha", 5, {})) + await blocking.query_started.wait() + await store.require_embedding_rebuild() + blocking.release_query.set() - await old_task - if store._embedding_backfill_task is not None: - await store._embedding_backfill_task - assert fake.node_embedding_calls == [["a"], ["a"]] - assert store.file_chunks["a"].embedding.tolist() == [1.0, 0.0] + assert await search_task == [] + await store.close() + + run(go()) + + +def test_embedding_reindex_retries_if_a_new_vector_space_is_required_while_finishing(): + """A new gate request cannot be cleared by an older reindex generation.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_local_store("t_reindex_new_embedding_generation") + await store.start() + await set_chunks_with_graph(store, {"a": chunk("a", "a.md", "alpha")}) + store.embedding_store = CountingFakeEmbeddingStore() + first_dump_started = asyncio.Event() + release_first_dump = asyncio.Event() + dump_calls = 0 + finalize_calls = 0 + real_dump = store._dump_owned_state + real_finalize = store._finalize_embedding_reindex + + async def blocking_dump(): + nonlocal dump_calls + dump_calls += 1 + if dump_calls == 1: + first_dump_started.set() + await release_first_dump.wait() + await real_dump() + + async def counting_finalize(): + nonlocal finalize_calls + finalize_calls += 1 + await real_finalize() + + store._dump_owned_state = blocking_dump + store._finalize_embedding_reindex = counting_finalize + reindex_task = asyncio.create_task(store.reindex("embedding")) + await first_dump_started.wait() + + await store.require_embedding_rebuild() + release_first_dump.set() + + assert await reindex_task == {"indexed": 1, "scope": "embedding"} + assert finalize_calls == 2 assert store._embedding_rebuild_pending is False await store.close() run(go()) -def test_unverified_rebuild_is_queued_behind_inflight_backfill(): - """An unverified rebuild request cannot be lost while another batch is running.""" +def test_upsert_discards_provider_result_from_previous_vector_space(): + """A late foreground embedding write is repaired in the current vector space.""" async def go(): with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): - store = _new_local_store("t_embedding_unverified_rebuild_race") + store = _new_local_store("t_upsert_previous_embedding_generation") await store.start() - await set_chunks_with_graph(store, {"a": chunk("a", "a.md", "alpha text")}) - fake = DelayedOldVectorStore() + old_space = DelayedOldVectorStore() + store.embedding_store = old_space + upsert_task = asyncio.create_task(store.upsert([(node("a.md"), [chunk("a", "a.md", "alpha")])])) + await old_space.first_batch_started.wait() + + current_space = CountingFakeEmbeddingStore() + store.embedding_store = current_space + await store.require_embedding_rebuild() + old_space.release_first_batch.set() + await upsert_task + await store.reindex("embedding") + + assert store.file_chunks["a"].embedding.tolist() == [1.0, 0.0] + assert current_space.node_embedding_calls == [["a"]] + await store.close() + + run(go()) + + +def test_embedding_requirement_cancels_backfill_without_clearing_gate(): + """Cancelling automatic repair must leave manual-reindex mode enabled.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_local_store("t_cancel_backfill_for_manual_reindex") + await store.start() + await set_chunks_with_graph(store, {"a": chunk("a", "a.md", "alpha")}) + fake = BlockingEmbeddingStore() store.embedding_store = fake store._start_embedding_backfill(skip_health_check=True) - old_task = store._embedding_backfill_task - await fake.first_batch_started.wait() + await fake.started.wait() - assert await store.resume_embedding(rebuild=True) is True - assert store._embedding_backfill_pending == (False, True) - fake.release_first_batch.set() + await store.require_embedding_rebuild() - await old_task - if store._embedding_backfill_task is not None: - await store._embedding_backfill_task - assert fake.node_embedding_calls == [["a"], ["a"]] - assert store.file_chunks["a"].embedding.tolist() == [1.0, 0.0] + assert store._embedding_backfill_task is None + assert store._embedding_rebuild_pending is True + assert store.file_chunks["a"].embedding is None + await store.close() + + run(go()) + + +def test_bm25_reindex_does_not_change_embedding_gate(): + """BM25 maintenance is independent from the embedding state machine.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_local_store("t_scoped_bm25_reindex") + await store.start() + await set_chunks_with_graph(store, {"a": chunk("a", "a.md", "uniquebm25word")}) + await store.keyword_index.clear() + await store.require_embedding_rebuild() + + result = await store.reindex("bm25") + + assert result == {"indexed": 1, "scope": "bm25"} + assert store._embedding_rebuild_pending is True + assert [item.id for item in await store.keyword_search("uniquebm25word", 5, {})] == ["a"] + await store.close() + + run(go()) + + +def test_all_reindex_composes_bm25_then_embedding_under_one_lock(): + """The all scope is exactly the ordered composition of both scopes.""" + + async def go(): + store = _new_local_store("t_all_reindex") + calls = [] + + async def rebuild(scope): + calls.append(scope) + return {"scope": scope, "indexed": 1} + + async def rebuild_bm25(): + return await rebuild("bm25") + + async def rebuild_embedding(): + return await rebuild("embedding") + + store._reindex_bm25 = rebuild_bm25 + store._reindex_embedding = rebuild_embedding + + result = await store.reindex("all") + + assert calls == ["bm25", "embedding"] + assert result["scope"] == "all" + + run(go()) + + +@pytest.mark.parametrize("store_factory", [_new_local_store, _new_faiss_store, _new_zvec_store]) +def test_embedding_reindex_clears_vectors_when_embedding_is_disabled(store_factory): + """Disabling embedding can explicitly discard the obsolete vectors.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = store_factory("t_disable_embedding_reindex") + await store.start() + stale = chunk("a", "a.md", "alpha") + stale.embedding = np.array([1.0, 0.0], dtype=np.float16) + await set_chunks_with_graph(store, {"a": stale}) + + result = await store.reindex("embedding") + + assert result == { + "indexed": 0, + "scope": "embedding", + } + assert stale.embedding is None assert store._embedding_rebuild_pending is False + if isinstance(store, ZvecLocalFileStore): + assert store._collection is None + assert not store.zvec_path.exists() + assert not store.zvec_sidecar_path.exists() + await store.close() + + run(go()) + + +@pytest.mark.parametrize("store_factory", [_new_local_store, _new_zvec_store]) +def test_clear_waits_for_in_flight_chunk_dump(store_factory): + """An older chunk snapshot cannot be published after clear() returns.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = store_factory("t_chunk_dump_clear_lock") + await store.start() + if isinstance(store, ZvecLocalFileStore): + store.embedding_store = FakeEmbeddingStore() + _ensure_zvec_collection(store) + await store.upsert([(node("a.md"), [chunk("a", "a.md", "alpha")])]) + + dump_started = threading.Event() + release_dump = threading.Event() + original_dump_chunks_sync = store._dump_chunks_sync + + def blocking_dump_chunks_sync(chunks): + dump_started.set() + release_dump.wait() + original_dump_chunks_sync(chunks) + + store._dump_chunks_sync = blocking_dump_chunks_sync + dump_task = asyncio.create_task(store.dump()) + assert await asyncio.to_thread(dump_started.wait, 1) + + clear_task = asyncio.create_task(store.clear()) + await asyncio.sleep(0.02) + assert not clear_task.done() + + release_dump.set() + await asyncio.gather(dump_task, clear_task) + assert store.file_chunks == {} + assert not store.chunks_path.exists() + + store._dump_chunks_sync = original_dump_chunks_sync + await store.close() + + run(go()) + + +@pytest.mark.parametrize("store_factory", [_new_faiss_store, _new_zvec_store]) +def test_embedding_reindex_keeps_gate_when_backend_checkpoint_fails(store_factory): + """A derived-index checkpoint failure must make the maintenance job fail.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = store_factory("t_embedding_reindex_checkpoint_failure") + await store.start() + await set_chunks_with_graph(store, {"a": chunk("a", "a.md", "alpha")}) + store.embedding_store = CountingFakeEmbeddingStore() + _ensure_zvec_collection(store) + if isinstance(store, FaissLocalFileStore): + store._faiss_index = store._new_index() + store._write_sidecar = AsyncMock(side_effect=OSError("disk full")) + + with pytest.raises(OSError, match="disk full"): + await store.reindex("embedding") + + assert store._embedding_rebuild_pending is True + store.embedding_store = None await store.close() run(go()) @pytest.mark.parametrize("store_factory", [_new_local_store, _new_faiss_store, _new_zvec_store]) -def test_clear_during_scheduled_rebuild_finishes_rebuild_state(store_factory): - """Clearing all chunks before the worker scan must not disable vector search forever.""" +def test_embedding_gate_discards_caller_supplied_vectors(store_factory): + """Upserts cannot inject vectors from an unverified space while gated.""" async def go(): with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): - store = store_factory("t_embedding_clear_during_rebuild") + store = store_factory("t_embedding_gate_supplied_vector") await store.start() - await set_chunks_with_graph(store, {"a": chunk("a", "a.md", "alpha text")}) store.embedding_store = CountingFakeEmbeddingStore() + _ensure_zvec_collection(store) + if isinstance(store, FaissLocalFileStore): + store._faiss_index = store._new_index() + await store.require_embedding_rebuild() + supplied = chunk("a", "a.md", "alpha") + supplied.embedding = np.array([1.0, 0.0], dtype=np.float16) - assert await store.resume_embedding(verified=True, rebuild=True) is True - task = store._embedding_backfill_task - await store.clear() - if task is not None: - await task + await store.upsert([(node("a.md"), [supplied])]) - assert store.file_chunks == {} - assert store._embedding_rebuild_pending is False + assert supplied.embedding is None + if isinstance(store, FaissLocalFileStore): + assert "a" not in store._id_to_row + elif isinstance(store, ZvecLocalFileStore): + assert "a" not in store._indexed_ids + store.embedding_store = None + await store.close() + + run(go()) + + +def test_embedding_reindex_retries_when_chunks_change_during_rebuild(): + """A rebuild only succeeds after one stable authoritative chunk generation.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_local_store("t_embedding_reindex_generation_retry") + await store.start() + await set_chunks_with_graph(store, {"a": chunk("a", "a.md", "alpha")}) + store.embedding_store = CountingFakeEmbeddingStore() + finalize_calls = 0 + + async def mutate_once(): + nonlocal finalize_calls + finalize_calls += 1 + if finalize_calls == 1: + store.file_chunks["b"] = chunk("b", "b.md", "beta") + store._mutation_generation += 1 + + store._finalize_embedding_reindex = mutate_once + + result = await store.reindex("embedding") + + assert finalize_calls == 2 + assert result == {"indexed": 2, "scope": "embedding"} + assert all(item.embedding is not None for item in store.file_chunks.values()) + await store.close() + + run(go()) + + +def test_upsert_waits_for_embedding_reindex(): + """A write cannot slip through the gate while reindex is finishing.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_local_store("t_upsert_waits_for_reindex") + await store.start() + await set_chunks_with_graph(store, {"a": chunk("a", "a.md", "alpha")}) + store.embedding_store = CountingFakeEmbeddingStore() + dump_started = asyncio.Event() + release_dump = asyncio.Event() + real_dump = store._dump_owned_state + + async def blocking_dump(): + dump_started.set() + await release_dump.wait() + await real_dump() + + store._dump_owned_state = blocking_dump + reindex_task = asyncio.create_task(store.reindex("embedding")) + await dump_started.wait() + upsert_task = asyncio.create_task(store.upsert([(node("b.md"), [chunk("b", "b.md", "beta")])])) + await asyncio.sleep(0) + assert "b" not in store.file_chunks + + release_dump.set() + await reindex_task + await upsert_task + + assert store.file_chunks["b"].embedding.tolist() == [0.0, 1.0] + await store.close() + + run(go()) + + +def test_faiss_explicit_finalizer_retries_stale_snapshot(): + """Explicit FAISS publication consumes the retry signal without a worker.""" + + async def go(): + store = _new_faiss_store("t_faiss_explicit_retry") + calls = 0 + + async def rebuild(): + nonlocal calls + calls += 1 + if calls == 1: + store._reindex_event.set() + + store._reindex_async = rebuild + await store._finalize_embedding_reindex() + assert calls == 2 + + run(go()) + + +def test_checkpoint_dump_keeps_event_loop_responsive(monkeypatch): + """Compression and file writes execute outside the request event loop.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_local_store("t_nonblocking_dump") + await store.start() + store.file_chunks["a"] = chunk("a", "a.md", "alpha") + entered = threading.Event() + release = threading.Event() + + def blocking_dump(_chunks): + assert _chunks[0] is not store.file_chunks["a"] + entered.set() + assert release.wait(timeout=2) + + monkeypatch.setattr(store, "_dump_chunks_sync", blocking_dump) + dump_task = asyncio.create_task(store._dump_owned_state()) + assert await asyncio.to_thread(entered.wait, 1) + ticks = 0 + for _ in range(5): + await asyncio.sleep(0) + ticks += 1 + store.file_chunks["a"].text = "changed while dumping" + assert ticks == 5 + assert not dump_task.done() + release.set() + await dump_task + await store.close() + + run(go()) + + +def test_checkpoint_dump_finishes_before_propagating_cancellation(monkeypatch): + """Cancellation cannot leave an old checkpoint writer running in the background.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_local_store("t_cancelled_dump") + await store.start() + store.file_chunks["a"] = chunk("a", "a.md", "alpha") + entered = threading.Event() + release = threading.Event() + + def blocking_dump(_chunks): + entered.set() + release.wait() + + monkeypatch.setattr(store, "_dump_chunks_sync", blocking_dump) + dump_task = asyncio.create_task(store._dump_owned_state()) + assert await asyncio.to_thread(entered.wait, 1) + dump_task.cancel() + await asyncio.sleep(0) + assert not dump_task.done() + + release.set() + with pytest.raises(asyncio.CancelledError): + await dump_task await store.close() run(go()) @@ -1107,8 +1496,14 @@ def test_search_filter_applies_to_vector_and_keyword_results(store_factory): await store.upsert( [ - (node("daily/a.md"), [chunk("a", "daily/a.md", "fresh topic", kind="daily")]), - (node("resource/b.md"), [chunk("b", "resource/b.md", "fresh topic", kind="resource")]), + ( + node("daily/a.md"), + [chunk("a", "daily/a.md", "fresh topic", kind="daily")], + ), + ( + node("resource/b.md"), + [chunk("b", "resource/b.md", "fresh topic", kind="resource")], + ), ], ) @@ -1280,9 +1675,18 @@ def test_date_filter_with_vector_and_keyword_search(store_factory): await store.upsert( [ - (node("daily/2026-01-10/a.md"), [chunk("a", "daily/2026-01-10/a.md", "alpha topic")]), - (node("daily/2026-02-15/b.md"), [chunk("b", "daily/2026-02-15/b.md", "alpha topic")]), - (node("daily/2026-03-20/c.md"), [chunk("c", "daily/2026-03-20/c.md", "alpha topic")]), + ( + node("daily/2026-01-10/a.md"), + [chunk("a", "daily/2026-01-10/a.md", "alpha topic")], + ), + ( + node("daily/2026-02-15/b.md"), + [chunk("b", "daily/2026-02-15/b.md", "alpha topic")], + ), + ( + node("daily/2026-03-20/c.md"), + [chunk("c", "daily/2026-03-20/c.md", "alpha topic")], + ), ], ) @@ -1684,7 +2088,10 @@ def test_faiss_async_reindex_no_lost_writes_during_build(): # After the follow-up rebuild the index reflects both concurrent writes; # the changed chunk now embeds as "beta". assert set(store._id_to_row) == {"a", "b"} - assert {c.id for c in await store.vector_search("beta", 5, {})} == {"a", "b"} + assert {c.id for c in await store.vector_search("beta", 5, {})} == { + "a", + "b", + } await store.close() run(go()) diff --git a/tests/unit/test_keyword_index.py b/tests/unit/test_keyword_index.py index b9554e5d..061c97b7 100644 --- a/tests/unit/test_keyword_index.py +++ b/tests/unit/test_keyword_index.py @@ -11,6 +11,7 @@ single-char ASCII words dropped). import asyncio import os import tempfile +import threading import warnings from reme.components.keyword_index import BM25Index @@ -840,6 +841,47 @@ def test_dump_failure_is_not_silent(): run(go()) +def test_concurrent_dumps_publish_in_invocation_order(): + """A newer dump must wait for and then supersede an older snapshot.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + bm25 = await create_bm25() + await bm25.add_docs({"d1": "alpha"}) + + first_started = threading.Event() + release_first = threading.Event() + dump_calls = 0 + original_dump_sync = bm25._dump_sync + + def blocking_dump_sync(snapshot): + nonlocal dump_calls + dump_calls += 1 + if dump_calls == 1: + first_started.set() + release_first.wait() + original_dump_sync(snapshot) + + bm25._dump_sync = blocking_dump_sync + first = asyncio.create_task(bm25.dump()) + assert await asyncio.to_thread(first_started.wait, 1) + + await bm25.add_docs({"d2": "beta"}) + second = asyncio.create_task(bm25.dump()) + await asyncio.sleep(0.02) + assert dump_calls == 1 + + release_first.set() + await asyncio.gather(first, second) + assert dump_calls == 2 + assert set(bm25._load_sync()["doc_id_to_idx"]) == {"d1", "d2"} + + bm25._dump_sync = original_dump_sync + await bm25.close() + + run(go()) + + # --------------------------------------------------------------------------- # # clear / optimize / reset_index # # --------------------------------------------------------------------------- # diff --git a/tests/unit/test_reindex.py b/tests/unit/test_reindex.py new file mode 100644 index 00000000..2dc428ba --- /dev/null +++ b/tests/unit/test_reindex.py @@ -0,0 +1,41 @@ +"""Tests for explicit scoped index rebuilds.""" + +from unittest.mock import AsyncMock + +import pytest + +from reme.components.file_store import LocalFileStore +from reme.components.runtime_context import RuntimeContext +from reme.steps.index import ReindexStep + + +@pytest.mark.asyncio +@pytest.mark.parametrize("scope", ["bm25", "embedding"]) +async def test_reindex_step_delegates_scope(scope): + """The step forwards each individual scope without clearing the store.""" + store = LocalFileStore(name=f"test_reindex_{scope}", embedding_store="") + store.reindex = AsyncMock(return_value={"indexed": 3, "scope": scope}) + store.clear = AsyncMock() + + response = await ReindexStep(file_store=store)(RuntimeContext(scope=scope)) + + store.reindex.assert_awaited_once_with(scope) + store.clear.assert_not_called() + assert response.metadata == {"indexed": 3, "scope": scope} + + +@pytest.mark.asyncio +async def test_reindex_step_delegates_all_once(): + """The step delegates the composite scope exactly once.""" + store = LocalFileStore(name="test_reindex_all", embedding_store="") + details = { + "scope": "all", + "bm25": {"scope": "bm25", "indexed": 3}, + "embedding": {"scope": "embedding", "indexed": 3}, + } + store.reindex = AsyncMock(return_value=details) + + response = await ReindexStep(file_store=store)(RuntimeContext(scope="all")) + + store.reindex.assert_awaited_once_with("all") + assert response.metadata == details