mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-07 08:26:06 +00:00
fix(index): make embedding rebuild explicit and scoped (#508)
* fix(index): make embedding rebuild explicit and scoped * fix(index): harden scoped reindex completion * fix(index): guard embedding space transitions * fix(index): serialize reindex with mutations * docs(index): clarify scoped reindex semantics * docs(index): explain synchronous checkpoint snapshots * fix(index): serialize checkpoint publication
This commit is contained in:
jinliyl
GitHub
parent
21f7757c80
commit
fc4a5398a8
33 changed files with 993 additions and 287 deletions
|
|
@ -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/<date>/<generated-name>.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/<date>/<resource-card>.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/<date>/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/<date>/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
|
||||
|
||||
|
|
|
|||
|
|
@ -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/<date>/<generated-name>.md` |
|
||||
| [`auto_resource`](docs/zh/auto_resource.md) | 资源监听或 `reme auto_resource` | 将 `resource/` 下的文件转为带来源链接、按内容命名的 daily 卡片。 | `daily/<date>/<resource-card>.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/<date>/interests.yaml` |
|
||||
| [`proactive`](docs/zh/proactive.md) | Agent 决定主动行动前调用 `reme proactive` | 读取 `auto_dream` 生成的 topics；是否以及如何提醒用户由宿主 Agent 决定。 | 来自 `daily/<date>/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 索引。 |
|
||||
|
||||
## 🤝 社区与贡献
|
||||
|
||||
|
|
|
|||
|
|
@ -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).
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
<p align="center">
|
||||
<img src="../figure/auto-index-and-memory-search.svg" alt="ReMe Auto Index and Memory Search indexing, recall, fusion, and link expansion" width="92%">
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
||||
|
|
|
|||
|
|
@ -73,7 +73,7 @@
|
|||
|
||||
<rect class="panel" x="142" y="484" width="916" height="76"/>
|
||||
<text class="note" x="190" y="514">Rebuild scope</text>
|
||||
<text class="chip-text" x="190" y="536">reindex adds resource + JSONL.</text>
|
||||
<text class="chip-text" x="190" y="536">BM25 + vectors from current chunks.</text>
|
||||
<line class="line" x1="410" y1="500" x2="410" y2="544"/>
|
||||
<text class="chip-title" x="454" y="514">BM25 is enabled by default</text>
|
||||
<text class="chip-text" x="454" y="536">Recall cap: 200; embeddings opt-in.</text>
|
||||
|
|
|
|||
|
Before Width: | Height: | Size: 7 KiB After Width: | Height: | Size: 7 KiB |
|
|
@ -1,8 +1,8 @@
|
|||
<svg xmlns="http://www.w3.org/2000/svg" width="1200" height="640" viewBox="0 0 1200 640" role="img"
|
||||
aria-labelledby="title desc">
|
||||
<title id="title">Memory Index</title>
|
||||
<desc id="desc">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.
|
||||
<desc id="desc">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.
|
||||
</desc>
|
||||
<defs>
|
||||
<style>.bg{fill:#fffdf8}.title{font:800 30px Arial,sans-serif;fill:#1f2430}.sub{font:14px Arial,sans-serif;fill:#667085}.panel{fill:#fff;stroke:#1f2430;stroke-width:2;rx:18}.head{font:700 17px Arial,sans-serif;fill:#1f2430}.text{font:13px Arial,sans-serif;fill:#5e6a7c}.tiny{font:12px Arial,sans-serif;fill:#667085}.mono{font:700 13px ui-monospace,SFMono-Regular,Menlo,monospace;fill:#1f2430}.chip{fill:#f8fbff;stroke:#1f2430;stroke-width:1.4;stroke-dasharray:6 5;rx:10}.orange{fill:#fff2e5}.blue{fill:#eef7ff}.green{fill:#effaf5}.arrow{fill:none;stroke:#7f8b9d;stroke-width:2;marker-end:url(#a)}.loop{fill:none;stroke:#ff963d;stroke-width:2;stroke-dasharray:7 6;marker-end:url(#o)}</style>
|
||||
|
|
@ -27,7 +27,7 @@
|
|||
<text class="mono" x="161" y="327" text-anchor="middle">digest/</text>
|
||||
<rect class="chip green" x="72" y="374" width="178" height="58"/>
|
||||
<text class="mono" x="161" y="401" text-anchor="middle">resource/</text>
|
||||
<text class="tiny" x="161" y="421" text-anchor="middle">via reindex</text>
|
||||
<text class="tiny" x="161" y="421" text-anchor="middle">via resource workflow</text>
|
||||
<text class="tiny" x="161" y="464" text-anchor="middle">Create · update · delete</text>
|
||||
|
||||
<path class="arrow" d="M280 312h58"/>
|
||||
|
|
|
|||
|
Before Width: | Height: | Size: 46 KiB After Width: | Height: | Size: 46 KiB |
|
|
@ -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)。
|
||||
|
|
|
|||
|
|
@ -219,7 +219,8 @@ FileLink
|
|||
|
||||
旧文档中的 `related:: [[path]]`、`- related:: [[path]]` 或
|
||||
`[related:: [[path]]]` 仍然可以读取。ReMe 会忽略外围文本，把内部 `[[path]]`
|
||||
作为普通链接建立索引。从曾存储 typed link 的版本升级后，应执行一次 `reme reindex`，用源文件重建不含旧关系字段的派生图索引。
|
||||
作为普通链接建立索引。源文件通过正常摄取路径时会应用图谱变更。`reme reindex` 只基于现有 chunks 重建 BM25 和 Embedding
|
||||
索引，不会重新解析文件或重建派生图谱。
|
||||
|
||||
### 来源和关系
|
||||
|
||||
|
|
|
|||
|
|
@ -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 图谱。
|
||||
|
||||
<p align="center">
|
||||
<img src="../figure/auto-index-and-memory-search.svg" alt="ReMe Auto Index and Memory Search 索引、召回、融合与链接展开流程" width="92%">
|
||||
|
|
@ -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` 快照。
|
||||
|
||||
## 怎么搜索
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -184,8 +184,8 @@ Auto Dream 默认查看以目标日期结尾的最近两天，只把相对上次
|
|||
</p>
|
||||
|
||||
Markdown 适合人读，但如果只是把文件堆进目录，Agent 仍然很难快速找到它们。默认实时索引持续监听 `daily/` 与 `digest/` 中的
|
||||
Markdown；`resource/` 由独立资源流程监听，转成 daily 卡片后进入同一索引。需要从现有文件完整重建时，`reme reindex` 还会扫描
|
||||
`resource/` 与 JSONL。
|
||||
Markdown；`resource/` 由独立资源流程监听，转成 daily 卡片后进入同一索引。手动 `reindex` 只基于这些摄取路径已经接受的
|
||||
chunks 重建 BM25 和 Embedding 索引，不会重新扫描文件或重建 Wikilink 图谱。
|
||||
|
||||
一份 Markdown 会被解析为：
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
"""ReMe CLI package."""
|
||||
|
||||
__version__ = "0.4.1.9"
|
||||
__version__ = "0.4.1.10"
|
||||
|
||||
from . import config
|
||||
from . import constants
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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] = []
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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 -----------------------------------------------------------
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
19
reme/steps/index/reindex.py
Normal file
19
reme/steps/index/reindex.py
Normal file
|
|
@ -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
|
||||
17
reme/utils/async_utils.py
Normal file
17
reme/utils/async_utils.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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 图谱。
|
||||
|
||||
## 检查
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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}"
|
||||
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
|
|
|||
|
|
@ -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 #
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
|
|
|||
41
tests/unit/test_reindex.py
Normal file
41
tests/unit/test_reindex.py
Normal file
|
|
@ -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
|
||||
Loading…
Add table
Reference in a new issue