From 2ed2e89e24ff4edf161a1e051191e3c37bda9c12 Mon Sep 17 00:00:00 2001 From: imrewce Date: Thu, 28 May 2026 18:10:04 +0800 Subject: [PATCH] port orthogonal steps (#262) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(file_io): port orthogonal crud_steps features onto upstream restructure * refactor(file_io): expose with_neighbors/max_neighbors_per_direction/max_bytes as step kwargs (not LLM params) Match search_step's convention: tuning knobs that are config-like (not part of the LLM-facing schema) live in the yaml steps: block and are read via self.kwargs.get(...) — not exposed under parameters.properties. Also simplify the write step metadata field description. --- reme4/config/default.yaml | 22 ++ reme4/constants.py | 4 + reme4/steps/__init__.py | 2 + reme4/steps/file_io/_file_io.py | 69 ++++++ reme4/steps/file_io/edit.py | 105 +++++---- reme4/steps/file_io/read.py | 47 +++- reme4/steps/file_io/read_image.py | 117 ++++++++++ reme4/steps/file_io/write.py | 61 +++-- tests4/unittest/test_read_image_steps.py | 233 ++++++++++++++++++++ tests4/unittest/test_read_with_neighbors.py | 181 +++++++++++++++ tests4/unittest/test_write_metadata_lock.py | 207 +++++++++++++++++ 11 files changed, 982 insertions(+), 66 deletions(-) create mode 100644 reme4/steps/file_io/read_image.py create mode 100644 tests4/unittest/test_read_image_steps.py create mode 100644 tests4/unittest/test_read_with_neighbors.py create mode 100644 tests4/unittest/test_write_metadata_lock.py diff --git a/reme4/config/default.yaml b/reme4/config/default.yaml index 5d69b2db..42ac251e 100644 --- a/reme4/config/default.yaml +++ b/reme4/config/default.yaml @@ -297,6 +297,25 @@ jobs: - path steps: - backend: read_step + with_neighbors: false + max_neighbors_per_direction: 10 + + read_image: + backend: base + description: "Read an image file as base64 (vault-relative path)." + parameters: + type: object + properties: + path: + type: string + description: >- + vault-relative path; common image formats supported + (png/jpg/jpeg/webp/gif/bmp/tiff/heic) + required: + - path + steps: + - backend: read_image_step + max_bytes: 5242880 write: backend: base @@ -316,6 +335,9 @@ jobs: content: type: string description: "body" + metadata: + type: object + description: "Optional extra frontmatter fields (md only)." required: - path - name diff --git a/reme4/constants.py b/reme4/constants.py index 819fbadd..ee4c1149 100644 --- a/reme4/constants.py +++ b/reme4/constants.py @@ -10,3 +10,7 @@ REME_DEFAULT_PORT = 2333 DEFAULT_MAX_BYTES = 50 * 1024 MAX_FILE_READ_BYTES = 200 * 1024 * 1024 TRUNCATION_NOTICE_MARKER = "<>" + +# read_image step: oversized images above this threshold return path & metadata +# only (no base64) to keep LLM context budgets safe. +DEFAULT_MAX_IMAGE_BYTES = 5 * 1024 * 1024 diff --git a/reme4/steps/__init__.py b/reme4/steps/__init__.py index 358ecd38..daa59838 100644 --- a/reme4/steps/__init__.py +++ b/reme4/steps/__init__.py @@ -17,6 +17,7 @@ from .file_io.frontmatter_update import FrontmatterUpdateStep from .file_io.list import ListStep from .file_io.move import MoveStep from .file_io.read import ReadStep +from .file_io.read_image import ReadImageStep from .file_io.stat import StatStep from .file_io.write import WriteStep from .index.clear_and_scan import ClearAndScanStep @@ -46,6 +47,7 @@ __all__ = [ "ListStep", "MoveStep", "ReadStep", + "ReadImageStep", "StatStep", "WriteStep", # file_io (daily) diff --git a/reme4/steps/file_io/_file_io.py b/reme4/steps/file_io/_file_io.py index 7a90acbd..fad27e10 100644 --- a/reme4/steps/file_io/_file_io.py +++ b/reme4/steps/file_io/_file_io.py @@ -10,6 +10,7 @@ Two related concerns, both private to the ``crud`` package: manual sections are preserved verbatim across refreshes. """ +import asyncio import re from pathlib import Path from typing import Iterable @@ -32,6 +33,24 @@ NON_MD_WARNING = ( "Operating in compatibility mode may carry risks of errors." ) +NON_IMAGE_WARNING = ( + "non-image file detected; CRUD image operations are recommended on standard image formats. " + "Operating in compatibility mode may carry risks of errors." +) + +# Image suffix → MIME mapping. SVG intentionally excluded (text format, base64 +# encoding has no benefit — caller should use read_step instead). +IMAGE_MIME_BY_EXT: dict[str, str] = { + ".png": "image/png", + ".jpg": "image/jpeg", + ".jpeg": "image/jpeg", + ".webp": "image/webp", + ".gif": "image/gif", + ".bmp": "image/bmp", + ".tiff": "image/tiff", + ".heic": "image/heic", +} + # Modern text formats — assume UTF-8 by convention. _STANDARD_TEXT_EXTS = { ".md", @@ -53,6 +72,39 @@ _STANDARD_TEXT_EXTS = { # Legacy formats that may use ANSI/GBK on Chinese Windows systems. _NON_STANDARD_EXTS = {".csv", ".bat", ".cmd", ".reg"} + +# --------------------------------------------------------------------------- +# In-process per-path write lock. +# +# Concurrent CRUD writes (write / edit) targeting the same path from the +# same process must be serialized so a read-modify-write cycle isn't +# interleaved by another coroutine. Different paths get different locks, +# so unrelated writes still run in parallel. +# +# NOTE: this is in-process only — multi-worker / multi-process deployments +# are NOT protected. That trade-off is acceptable for the current single- +# process reme server; cross-process protection would need flock or OCC. +# --------------------------------------------------------------------------- +_PATH_LOCKS: dict[str, asyncio.Lock] = {} +_PATH_LOCKS_REGISTRY = asyncio.Lock() + + +async def get_path_lock(target: Path) -> asyncio.Lock: + """Return the asyncio.Lock for ``target``; created lazily on first request. + + The lock is keyed by the string form of ``target`` — callers should pass + a path that has already been normalized by :func:`resolve_path` so two + equivalent paths share one lock. + """ + key = str(target) + async with _PATH_LOCKS_REGISTRY: + lock = _PATH_LOCKS.get(key) + if lock is None: + lock = asyncio.Lock() + _PATH_LOCKS[key] = lock + return lock + + # Path helpers # ------------ @@ -149,6 +201,23 @@ def gate_md(target: Path) -> tuple[Path, bool]: return target, True +def gate_image(target: Path) -> tuple[Path, bool, str | None]: + """Image gate with compatibility fallback. Returns ``(path, is_image, mime)``. + + Behavior diverges from :func:`gate_md` in two ways: + - **No suffix is NOT auto-appended.** Image formats have no single + reasonable default; guessing would mislead. + - **No path mutation** — ``target`` is returned unchanged. + + Suffix routing: + - Known image suffix (see ``IMAGE_MIME_BY_EXT``) → ``(target, True, mime)`` + - Empty suffix or unknown suffix → ``(target, False, None)`` + Caller may still read the file and surface a ``NON_IMAGE_WARNING``. + """ + mime = IMAGE_MIME_BY_EXT.get(target.suffix.lower()) + return target, mime is not None, mime + + # Encoding detection (private) # ---------------------------- diff --git a/reme4/steps/file_io/edit.py b/reme4/steps/file_io/edit.py index fc0462cf..4e6809aa 100644 --- a/reme4/steps/file_io/edit.py +++ b/reme4/steps/file_io/edit.py @@ -2,7 +2,14 @@ import frontmatter -from ._file_io import NON_MD_WARNING, gate_md, read_file_safe, resolve_path, write_file_safe +from ._file_io import ( + NON_MD_WARNING, + gate_md, + get_path_lock, + read_file_safe, + resolve_path, + write_file_safe, +) from ..base_step import BaseStep from ...components import R @@ -13,7 +20,11 @@ class EditStep(BaseStep): The YAML front matter block (if any) is parsed out, kept verbatim and re-emitted unchanged — matches that fall inside front matter are ignored, - so a typo in `old` cannot corrupt structured metadata.""" + so a typo in `old` cannot corrupt structured metadata. + + Concurrency: in-process per-path ``asyncio.Lock`` serializes the + read-modify-write cycle against the same file (multi-worker / multi- + process safety is out of scope).""" def _fail(self, message: str, **meta) -> None: assert self.context is not None @@ -44,56 +55,58 @@ class EditStep(BaseStep): target, is_md = gate_md(target) - if not target.exists(): - self._fail(f"file {target} does not exist", path=str(target)) - return None - if not target.is_file(): - self._fail(f"path {target} is not a file", path=str(target)) - return None + lock = await get_path_lock(target) + async with lock: + if not target.exists(): + self._fail(f"file {target} does not exist", path=str(target)) + return None + if not target.is_file(): + self._fail(f"path {target} is not a file", path=str(target)) + return None - try: - raw_text, encoding = await read_file_safe(target) - except Exception as e: # pylint: disable=broad-except - self._fail(f"read failed: {e}", path=str(target)) - return None + try: + raw_text, encoding = await read_file_safe(target) + except Exception as e: # pylint: disable=broad-except + self._fail(f"read failed: {e}", path=str(target)) + return None - # Markdown: parse frontmatter and operate on body only. Non-markdown: - # there's no frontmatter convention, so operate on the full text. - if is_md: - post = frontmatter.loads(raw_text) - body = post.content - not_found_msg = ( - f"text to replace was not found in the body of {target} (front matter is excluded from edit)" - ) - else: - post = None - body = raw_text - not_found_msg = f"text to replace was not found in {target}" + # Markdown: parse frontmatter and operate on body only. Non-markdown: + # there's no frontmatter convention, so operate on the full text. + if is_md: + post = frontmatter.loads(raw_text) + body = post.content + not_found_msg = ( + f"text to replace was not found in the body of {target} (front matter is excluded from edit)" + ) + else: + post = None + body = raw_text + not_found_msg = f"text to replace was not found in {target}" - if old_str not in body: - self._fail(not_found_msg, path=str(target)) - return None + if old_str not in body: + self._fail(not_found_msg, path=str(target)) + return None - count = body.count(old_str) - new_body = body.replace(old_str, new_str) + count = body.count(old_str) + new_body = body.replace(old_str, new_str) - if is_md and post is not None: - post.content = new_body - # Re-serialize: keep front matter when present, otherwise emit body alone - # so we don't introduce an empty `---\n---\n` block. - new_text = frontmatter.dumps(post) if post.metadata else post.content - if not new_text.endswith("\n"): - new_text += "\n" - else: - new_text = new_body + if is_md and post is not None: + post.content = new_body + # Re-serialize: keep front matter when present, otherwise emit body alone + # so we don't introduce an empty `---\n---\n` block. + new_text = frontmatter.dumps(post) if post.metadata else post.content + if not new_text.endswith("\n"): + new_text += "\n" + else: + new_text = new_body - # Preserve the file's original encoding (returned by read_file_safe above) - # so edits don't silently re-encode non-UTF-8 files (e.g. GBK CSV) to UTF-8. - try: - await write_file_safe(target, new_text, encoding=encoding) - except Exception as e: # pylint: disable=broad-except - self._fail(f"write failed: {e}", path=str(target)) - return None + # Preserve the file's original encoding (returned by read_file_safe above) + # so edits don't silently re-encode non-UTF-8 files (e.g. GBK CSV) to UTF-8. + try: + await write_file_safe(target, new_text, encoding=encoding) + except Exception as e: # pylint: disable=broad-except + self._fail(f"write failed: {e}", path=str(target)) + return None self.context.response.success = True answer = f"Replaced {count} occurrence(s) in {target}" diff --git a/reme4/steps/file_io/read.py b/reme4/steps/file_io/read.py index 3263677f..1114e5f3 100644 --- a/reme4/steps/file_io/read.py +++ b/reme4/steps/file_io/read.py @@ -5,11 +5,21 @@ from pathlib import Path from ._file_io import NON_MD_WARNING, gate_md, read_file_safe, resolve_path, truncate_text_output from ..base_step import BaseStep from ...components import R +from ...utils import expand_links, render_expansion_lines @R.register("read_step") class ReadStep(BaseStep): - """Read a markdown file. Optional `start_line`/`end_line` for ranged reads.""" + """Read a markdown file. Optional `start_line`/`end_line` for ranged reads. + + Step-level attributes (``kwargs``, configured in yaml under ``steps:`` — + not exposed to LLM): + with_neighbors (bool, default False): when true and the file is + markdown, append a block listing first-order bidirectional + neighbors (out/in link targets) with name/description meta, + fetched via the file_store. Same rendering as SearchStep. + max_neighbors_per_direction (int, default 10): cap per direction. + """ def _fail(self, message: str, **meta) -> None: """Mark the response failed and stash a human-readable error.""" @@ -81,6 +91,8 @@ class ReadStep(BaseStep): assert self.context is not None raw = str(self.context.get("path") or "") start_line, end_line = self.context.get("start_line"), self.context.get("end_line") + with_neighbors: bool = bool(self.kwargs.get("with_neighbors", False)) + max_neighbors_per_direction: int = int(self.kwargs.get("max_neighbors_per_direction", 10)) # Validate inputs and target before touching the filesystem twice. target = self._resolve_target(raw) @@ -112,4 +124,37 @@ class ReadStep(BaseStep): self.context.response.success = True self.context.response.answer = text self.logger.info(f"[{self.name}] read path={target} lines={s}-{e}/{total} bytes={len(text.encode('utf-8'))}") + + if with_neighbors and target.suffix.lower() == ".md": + await self._maybe_inject_neighbors(target, text, max_neighbors_per_direction) + return self.context.response + + # -- neighbor injection (opt-in) ----------------------------------------- + + async def _maybe_inject_neighbors(self, target: Path, text: str, max_per_direction: int) -> None: + """Append the rendered neighbor block + stash raw expansion in metadata.""" + assert self.context is not None + try: + rel_path = str(target.relative_to(self.vault_path)) + except ValueError: + self.logger.info(f"[{self.name}] skip neighbors: path outside vault_path path={target}") + return + + try: + expansion = await expand_links(self.file_store, [rel_path], max_per_direction) + except Exception as exc: # noqa: BLE001 + self.logger.warning(f"[{self.name}] neighbor fetch failed: {type(exc).__name__}: {exc}") + return + + per_path = expansion.get(rel_path, {}) + lines = render_expansion_lines(per_path) + if not lines: + return + + out_n = len(per_path.get("outlinks") or []) + in_n = len(per_path.get("inlinks") or []) + header = f"========== Related neighbors (outlinks={out_n}, inlinks={in_n}) ==========" + block = "\n".join([header, *lines]) + self.context.response.answer = f"{text}\n\n{block}" + self.context.response.metadata["link_expansion"] = expansion diff --git a/reme4/steps/file_io/read_image.py b/reme4/steps/file_io/read_image.py new file mode 100644 index 00000000..d5a60f67 --- /dev/null +++ b/reme4/steps/file_io/read_image.py @@ -0,0 +1,117 @@ +"""Read an image file as base64; oversized images degrade to path + metadata.""" + +import base64 + +import aiofiles +import aiofiles.os + +from ._file_io import NON_IMAGE_WARNING, gate_image, resolve_path +from ..base_step import BaseStep +from ...components import R +from ...constants import DEFAULT_MAX_IMAGE_BYTES + + +@R.register("read_image_step") +class ReadImageStep(BaseStep): + """Read an image file under ``vault_dir`` and return base64 in ``answer``. + + Step-level attributes (``kwargs``, configured in yaml under ``steps:`` — + not exposed to LLM): + max_bytes (int, default ``DEFAULT_MAX_IMAGE_BYTES``): cap for the + base64 path. Above this, ``answer`` carries a notice and + ``metadata.oversized=True`` (no base64). + + Invariants (callers depend on these): + - Normal branch: ``answer`` is pure base64 (no ``data:`` prefix, no + notice suffix). Use ``f"data:{mime};base64,{answer}"`` if you need a + data URL. + - Oversized branch (``metadata.oversized=True``): ``answer`` is a + human-readable notice, **not** base64. Inspect ``metadata.oversized`` + before decoding. + - Unknown / missing suffix (compatibility mode): ``answer`` stays pure + base64; ``metadata.non_image_warning=True`` flags it. + """ + + def _fail(self, message: str, **meta) -> None: + assert self.context is not None + self.context.response.success = False + self.context.response.answer = f"Error: {message}" + if meta: + self.context.response.metadata.update(meta) + + async def execute(self): # pylint: disable=too-many-return-statements + assert self.context is not None + raw = str(self.context.get("path") or "") + max_bytes_raw = self.kwargs.get("max_bytes", DEFAULT_MAX_IMAGE_BYTES) + + try: + max_bytes = int(max_bytes_raw) + if max_bytes <= 0: + raise ValueError() + except (TypeError, ValueError): + self._fail(f"`max_bytes` must be a positive integer, got {max_bytes_raw!r}") + return None + + target, err = resolve_path(self.vault_path, raw) + if err: + self._fail(err) + return None + + target, is_image, mime = gate_image(target) + if not is_image: + self.logger.info(f"[{self.name}] {NON_IMAGE_WARNING} path={target}") + + if not target.exists(): + self._fail(f"file {target} does not exist", path=str(target)) + return None + if not target.is_file(): + self._fail(f"path {target} is not a file", path=str(target)) + return None + + try: + stat = await aiofiles.os.stat(str(target)) + except Exception as e: # pylint: disable=broad-except + self._fail(f"read failed: {e}", path=str(target)) + return None + size_bytes = stat.st_size + + if size_bytes > max_bytes: + self.context.response.success = True + self.context.response.answer = ( + f"image exceeds max_bytes ({size_bytes} > {max_bytes}), " + f"base64 omitted; use path={target} to access directly" + ) + md = { + "path": str(target), + "size_bytes": size_bytes, + "mime": mime, + "oversized": True, + "max_bytes": max_bytes, + } + if not is_image: + md["non_image_warning"] = True + self.context.response.metadata.update(md) + self.logger.info( + f"[{self.name}] read path={target} size={size_bytes} mime={mime} " + f"oversized=True is_image={is_image}", + ) + return self.context.response + + try: + async with aiofiles.open(str(target), "rb") as f: + data = await f.read() + except Exception as e: # pylint: disable=broad-except + self._fail(f"read failed: {e}", path=str(target)) + return None + + b64 = base64.b64encode(data).decode("ascii") + self.context.response.success = True + self.context.response.answer = b64 + md = {"path": str(target), "size_bytes": size_bytes, "mime": mime} + if not is_image: + md["non_image_warning"] = True + self.context.response.metadata.update(md) + self.logger.info( + f"[{self.name}] read path={target} size={size_bytes} mime={mime} oversized=False is_image={is_image}", + ) + return self.context.response diff --git a/reme4/steps/file_io/write.py b/reme4/steps/file_io/write.py index ba0f478f..26aee115 100644 --- a/reme4/steps/file_io/write.py +++ b/reme4/steps/file_io/write.py @@ -2,20 +2,32 @@ import frontmatter -from ._file_io import NON_MD_WARNING, detect_file_encoding, gate_md, resolve_path, write_file_safe +from ._file_io import ( + NON_MD_WARNING, + detect_file_encoding, + gate_md, + get_path_lock, + resolve_path, + write_file_safe, +) from ..base_step import BaseStep from ...components import R @R.register("write_step") class WriteStep(BaseStep): - """Write (create or overwrite) a markdown file. When the target already exists, - its contents are replaced and a system notice is appended to the answer. + """Write (create or overwrite) a markdown file. - Front matter is restricted to two string fields: ``name`` and ``description``. - The CLI schema declares them as required, but the step itself is lenient — - missing or empty values are silently skipped so manual invocations don't - fail catastrophically.""" + Frontmatter accepts two reserved string fields (``name`` / ``description``) + plus an optional free-form ``metadata`` dict whose entries are expanded + into the frontmatter as-is. ``name`` / ``description`` keys inside + ``metadata`` are ignored — only the top-level explicit parameters are + honored for those two reserved fields. + + Concurrency: in-process per-path ``asyncio.Lock`` serializes concurrent + writes to the same file (multi-worker / multi-process safety is out of + scope). + """ def _fail(self, message: str, **meta) -> None: assert self.context is not None @@ -29,6 +41,7 @@ class WriteStep(BaseStep): raw = str(self.context.get("path") or "") content = self.context.get("content") content = "" if content is None else str(content) + metadata_raw = self.context.get("metadata") target, err = resolve_path(self.vault_path, raw) if err: @@ -37,12 +50,19 @@ class WriteStep(BaseStep): target, is_md = gate_md(target) - existed = target.exists() - # Non-markdown files have no frontmatter convention: name/description - # are silently dropped and the body is written verbatim. + # and metadata are silently dropped and the body is written verbatim. if is_md: meta: dict = {} + # Expand `metadata` dict first, skipping the two reserved keys. + if isinstance(metadata_raw, dict): + for k, v in metadata_raw.items(): + if k in ("name", "description"): + continue + if v is None: + continue + meta[str(k)] = v + # Layer explicit name/description on top (always wins). for key in ("name", "description"): value = self.context.get(key) if value is None: @@ -62,14 +82,17 @@ class WriteStep(BaseStep): else: body = content - # Preserve the existing file's encoding when overwriting (e.g. GBK CSV - # stays GBK). New files are written as UTF-8. - encoding = await detect_file_encoding(target) if existed else "utf-8" - try: - await write_file_safe(target, body, encoding=encoding) - except Exception as e: # pylint: disable=broad-except - self._fail(f"write failed: {e}", path=str(target)) - return None + lock = await get_path_lock(target) + async with lock: + existed = target.exists() + # Preserve the existing file's encoding when overwriting (e.g. GBK + # CSV stays GBK). New files are written as UTF-8. + encoding = await detect_file_encoding(target) if existed else "utf-8" + try: + await write_file_safe(target, body, encoding=encoding) + except Exception as e: # pylint: disable=broad-except + self._fail(f"write failed: {e}", path=str(target)) + return None try: nbytes = len(body.encode(encoding)) @@ -77,7 +100,7 @@ class WriteStep(BaseStep): nbytes = len(body.encode("utf-8")) self.context.response.success = True if existed: - answer = f"Wrote {target} ({nbytes} bytes) " f"[system notice: target already existed and was overwritten]" + answer = f"Wrote {target} ({nbytes} bytes) [system notice: target already existed and was overwritten]" else: answer = f"Wrote {target} ({nbytes} bytes)" if not is_md: diff --git a/tests4/unittest/test_read_image_steps.py b/tests4/unittest/test_read_image_steps.py new file mode 100644 index 00000000..849d6649 --- /dev/null +++ b/tests4/unittest/test_read_image_steps.py @@ -0,0 +1,233 @@ +"""Tests for the ``read_image`` step. + +Style mirrors ``test_crud_steps.py`` — direct step invocation against a +freshly built ``LocalFileStore`` with the step's ``vault_path`` rooted at +``cwd()`` via chdir. + +Coverage: + - happy paths for known suffixes (png/jpeg) + - oversized branch (``answer`` is a notice, ``metadata.oversized=True``) + - unknown / missing suffix (compatibility mode; ``non_image_warning=True``) + - error branches: missing file, directory, empty path, bad ``max_bytes`` +""" + +# pylint: disable=protected-access + +import asyncio +import base64 +import os +import tempfile +import warnings +from pathlib import Path + +from reme4.components.file_store import LocalFileStore +from reme4.steps.file_io import read_image as crud_read_image + +warnings.filterwarnings("ignore", category=DeprecationWarning, module="jieba") +warnings.filterwarnings("ignore", category=DeprecationWarning, module="pkg_resources") + + +class temp_chdir: + """Context manager to temporarily chdir into a path and restore on exit.""" + + def __init__(self, path): + self.path = path + self.old = None + + def __enter__(self): + self.old = os.getcwd() + os.chdir(self.path) + return self + + def __exit__(self, *exc): + os.chdir(self.old) + + +def _run(coro): + asyncio.run(coro) + + +async def _make_store() -> LocalFileStore: + store = LocalFileStore(name="t_img", embedding_model="") + await store.start() + return store + + +def _seed_bytes(rel: str, data: bytes) -> Path: + """Drop raw bytes at ``cwd/rel``. Step is byte-level — no real PNG needed.""" + target = Path.cwd() / rel + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(data) + return target + + +async def _read_image(store: LocalFileStore, *, step_kwargs: dict | None = None, **call_kwargs): + """Run ReadImageStep; ``step_kwargs`` go to step init (kwargs/attrs).""" + step = crud_read_image.ReadImageStep(file_store=store, **(step_kwargs or {})) + await step(**call_kwargs) + return step.context.response + + +def test_read_image_png(): + """``read_image path=img/cat.png`` returns base64 + ``image/png`` mime.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + payload = b"\x89PNG\r\n\x1a\n" + b"fake-png-body-bytes" + _seed_bytes("img/cat.png", payload) + store = await _make_store() + resp = await _read_image(store, path="img/cat.png") + assert resp.success is True, resp + assert base64.b64decode(resp.answer) == payload, "base64 round-trip mismatch" + assert resp.metadata["mime"] == "image/png", resp.metadata + assert resp.metadata["size_bytes"] == len(payload), resp.metadata + assert "oversized" not in resp.metadata, resp.metadata + await store.close() + print("✓ test_read_image_png passed") + + _run(run()) + + +def test_read_image_jpeg(): + """``.jpg`` suffix maps to ``image/jpeg``.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + payload = b"\xff\xd8\xff\xe0" + b"jpeg-body" + _seed_bytes("dog.jpg", payload) + store = await _make_store() + resp = await _read_image(store, path="dog.jpg") + assert resp.success is True, resp + assert base64.b64decode(resp.answer) == payload + assert resp.metadata["mime"] == "image/jpeg", resp.metadata + await store.close() + print("✓ test_read_image_jpeg passed") + + _run(run()) + + +def test_read_image_oversized(): + """Above ``max_bytes`` → ``answer`` is a notice, ``metadata.oversized=True``.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + payload = b"\x89PNG\r\n\x1a\n" + b"x" * 2048 + _seed_bytes("big.png", payload) + store = await _make_store() + resp = await _read_image(store, step_kwargs={"max_bytes": 1024}, path="big.png") + assert resp.success is True, resp + assert resp.metadata["oversized"] is True, resp.metadata + assert resp.metadata["max_bytes"] == 1024, resp.metadata + assert resp.metadata["size_bytes"] == len(payload), resp.metadata + assert resp.metadata["mime"] == "image/png", resp.metadata + try: + decoded = base64.b64decode(resp.answer, validate=True) + assert decoded != payload, "oversized branch must not return real base64" + except Exception: + pass # expected — answer is notice text, not base64 + assert "exceeds max_bytes" in resp.answer + await store.close() + print("✓ test_read_image_oversized passed") + + _run(run()) + + +def test_read_image_unknown_suffix(): + """Unknown suffix → still returns base64, ``metadata.non_image_warning=True``.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + payload = b"any-bytes-here-for-blob" + _seed_bytes("blob.xyz", payload) + store = await _make_store() + resp = await _read_image(store, path="blob.xyz") + assert resp.success is True, resp + assert base64.b64decode(resp.answer) == payload + assert resp.metadata["non_image_warning"] is True, resp.metadata + assert resp.metadata["mime"] is None, resp.metadata + await store.close() + print("✓ test_read_image_unknown_suffix passed") + + _run(run()) + + +def test_read_image_no_suffix(): + """No suffix → compatibility mode (no auto-append), still reads as base64.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + payload = b"\x89PNG\r\n\x1a\nbody" + _seed_bytes("no_suffix_blob", payload) + store = await _make_store() + resp = await _read_image(store, path="no_suffix_blob") + assert resp.success is True, resp + assert resp.metadata["non_image_warning"] is True, resp.metadata + assert resp.metadata["mime"] is None, resp.metadata + assert base64.b64decode(resp.answer) == payload + await store.close() + print("✓ test_read_image_no_suffix passed") + + _run(run()) + + +def test_read_image_missing(): + """Non-existent path → ``success=False`` with ``does not exist`` message.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _make_store() + resp = await _read_image(store, path="never_existed.png") + assert resp.success is False, resp + assert resp.answer.startswith("Error:"), resp + assert "does not exist" in resp.answer + await store.close() + print("✓ test_read_image_missing passed") + + _run(run()) + + +def test_read_image_is_directory(): + """Path pointing to a directory → ``success=False`` with ``is not a file``.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + (Path(tmp) / "subdir").mkdir(parents=True, exist_ok=True) + store = await _make_store() + resp = await _read_image(store, path="subdir") + assert resp.success is False, resp + assert "is not a file" in resp.answer, resp + await store.close() + print("✓ test_read_image_is_directory passed") + + _run(run()) + + +def test_read_image_path_required(): + """Empty ``path`` → ``success=False`` with ``path is required`` message.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _make_store() + resp = await _read_image(store, path="") + assert resp.success is False, resp + assert "`path` is required" in resp.answer, resp + await store.close() + print("✓ test_read_image_path_required passed") + + _run(run()) + + +def test_read_image_invalid_max_bytes(): + """``max_bytes=-1`` → ``success=False`` with positive-integer error.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + _seed_bytes("a.png", b"\x89PNG\r\n\x1a\nx") + store = await _make_store() + resp = await _read_image(store, step_kwargs={"max_bytes": -1}, path="a.png") + assert resp.success is False, resp + assert "positive integer" in resp.answer, resp + await store.close() + print("✓ test_read_image_invalid_max_bytes passed") + + _run(run()) diff --git a/tests4/unittest/test_read_with_neighbors.py b/tests4/unittest/test_read_with_neighbors.py new file mode 100644 index 00000000..45ae476d --- /dev/null +++ b/tests4/unittest/test_read_with_neighbors.py @@ -0,0 +1,181 @@ +"""Tests for ``ReadStep`` with ``with_neighbors=True`` opt-in injection. + +The opt-in re-uses ``reme4.utils.link_expansion.{expand_links, +render_expansion_lines}`` (already covered by ``test_link_expansion.py``), +so these tests focus on the *integration* layer: + +* off by default — no extra answer suffix, no ``link_expansion`` metadata +* on + has neighbors — answer ends with a ``Related neighbors`` block, + metadata carries the raw expansion dict +* on + no neighbors — no block appended, no metadata key +* on + non-markdown — falls through (no neighbor injection) +""" + +# pylint: disable=protected-access + +import asyncio +import os +import tempfile +import warnings +from pathlib import Path + +from reme4.components.file_store import LocalFileStore +from reme4.schema import FileFrontMatter, FileNode +from reme4.steps.file_io import read as crud_read +from reme4.utils.wikilink_handler import WikilinkHandler + +warnings.filterwarnings("ignore", category=DeprecationWarning, module="jieba") +warnings.filterwarnings("ignore", category=DeprecationWarning, module="pkg_resources") + + +class temp_chdir: + """Context manager to temporarily chdir into a path and restore on exit.""" + + def __init__(self, path): + self.path = path + self.old = None + + def __enter__(self): + self.old = os.getcwd() + os.chdir(self.path) + return self + + def __exit__(self, *exc): + os.chdir(self.old) + + +def _run(coro): + asyncio.run(coro) + + +async def _store_with(files: dict[str, dict]) -> LocalFileStore: + """LocalFileStore seeded with files + parsed wikilinks + optional frontmatter.""" + store = LocalFileStore(name="t_read_neighbors", embedding_model="") + await store.start() + nodes: list[FileNode] = [] + root = Path.cwd() + for rel, spec in files.items(): + body = spec["body"] + abs_path = root / rel + abs_path.parent.mkdir(parents=True, exist_ok=True) + abs_path.write_text(body, encoding="utf-8") + fm = FileFrontMatter( + name=spec.get("name", ""), + description=spec.get("description", ""), + ) + nodes.append( + FileNode( + path=rel, + st_mtime=abs_path.stat().st_mtime, + links=WikilinkHandler.extract_links(body, rel), + front_matter=fm, + ), + ) + if nodes: + await store.file_graph.upsert_nodes(nodes) + return store + + +async def _read(store: LocalFileStore, *, step_kwargs: dict | None = None, **call_kwargs): + """Run a ReadStep against ``store``; ``step_kwargs`` go to step init (kwargs/attrs).""" + step = crud_read.ReadStep(file_store=store, **(step_kwargs or {})) + await step(**call_kwargs) + return step.context.response + + +# -- off (default) ----------------------------------------------------------- + + +def test_read_without_neighbors_default(): + """``with_neighbors`` defaults to off — no block, no metadata key.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _store_with( + { + "A.md": {"body": "See [[B.md]] for details.", "name": "A Doc"}, + "B.md": {"body": "End node.", "name": "B Doc"}, + }, + ) + resp = await _read(store, path="A.md") + assert resp.success is True + assert "Related neighbors" not in str(resp.answer) + assert "link_expansion" not in resp.metadata + await store.close() + print("✓ test_read_without_neighbors_default passed") + + _run(run()) + + +# -- on + has neighbors ------------------------------------------------------ + + +def test_read_with_neighbors_injects_block_and_metadata(): + """Block is appended; metadata carries the raw expansion dict.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _store_with( + { + "A.md": {"body": "See [[B.md]] for details.", "name": "A Doc", "description": "alpha"}, + "B.md": {"body": "End node.", "name": "B Doc", "description": "beta"}, + }, + ) + resp = await _read(store, step_kwargs={"with_neighbors": True}, path="A.md") + assert resp.success is True + answer = str(resp.answer) + assert "Related neighbors" in answer + assert "outlinks=1" in answer + assert "inlinks=0" in answer + assert "B.md" in answer + assert 'name="B Doc"' in answer + + expansion = resp.metadata.get("link_expansion") + assert expansion is not None + assert "A.md" in expansion + assert expansion["A.md"]["outlinks"][0]["path"] == "B.md" + assert expansion["A.md"]["outlinks"][0]["meta"] == {"name": "B Doc", "description": "beta"} + await store.close() + print("✓ test_read_with_neighbors_injects_block_and_metadata passed") + + _run(run()) + + +# -- on + zero neighbors ----------------------------------------------------- + + +def test_read_with_neighbors_no_links_no_block(): + """File with no out/in links → no block appended, no metadata key.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _store_with({"lonely.md": {"body": "no links here", "name": "Lonely"}}) + resp = await _read(store, step_kwargs={"with_neighbors": True}, path="lonely.md") + assert resp.success is True + assert "Related neighbors" not in str(resp.answer) + assert "link_expansion" not in resp.metadata + await store.close() + print("✓ test_read_with_neighbors_no_links_no_block passed") + + _run(run()) + + +# -- on + non-md ------------------------------------------------------------- + + +def test_read_with_neighbors_non_md_falls_through(): + """Non-md target: neighbor injection is skipped even when ``with_neighbors=True``.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + (Path(tmp) / "notes.txt").write_text("plain text body", encoding="utf-8") + store = LocalFileStore(name="t_read_neighbors_nonmd", embedding_model="") + await store.start() + resp = await _read(store, step_kwargs={"with_neighbors": True}, path="notes.txt") + assert resp.success is True + assert "Related neighbors" not in str(resp.answer) + assert "link_expansion" not in resp.metadata + await store.close() + print("✓ test_read_with_neighbors_non_md_falls_through passed") + + _run(run()) diff --git a/tests4/unittest/test_write_metadata_lock.py b/tests4/unittest/test_write_metadata_lock.py new file mode 100644 index 00000000..1cfb98f2 --- /dev/null +++ b/tests4/unittest/test_write_metadata_lock.py @@ -0,0 +1,207 @@ +"""Tests for the WriteStep ``metadata`` param and the per-path write lock. + +The ``metadata`` feature lets callers extend the on-disk frontmatter beyond +the two reserved fields (``name`` / ``description``) without touching the +step interface. Reserved keys inside the dict are ignored — explicit +top-level parameters always win. + +The per-path lock serializes concurrent write_step invocations targeting +the same path within a single process. We exercise it by firing many +concurrent writes at one path and asserting the final state is consistent +(the lock guarantees the last write's bytes land intact, not a torn +interleaving). +""" + +# pylint: disable=protected-access + +import asyncio +import os +import tempfile +import warnings +from pathlib import Path + +import frontmatter + +from reme4.components.file_store import LocalFileStore +from reme4.steps.file_io import write as crud_write + +warnings.filterwarnings("ignore", category=DeprecationWarning, module="jieba") +warnings.filterwarnings("ignore", category=DeprecationWarning, module="pkg_resources") + + +class temp_chdir: + """Context manager to temporarily chdir into a path and restore on exit.""" + + def __init__(self, path): + self.path = path + self.old = None + + def __enter__(self): + self.old = os.getcwd() + os.chdir(self.path) + return self + + def __exit__(self, *exc): + os.chdir(self.old) + + +def _run(coro): + asyncio.run(coro) + + +async def _make_store() -> LocalFileStore: + store = LocalFileStore(name="t_write_meta", embedding_model="") + await store.start() + return store + + +async def _write(store: LocalFileStore, **kwargs): + step = crud_write.WriteStep(file_store=store) + await step(**kwargs) + return step.context.response + + +# -- metadata expansion ------------------------------------------------------ + + +def test_write_metadata_extends_frontmatter(): + """``metadata={"tags": [...]}`` ends up as a frontmatter field on disk.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _make_store() + resp = await _write( + store, + path="note.md", + name="Hello", + description="A note", + content="body text", + metadata={"tags": ["alpha", "beta"], "priority": 3}, + ) + assert resp.success is True, resp + on_disk = (Path(tmp) / "note.md").read_text(encoding="utf-8") + post = frontmatter.loads(on_disk) + assert post.metadata["name"] == "Hello" + assert post.metadata["description"] == "A note" + assert post.metadata["tags"] == ["alpha", "beta"] + assert post.metadata["priority"] == 3 + assert post.content.strip() == "body text" + await store.close() + print("✓ test_write_metadata_extends_frontmatter passed") + + _run(run()) + + +def test_write_metadata_reserved_keys_ignored(): + """``name`` / ``description`` inside ``metadata`` are dropped; explicit args win.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _make_store() + resp = await _write( + store, + path="note.md", + name="ExplicitName", + description="ExplicitDesc", + content="body", + metadata={"name": "DroppedName", "description": "DroppedDesc", "tag": "kept"}, + ) + assert resp.success is True, resp + post = frontmatter.loads((Path(tmp) / "note.md").read_text(encoding="utf-8")) + assert post.metadata["name"] == "ExplicitName" + assert post.metadata["description"] == "ExplicitDesc" + assert post.metadata["tag"] == "kept" + await store.close() + print("✓ test_write_metadata_reserved_keys_ignored passed") + + _run(run()) + + +def test_write_no_metadata_preserves_legacy_shape(): + """No ``metadata`` arg → frontmatter still contains only name/description.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _make_store() + resp = await _write( + store, + path="note.md", + name="N", + description="D", + content="body", + ) + assert resp.success is True, resp + post = frontmatter.loads((Path(tmp) / "note.md").read_text(encoding="utf-8")) + assert set(post.metadata.keys()) == {"name", "description"} + await store.close() + print("✓ test_write_no_metadata_preserves_legacy_shape passed") + + _run(run()) + + +def test_write_non_md_drops_metadata(): + """Non-markdown target: metadata silently dropped, body written verbatim.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _make_store() + resp = await _write( + store, + path="note.txt", + name="N", + description="D", + content="plain body", + metadata={"tag": "ignored"}, + ) + assert resp.success is True, resp + on_disk = (Path(tmp) / "note.txt").read_text(encoding="utf-8") + assert on_disk == "plain body" + await store.close() + print("✓ test_write_non_md_drops_metadata passed") + + _run(run()) + + +# -- per-path lock ----------------------------------------------------------- + + +def test_write_lock_serializes_concurrent_writes(): + """Many concurrent writes at one path land cleanly — no torn frontmatter. + + Without the lock, concurrent writers can interleave reads-of-existence + and writes-of-bytes; with it, each write either runs before or after + every other write. The on-disk file at the end must parse as valid + frontmatter with name matching exactly one of the writers. + """ + + async def run(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _make_store() + n = 16 + await asyncio.gather( + *( + _write( + store, + path="shared.md", + name=f"writer-{i}", + description=f"desc-{i}", + content=f"body-{i}", + ) + for i in range(n) + ), + ) + on_disk = (Path(tmp) / "shared.md").read_text(encoding="utf-8") + post = frontmatter.loads(on_disk) + assert post.metadata.get("name", "").startswith("writer-"), post.metadata + assert post.metadata.get("description", "").startswith("desc-"), post.metadata + assert post.content.strip().startswith("body-"), post.content + # The body, name, and description must all come from the SAME write + # (no interleaving). Extract the index from each. + idx_name = post.metadata["name"].split("-", 1)[1] + idx_desc = post.metadata["description"].split("-", 1)[1] + idx_body = post.content.strip().split("-", 1)[1] + assert idx_name == idx_desc == idx_body, (idx_name, idx_desc, idx_body) + await store.close() + print("✓ test_write_lock_serializes_concurrent_writes passed") + + _run(run())