mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-10-13 04:01:35 +00:00
fix(index): move file parsing off the service event loop
Signed-off-by: mikemikimike <186855910+mikemikimike@users.noreply.github.com>
This commit is contained in:
mikemikimike
parent
c1de31ab2c
commit
55a3b6632c
8 changed files with 1126 additions and 17 deletions
|
|
@ -1,12 +1,36 @@
|
|||
"""Abstract base for file chunkers."""
|
||||
|
||||
import asyncio
|
||||
import contextvars
|
||||
import os
|
||||
from abc import abstractmethod
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
from typing import Callable, TypeVar
|
||||
|
||||
from ..base_component import BaseComponent
|
||||
from ...enumeration import ComponentEnum
|
||||
from ...schema import FileChunk, FileNode
|
||||
|
||||
WorkResult = TypeVar("WorkResult")
|
||||
|
||||
|
||||
async def _finish_work(future: asyncio.Future[WorkResult]) -> WorkResult:
|
||||
"""Defer cancellation until submitted work finishes, then discard its result."""
|
||||
cancellation: asyncio.CancelledError | None = None
|
||||
while not future.done():
|
||||
try:
|
||||
await asyncio.shield(future)
|
||||
except asyncio.CancelledError as exc:
|
||||
cancellation = exc
|
||||
except Exception:
|
||||
break
|
||||
if cancellation is not None:
|
||||
if not future.cancelled():
|
||||
future.exception()
|
||||
raise cancellation
|
||||
return future.result()
|
||||
|
||||
|
||||
class BaseFileChunker(BaseComponent):
|
||||
"""Abstract base for file chunkers. Subclasses implement `chunk`."""
|
||||
|
|
@ -16,6 +40,81 @@ class BaseFileChunker(BaseComponent):
|
|||
def __init__(self, supported_extensions: list[str] | None = None, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.supported_extensions: list[str] = supported_extensions or []
|
||||
self._worker: ThreadPoolExecutor | None = None
|
||||
self._worker_lock = asyncio.Lock()
|
||||
self._worker_generation = 0
|
||||
self._worker_closing = False
|
||||
|
||||
async def _start(self) -> None:
|
||||
self._worker = ThreadPoolExecutor(max_workers=1, thread_name_prefix="reme-chunker")
|
||||
self._worker_generation += 1
|
||||
self._worker_closing = False
|
||||
|
||||
async def _close(self) -> None:
|
||||
self._worker_closing = True
|
||||
self._worker_generation += 1
|
||||
await _finish_work(asyncio.create_task(self._close_worker()))
|
||||
|
||||
async def _close_worker(self) -> None:
|
||||
async with self._worker_lock:
|
||||
if self._worker is not None:
|
||||
# The admission lock is held through completion, so shutdown has no running work.
|
||||
self._worker.shutdown(wait=True)
|
||||
self._worker = None
|
||||
|
||||
async def _chunk_in_worker(
|
||||
self,
|
||||
path: str | Path,
|
||||
chunk: Callable[[Path, str], tuple[FileNode, list[FileChunk]]],
|
||||
) -> tuple[FileNode, list[FileChunk]]:
|
||||
"""Admit one parse at a time without accumulating an executor work queue."""
|
||||
generation = self._worker_generation
|
||||
file_path = Path(path)
|
||||
rel_path = self.to_workspace_relative(path)
|
||||
async with self._worker_lock:
|
||||
if self._worker_closing or generation != self._worker_generation:
|
||||
raise RuntimeError("File chunker closed while waiting to parse")
|
||||
worker = self._worker
|
||||
temporary_worker = worker is None
|
||||
if worker is None:
|
||||
# Standalone chunk() remains usable without start(), with no persistent threads.
|
||||
worker = ThreadPoolExecutor(max_workers=1, thread_name_prefix="reme-chunker")
|
||||
try:
|
||||
context = contextvars.copy_context()
|
||||
future = asyncio.get_running_loop().run_in_executor(
|
||||
worker,
|
||||
context.run,
|
||||
self._chunk_unchanged_file,
|
||||
file_path,
|
||||
rel_path,
|
||||
chunk,
|
||||
)
|
||||
result, version = await _finish_work(future)
|
||||
if self._worker_closing or generation != self._worker_generation:
|
||||
raise RuntimeError("File chunker closed while parsing")
|
||||
self._check_file_version(file_path, version, file_path.stat())
|
||||
return result
|
||||
finally:
|
||||
if temporary_worker:
|
||||
worker.shutdown(wait=True)
|
||||
|
||||
@staticmethod
|
||||
def _chunk_unchanged_file(
|
||||
path: Path,
|
||||
rel_path: str,
|
||||
chunk: Callable[[Path, str], tuple[FileNode, list[FileChunk]]],
|
||||
) -> tuple[tuple[FileNode, list[FileChunk]], os.stat_result]:
|
||||
before = path.stat()
|
||||
result = chunk(path, rel_path)
|
||||
after = path.stat()
|
||||
BaseFileChunker._check_file_version(path, before, after)
|
||||
return result, after
|
||||
|
||||
@staticmethod
|
||||
def _check_file_version(path: Path, before: os.stat_result, after: os.stat_result) -> None:
|
||||
fields = ("st_dev", "st_ino", "st_size", "st_mtime_ns", "st_ctime_ns")
|
||||
if any(getattr(before, field) != getattr(after, field) for field in fields):
|
||||
raise RuntimeError(f"File changed while chunking: {path}")
|
||||
|
||||
@abstractmethod
|
||||
async def chunk(self, path: str | Path) -> tuple[FileNode, list[FileChunk]]:
|
||||
|
|
|
|||
|
|
@ -4,7 +4,6 @@ from bisect import bisect_right
|
|||
from pathlib import Path
|
||||
from typing import Literal
|
||||
|
||||
import aiofiles
|
||||
import yaml
|
||||
|
||||
from .base_file_chunker import BaseFileChunker
|
||||
|
|
@ -53,10 +52,9 @@ class DefaultFileChunker(BaseFileChunker):
|
|||
front_matter = FileFrontMatter()
|
||||
return front_matter, text[end_idx + 4 :].lstrip("\n")
|
||||
|
||||
async def _read_text_for_indexing(self, file_path: Path) -> str:
|
||||
def _read_text_for_indexing(self, file_path: Path) -> str:
|
||||
"""Decode text for a derived index without modifying the source file."""
|
||||
async with aiofiles.open(file_path, "rb") as f:
|
||||
data = await f.read()
|
||||
data = file_path.read_bytes()
|
||||
try:
|
||||
text = data.decode(self.encoding)
|
||||
except UnicodeDecodeError as exc:
|
||||
|
|
@ -77,11 +75,12 @@ class DefaultFileChunker(BaseFileChunker):
|
|||
return text.replace("\r\n", "\n").replace("\r", "\n")
|
||||
|
||||
async def chunk(self, path: str | Path) -> tuple[FileNode, list[FileChunk]]:
|
||||
file_path = Path(path)
|
||||
stat = file_path.stat()
|
||||
rel_path = self.to_workspace_relative(path)
|
||||
return await self._chunk_in_worker(path, self._chunk_sync)
|
||||
|
||||
text = await self._read_text_for_indexing(file_path)
|
||||
def _chunk_sync(self, file_path: Path, rel_path: str) -> tuple[FileNode, list[FileChunk]]:
|
||||
stat = file_path.stat()
|
||||
|
||||
text = self._read_text_for_indexing(file_path)
|
||||
|
||||
if not text:
|
||||
return FileNode(path=rel_path, st_mtime=stat.st_mtime), []
|
||||
|
|
|
|||
|
|
@ -112,9 +112,10 @@ class JsonFileChunker(BaseFileChunker):
|
|||
|
||||
async def chunk(self, path: str | Path) -> tuple[FileNode, list[FileChunk]]:
|
||||
"""Read and chunk a JSON file at *path*."""
|
||||
file_path = Path(path)
|
||||
return await self._chunk_in_worker(path, self._chunk_sync)
|
||||
|
||||
def _chunk_sync(self, file_path: Path, rel_path: str) -> tuple[FileNode, list[FileChunk]]:
|
||||
stat = file_path.stat()
|
||||
rel_path = self.to_workspace_relative(path)
|
||||
|
||||
raw_text = file_path.read_text(encoding=self.encoding)
|
||||
if not raw_text.strip():
|
||||
|
|
|
|||
|
|
@ -124,9 +124,10 @@ class JsonlFileChunker(BaseFileChunker):
|
|||
|
||||
async def chunk(self, path: str | Path) -> tuple[FileNode, list[FileChunk]]:
|
||||
"""Read and chunk a JSONL file at *path*."""
|
||||
file_path = Path(path)
|
||||
return await self._chunk_in_worker(path, self._chunk_sync)
|
||||
|
||||
def _chunk_sync(self, file_path: Path, rel_path: str) -> tuple[FileNode, list[FileChunk]]:
|
||||
stat = file_path.stat()
|
||||
rel_path = self.to_workspace_relative(path)
|
||||
|
||||
text = file_path.read_text(encoding=self.encoding)
|
||||
if not text.strip():
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ the single source of truth for ``[[...]]`` syntax.
|
|||
"""
|
||||
|
||||
import re
|
||||
import threading
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
|
@ -39,6 +40,9 @@ _ATX_HEADING_RE = re.compile(r"^ {0,3}#{1,6}(?:[ \t]+|$)")
|
|||
_SETEXT_HEADING_RE = re.compile(r"^ {0,3}(?:=+|-+)[ \t]*$")
|
||||
_FENCE_RE = re.compile(r"^ {0,3}(`{3,}|~{3,})")
|
||||
|
||||
# MarkdownRenderer changes mistletoe's process-wide token registry.
|
||||
_RENDERER_LOCK = threading.Lock()
|
||||
|
||||
|
||||
@dataclass
|
||||
class MdNode:
|
||||
|
|
@ -124,9 +128,10 @@ class MarkdownFileChunker(DefaultFileChunker):
|
|||
self.include_frontmatter_keys_in_metadata = list(include_frontmatter_keys_in_metadata or [])
|
||||
|
||||
async def chunk(self, path: str | Path) -> tuple[FileNode, list[FileChunk]]:
|
||||
file_path = Path(path)
|
||||
rel_path = self.to_workspace_relative(path)
|
||||
text = await self._read_text_for_indexing(file_path)
|
||||
return await self._chunk_in_worker(path, self._chunk_sync)
|
||||
|
||||
def _chunk_sync(self, file_path: Path, rel_path: str) -> tuple[FileNode, list[FileChunk]]:
|
||||
text = self._read_text_for_indexing(file_path)
|
||||
front_matter, content, line_offset = self._parse_front_matter(text)
|
||||
|
||||
chunks: list[FileChunk] = []
|
||||
|
|
@ -142,7 +147,7 @@ class MarkdownFileChunker(DefaultFileChunker):
|
|||
from mistletoe.markdown_renderer import MarkdownRenderer
|
||||
from mistletoe.block_token import Document
|
||||
|
||||
with MarkdownRenderer() as renderer:
|
||||
with _RENDERER_LOCK, MarkdownRenderer() as renderer:
|
||||
tree = self._build_tree(Document(content), renderer, line_offset=line_offset)
|
||||
chunks = self._chunk_node(tree, (), rel_path, renderer)
|
||||
if self.include_frontmatter_in_metadata:
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
|
||||
import asyncio
|
||||
import math
|
||||
import os
|
||||
from abc import abstractmethod
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable
|
||||
|
|
@ -46,6 +47,7 @@ class ChangeApplyStep(BaseStep):
|
|||
self.chunk_memory_overhead_bytes = int(chunk_memory_overhead_bytes)
|
||||
self.estimated_chunk_bytes = int(estimated_chunk_bytes)
|
||||
self.float16_bytes = int(float16_bytes)
|
||||
self._source_versions: dict[int, os.stat_result] = {}
|
||||
if self.batch_max_files <= 0:
|
||||
raise ValueError("batch_max_files must be greater than zero")
|
||||
if not math.isfinite(self.batch_available_memory_ratio) or not 0 < self.batch_available_memory_ratio <= 1:
|
||||
|
|
@ -323,8 +325,51 @@ class UpdateIndexStep(ChangeApplyStep):
|
|||
target_name = "file_store"
|
||||
reads_file_content = True
|
||||
|
||||
async def _apply_existing(self, buckets: dict[Change, list[str]]) -> list[dict]:
|
||||
try:
|
||||
return await super()._apply_existing(buckets)
|
||||
finally:
|
||||
# Steps are invocation-scoped; cancelled batches must release their snapshots too.
|
||||
self._source_versions.clear()
|
||||
|
||||
async def build_item(self, path: Path) -> tuple[FileNode, list[FileChunk]]:
|
||||
return await self.chunk_file(path)
|
||||
before = path.stat()
|
||||
item = await self.chunk_file(path)
|
||||
after = path.stat()
|
||||
BaseFileChunker._check_file_version(path, before, after) # pylint: disable=protected-access
|
||||
self._source_versions[id(item[0])] = after
|
||||
return item
|
||||
|
||||
async def _try_upsert(
|
||||
self,
|
||||
change: Change,
|
||||
items: list[tuple[FileNode, list[FileChunk]]],
|
||||
ok_paths: list[str],
|
||||
) -> list[dict]:
|
||||
# Recheck after waiting for the same guard used by every store mutation.
|
||||
async with self.file_store._maintenance_guard(): # pylint: disable=protected-access
|
||||
current_items: list[tuple[FileNode, list[FileChunk]]] = []
|
||||
current_paths: list[str] = []
|
||||
results: list[dict] = []
|
||||
for item, path in zip(items, ok_paths):
|
||||
node = item[0]
|
||||
try:
|
||||
BaseFileChunker._check_file_version( # pylint: disable=protected-access
|
||||
self._to_abs_path(node.path),
|
||||
self._source_versions[id(node)],
|
||||
self._to_abs_path(node.path).stat(),
|
||||
)
|
||||
except (OSError, RuntimeError) as exc:
|
||||
results.append({"change": change.name, "path": path, "success": False, "error": str(exc)})
|
||||
else:
|
||||
current_items.append(item)
|
||||
current_paths.append(path)
|
||||
finally:
|
||||
self._source_versions.pop(id(node))
|
||||
if current_items:
|
||||
results.extend(await super()._try_upsert(change, current_items, current_paths))
|
||||
by_path = {result["path"]: result for result in results}
|
||||
return [by_path[path] for path in ok_paths]
|
||||
|
||||
async def upsert_items(self, items: list[tuple[FileNode, list[FileChunk]]]) -> None:
|
||||
await self.file_store.upsert(items)
|
||||
|
|
|
|||
412
tests/unit/fixtures/file_chunker_outputs.json
Normal file
412
tests/unit/fixtures/file_chunker_outputs.json
Normal file
|
|
@ -0,0 +1,412 @@
|
|||
{
|
||||
"plain.txt": {
|
||||
"node": {
|
||||
"path": "plain.txt",
|
||||
"links": [],
|
||||
"chunk_ids": [
|
||||
"27beade9f75899604a5c33dc2a2cad48c68782309d1961d34aa416fa4a9f44e8",
|
||||
"339017807ea80235b88b9a55e5f38a0eef93d0e5515ef7b5ae5d5586551f38c7",
|
||||
"339017807ea80235b88b9a55e5f38a0eef93d0e5515ef7b5ae5d5586551f38c7",
|
||||
"339017807ea80235b88b9a55e5f38a0eef93d0e5515ef7b5ae5d5586551f38c7",
|
||||
"5fec11d28b94df329491529d38cc187104d52ea3500622503103a89feae73a23"
|
||||
],
|
||||
"front_matter": {
|
||||
"name": "",
|
||||
"description": ""
|
||||
}
|
||||
},
|
||||
"chunks": [
|
||||
{
|
||||
"id": "27beade9f75899604a5c33dc2a2cad48c68782309d1961d34aa416fa4a9f44e8",
|
||||
"text": "first\nabc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc ab",
|
||||
"embedding": null,
|
||||
"metadata": {},
|
||||
"path": "plain.txt",
|
||||
"start_line": 1,
|
||||
"end_line": 2,
|
||||
"scores": {}
|
||||
},
|
||||
{
|
||||
"id": "339017807ea80235b88b9a55e5f38a0eef93d0e5515ef7b5ae5d5586551f38c7",
|
||||
"text": "c abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc ab",
|
||||
"embedding": null,
|
||||
"metadata": {},
|
||||
"path": "plain.txt",
|
||||
"start_line": 2,
|
||||
"end_line": 2,
|
||||
"scores": {}
|
||||
},
|
||||
{
|
||||
"id": "339017807ea80235b88b9a55e5f38a0eef93d0e5515ef7b5ae5d5586551f38c7",
|
||||
"text": "c abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc ab",
|
||||
"embedding": null,
|
||||
"metadata": {},
|
||||
"path": "plain.txt",
|
||||
"start_line": 2,
|
||||
"end_line": 2,
|
||||
"scores": {}
|
||||
},
|
||||
{
|
||||
"id": "339017807ea80235b88b9a55e5f38a0eef93d0e5515ef7b5ae5d5586551f38c7",
|
||||
"text": "c abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc ab",
|
||||
"embedding": null,
|
||||
"metadata": {},
|
||||
"path": "plain.txt",
|
||||
"start_line": 2,
|
||||
"end_line": 2,
|
||||
"scores": {}
|
||||
},
|
||||
{
|
||||
"id": "5fec11d28b94df329491529d38cc187104d52ea3500622503103a89feae73a23",
|
||||
"text": "c abc ",
|
||||
"embedding": null,
|
||||
"metadata": {},
|
||||
"path": "plain.txt",
|
||||
"start_line": 2,
|
||||
"end_line": 2,
|
||||
"scores": {}
|
||||
}
|
||||
]
|
||||
},
|
||||
"note.md": {
|
||||
"node": {
|
||||
"path": "note.md",
|
||||
"links": [
|
||||
{
|
||||
"source_path": "note.md",
|
||||
"target_path": "other.md",
|
||||
"target_anchor": null
|
||||
}
|
||||
],
|
||||
"chunk_ids": [
|
||||
"c135f79d777949fa45848ad6302adcf6368a348910b2bcc4d99369b1acee75c3",
|
||||
"03ee7e6359843ec17c78ecc037d232e473d1b6b750ea544a0b1333a479a35f9c",
|
||||
"f74ee6f8929ad297e613b0341f0307ff1c9ff9656ec0b28b53de9ae58373c2d7"
|
||||
],
|
||||
"front_matter": {
|
||||
"name": "",
|
||||
"description": "",
|
||||
"title": "Example",
|
||||
"tags": [
|
||||
"one",
|
||||
"two"
|
||||
]
|
||||
}
|
||||
},
|
||||
"chunks": [
|
||||
{
|
||||
"id": "c135f79d777949fa45848ad6302adcf6368a348910b2bcc4d99369b1acee75c3",
|
||||
"text": "# First\n\nbody [[other.md]]",
|
||||
"embedding": null,
|
||||
"metadata": {
|
||||
"title": "Example",
|
||||
"tags": [
|
||||
"one",
|
||||
"two"
|
||||
]
|
||||
},
|
||||
"path": "note.md",
|
||||
"start_line": 5,
|
||||
"end_line": 7,
|
||||
"scores": {}
|
||||
},
|
||||
{
|
||||
"id": "03ee7e6359843ec17c78ecc037d232e473d1b6b750ea544a0b1333a479a35f9c",
|
||||
"text": "# First\n\n## Second",
|
||||
"embedding": null,
|
||||
"metadata": {
|
||||
"title": "Example",
|
||||
"tags": [
|
||||
"one",
|
||||
"two"
|
||||
]
|
||||
},
|
||||
"path": "note.md",
|
||||
"start_line": 9,
|
||||
"end_line": 9,
|
||||
"scores": {}
|
||||
},
|
||||
{
|
||||
"id": "f74ee6f8929ad297e613b0341f0307ff1c9ff9656ec0b28b53de9ae58373c2d7",
|
||||
"text": "text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text",
|
||||
"embedding": null,
|
||||
"metadata": {
|
||||
"title": "Example",
|
||||
"tags": [
|
||||
"one",
|
||||
"two"
|
||||
]
|
||||
},
|
||||
"path": "note.md",
|
||||
"start_line": 11,
|
||||
"end_line": 11,
|
||||
"scores": {}
|
||||
}
|
||||
]
|
||||
},
|
||||
"fallback.md": {
|
||||
"node": {
|
||||
"path": "fallback.md",
|
||||
"links": [],
|
||||
"chunk_ids": [
|
||||
"22b74c7a4dcd76d21c6aa9b3717065fb8328a8d77215fb7cb37003a9ecafbbab",
|
||||
"adfc277818b38936075c7aaf1fbc53d397ac6f9c2818dfa507d4598d67c2f203",
|
||||
"adfc277818b38936075c7aaf1fbc53d397ac6f9c2818dfa507d4598d67c2f203",
|
||||
"adfc277818b38936075c7aaf1fbc53d397ac6f9c2818dfa507d4598d67c2f203",
|
||||
"adfc277818b38936075c7aaf1fbc53d397ac6f9c2818dfa507d4598d67c2f203",
|
||||
"2560053ffa2a0c6decb9265aa1e33c0b264792535d3c9cf2a7b3a992084ccd19"
|
||||
],
|
||||
"front_matter": {
|
||||
"name": "",
|
||||
"description": ""
|
||||
}
|
||||
},
|
||||
"chunks": [
|
||||
{
|
||||
"id": "22b74c7a4dcd76d21c6aa9b3717065fb8328a8d77215fb7cb37003a9ecafbbab",
|
||||
"text": "# First\ntext text text text text text text text text text text text text text text text text text te",
|
||||
"embedding": null,
|
||||
"metadata": {},
|
||||
"path": "fallback.md",
|
||||
"start_line": 1,
|
||||
"end_line": 2,
|
||||
"scores": {}
|
||||
},
|
||||
{
|
||||
"id": "adfc277818b38936075c7aaf1fbc53d397ac6f9c2818dfa507d4598d67c2f203",
|
||||
"text": "xt text text text text text text text text text text text text text text text text text text text te",
|
||||
"embedding": null,
|
||||
"metadata": {},
|
||||
"path": "fallback.md",
|
||||
"start_line": 2,
|
||||
"end_line": 2,
|
||||
"scores": {}
|
||||
},
|
||||
{
|
||||
"id": "adfc277818b38936075c7aaf1fbc53d397ac6f9c2818dfa507d4598d67c2f203",
|
||||
"text": "xt text text text text text text text text text text text text text text text text text text text te",
|
||||
"embedding": null,
|
||||
"metadata": {},
|
||||
"path": "fallback.md",
|
||||
"start_line": 2,
|
||||
"end_line": 2,
|
||||
"scores": {}
|
||||
},
|
||||
{
|
||||
"id": "adfc277818b38936075c7aaf1fbc53d397ac6f9c2818dfa507d4598d67c2f203",
|
||||
"text": "xt text text text text text text text text text text text text text text text text text text text te",
|
||||
"embedding": null,
|
||||
"metadata": {},
|
||||
"path": "fallback.md",
|
||||
"start_line": 2,
|
||||
"end_line": 2,
|
||||
"scores": {}
|
||||
},
|
||||
{
|
||||
"id": "adfc277818b38936075c7aaf1fbc53d397ac6f9c2818dfa507d4598d67c2f203",
|
||||
"text": "xt text text text text text text text text text text text text text text text text text text text te",
|
||||
"embedding": null,
|
||||
"metadata": {},
|
||||
"path": "fallback.md",
|
||||
"start_line": 2,
|
||||
"end_line": 2,
|
||||
"scores": {}
|
||||
},
|
||||
{
|
||||
"id": "2560053ffa2a0c6decb9265aa1e33c0b264792535d3c9cf2a7b3a992084ccd19",
|
||||
"text": "xt text ",
|
||||
"embedding": null,
|
||||
"metadata": {},
|
||||
"path": "fallback.md",
|
||||
"start_line": 2,
|
||||
"end_line": 2,
|
||||
"scores": {}
|
||||
}
|
||||
]
|
||||
},
|
||||
"data.json": {
|
||||
"node": {
|
||||
"path": "data.json",
|
||||
"links": [],
|
||||
"chunk_ids": [
|
||||
"cacaea11e2cb2ee3a14b64163ef61cb4033f937a9e1926b8e5c59f938ab883a1",
|
||||
"562530d77b54d6f7ea397db77ca543023f6f082bc9a2d17d0693a9057f6f9c76",
|
||||
"147aa7fe035030c8a7632aefa473adbbe447bcb4ded262234c53c65c120eb799",
|
||||
"79cbd086821831544adee04ddc6e61d4ea8a052c21941710c754f8a8248eed62",
|
||||
"352cc7ee59226ac87644680d8c33f239bc5a8a30b0d088383ebfbf4559ba5bba",
|
||||
"04c5f9658b948d4f9f8eb9cb98f2096e1dd6507b4429211925a7eb5f4aed0f7e"
|
||||
],
|
||||
"front_matter": {
|
||||
"name": "",
|
||||
"description": ""
|
||||
}
|
||||
},
|
||||
"chunks": [
|
||||
{
|
||||
"id": "cacaea11e2cb2ee3a14b64163ef61cb4033f937a9e1926b8e5c59f938ab883a1",
|
||||
"text": "{\"items\": [{\"name\": \"abcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabc\", \"n\": 0}]}",
|
||||
"embedding": null,
|
||||
"metadata": {},
|
||||
"path": "data.json",
|
||||
"start_line": 4,
|
||||
"end_line": 5,
|
||||
"scores": {}
|
||||
},
|
||||
{
|
||||
"id": "562530d77b54d6f7ea397db77ca543023f6f082bc9a2d17d0693a9057f6f9c76",
|
||||
"text": "{\"items\": [{\"name\": \"abcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabc\", \"n\": 1}]}",
|
||||
"embedding": null,
|
||||
"metadata": {},
|
||||
"path": "data.json",
|
||||
"start_line": 8,
|
||||
"end_line": 9,
|
||||
"scores": {}
|
||||
},
|
||||
{
|
||||
"id": "147aa7fe035030c8a7632aefa473adbbe447bcb4ded262234c53c65c120eb799",
|
||||
"text": "{\"items\": [{\"name\": \"abcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabc\", \"n\": 2}]}",
|
||||
"embedding": null,
|
||||
"metadata": {},
|
||||
"path": "data.json",
|
||||
"start_line": 12,
|
||||
"end_line": 13,
|
||||
"scores": {}
|
||||
},
|
||||
{
|
||||
"id": "79cbd086821831544adee04ddc6e61d4ea8a052c21941710c754f8a8248eed62",
|
||||
"text": "{\"items\": [{\"name\": \"abcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabc\", \"n\": 3}]}",
|
||||
"embedding": null,
|
||||
"metadata": {},
|
||||
"path": "data.json",
|
||||
"start_line": 16,
|
||||
"end_line": 17,
|
||||
"scores": {}
|
||||
},
|
||||
{
|
||||
"id": "352cc7ee59226ac87644680d8c33f239bc5a8a30b0d088383ebfbf4559ba5bba",
|
||||
"text": "{\"items\": [{\"name\": \"abcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabc\", \"n\": 4}]}",
|
||||
"embedding": null,
|
||||
"metadata": {},
|
||||
"path": "data.json",
|
||||
"start_line": 20,
|
||||
"end_line": 21,
|
||||
"scores": {}
|
||||
},
|
||||
{
|
||||
"id": "04c5f9658b948d4f9f8eb9cb98f2096e1dd6507b4429211925a7eb5f4aed0f7e",
|
||||
"text": "{\"items\": [{\"name\": \"abcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabc\", \"n\": 5}]}",
|
||||
"embedding": null,
|
||||
"metadata": {},
|
||||
"path": "data.json",
|
||||
"start_line": 24,
|
||||
"end_line": 25,
|
||||
"scores": {}
|
||||
}
|
||||
]
|
||||
},
|
||||
"bad.json": {
|
||||
"node": {
|
||||
"path": "bad.json",
|
||||
"links": [],
|
||||
"chunk_ids": [
|
||||
"e974bca33c3651cf36a114da6f575a5af54134648f0368b34b9e1d2bc52acbb3"
|
||||
],
|
||||
"front_matter": {
|
||||
"name": "",
|
||||
"description": ""
|
||||
}
|
||||
},
|
||||
"chunks": [
|
||||
{
|
||||
"id": "e974bca33c3651cf36a114da6f575a5af54134648f0368b34b9e1d2bc52acbb3",
|
||||
"text": "{invalid\ntext\n",
|
||||
"embedding": null,
|
||||
"metadata": {},
|
||||
"path": "bad.json",
|
||||
"start_line": 1,
|
||||
"end_line": 3,
|
||||
"scores": {}
|
||||
}
|
||||
]
|
||||
},
|
||||
"events.jsonl": {
|
||||
"node": {
|
||||
"path": "events.jsonl",
|
||||
"links": [],
|
||||
"chunk_ids": [
|
||||
"fe4880a04aedd78ae5743cb495b0d9d8500ab9d3ba79515f4230e1b2f5f7889f",
|
||||
"7725431a86471c427c1a123803ccb0218fec1fd05877a6a70902380d50999fc6",
|
||||
"48c0edd450850c2382a954df3a907b6ee87f210de7a21a6341797438a30dfd3b",
|
||||
"116a24007dabf192a80d4d011a8cbbcf29d6831cc4a0c5941233fadac054e634",
|
||||
"c91f51a356df943b80ba09cf654f1a042aa08bbe8912f393f13ac99211b19b93",
|
||||
"489e1d724b739ddc864755e3d8e9a8d0c669dd1962102d525056ea265f8915a3"
|
||||
],
|
||||
"front_matter": {
|
||||
"name": "",
|
||||
"description": ""
|
||||
}
|
||||
},
|
||||
"chunks": [
|
||||
{
|
||||
"id": "fe4880a04aedd78ae5743cb495b0d9d8500ab9d3ba79515f4230e1b2f5f7889f",
|
||||
"text": "{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n",
|
||||
"embedding": null,
|
||||
"metadata": {},
|
||||
"path": "events.jsonl",
|
||||
"start_line": 1,
|
||||
"end_line": 3,
|
||||
"scores": {}
|
||||
},
|
||||
{
|
||||
"id": "7725431a86471c427c1a123803ccb0218fec1fd05877a6a70902380d50999fc6",
|
||||
"text": "{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n",
|
||||
"embedding": null,
|
||||
"metadata": {},
|
||||
"path": "events.jsonl",
|
||||
"start_line": 3,
|
||||
"end_line": 5,
|
||||
"scores": {}
|
||||
},
|
||||
{
|
||||
"id": "48c0edd450850c2382a954df3a907b6ee87f210de7a21a6341797438a30dfd3b",
|
||||
"text": "{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n",
|
||||
"embedding": null,
|
||||
"metadata": {},
|
||||
"path": "events.jsonl",
|
||||
"start_line": 5,
|
||||
"end_line": 7,
|
||||
"scores": {}
|
||||
},
|
||||
{
|
||||
"id": "116a24007dabf192a80d4d011a8cbbcf29d6831cc4a0c5941233fadac054e634",
|
||||
"text": "{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n",
|
||||
"embedding": null,
|
||||
"metadata": {},
|
||||
"path": "events.jsonl",
|
||||
"start_line": 7,
|
||||
"end_line": 9,
|
||||
"scores": {}
|
||||
},
|
||||
{
|
||||
"id": "c91f51a356df943b80ba09cf654f1a042aa08bbe8912f393f13ac99211b19b93",
|
||||
"text": "{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n",
|
||||
"embedding": null,
|
||||
"metadata": {},
|
||||
"path": "events.jsonl",
|
||||
"start_line": 9,
|
||||
"end_line": 11,
|
||||
"scores": {}
|
||||
},
|
||||
{
|
||||
"id": "489e1d724b739ddc864755e3d8e9a8d0c669dd1962102d525056ea265f8915a3",
|
||||
"text": "{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n",
|
||||
"embedding": null,
|
||||
"metadata": {},
|
||||
"path": "events.jsonl",
|
||||
"start_line": 11,
|
||||
"end_line": 12,
|
||||
"scores": {}
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
547
tests/unit/test_file_chunker_workers.py
Normal file
547
tests/unit/test_file_chunker_workers.py
Normal file
|
|
@ -0,0 +1,547 @@
|
|||
"""Worker admission, output compatibility, and real indexing/search regressions."""
|
||||
|
||||
import asyncio
|
||||
import contextvars
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from typing import Callable
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from httpx import ASGITransport, AsyncClient
|
||||
|
||||
from reme.application import Application
|
||||
from reme.components.application_context import ApplicationContext
|
||||
from reme.components.file_chunker import DefaultFileChunker, JsonFileChunker, JsonlFileChunker, MarkdownFileChunker
|
||||
from reme.components.file_store import LocalFileStore
|
||||
from reme.components.job import BaseJob, StreamJob
|
||||
from reme.components.runtime_context import RuntimeContext
|
||||
from reme.components.service import HttpService
|
||||
from reme.enumeration import ChunkEnum, ComponentEnum
|
||||
from reme.schema import FileChunk, FileNode, Response, StreamChunk
|
||||
from reme.steps.index import UpdateIndexStep
|
||||
|
||||
# Worker tests intentionally inspect admission and teardown at the internal boundary.
|
||||
# pylint: disable=protected-access
|
||||
|
||||
|
||||
class _StoreSearchJob(BaseJob):
|
||||
"""Use the actual keyword store behind the HTTP route and MCP tool."""
|
||||
|
||||
def __init__(self, store: LocalFileStore) -> None:
|
||||
super().__init__(name="search", parameters={"type": "object", "properties": {"query": {"type": "string"}}})
|
||||
self.store = store
|
||||
|
||||
async def __call__(self, query: str, **kwargs) -> Response:
|
||||
del kwargs
|
||||
chunks = await self.store.keyword_search(query, limit=10, search_filter={})
|
||||
return Response(answer=[chunk.text for chunk in chunks])
|
||||
|
||||
|
||||
class _StoreSearchStreamJob(StreamJob):
|
||||
"""Stream real keyword matches through the existing SSE endpoint."""
|
||||
|
||||
def __init__(self, store: LocalFileStore) -> None:
|
||||
super().__init__(name="search_stream")
|
||||
self.store = store
|
||||
|
||||
async def __call__(self, stream_queue: asyncio.Queue[StreamChunk], query: str, **kwargs) -> None:
|
||||
del kwargs
|
||||
chunks = await self.store.keyword_search(query, limit=10, search_filter={})
|
||||
for chunk in chunks:
|
||||
await stream_queue.put(StreamChunk(chunk=chunk.text))
|
||||
await asyncio.sleep(0)
|
||||
await stream_queue.put(StreamChunk(chunk_type=ChunkEnum.DONE, done=True))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_outputs_equal_pre_worker_fixtures(tmp_path: Path) -> None:
|
||||
"""Compare every serialized field, including hashes, links, metadata and ranges."""
|
||||
context = ApplicationContext(workspace_dir=str(tmp_path))
|
||||
cases = [
|
||||
(DefaultFileChunker(chunk_byte_size=100), "plain.txt", "first\r\n" + "abc " * 100),
|
||||
(
|
||||
MarkdownFileChunker(chunk_byte_size=100, include_frontmatter_in_metadata=True),
|
||||
"note.md",
|
||||
"---\ntitle: Example\ntags: [one, two]\n---\n# First\n\nbody [[other.md]]\n\n## Second\n\n" + "text " * 70,
|
||||
),
|
||||
(MarkdownFileChunker(chunk_byte_size=100, max_ast_sections=0), "fallback.md", "# First\n" + "text " * 100),
|
||||
(
|
||||
JsonFileChunker(chunk_chars=256),
|
||||
"data.json",
|
||||
json.dumps({"items": [{"name": "abc" * 35, "n": i} for i in range(6)]}, indent=2),
|
||||
),
|
||||
(JsonFileChunker(chunk_chars=256), "bad.json", "{invalid\ntext\n"),
|
||||
(JsonlFileChunker(max_chars=64, max_overlap_chars=30), "events.jsonl", '{"value": "alpha"}\n' * 12),
|
||||
]
|
||||
expected = json.loads((Path(__file__).parent / "fixtures" / "file_chunker_outputs.json").read_text())
|
||||
for chunker, name, text in cases:
|
||||
chunker.app_context = context
|
||||
path = tmp_path / name
|
||||
path.write_bytes(text.encode())
|
||||
node, chunks = await chunker.chunk(path)
|
||||
output = {"node": node.model_dump(mode="json"), "chunks": [c.model_dump(mode="json") for c in chunks]}
|
||||
assert output["node"].pop("st_mtime") == path.stat().st_mtime
|
||||
assert output == expected[name]
|
||||
|
||||
|
||||
async def _wait_started(event: threading.Event) -> None:
|
||||
for _ in range(500):
|
||||
if event.is_set():
|
||||
return
|
||||
await asyncio.sleep(0.01)
|
||||
raise AssertionError("worker did not start")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("chunker_class", [DefaultFileChunker, MarkdownFileChunker, JsonFileChunker, JsonlFileChunker])
|
||||
async def test_parsing_leaves_event_loop_responsive(
|
||||
tmp_path: Path,
|
||||
chunker_class: type[DefaultFileChunker | MarkdownFileChunker | JsonFileChunker | JsonlFileChunker],
|
||||
) -> None:
|
||||
"""A parser can block while unrelated async requests keep progressing."""
|
||||
chunker = chunker_class()
|
||||
path = tmp_path / "source"
|
||||
path.write_text("{}")
|
||||
entered, release = threading.Event(), threading.Event()
|
||||
main_thread = threading.get_ident()
|
||||
original = chunker._chunk_sync
|
||||
|
||||
def blocking_parse(file_path: Path, rel_path: str) -> tuple[FileNode, list[FileChunk]]:
|
||||
assert threading.get_ident() != main_thread
|
||||
entered.set()
|
||||
assert release.wait(5)
|
||||
return original(file_path, rel_path)
|
||||
|
||||
with patch.object(chunker, "_chunk_sync", side_effect=blocking_parse):
|
||||
task = asyncio.create_task(chunker.chunk(path))
|
||||
try:
|
||||
await _wait_started(entered)
|
||||
for _ in range(10):
|
||||
await asyncio.sleep(0)
|
||||
assert not task.done()
|
||||
finally:
|
||||
release.set()
|
||||
await task
|
||||
assert chunker._worker is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancellation_drains_worker_and_preserves_admission(tmp_path: Path) -> None:
|
||||
"""Repeated cancellation cannot free admission while parsing is still running."""
|
||||
path = tmp_path / "source.txt"
|
||||
path.write_text("source")
|
||||
chunker = DefaultFileChunker()
|
||||
entered, release = threading.Event(), threading.Event()
|
||||
calls = []
|
||||
original = chunker._chunk_sync
|
||||
|
||||
def blocking_parse(file_path: Path, rel_path: str) -> tuple[FileNode, list[FileChunk]]:
|
||||
calls.append(rel_path)
|
||||
entered.set()
|
||||
assert release.wait(5)
|
||||
return original(file_path, rel_path)
|
||||
|
||||
async with chunker:
|
||||
with patch.object(chunker, "_chunk_sync", side_effect=blocking_parse):
|
||||
first = asyncio.create_task(chunker.chunk(path))
|
||||
await _wait_started(entered)
|
||||
first.cancel()
|
||||
second = asyncio.create_task(chunker.chunk(path))
|
||||
await asyncio.sleep(0.01)
|
||||
first.cancel()
|
||||
await asyncio.sleep(0.01)
|
||||
try:
|
||||
assert not first.done()
|
||||
assert len(calls) == 1
|
||||
finally:
|
||||
release.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await first
|
||||
await second
|
||||
assert chunker._worker is None
|
||||
assert len(calls) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_shutdown_rejects_queued_work_and_drains_despite_cancellation(tmp_path: Path) -> None:
|
||||
"""Closing waits for running work, rejects queued parsing, and permits restart."""
|
||||
path = tmp_path / "source.txt"
|
||||
path.write_text("source")
|
||||
chunker = DefaultFileChunker()
|
||||
entered, release = threading.Event(), threading.Event()
|
||||
original = chunker._chunk_sync
|
||||
|
||||
def blocking_parse(file_path: Path, rel_path: str) -> tuple[FileNode, list[FileChunk]]:
|
||||
entered.set()
|
||||
assert release.wait(5)
|
||||
return original(file_path, rel_path)
|
||||
|
||||
await chunker.start()
|
||||
with patch.object(chunker, "_chunk_sync", side_effect=blocking_parse):
|
||||
running = asyncio.create_task(chunker.chunk(path))
|
||||
await _wait_started(entered)
|
||||
queued = asyncio.create_task(chunker.chunk(path))
|
||||
await asyncio.sleep(0)
|
||||
closing = asyncio.create_task(chunker.close())
|
||||
await asyncio.sleep(0.01)
|
||||
closing.cancel()
|
||||
try:
|
||||
assert not closing.done()
|
||||
finally:
|
||||
release.set()
|
||||
with pytest.raises(RuntimeError, match="closed"):
|
||||
await running
|
||||
with pytest.raises(RuntimeError, match="closed"):
|
||||
await queued
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await closing
|
||||
assert chunker._worker is None
|
||||
assert not chunker.is_started
|
||||
async with chunker:
|
||||
assert (await chunker.chunk(path))[1][0].text == "source"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("change", ["modify", "replace", "delete"])
|
||||
async def test_source_changes_do_not_return_stale_chunks(tmp_path: Path, change: str) -> None:
|
||||
"""Changes after reading cannot be mistaken for successfully parsed current content."""
|
||||
path = tmp_path / "source.txt"
|
||||
path.write_text("old source")
|
||||
chunker = DefaultFileChunker()
|
||||
original = chunker._chunk_sync
|
||||
entered, release = threading.Event(), threading.Event()
|
||||
|
||||
def blocking_parse(file_path: Path, rel_path: str) -> tuple[FileNode, list[FileChunk]]:
|
||||
result = original(file_path, rel_path)
|
||||
entered.set()
|
||||
assert release.wait(5)
|
||||
return result
|
||||
|
||||
with patch.object(chunker, "_chunk_sync", side_effect=blocking_parse):
|
||||
task = asyncio.create_task(chunker.chunk(path))
|
||||
await _wait_started(entered)
|
||||
try:
|
||||
if change == "modify":
|
||||
path.write_text("new source with different size")
|
||||
elif change == "replace":
|
||||
replacement = tmp_path / "replacement"
|
||||
replacement.write_text("replacement source")
|
||||
replacement.replace(path)
|
||||
else:
|
||||
path.unlink()
|
||||
finally:
|
||||
release.set()
|
||||
with pytest.raises((RuntimeError, FileNotFoundError)):
|
||||
await task
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_change_before_event_loop_receives_result_is_rejected(tmp_path: Path) -> None:
|
||||
"""Validate the worker snapshot again when its result reaches the calling loop."""
|
||||
path = tmp_path / "source.txt"
|
||||
path.write_text("old source")
|
||||
chunker = DefaultFileChunker()
|
||||
original = chunker._chunk_unchanged_file
|
||||
loop = asyncio.get_running_loop()
|
||||
changed = threading.Event()
|
||||
|
||||
def change_on_loop() -> None:
|
||||
path.write_text("changed after the worker snapshot")
|
||||
changed.set()
|
||||
|
||||
def change_after_parse(
|
||||
file_path: Path,
|
||||
rel_path: str,
|
||||
chunk: Callable[[Path, str], tuple[FileNode, list[FileChunk]]],
|
||||
) -> tuple[tuple[FileNode, list[FileChunk]], os.stat_result]:
|
||||
result = original(file_path, rel_path, chunk)
|
||||
loop.call_soon_threadsafe(change_on_loop)
|
||||
assert changed.wait(5)
|
||||
return result
|
||||
|
||||
with patch.object(chunker, "_chunk_unchanged_file", side_effect=change_after_parse):
|
||||
with pytest.raises(RuntimeError, match="File changed"):
|
||||
await chunker.chunk(path)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_queued_cancellation_does_not_submit_and_context_is_preserved(tmp_path: Path) -> None:
|
||||
"""A cancelled waiter uses no worker slot; caller context is copied into the worker."""
|
||||
path = tmp_path / "source.txt"
|
||||
path.write_text("source")
|
||||
chunker = DefaultFileChunker()
|
||||
entered, release = threading.Event(), threading.Event()
|
||||
marker = contextvars.ContextVar("chunk_test_marker", default="unset")
|
||||
original = chunker._chunk_sync
|
||||
calls = []
|
||||
|
||||
def blocking_parse(file_path: Path, rel_path: str) -> tuple[FileNode, list[FileChunk]]:
|
||||
calls.append(marker.get())
|
||||
entered.set()
|
||||
assert release.wait(5)
|
||||
return original(file_path, rel_path)
|
||||
|
||||
token = marker.set("caller")
|
||||
try:
|
||||
async with chunker:
|
||||
with patch.object(chunker, "_chunk_sync", side_effect=blocking_parse):
|
||||
first = asyncio.create_task(chunker.chunk(path))
|
||||
await _wait_started(entered)
|
||||
queued = asyncio.create_task(chunker.chunk(path))
|
||||
await asyncio.sleep(0)
|
||||
queued.cancel()
|
||||
try:
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await queued
|
||||
finally:
|
||||
release.set()
|
||||
await first
|
||||
finally:
|
||||
marker.reset(token)
|
||||
assert calls == ["caller"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_markdown_renderer_is_serialized_across_components(tmp_path: Path) -> None:
|
||||
"""Independent chunker instances cannot overlap renderer registry changes."""
|
||||
from mistletoe.markdown_renderer import MarkdownRenderer
|
||||
|
||||
path = tmp_path / "source.md"
|
||||
path.write_text("# Title\n\nparagraph **bold** [[other.md]]\n")
|
||||
entered, release = threading.Event(), threading.Event()
|
||||
calls = []
|
||||
original_enter = MarkdownRenderer.__enter__
|
||||
|
||||
def blocking_enter(renderer: MarkdownRenderer) -> MarkdownRenderer:
|
||||
calls.append(threading.get_ident())
|
||||
entered.set()
|
||||
assert release.wait(5)
|
||||
return original_enter(renderer)
|
||||
|
||||
with patch.object(MarkdownRenderer, "__enter__", blocking_enter):
|
||||
first = asyncio.create_task(MarkdownFileChunker().chunk(path))
|
||||
await _wait_started(entered)
|
||||
second = asyncio.create_task(MarkdownFileChunker().chunk(path))
|
||||
try:
|
||||
await asyncio.sleep(0.05)
|
||||
assert len(calls) == 1
|
||||
finally:
|
||||
release.set()
|
||||
first_output, second_output = await asyncio.gather(first, second)
|
||||
assert first_output == second_output
|
||||
assert len(calls) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("outcome", ["complete", "cancel", "change", "error"])
|
||||
async def test_existing_content_searchable_during_indexing(
|
||||
tmp_path: Path,
|
||||
outcome: str,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Exercise UpdateIndexStep and LocalFileStore while parsing is suspended."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
path = tmp_path / "source.txt"
|
||||
path.write_text("old searchable content")
|
||||
context = ApplicationContext(workspace_dir=str(tmp_path))
|
||||
chunker = DefaultFileChunker(app_context=context)
|
||||
context.components = {ComponentEnum.FILE_CHUNKER: {"default": chunker}}
|
||||
store = LocalFileStore(embedding_store="")
|
||||
async with store, chunker:
|
||||
await store.upsert([await chunker.chunk(path)])
|
||||
path.write_text("new searchable content")
|
||||
entered, release = threading.Event(), threading.Event()
|
||||
original = chunker._chunk_sync
|
||||
|
||||
def blocking_parse(file_path: Path, rel_path: str) -> tuple[FileNode, list[FileChunk]]:
|
||||
result = original(file_path, rel_path)
|
||||
entered.set()
|
||||
assert release.wait(5)
|
||||
if outcome == "error":
|
||||
raise ValueError("parser failed")
|
||||
return result
|
||||
|
||||
step = UpdateIndexStep(file_store=store, persist=False, app_context=context)
|
||||
app = Application(
|
||||
workspace_dir=str(tmp_path / "service"),
|
||||
service={"backend": "http", "web_enabled": False},
|
||||
components={},
|
||||
jobs={},
|
||||
enable_logo=False,
|
||||
log_to_console=False,
|
||||
log_to_file=False,
|
||||
)
|
||||
service = app.context.service
|
||||
assert isinstance(service, HttpService)
|
||||
service.build_service(app)
|
||||
service.add_job(_StoreSearchJob(store))
|
||||
service.add_job(_StoreSearchStreamJob(store))
|
||||
with patch.object(chunker, "_chunk_sync", side_effect=blocking_parse):
|
||||
updating = asyncio.create_task(step(RuntimeContext(changes=[{"change": "modified", "path": str(path)}])))
|
||||
await _wait_started(entered)
|
||||
try:
|
||||
nodes = await store.get_nodes()
|
||||
assert len(nodes) == 1
|
||||
matches = await store.keyword_search("searchable", limit=10, search_filter={})
|
||||
assert matches[0].text == "old searchable content"
|
||||
async with AsyncClient(transport=ASGITransport(app=service.service), base_url="http://test") as client:
|
||||
response = await client.post("/search", json={"query": "searchable"})
|
||||
assert response.status_code == 200
|
||||
assert response.json()["answer"] == ["old searchable content"]
|
||||
streaming = await client.post("/search_stream", json={"query": "searchable"})
|
||||
assert streaming.status_code == 200
|
||||
assert "old searchable content" in streaming.text
|
||||
assert streaming.text.endswith("data:[DONE]\n\n")
|
||||
assert service.mcp_server is not None
|
||||
tool = await service.mcp_server.get_tool("search")
|
||||
assert tool is not None
|
||||
result = await tool.run({"query": "searchable"})
|
||||
assert "old searchable content" in str(result.content)
|
||||
if outcome == "cancel":
|
||||
updating.cancel()
|
||||
elif outcome == "change":
|
||||
path.write_text("latest source content")
|
||||
finally:
|
||||
release.set()
|
||||
if outcome == "cancel":
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await updating
|
||||
else:
|
||||
response = await updating
|
||||
assert response.success == (outcome == "complete")
|
||||
nodes = await store.get_nodes()
|
||||
chunks = [store.file_chunks[chunk_id] for chunk_id in nodes[0].chunk_ids]
|
||||
assert chunks[0].text == ("new searchable content" if outcome == "complete" else "old searchable content")
|
||||
if outcome != "complete":
|
||||
recovered = await step(RuntimeContext(changes=[{"change": "modified", "path": str(path)}]))
|
||||
assert recovered.success is True
|
||||
nodes = await store.get_nodes()
|
||||
assert store.file_chunks[nodes[0].chunk_ids[0]].text == path.read_text()
|
||||
assert not step._source_versions
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("change", ["modify", "delete", "replace", "cancel"])
|
||||
async def test_changed_batched_file_is_not_published_and_valid_files_update(
|
||||
tmp_path: Path,
|
||||
change: str,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""An earlier parsed file can change while the batch waits for the next worker."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
first_path, second_path = tmp_path / "first.txt", tmp_path / "second.txt"
|
||||
first_path.write_text("original first")
|
||||
second_path.write_text("original second")
|
||||
context = ApplicationContext(workspace_dir=str(tmp_path))
|
||||
chunker = DefaultFileChunker(app_context=context)
|
||||
context.components = {ComponentEnum.FILE_CHUNKER: {"default": chunker}}
|
||||
async with LocalFileStore(embedding_store="") as store, chunker:
|
||||
await store.upsert([await chunker.chunk(first_path), await chunker.chunk(second_path)])
|
||||
first_path.write_text("updated first")
|
||||
second_path.write_text("updated second")
|
||||
entered, release = threading.Event(), threading.Event()
|
||||
original = chunker.chunk
|
||||
|
||||
async def blocking_second(file_path: Path) -> tuple[FileNode, list[FileChunk]]:
|
||||
if file_path == second_path:
|
||||
entered.set()
|
||||
while not release.is_set():
|
||||
await asyncio.sleep(0.01)
|
||||
return await original(file_path)
|
||||
|
||||
changes = [{"change": "modified", "path": str(path)} for path in (first_path, second_path)]
|
||||
step = UpdateIndexStep(file_store=store, persist=False, app_context=context)
|
||||
with patch.object(chunker, "chunk", side_effect=blocking_second):
|
||||
task = asyncio.create_task(step(RuntimeContext(changes=changes)))
|
||||
await _wait_started(entered)
|
||||
try:
|
||||
if change == "modify":
|
||||
first_path.write_text("latest first")
|
||||
elif change == "delete":
|
||||
first_path.unlink()
|
||||
elif change == "cancel":
|
||||
task.cancel()
|
||||
else:
|
||||
replacement = tmp_path / "replacement"
|
||||
replacement.write_text("replacement first")
|
||||
replacement.replace(first_path)
|
||||
finally:
|
||||
release.set()
|
||||
if change == "cancel":
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
else:
|
||||
response = await task
|
||||
assert response.success is False
|
||||
assert [result["success"] for result in response.answer] == [False, True]
|
||||
assert not step._source_versions
|
||||
nodes = {node.path: node for node in await store.get_nodes()}
|
||||
assert store.file_chunks[nodes["first.txt"].chunk_ids[0]].text == "original first"
|
||||
assert store.file_chunks[nodes["second.txt"].chunk_ids[0]].text == (
|
||||
"original second" if change == "cancel" else "updated second"
|
||||
)
|
||||
recovered = await step(
|
||||
RuntimeContext(
|
||||
changes=[{"change": "deleted" if change == "delete" else "modified", "path": str(first_path)}],
|
||||
),
|
||||
)
|
||||
assert recovered.success is True
|
||||
nodes = {node.path: node for node in await store.get_nodes()}
|
||||
if change == "delete":
|
||||
assert "first.txt" not in nodes
|
||||
else:
|
||||
assert store.file_chunks[nodes["first.txt"].chunk_ids[0]].text == first_path.read_text()
|
||||
assert not step._source_versions
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("outcome", ["change", "cancel"])
|
||||
async def test_version_is_checked_after_waiting_for_store_lock(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
outcome: str,
|
||||
) -> None:
|
||||
"""A delayed mutation lock cannot admit a result that became stale while waiting."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
path = tmp_path / "source.txt"
|
||||
path.write_text("original source")
|
||||
context = ApplicationContext(workspace_dir=str(tmp_path))
|
||||
chunker = DefaultFileChunker(app_context=context)
|
||||
context.components = {ComponentEnum.FILE_CHUNKER: {"default": chunker}}
|
||||
async with LocalFileStore(embedding_store="") as store, chunker:
|
||||
await store.upsert([await chunker.chunk(path)])
|
||||
path.write_text("updated source")
|
||||
built = asyncio.Event()
|
||||
step = UpdateIndexStep(file_store=store, persist=False, app_context=context)
|
||||
original = step.build_item
|
||||
|
||||
async def record_build(file_path: Path) -> tuple[FileNode, list[FileChunk]]:
|
||||
item = await original(file_path)
|
||||
built.set()
|
||||
return item
|
||||
|
||||
with patch.object(step, "build_item", side_effect=record_build):
|
||||
async with store._maintenance_guard():
|
||||
task = asyncio.create_task(step(RuntimeContext(changes=[{"change": "modified", "path": str(path)}])))
|
||||
await asyncio.wait_for(built.wait(), 5)
|
||||
await asyncio.sleep(0)
|
||||
if outcome == "cancel":
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
else:
|
||||
path.write_text("latest source")
|
||||
if outcome == "change":
|
||||
response = await task
|
||||
assert response.success is False
|
||||
assert not step._source_versions
|
||||
nodes = await store.get_nodes()
|
||||
assert store.file_chunks[nodes[0].chunk_ids[0]].text == "original source"
|
||||
recovered = await step(RuntimeContext(changes=[{"change": "modified", "path": str(path)}]))
|
||||
assert recovered.success is True
|
||||
nodes = await store.get_nodes()
|
||||
assert store.file_chunks[nodes[0].chunk_ids[0]].text == path.read_text()
|
||||
assert not step._source_versions
|
||||
Loading…
Add table
Reference in a new issue