From ad4f23e4dc00bc77418ae5d5c5e4e91381c17743 Mon Sep 17 00:00:00 2001 From: lichen2015 Date: Wed, 5 Aug 2026 22:08:30 +0800 Subject: [PATCH 1/3] feat(file_store): add ZvecLocalFileStore backend (#410) * feat(file_store): add ZvecLocalFileStore backend - Implement ZvecLocalFileStore with native zvec collection for ANN search. - Keep JSONL chunks as the source of truth; rebuild collection from chunks when sidecar digest/dimension/HNSW M mismatch is detected. - Add dedicated unit tests in tests/unit/test_zvec_file_store.py. - Parametrize existing file_store consistency tests to cover both LocalFileStore and ZvecLocalFileStore. - Register the new backend in reme/components/file_store/__init__.py. * fix(file_store): fix zvec collection sync and content validation, declare zvec dependency --- pyproject.toml | 1 + reme/components/file_store/__init__.py | 2 + .../file_store/zvec_local_file_store.py | 476 ++++++++++++++ tests/unit/test_file_store_consistency.py | 62 +- tests/unit/test_zvec_file_store.py | 614 ++++++++++++++++++ 5 files changed, 1141 insertions(+), 14 deletions(-) create mode 100644 reme/components/file_store/zvec_local_file_store.py create mode 100644 tests/unit/test_zvec_file_store.py diff --git a/pyproject.toml b/pyproject.toml index 2e1efc6f..cba14f77 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,7 @@ core = [ "openai-codex>=0.144.4", "pproxy>=2.7.9,<3", "faiss-cpu>=1.13.2", + "zvec>=0.6.0", "jieba>=0.42.1", "rjieba>=0.2.1", "neo4j>=6.2.0", diff --git a/reme/components/file_store/__init__.py b/reme/components/file_store/__init__.py index 61e9b14f..116330f8 100644 --- a/reme/components/file_store/__init__.py +++ b/reme/components/file_store/__init__.py @@ -8,9 +8,11 @@ shipped today. from .base_file_store import BaseFileStore from .faiss_local_file_store import FaissLocalFileStore from .local_file_store import LocalFileStore +from .zvec_local_file_store import ZvecLocalFileStore __all__ = [ "BaseFileStore", "FaissLocalFileStore", "LocalFileStore", + "ZvecLocalFileStore", ] diff --git a/reme/components/file_store/zvec_local_file_store.py b/reme/components/file_store/zvec_local_file_store.py new file mode 100644 index 00000000..80129eed --- /dev/null +++ b/reme/components/file_store/zvec_local_file_store.py @@ -0,0 +1,476 @@ +"""Zvec-backed file store: chunk JSONL stays authoritative; a zvec collection replaces the linear vector scan.""" + +import hashlib +import json +import shutil +import time +from uuid import uuid4 + +import aiofiles +import numpy as np + +from .local_file_store import LocalFileStore +from ..component_registry import R +from ...schema import FileChunk, FileNode + +# Batch size for bulk inserts during a rebuild. +_ZVEC_INSERT_BATCH_SIZE = 1024 +# Batch size for the startup fetch that verifies persisted collection contents. +_ZVEC_VERIFY_BATCH_SIZE = 1024 + + +@R.register("zvec") +class ZvecLocalFileStore(LocalFileStore): + """LocalFileStore variant whose vector_search is backed by a zvec collection. + + Chunk persistence is unchanged (JSONL, owned by the parent); the zvec + collection only stores ``(chunk_id, embedding)`` pairs and serves ANN + queries. ``self.file_chunks`` remains the source of truth: if the + collection directory or its digest sidecar is missing or stale, the + collection is rebuilt from the chunks. + + zvec (https://zvec.org) is an in-process vector database, so unlike the + FAISS backend there is no tombstone bookkeeping: documents are updated and + removed natively via ``Collection.upsert`` / ``Collection.delete``, and + the collection persists itself inside its own directory. A small JSON + sidecar records an order-independent digest of the live + ``(chunk_id, embedding)`` set at dump time; on load a digest mismatch + (crash between the chunk dump and the collection flush, externally mixed + files, model change) triggers a clean rebuild, and the opened collection is + additionally verified against the ids and vectors it should hold so a + damaged or externally modified collection is rebuilt instead of served. + + HNSW parameters (``hnsw_m``, ``hnsw_ef_construction``) are baked into the + collection schema at creation time; changing ``hnsw_m`` on an existing + store triggers a rebuild so the graph topology matches the configuration. + + ``optimize_index`` maps directly onto ``Collection.optimize()``, zvec's + idle-time index compaction hook. + + zvec is imported lazily inside ``__init__`` so that merely importing this + module does not require the optional dependency; the backend targets the + zvec version declared in ``pyproject.toml`` (``zvec>=0.6.0``). + """ + + def __init__( + self, + hnsw_m: int = 32, + hnsw_ef_construction: int = 64, + **kwargs, + ): + super().__init__(**kwargs) + self._zvec = self._import_zvec() + self.hnsw_m = hnsw_m + self.hnsw_ef_construction = hnsw_ef_construction + self.zvec_path = self.component_metadata_path / f"zvec_index_{self.name}_{self.store_version}" + self.zvec_sidecar_path = self.component_metadata_path / f"zvec_sidecar_{self.name}_{self.store_version}.json" + self._collection = None # zvec.Collection | None + self._indexed_ids: set[str] = set() # chunk ids currently present in the collection + + @staticmethod + def _import_zvec(): + try: + import zvec + except ImportError as e: + raise ImportError( + "zvec is required for ZvecLocalFileStore. Install with `pip install 'zvec>=0.6.0'`.", + ) from e + return zvec + + # -- helpers ---------------------------------------------------------- + + @property + def _dim(self) -> int: + return self.embedding_store.dimensions if self.embedding_store is not None else 0 + + def _collection_schema(self): + """Vector-only schema: chunk data lives in the JSONL, zvec is pure ANN.""" + zvec = self._zvec + return zvec.CollectionSchema( + name=self.name, + fields=[], + vectors=[ + zvec.VectorSchema( + name="embedding", + data_type=zvec.DataType.VECTOR_FP32, + dimension=self._dim, + index_param=zvec.HnswIndexParam( + metric_type=zvec.MetricType.COSINE, + m=self.hnsw_m, + ef_construction=self.hnsw_ef_construction, + ), + ), + ], + ) + + def _create_collection(self): + """Create a fresh (empty) collection, replacing any directory on disk.""" + # Release any open handle first: zvec holds an in-process lock on the + # collection directory, so the old object must be dropped before the + # directory is wiped and re-created. + self._collection = None + if self.zvec_path.exists(): + shutil.rmtree(self.zvec_path, ignore_errors=True) + self._indexed_ids = set() + return self._zvec.create_and_open(path=str(self.zvec_path), schema=self._collection_schema()) + + def _to_doc(self, chunk: FileChunk): + """Build a zvec Doc carrying only the id and the float32 vector.""" + vector = np.asarray(chunk.embedding, dtype=np.float32).tolist() + return self._zvec.Doc(id=chunk.id, vectors={"embedding": vector}) + + @staticmethod + def _vector_fingerprint(vector) -> bytes: + """Canonical bytes of a vector, comparable across an insert/fetch round trip. + + Vectors are inserted as float32 and hashed as float16 (the JSONL + serialization dtype), so normalizing through float32 then float16 makes + an in-memory chunk embedding and the vector read back from the + collection byte-identical whenever they represent the same value. + """ + if vector is None: + return b"" + return np.asarray(np.asarray(vector, dtype=np.float32), dtype=np.float16).tobytes() + + def _upsert_docs(self, chunks: list[FileChunk]) -> None: + if not chunks or self._collection is None: + return + self._collection.upsert([self._to_doc(c) for c in chunks]) + self._indexed_ids.update(c.id for c in chunks) + + def _delete_docs(self, chunk_ids: list[str]) -> None: + if self._collection is None: + return + stale = [cid for cid in chunk_ids if cid in self._indexed_ids] + if not stale: + return + self._collection.delete(ids=stale) + self._indexed_ids.difference_update(stale) + + def _rebuild_collection(self) -> None: + """Rebuild the zvec collection from self.file_chunks (the source of truth).""" + self._collection = self._create_collection() + chunks = [c for c in self.file_chunks.values() if self._embedding_dim_matches(c.embedding)] + for start in range(0, len(chunks), _ZVEC_INSERT_BATCH_SIZE): + batch = chunks[start : start + _ZVEC_INSERT_BATCH_SIZE] + self._collection.insert([self._to_doc(c) for c in batch]) + self._indexed_ids.update(c.id for c in batch) + self.logger.info(f"{self.name}: rebuilt zvec collection with {len(chunks)} vectors at {self.zvec_path}") + + async def _after_embedding_backfill(self) -> None: + """Make newly backfilled vectors visible to zvec. + + zvec upserts are incremental by nature, so backfilled vectors are added + directly to the live collection; no rebuild or tombstone accounting is + needed. + """ + if self.embedding_store is None or self._dim == 0: + return + if self._collection is None: + self._rebuild_collection() + return + to_add = [ + chunk + for cid, chunk in self.file_chunks.items() + if cid not in self._indexed_ids and self._embedding_dim_matches(chunk.embedding) + ] + self._upsert_docs(to_add) + + # -- maintenance ------------------------------------------------------ + + async def optimize_index(self) -> None: + """Idle-time maintenance: delegate to zvec's native index optimization. + + ``Collection.optimize()`` merges staged segments into the persistent + HNSW index, the direct analogue of a tombstone compaction pass. + """ + await super().optimize_index() + if self._collection is None: + return + try: + self._collection.optimize() + self.logger.info(f"{self.name}: zvec collection optimized") + except Exception as e: + self.logger.exception(f"{self.name}: zvec optimize failed: {e}") + + # -- lifecycle ---------------------------------------------------------- + + async def _close(self) -> None: + """Persist via the parent, then release the collection handle. + + zvec holds an in-process lock on the collection directory for the + lifetime of the ``Collection`` object; dropping the reference releases + it so another store instance can reopen the same directory. + """ + await super()._close() # parent dump() flushes the live collection + self._collection = None + self._indexed_ids = set() + + # -- persistence ------------------------------------------------------ + + def _chunks_embedding_digest(self) -> str: + """Order-independent digest of the live (chunk_id, embedding) set. + + Hashes float16 canonical bytes (the chunk JSONL serialization dtype) so + the value is identical whether an embedding sits fresh in memory or has + round-tripped through the JSONL. Written into the sidecar at dump time + and recomputed from ``self.file_chunks`` at load time: a mismatch means + the collection belongs to a different chunk generation than the + authoritative JSONL. + """ + digest = hashlib.sha256() + eligible = sorted( + cid for cid, chunk in self.file_chunks.items() if self._embedding_dim_matches(chunk.embedding) + ) + for cid in eligible: + digest.update(cid.encode("utf-8")) + digest.update(b"\x00") + digest.update(np.asarray(self.file_chunks[cid].embedding, dtype=np.float16).tobytes()) + return digest.hexdigest() + + async def load(self) -> None: + """Load chunks via the parent, then attach the zvec collection (open or rebuild).""" + await super().load() + if self.embedding_store is None or self._dim == 0: + self._collection = None + return + if not await self._try_open_collection(): + self._rebuild_collection() + + async def _try_open_collection(self) -> bool: + """Open the persisted collection and validate it against the chunks. + + On any mismatch or open error the collection directory and sidecar are + wiped so the caller can rebuild from chunks cleanly. Validated: + + - vector dimension against the active embedding model; + - HNSW ``M`` against the active config (baked into the graph topology); + - the sidecar embedding digest against the authoritative JSONL; + - the collection contents against the ids and vectors it should hold. + """ + if not (self.zvec_path.exists() and self.zvec_sidecar_path.exists()): + return False + collection = None + try: + async with aiofiles.open(self.zvec_sidecar_path, encoding=self.encoding) as f: + sidecar = json.loads(await f.read()) + if sidecar.get("digest") != self._chunks_embedding_digest(): + raise ValueError("zvec sidecar embedding digest does not match persisted chunks") + indexed_ids = set(sidecar.get("ids", [])) + expected_ids = { + cid for cid, chunk in self.file_chunks.items() if self._embedding_dim_matches(chunk.embedding) + } + if indexed_ids != expected_ids: + raise ValueError("zvec sidecar ids do not match persisted chunks") + + collection = self._zvec.open(path=str(self.zvec_path)) + vector_schema = collection.schema.vectors[0] + if vector_schema.dimension != self._dim: + raise ValueError(f"zvec dim {vector_schema.dimension} != embedding dim {self._dim}") + persisted_m = vector_schema.index_param.m + if persisted_m != self.hnsw_m: + raise ValueError(f"zvec HNSW M mismatch: persisted={persisted_m}, configured={self.hnsw_m}") + self._verify_collection_contents(collection, indexed_ids) + + self._collection = collection + self._indexed_ids = indexed_ids + self.logger.info(f"Opened zvec collection: {len(indexed_ids)} vectors from {self.zvec_path}") + return True + except Exception as e: + self.logger.warning(f"Failed to open zvec collection, will rebuild: {e}") + collection = None # release the handle before wiping the directory + if self.zvec_path.exists(): + shutil.rmtree(self.zvec_path, ignore_errors=True) + self.zvec_sidecar_path.unlink(missing_ok=True) + return False + + def _verify_collection_contents(self, collection, expected_ids: set[str]) -> None: + """Check that the opened collection really holds the expected (id, vector) pairs. + + The sidecar only binds the collection to a chunk generation; it cannot + show that the collection itself lost, gained, or corrupted documents + (a crash between flushes, an external write, a partial copy). The + document count catches missing or extra ids cheaply, and the fetched + vectors catch same-id corruption. Any failure raises so the caller + rebuilds from the authoritative chunks. + """ + started_at = time.monotonic() + doc_count = collection.stats.doc_count + if doc_count != len(expected_ids): + raise ValueError(f"zvec collection holds {doc_count} documents, expected {len(expected_ids)}") + + ids = sorted(expected_ids) + for start in range(0, len(ids), _ZVEC_VERIFY_BATCH_SIZE): + batch = ids[start : start + _ZVEC_VERIFY_BATCH_SIZE] + docs = collection.fetch(batch, include_vector=True) + for cid in batch: + doc = docs.get(cid) + if doc is None: + raise ValueError(f"zvec collection is missing chunk {cid}") + indexed = self._vector_fingerprint(doc.vectors.get("embedding")) + if indexed != self._vector_fingerprint(self.file_chunks[cid].embedding): + raise ValueError(f"zvec vector for chunk {cid} does not match the persisted chunk") + self.logger.info( + f"{self.name}: verified zvec collection contents: docs={doc_count}, " + f"elapsed={time.monotonic() - started_at:.3f}s", + ) + + async def dump(self) -> None: + """Persist chunks JSONL via the parent, then flush zvec and write the sidecar.""" + await super().dump() + if self._collection is None or self.embedding_store is None: + return + try: + self._collection.flush() + await self._write_sidecar() + self.logger.info(f"Saved zvec collection: {len(self._indexed_ids)} vectors to {self.zvec_path}") + except Exception as e: + self.logger.exception(f"Failed to persist zvec collection: {e}") + + async def _write_sidecar(self) -> None: + """Atomically write the digest sidecar binding the collection to the chunk generation.""" + tmp = self.zvec_sidecar_path.with_name(f".{self.zvec_sidecar_path.name}.{uuid4().hex}.tmp") + payload = json.dumps( + { + "ids": sorted(self._indexed_ids), + "digest": self._chunks_embedding_digest(), + }, + ) + try: + async with aiofiles.open(tmp, "w", encoding=self.encoding) as f: + await f.write(payload) + tmp.replace(self.zvec_sidecar_path) + finally: + tmp.unlink(missing_ok=True) + + # -- CRUD overrides --------------------------------------------------- + + async def upsert(self, files: list[tuple[FileNode, list[FileChunk]]]) -> None: + if not files: + return + assert self.file_graph is not None + + # Snapshot pre-upsert chunk ids so we can drop chunks the new revision removed. + old_nodes = await self.file_graph.get_nodes([node.path for node, _ in files]) + old_ids_by_path = {n.path: set(n.chunk_ids) for n in old_nodes} + await super().upsert(files) + + if self._collection is None or self.embedding_store is None: + return + self._sync_collection_after_upsert(files, old_ids_by_path) + + def _sync_collection_after_upsert( + self, + files: list[tuple[FileNode, list[FileChunk]]], + old_ids_by_path: dict[str, set[str]], + ) -> None: + """Apply add / delete deltas to the zvec collection natively. + + Every eligible chunk of the request is re-upserted instead of being + diffed against its previous text: the parent accepts a caller-provided + embedding, so unchanged text does not imply an unchanged vector. zvec + upsert is idempotent, so re-sending an identical vector is cheap and + keeps the collection consistent with ``self.file_chunks``. Chunks that + no longer carry a usable vector are removed so the indexed id set stays + exactly the set of embeddable chunks. + """ + to_delete: list[str] = [] + to_upsert: list[FileChunk] = [] + for node, _ in files: + new_ids = set(node.chunk_ids) + to_delete.extend(old_ids_by_path.get(node.path, set()) - new_ids) + for cid in new_ids: + chunk = self.file_chunks.get(cid) + if chunk is None or not self._embedding_dim_matches(chunk.embedding): + to_delete.append(cid) # _delete_docs ignores ids that were never indexed + continue + to_upsert.append(chunk) + self._delete_docs(to_delete) + self._upsert_docs(to_upsert) + + async def delete(self, path: str | list[str]) -> None: + assert self.file_graph is not None + paths = [path] if isinstance(path, str) else path + nodes = await self.file_graph.get_nodes(paths) + deleted_ids = [cid for n in nodes for cid in n.chunk_ids] + await self._delete_nodes(nodes) # reuse resolved nodes; avoids a second get_nodes + self._delete_docs(deleted_ids) + + async def clear(self) -> None: + await super().clear() + if self._collection is not None: + try: + self._collection.destroy() # drops the collection directory + except Exception: # pragma: no cover - defensive + shutil.rmtree(self.zvec_path, ignore_errors=True) + elif self.zvec_path.exists(): + shutil.rmtree(self.zvec_path, ignore_errors=True) + self._indexed_ids = set() + self.zvec_sidecar_path.unlink(missing_ok=True) + self._collection = self._create_collection() if self.embedding_store is not None else None + + # -- search ----------------------------------------------------------- + + async def vector_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]: + if self.embedding_store is None or not query or limit <= 0 or self._collection is None or not self._indexed_ids: + return [] + + query_embedding = None + try: + query_embedding = await self.embedding_store.get_embedding(query) + except Exception as e: + self._disable_embedding(f"search: {type(e).__name__}: {e}") + if query_embedding is None or not self._embedding_dim_matches(query_embedding): + if query_embedding is not None: + self._disable_embedding( + f"search: query embedding dimension {len(query_embedding)} != {self.embedding_store.dimensions}", + ) + return [] + + # get_embedding above yielded control; a concurrent clear() may have + # swapped or dropped the collection. Re-read before dereferencing. + collection = self._collection + if collection is None or not self._indexed_ids: + return [] + + vector = np.asarray(query_embedding, dtype=np.float32).tolist() + ntotal = len(self._indexed_ids) + + if not search_filter: + hits = self._query_collection(collection, vector, min(limit, ntotal)) + return self._collect_hits(hits, limit, search_filter) + + # With a post-filter: progressively increase k until we collect enough + # results or exhaust the reachable index. + k = min(ntotal, 3 * limit) + while True: + hits = self._query_collection(collection, vector, k) + results = self._collect_hits(hits, limit, search_filter) + if len(results) >= limit or k >= ntotal: + return results + k = min(ntotal, k * 2) + + def _query_collection(self, collection, vector: list[float], topk: int) -> list[tuple[str, float]]: + """Run an ANN query and convert cosine distance to similarity (higher = closer).""" + docs = collection.query( + self._zvec.Query(field_name="embedding", vector=vector), + topk=max(1, topk), + ) + return [(doc.id, 1.0 - float(doc.score)) for doc in docs] + + def _collect_hits( + self, + hits: list[tuple[str, float]], + limit: int, + search_filter: dict | None = None, + ) -> list[FileChunk]: + """Map zvec hits back to chunks, skipping stale ids and filtered-out chunks.""" + results: list[FileChunk] = [] + for chunk_id, score in hits: + chunk = self.file_chunks.get(chunk_id) + if chunk is None or not self._matches_search_filter(chunk, search_filter): + continue + results.append(chunk.model_copy(update={"scores": {"vector": score, "score": score}})) + if len(results) >= limit: + break + return results diff --git a/tests/unit/test_file_store_consistency.py b/tests/unit/test_file_store_consistency.py index e8ab7cdb..cd61cbdd 100644 --- a/tests/unit/test_file_store_consistency.py +++ b/tests/unit/test_file_store_consistency.py @@ -14,7 +14,7 @@ import time import numpy as np import pytest -from reme.components.file_store import FaissLocalFileStore, LocalFileStore +from reme.components.file_store import FaissLocalFileStore, LocalFileStore, ZvecLocalFileStore from reme.components.file_store import local_file_store as local_file_store_module from reme.schema import FileChunk, FileNode from reme.utils.jsonl_zst import read_jsonl_zst, write_jsonl_zst @@ -148,6 +148,32 @@ def chunk(chunk_id: str, path: str, text: str, **metadata) -> FileChunk: return FileChunk(id=chunk_id, path=path, text=text, start_line=1, end_line=1, metadata=metadata) +def _new_local_store(name, **kwargs): + """Construct a LocalFileStore with embedding disabled at bind time.""" + return LocalFileStore(name=name, embedding_store="", **kwargs) + + +def _new_zvec_store(name, **kwargs): + """Construct a zvec store with embedding disabled at bind time.""" + try: + store = ZvecLocalFileStore(name=name, embedding_store="", **kwargs) + except ImportError: + pytest.skip("zvec is not installed") + return store + + +def _ensure_zvec_collection(store): + """Materialize the zvec collection once an embedding backend is attached. + + Fresh zvec stores start with no collection because ``embedding_store=""``. + Tests that attach a fake provider after ``start()`` must explicitly create + the collection before the first upsert, otherwise vectors are accepted by + the parent but never synced into zvec. + """ + if isinstance(store, ZvecLocalFileStore) and store._collection is None and store.embedding_store is not None: + store._collection = store._create_collection() + + async def set_chunks_with_graph(store: LocalFileStore, chunks: dict[str, FileChunk]) -> None: """Seed a graph/chunk snapshot that satisfies the persistence invariant.""" store.file_chunks = chunks @@ -454,14 +480,16 @@ def test_chunk_persistence_loads_legacy_json_embedding_list(): run(go()) -def test_same_chunk_id_with_changed_text_gets_new_embedding(): +@pytest.mark.parametrize("store_factory", [_new_local_store, _new_zvec_store]) +def test_same_chunk_id_with_changed_text_gets_new_embedding(store_factory): """Changing a chunk text refreshes its embedding.""" async def go(): with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): - store = LocalFileStore(name="t_embedding_reuse", embedding_store="") + store = store_factory(name="t_embedding_reuse") await store.start() store.embedding_store = FakeEmbeddingStore() + _ensure_zvec_collection(store) await store.upsert([(node("note.md"), [chunk("same", "note.md", "alpha text")])]) assert store.file_chunks["same"].embedding.tolist() == [1.0, 0.0] @@ -474,12 +502,13 @@ def test_same_chunk_id_with_changed_text_gets_new_embedding(): run(go()) -def test_load_backfills_missing_embeddings_from_persisted_chunks(): +@pytest.mark.parametrize("store_factory", [_new_local_store, _new_zvec_store]) +def test_load_backfills_missing_embeddings_from_persisted_chunks(store_factory): """Startup backfills old chunks in the background and persists vectors.""" async def go(): with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): - store = LocalFileStore(name="t_embedding_backfill", embedding_store="") + store = store_factory(name="t_embedding_backfill") await store.start() await store.upsert( [ @@ -489,7 +518,7 @@ def test_load_backfills_missing_embeddings_from_persisted_chunks(): ) await store.close() - store = LocalFileStore(name="t_embedding_backfill", embedding_store="") + store = store_factory(name="t_embedding_backfill") store.embedding_store = FakeEmbeddingStore() await store.start() await store._embedding_backfill_task @@ -498,7 +527,7 @@ def test_load_backfills_missing_embeddings_from_persisted_chunks(): assert store.file_chunks["b"].embedding.tolist() == [0.0, 1.0] await store.close() - store = LocalFileStore(name="t_embedding_backfill", embedding_store="") + store = store_factory(name="t_embedding_backfill") await store.start() assert store.file_chunks["a"].embedding.tolist() == [1.0, 0.0] assert store.file_chunks["b"].embedding.tolist() == [0.0, 1.0] @@ -588,12 +617,13 @@ def test_load_skips_backfill_when_embedding_health_check_fails(): run(go()) -def test_load_reembeds_persisted_chunks_with_stale_embedding_dimensions(): +@pytest.mark.parametrize("store_factory", [_new_local_store, _new_zvec_store]) +def test_load_reembeds_persisted_chunks_with_stale_embedding_dimensions(store_factory): """Loading persisted chunks re-embeds vectors that do not match current dimensions.""" async def go(): with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): - store = LocalFileStore(name="t_embedding_stale_dim", embedding_store="") + store = store_factory(name="t_embedding_stale_dim") await store.start() stale = chunk("a", "a.md", "alpha text") stale.embedding = np.array([1.0], dtype=np.float16) @@ -601,7 +631,7 @@ def test_load_reembeds_persisted_chunks_with_stale_embedding_dimensions(): await store.dump() await store.close() - store = LocalFileStore(name="t_embedding_stale_dim", embedding_store="") + store = store_factory(name="t_embedding_stale_dim") fake = CountingFakeEmbeddingStore() store.embedding_store = fake await store.start() @@ -689,14 +719,16 @@ def test_upsert_reembeds_prefilled_chunk_with_stale_dimension(): run(go()) -def test_search_filter_applies_to_vector_and_keyword_results(): +@pytest.mark.parametrize("store_factory", [_new_local_store, _new_zvec_store]) +def test_search_filter_applies_to_vector_and_keyword_results(store_factory): """Search filters apply consistently to vector and keyword results.""" async def go(): with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): - store = LocalFileStore(name="t_filter", embedding_store="") + store = store_factory(name="t_filter") await store.start() store.embedding_store = FakeEmbeddingStore() + _ensure_zvec_collection(store) await store.upsert( [ @@ -860,14 +892,16 @@ def test_date_filter_extract_and_match(): assert LocalFileStore._matches_search_filter(c, {"start_date": "2026-01-01", "end_date": "2026-12-31"}) is True -def test_date_filter_with_vector_and_keyword_search(): +@pytest.mark.parametrize("store_factory", [_new_local_store, _new_zvec_store]) +def test_date_filter_with_vector_and_keyword_search(store_factory): """vector_search and keyword_search respect start_date/end_date filters.""" async def go(): with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): - store = LocalFileStore(name="t_date_search", embedding_store="") + store = store_factory(name="t_date_search") await store.start() store.embedding_store = FakeEmbeddingStore() + _ensure_zvec_collection(store) await store.upsert( [ diff --git a/tests/unit/test_zvec_file_store.py b/tests/unit/test_zvec_file_store.py new file mode 100644 index 00000000..ba1566f4 --- /dev/null +++ b/tests/unit/test_zvec_file_store.py @@ -0,0 +1,614 @@ +"""Tests for ZvecLocalFileStore: native upsert/delete sync, persistence, and rebuild triggers.""" + +# pylint: disable=protected-access + +import asyncio +import json +import os +import tempfile +import tomllib +import warnings +from pathlib import Path + +import numpy as np +import pytest + +from reme.components.file_store import ZvecLocalFileStore +from reme.schema import FileChunk, FileNode + + +class temp_chdir: + """Temporarily chdir into a test workspace.""" + + 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) + + +class FakeEmbeddingStore: + """Small deterministic embedding provider used by file-store tests.""" + + dimensions = 2 + max_batch_size = 10 + + def _embed(self, text: str) -> np.ndarray: + if "beta" in text or "fresh" in text: + return np.array([0.0, 1.0], dtype=np.float16) + return np.array([1.0, 0.0], dtype=np.float16) + + async def health_check(self, _timeout: float = 2.0) -> bool: + """Report the fake embedding service as healthy.""" + return True + + async def get_embedding(self, input_text: str, **_kwargs) -> np.ndarray: + """Return a deterministic embedding for a single text.""" + return self._embed(input_text) + + async def get_node_embeddings(self, nodes: list[FileChunk], **_kwargs) -> list[FileChunk]: + """Attach deterministic embeddings to file chunks.""" + for chunk_node in nodes: + chunk_node.embedding = self._embed(chunk_node.text) + return nodes + + +def run(coro): + """Run an async test body.""" + return asyncio.run(coro) + + +def node(path: str) -> FileNode: + """Build a minimal file node.""" + return FileNode(path=path, st_mtime=1.0) + + +def chunk(chunk_id: str, path: str, text: str, **metadata) -> FileChunk: + """Build a minimal file chunk.""" + return FileChunk(id=chunk_id, path=path, text=text, start_line=1, end_line=1, metadata=metadata) + + +def embedded_chunk(chunk_id: str, path: str, text: str, embedding: list[float]) -> FileChunk: + """Build a chunk carrying an explicit, caller-provided embedding.""" + file_chunk = chunk(chunk_id, path, text) + file_chunk.embedding = np.asarray(embedding, dtype=np.float16) + return file_chunk + + +def _track_rebuilds(store: ZvecLocalFileStore) -> list[bool]: + """Record calls to the collection rebuild path.""" + rebuilds: list[bool] = [] + original_rebuild = store._rebuild_collection + store._rebuild_collection = lambda: rebuilds.append(True) or original_rebuild() + return rebuilds + + +def _tamper_with_collection(store: ZvecLocalFileStore, mutate) -> None: + """Mutate the persisted collection behind the store's back, then release the handle. + + zvec holds an in-process lock on the collection directory, so the store must + already be closed and the local handle must be dropped before another store + reopens the same path. + """ + zvec = pytest.importorskip("zvec") + collection = zvec.open(path=str(store.zvec_path)) + mutate(zvec, collection) + collection.flush() + del collection + + +def _new_zvec_store(name, **kwargs): + """Construct a zvec store with embedding disabled at bind time.""" + try: + store = ZvecLocalFileStore(name=name, embedding_store="", **kwargs) + except ImportError: + pytest.skip("zvec is not installed") + return store + + +async def _started_store(name, **kwargs) -> ZvecLocalFileStore: + """Start a store with a fake embedding provider and a live collection.""" + store = _new_zvec_store(name, **kwargs) + store.embedding_store = FakeEmbeddingStore() + await store.start() + if store._collection is None: + store._collection = store._create_collection() + return store + + +async def _seed_unembedded_chunk(store: ZvecLocalFileStore, chunk_id: str, path: str, text: str) -> None: + """Attach one chunk without a vector, keeping the graph invariant intact.""" + store.file_chunks[chunk_id] = chunk(chunk_id, path, text) + file_node = node(path) + file_node.chunk_ids = [chunk_id] + await store.file_graph.upsert_nodes([file_node]) + + +# -- CRUD / search ------------------------------------------------------------ + + +def test_zvec_upsert_and_vector_search_ranks_by_similarity(): + """Upserted chunks are searchable; results are ranked by cosine similarity.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _started_store("t_zvec_basic") + + await store.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + await store.upsert([(node("b.md"), [chunk("b", "b.md", "beta text")])]) + + results = await store.vector_search("alpha", 10, {}) + assert [c.id for c in results] == ["a", "b"] + # Identical vector -> cosine distance 0 -> similarity score 1. + assert results[0].scores["vector"] == pytest.approx(1.0, abs=1e-5) + assert results[0].scores["score"] == results[0].scores["vector"] + await store.close() + + run(go()) + + +def test_zvec_same_id_text_change_updates_vector_in_place(): + """A same-id text change replaces the vector via native zvec upsert.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _started_store("t_zvec_update") + + await store.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + assert [c.id for c in await store.vector_search("alpha", 5, {})] == ["a"] + + await store.upsert([(node("a.md"), [chunk("a", "a.md", "beta text")])]) + + assert store._indexed_ids == {"a"} + results = await store.vector_search("beta", 5, {}) + assert [c.id for c in results] == ["a"] + assert results[0].scores["vector"] == pytest.approx(1.0, abs=1e-5) + await store.close() + + run(go()) + + +def test_zvec_explicit_embedding_change_updates_collection(): + """A same-id, same-text update with a new explicit embedding replaces the vector.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _started_store("t_zvec_explicit_embedding") + + await store.upsert([(node("a.md"), [embedded_chunk("same", "a.md", "same text", [1.0, 0.0])])]) + await store.upsert([(node("a.md"), [embedded_chunk("same", "a.md", "same text", [0.0, 1.0])])]) + + # "beta" embeds to [0, 1]; the collection must hold the new vector. + results = await store.vector_search("beta", 5, {}) + assert [c.id for c in results] == ["same"] + assert results[0].scores["vector"] == pytest.approx(1.0, abs=1e-5) + await store.close() + + # The stale vector must not survive a restart either. + reopened = _new_zvec_store("t_zvec_explicit_embedding") + reopened.embedding_store = FakeEmbeddingStore() + rebuilds = _track_rebuilds(reopened) + await reopened.start() + + assert not rebuilds # the collection was already in sync, no repair needed + results = await reopened.vector_search("beta", 5, {}) + assert [c.id for c in results] == ["same"] + assert results[0].scores["vector"] == pytest.approx(1.0, abs=1e-5) + await reopened.close() + + run(go()) + + +def test_zvec_upsert_removes_stale_chunks_from_collection(): + """Re-upserting a path deletes vectors of chunks that no longer exist.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _started_store("t_zvec_stale") + + await store.upsert([(node("a.md"), [chunk("old", "a.md", "alpha old")])]) + await store.upsert([(node("a.md"), [chunk("new", "a.md", "alpha new")])]) + + assert store._indexed_ids == {"new"} + assert [c.id for c in await store.vector_search("alpha", 10, {})] == ["new"] + await store.close() + + run(go()) + + +def test_zvec_delete_removes_vectors(): + """Deleting a path removes its chunk vectors from the collection.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _started_store("t_zvec_delete") + + await store.upsert( + [ + (node("a.md"), [chunk("a", "a.md", "alpha text")]), + (node("b.md"), [chunk("b", "b.md", "beta text")]), + ], + ) + await store.delete("a.md") + + assert store._indexed_ids == {"b"} + assert [c.id for c in await store.vector_search("alpha", 10, {})] == ["b"] + await store.close() + + run(go()) + + +def test_zvec_vector_search_applies_post_filter(): + """The shared search_filter semantics apply on top of ANN results.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _started_store("t_zvec_filter") + + await store.upsert( + [ + (node("a.md"), [chunk("a", "a.md", "alpha one", kind="x")]), + (node("b.md"), [chunk("b", "b.md", "alpha two", kind="y")]), + ], + ) + + assert [c.id for c in await store.vector_search("alpha", 10, {"path": "b.md"})] == ["b"] + assert [c.id for c in await store.vector_search("alpha", 10, {"kind": "x"})] == ["a"] + assert await store.vector_search("alpha", 10, {"path": "c.md"}) == [] + await store.close() + + run(go()) + + +def test_zvec_vector_search_uses_current_query_api(): + """The ANN query path must not rely on zvec's deprecated VectorQuery alias.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _started_store("t_zvec_query_api") + await store.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + + with warnings.catch_warnings(): + warnings.simplefilter("error", DeprecationWarning) + assert [c.id for c in await store.vector_search("alpha", 5, {})] == ["a"] + await store.close() + + run(go()) + + +def test_zvec_clear_resets_collection(): + """clear() drops the collection, sidecar, and indexed-id state.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _started_store("t_zvec_clear") + + await store.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + await store.dump() + assert store.zvec_sidecar_path.exists() + + await store.clear() + + assert store._indexed_ids == set() + assert not store.zvec_sidecar_path.exists() + assert await store.vector_search("alpha", 10, {}) == [] + # The store stays usable after clear. + await store.upsert([(node("b.md"), [chunk("b", "b.md", "beta text")])]) + assert [c.id for c in await store.vector_search("beta", 10, {})] == ["b"] + await store.close() + + run(go()) + + +def test_zvec_keyword_only_mode_keeps_keyword_search_working(): + """Without an embedding store, vector search is empty but keyword search works.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_zvec_store("t_zvec_keyword_only") + await store.start() + + await store.upsert([(node("a.md"), [chunk("a", "a.md", "uniquezvecword only")])]) + + assert store._collection is None + assert await store.vector_search("uniquezvecword", 5, {}) == [] + assert [c.id for c in await store.keyword_search("uniquezvecword", 5, {})] == ["a"] + await store.close() + + run(go()) + + +# -- embedding backfill --------------------------------------------------------- + + +def test_zvec_backfill_adds_incrementally(): + """Backfilled vectors are upserted into the live collection without a rebuild.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _started_store("t_zvec_backfill") + + await store.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + await _seed_unembedded_chunk(store, "c", "c.md", "alpha extra") + + rebuilds = [] + original_rebuild = store._rebuild_collection + store._rebuild_collection = lambda: rebuilds.append(True) or original_rebuild() + + await store._backfill_missing_embeddings() + + assert not rebuilds + assert store._indexed_ids == {"a", "c"} + assert {c.id for c in await store.vector_search("alpha", 10, {})} == {"a", "c"} + await store.close() + + run(go()) + + +# -- persistence ---------------------------------------------------------------- + + +def test_zvec_persistence_round_trip_reopens_without_rebuild(): + """dump() + a fresh store reattaches the persisted collection via the sidecar.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + seed = await _started_store("t_zvec_persist") + await seed.upsert( + [ + (node("a.md"), [chunk("a", "a.md", "alpha text")]), + (node("b.md"), [chunk("b", "b.md", "beta text")]), + ], + ) + await seed.close() + + store = _new_zvec_store("t_zvec_persist") + store.embedding_store = FakeEmbeddingStore() + + rebuilds = [] + original_rebuild = store._rebuild_collection + store._rebuild_collection = lambda: rebuilds.append(True) or original_rebuild() + + await store.start() + + assert not rebuilds + assert store._indexed_ids == {"a", "b"} + assert [c.id for c in await store.vector_search("alpha", 5, {})][0] == "a" + assert [c.id for c in await store.vector_search("beta", 5, {})][0] == "b" + await store.close() + + run(go()) + + +def test_zvec_digest_mismatch_triggers_rebuild(): + """A sidecar whose digest fell behind the chunk JSONL forces a clean rebuild.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + seed = await _started_store("t_zvec_digest") + await seed.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + await seed.close() + + sidecar = json.loads(seed.zvec_sidecar_path.read_text()) + sidecar["digest"] = "0" * 64 + seed.zvec_sidecar_path.write_text(json.dumps(sidecar)) + + store = _new_zvec_store("t_zvec_digest") + store.embedding_store = FakeEmbeddingStore() + + rebuilds = [] + original_rebuild = store._rebuild_collection + store._rebuild_collection = lambda: rebuilds.append(True) or original_rebuild() + + await store.start() + + assert rebuilds == [True] + assert store._indexed_ids == {"a"} + assert [c.id for c in await store.vector_search("alpha", 5, {})] == ["a"] + await store.close() + + run(go()) + + +def test_zvec_missing_sidecar_triggers_rebuild(): + """A collection directory without its sidecar cannot be trusted and is rebuilt.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + seed = await _started_store("t_zvec_no_sidecar") + await seed.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + await seed.close() + + seed.zvec_sidecar_path.unlink() + + store = _new_zvec_store("t_zvec_no_sidecar") + store.embedding_store = FakeEmbeddingStore() + await store.start() + + assert store._indexed_ids == {"a"} + assert [c.id for c in await store.vector_search("alpha", 5, {})] == ["a"] + await store.close() + + run(go()) + + +def test_zvec_hnsw_m_mismatch_triggers_rebuild(): + """A persisted collection built with a different HNSW M is rebuilt.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + seed = await _started_store("t_zvec_m_mismatch", hnsw_m=16) + await seed.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + await seed.close() + + store = _new_zvec_store("t_zvec_m_mismatch", hnsw_m=48) + store.embedding_store = FakeEmbeddingStore() + + rebuilds = [] + original_rebuild = store._rebuild_collection + store._rebuild_collection = lambda: rebuilds.append(True) or original_rebuild() + + await store.start() + + assert rebuilds == [True] + assert [c.id for c in await store.vector_search("alpha", 5, {})] == ["a"] + await store.close() + + run(go()) + + +def test_zvec_missing_collection_document_triggers_rebuild(): + """A collection that lost a document is rebuilt even though the sidecar matches.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + seed = await _started_store("t_zvec_missing_doc") + await seed.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + await seed.close() + + _tamper_with_collection(seed, lambda _zvec, collection: collection.delete(ids=["a"])) + # The sidecar still claims the chunk is indexed. + assert json.loads(seed.zvec_sidecar_path.read_text())["ids"] == ["a"] + + store = _new_zvec_store("t_zvec_missing_doc") + store.embedding_store = FakeEmbeddingStore() + rebuilds = _track_rebuilds(store) + await store.start() + + assert rebuilds == [True] + assert store._indexed_ids == {"a"} + assert store._collection.stats.doc_count == 1 + assert [c.id for c in await store.vector_search("alpha", 5, {})] == ["a"] + await store.close() + + run(go()) + + +def test_zvec_unexpected_collection_document_triggers_rebuild(): + """A collection holding a document no chunk owns is rebuilt.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + seed = await _started_store("t_zvec_extra_doc") + await seed.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + await seed.close() + + _tamper_with_collection( + seed, + lambda zvec, collection: collection.upsert( + [zvec.Doc(id="ghost", vectors={"embedding": [0.0, 1.0]})], + ), + ) + + store = _new_zvec_store("t_zvec_extra_doc") + store.embedding_store = FakeEmbeddingStore() + rebuilds = _track_rebuilds(store) + await store.start() + + assert rebuilds == [True] + assert store._indexed_ids == {"a"} + assert store._collection.stats.doc_count == 1 + assert [c.id for c in await store.vector_search("beta", 5, {})] == ["a"] + await store.close() + + run(go()) + + +def test_zvec_corrupted_collection_vector_triggers_rebuild(): + """An expected id whose stored vector was replaced is rebuilt from the chunks.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + seed = await _started_store("t_zvec_bad_vector") + await seed.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + await seed.close() + + _tamper_with_collection( + seed, + lambda zvec, collection: collection.upsert( + [zvec.Doc(id="a", vectors={"embedding": [0.0, 1.0]})], + ), + ) + + store = _new_zvec_store("t_zvec_bad_vector") + store.embedding_store = FakeEmbeddingStore() + rebuilds = _track_rebuilds(store) + await store.start() + + assert rebuilds == [True] + # The rebuilt collection carries the vector of the authoritative chunk. + results = await store.vector_search("alpha", 5, {}) + assert [c.id for c in results] == ["a"] + assert results[0].scores["vector"] == pytest.approx(1.0, abs=1e-5) + await store.close() + + run(go()) + + +def test_zvec_stale_embedding_dimension_rebuilds_from_backfill(): + """Persisted vectors with a stale dimension are dropped and re-embedded.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + seed = await _started_store("t_zvec_stale_dim") + await seed.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + await seed.close() + + class WideEmbeddingStore(FakeEmbeddingStore): + """Same provider with a different dimension.""" + + dimensions = 4 + + def _embed(self, text: str) -> np.ndarray: + base = super()._embed(text) + return np.concatenate([base, base]).astype(np.float16) + + store = _new_zvec_store("t_zvec_stale_dim") + store.embedding_store = WideEmbeddingStore() + await store.start() + # Startup backfill re-embeds the stale chunk in the background. + await store._embedding_backfill_task + + assert store._collection.schema.vectors[0].dimension == 4 + assert store._indexed_ids == {"a"} + assert [c.id for c in await store.vector_search("alpha", 5, {})] == ["a"] + await store.close() + + run(go()) + + +# -- packaging ------------------------------------------------------------ + + +def test_zvec_declared_as_installable_dependency(): + """The registered backend needs a supported installation path in pyproject.""" + pyproject = Path(__file__).resolve().parents[2] / "pyproject.toml" + if not pyproject.exists(): # running against an installed package, not the repo + pytest.skip("pyproject.toml is not available") + optional = tomllib.loads(pyproject.read_text(encoding="utf-8"))["project"]["optional-dependencies"] + assert any(dep.replace(" ", "").startswith("zvec") for dep in optional["core"]) + + +# -- maintenance ------------------------------------------------------------ + + +def test_zvec_optimize_index_runs_native_optimize(): + """optimize_index() delegates to Collection.optimize() and keeps search intact.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _started_store("t_zvec_optimize") + + await store.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + await store.optimize_index() + + assert [c.id for c in await store.vector_search("alpha", 5, {})] == ["a"] + await store.close() + + run(go()) From b00eb0a9ea292924d57915fc42477f0697525b54 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Wed, 5 Aug 2026 23:07:05 +0800 Subject: [PATCH 2/3] feat(backend): improve workspace support for web clients (#417) --- .../agent_wrapper/as_agent_wrapper.py | 86 ++++++++++++++++++- reme/config/default.yaml | 12 +++ reme/schema/application_config.py | 15 +++- reme/steps/common/chat.py | 29 ++++++- reme/steps/file_io/list.py | 38 ++++++-- tests/unit/test_base_agent_wrapper.py | 36 ++++++++ tests/unit/test_crud_steps.py | 23 +++++ tests/unit/test_workspace_web_steps.py | 42 +++++++-- 8 files changed, 259 insertions(+), 22 deletions(-) diff --git a/reme/components/agent_wrapper/as_agent_wrapper.py b/reme/components/agent_wrapper/as_agent_wrapper.py index 3b539bf5..7360543c 100644 --- a/reme/components/agent_wrapper/as_agent_wrapper.py +++ b/reme/components/agent_wrapper/as_agent_wrapper.py @@ -123,6 +123,86 @@ class WorkspaceBackend(LocalBackend): return ExecResult(exit_code=process.returncode or 0, stdout=stdout, stderr=stderr) +class WorkspaceRead(Read): + """AgentScope Read accepting paths relative to the configured agent workspace.""" + + description = Read.description.replace( + "must be an absolute path, not a relative path", + "may be absolute or relative to the workspace", + ) + input_schema = { + **Read.input_schema, + "properties": { + **Read.input_schema["properties"], + "file_path": { + "type": "string", + "description": "Absolute path or path relative to the workspace.", + }, + }, + } + + async def call( + self, + file_path: str, + offset: int = 1, + limit: int = 2000, + _agent_state: AgentState | None = None, + ) -> ToolChunk: + """Resolve a workspace-relative path before delegating to AgentScope.""" + cwd = await self._backend.getcwd() # pylint: disable=protected-access + target = self._backend.abspath(file_path, cwd=cwd) # pylint: disable=protected-access + return await super().call(target, offset, limit, _agent_state) + + +class WorkspaceWrite(Write): + """AgentScope Write accepting paths relative to the configured agent workspace.""" + + input_schema = { + **Write.input_schema, + "properties": { + **Write.input_schema["properties"], + "file_path": { + "type": "string", + "description": "Absolute path or path relative to the workspace.", + }, + }, + } + + async def call(self, file_path: str, content: str, _agent_state: AgentState | None = None) -> ToolChunk: + """Resolve a workspace-relative path before delegating to AgentScope.""" + cwd = await self._backend.getcwd() # pylint: disable=protected-access + target = self._backend.abspath(file_path, cwd=cwd) # pylint: disable=protected-access + return await super().call(target, content, _agent_state) + + +class WorkspaceEdit(Edit): + """AgentScope Edit accepting paths relative to the configured agent workspace.""" + + input_schema = { + **Edit.input_schema, + "properties": { + **Edit.input_schema["properties"], + "file_path": { + "type": "string", + "description": "Absolute path or path relative to the workspace.", + }, + }, + } + + async def call( + self, + file_path: str, + old_string: str, + new_string: str, + replace_all: bool = False, + _agent_state: AgentState | None = None, + ) -> ToolChunk: + """Resolve a workspace-relative path before delegating to AgentScope.""" + cwd = await self._backend.getcwd() # pylint: disable=protected-access + target = self._backend.abspath(file_path, cwd=cwd) # pylint: disable=protected-access + return await super().call(target, old_string, new_string, replace_all, _agent_state) + + class BypassAnalysisBash(Bash): """Bash variant that delegates permission decisions to PermissionEngine. @@ -192,11 +272,11 @@ class AsAgentWrapper(BaseAgentWrapper): backend = WorkspaceBackend(cwd, self.bash_environment) factories = { "bash": lambda: BypassAnalysisBash(cwd=cwd, backend=backend), - "edit": lambda: Edit(backend=backend), + "edit": lambda: WorkspaceEdit(backend=backend), "glob": lambda: Glob(backend=backend), "grep": lambda: Grep(backend=backend), - "read": lambda: Read(backend=backend), - "write": lambda: Write(backend=backend), + "read": lambda: WorkspaceRead(backend=backend), + "write": lambda: WorkspaceWrite(backend=backend), } if names is False: selected_names = [] diff --git a/reme/config/default.yaml b/reme/config/default.yaml index 30b29f1a..36036b65 100644 --- a/reme/config/default.yaml +++ b/reme/config/default.yaml @@ -1,3 +1,5 @@ +workspace_dir: ~/.copaw/workspaces/default/ + service: backend: http @@ -503,6 +505,16 @@ jobs: type: boolean description: "recurse" default: false + sort_by: + type: string + enum: + - mtime + description: "optional ordering; mtime returns most recently modified files first" + extensions: + type: array + items: + type: string + description: "optional extension allowlist applied before sorting and limiting" limit: type: integer description: "max results" diff --git a/reme/schema/application_config.py b/reme/schema/application_config.py index 00a0b286..ff1b8697 100644 --- a/reme/schema/application_config.py +++ b/reme/schema/application_config.py @@ -1,8 +1,9 @@ """Application configuration models.""" import os +from pathlib import Path -from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel, ConfigDict, Field, field_validator from ..enumeration import ComponentEnum @@ -32,7 +33,11 @@ class ApplicationConfig(BaseModel): default_factory=dict, description="Environment variables loaded once at startup and passed to agent subprocesses", ) - workspace_dir: str = Field(default=".reme", description="Workspace root directory for runtime files") + workspace_dir: str = Field( + default=".reme", + description="Workspace root directory for runtime files", + validate_default=True, + ) metadata_dir: str = Field(default="metadata", description="Subdirectory for ReMe persistent state") session_dir: str = Field(default="session", description="Subdirectory for persisted agent sessions") mem_session_dir: str = Field(default="mem_session", description="Subdirectory for persisted agent sessions") @@ -53,3 +58,9 @@ class ApplicationConfig(BaseModel): default_factory=dict, description="Component registry keyed by type then name", ) + + @field_validator("workspace_dir", mode="before") + @classmethod + def normalize_workspace_dir(cls, value) -> str: + """Expand home-relative paths once so every component sees the same absolute workspace.""" + return str(Path(value).expanduser().resolve(strict=False)) diff --git a/reme/steps/common/chat.py b/reme/steps/common/chat.py index 56f18de1..bf8002de 100644 --- a/reme/steps/common/chat.py +++ b/reme/steps/common/chat.py @@ -1,5 +1,8 @@ """Workspace-aware streaming chat step for the ReMe web interface.""" +import datetime +import zoneinfo + from ..base_step import BaseStep from ...components import R from ...enumeration import ChunkEnum @@ -13,7 +16,25 @@ class ChatStep(BaseStep): Use the available ReMe tools when workspace facts are needed. Cite workspace-relative file paths when referring to notes. Never invent file contents, and do not claim to have changed files because this chat intentionally provides read-only tools. Reply in the user's language.""" - READ_ONLY_TOOLS = ["search", "list", "read", "stat", "traverse"] + READ_ONLY_TOOLS = ["search", "list", "stat", "traverse"] + + def _system_prompt(self) -> str: + """Append request-time environment facts to the configured prompt.""" + timezone = self.app_context.app_config.timezone if self.app_context is not None else None + try: + current = datetime.datetime.now(zoneinfo.ZoneInfo(timezone)) if timezone else datetime.datetime.now() + except (zoneinfo.ZoneInfoNotFoundError, ValueError): + current = datetime.datetime.now() + + assert self.context is not None + base_prompt = str(self.context.get("system_prompt") or self.DEFAULT_SYSTEM_PROMPT).rstrip() + return ( + f"{base_prompt}\n\n" + "\n" + f"Current date: {current.date().isoformat()}\n" + f"Current working directory: {self.agent_wrapper.cwd.resolve(strict=False)}\n" + "" + ) async def execute(self): assert self.context is not None @@ -26,7 +47,7 @@ this chat intentionally provides read-only tools. Reply in the user's language." raise RuntimeError("chat_step requires an agent_wrapper") wrapper_kwargs = { - "system_prompt": self.context.get("system_prompt") or self.DEFAULT_SYSTEM_PROMPT, + "system_prompt": self._system_prompt(), "job_tools": self.READ_ONLY_TOOLS, } if session_id := self.context.get("session_id"): @@ -34,6 +55,10 @@ this chat intentionally provides read-only tools. Reply in the user's language." parts: list[str] = [] async for chunk in self.agent_wrapper.reply_stream(query, **wrapper_kwargs): + if chunk.chunk_type == ChunkEnum.REPLY_END: + answer = "".join(parts).strip() + if answer: + chunk.metadata["answer"] = answer await self.context.add_stream_chunk(chunk) if chunk.chunk_type == ChunkEnum.CONTENT and isinstance(chunk.chunk, str): parts.append(chunk.chunk) diff --git a/reme/steps/file_io/list.py b/reme/steps/file_io/list.py index ebacc967..cee573ba 100644 --- a/reme/steps/file_io/list.py +++ b/reme/steps/file_io/list.py @@ -8,6 +8,8 @@ Parameters: path — dir to list under (relative to the workspace or absolute). Empty = workspace root. limit — cap on the number of returned items (default 100, must be > 0). recursive — descend into subdirectories. Default False = direct children only. + sort_by — optional ordering; ``mtime`` returns most recently modified files first. + extensions — optional extension allowlist applied before sorting and limiting. No frontmatter is read. Callers needing frontmatter-based filtering should iterate the result and call ``frontmatter_read`` per candidate. @@ -36,7 +38,7 @@ class ListStep(BaseStep): if meta: self.context.response.metadata.update(meta) - def _collect_params(self) -> tuple[str, bool, int]: + def _collect_params(self) -> tuple[str, bool, int, str, frozenset[str]]: """Read ``path`` / ``recursive`` / ``limit`` from context; coerce permissively.""" assert self.context is not None path = str(self.context.get("path") or "") @@ -47,20 +49,37 @@ class ListStep(BaseStep): limit = int(raw_limit) if raw_limit is not None else DEFAULT_LIMIT except (TypeError, ValueError): limit = DEFAULT_LIMIT - return path, recursive, limit if limit > 0 else DEFAULT_LIMIT + sort_by = str(self.context.get("sort_by") or "") + raw_extensions = self.context.get("extensions") or [] + if isinstance(raw_extensions, str): + raw_extensions = raw_extensions.split(",") + extensions = frozenset( + normalized for value in raw_extensions if (normalized := str(value).strip().lower().lstrip(".")) + ) + return path, recursive, limit if limit > 0 else DEFAULT_LIMIT, sort_by, extensions @staticmethod - def _walk_files(target_dir: Path, recursive: bool, limit: int) -> list[Path]: - """Return up to ``limit`` regular files under ``target_dir``; short-circuits at the cap.""" + def _walk_files( + target_dir: Path, + recursive: bool, + limit: int, + sort_by: str = "", + extensions: frozenset[str] = frozenset(), + ) -> list[Path]: + """Return up to ``limit`` regular files, scanning all entries only when sorting.""" entries: Iterable[Path] = target_dir.rglob("*") if recursive else target_dir.iterdir() files: list[Path] = [] for entry in entries: if not entry.is_file(): # skip dirs, sockets, broken links, etc. continue + if extensions and entry.suffix.lower().lstrip(".") not in extensions: + continue files.append(entry) - if len(files) >= limit: + if sort_by != "mtime" and len(files) >= limit: break - return files + if sort_by == "mtime": + files.sort(key=lambda entry: (-entry.stat().st_mtime, entry.as_posix())) + return files[:limit] @staticmethod def _format_relative(files: list[Path], workspace_dir: Path) -> list[str]: @@ -75,7 +94,7 @@ class ListStep(BaseStep): async def execute(self): assert self.context is not None - path, recursive, limit = self._collect_params() + path, recursive, limit, sort_by, extensions = self._collect_params() workspace_dir = Path(self.file_store.workspace_path or ".").resolve() target_dir, err = resolve_path(workspace_dir, path, allow_empty=True) if err or target_dir is None: @@ -89,7 +108,10 @@ class ListStep(BaseStep): self._fail(f"path {target_dir} is not a directory", path=str(target_dir)) return None - items = self._format_relative(self._walk_files(target_dir, recursive, limit), workspace_dir) + items = self._format_relative( + self._walk_files(target_dir, recursive, limit, sort_by, extensions), + workspace_dir, + ) self.context.response.success = True location = path or "." diff --git a/tests/unit/test_base_agent_wrapper.py b/tests/unit/test_base_agent_wrapper.py index 17a1b1fa..53fda2bb 100644 --- a/tests/unit/test_base_agent_wrapper.py +++ b/tests/unit/test_base_agent_wrapper.py @@ -139,6 +139,42 @@ async def test_agentscope_backend_passes_configured_environment_to_bash(tmp_path assert result.stdout == b"configured\n" +def test_workspace_dir_expands_user_home(monkeypatch, tmp_path): + """Home-relative workspace config is normalized before components consume it.""" + monkeypatch.setenv("HOME", str(tmp_path)) + + context = ApplicationContext(workspace_dir="~/.copaw/workspaces/default") + + assert context.app_config.workspace_dir == str(tmp_path / ".copaw/workspaces/default") + + +@pytest.mark.asyncio +async def test_agentscope_file_tools_accept_workspace_relative_paths(tmp_path): + """Built-in Read/Write/Edit share the workspace-relative path convention used by ReMe jobs.""" + wrapper = AsAgentWrapper( + app_context=ApplicationContext(workspace_dir=str(tmp_path)), + as_llm="", + ) + read, write, edit = wrapper._builtin_tools(["read", "write", "edit"]) # pylint: disable=protected-access + source = tmp_path / ".reme/daily/note.md" + source.parent.mkdir(parents=True) + source.write_text("before", encoding="utf-8") + + read_result = await read.call(file_path=".reme/daily/note.md") + write_result = await write.call(file_path=".reme/daily/new.md", content="new") + edit_result = await edit.call( + file_path=".reme/daily/note.md", + old_string="before", + new_string="after", + ) + + assert "absolute path" not in str(read_result.content) + assert "absolute path" not in str(write_result.content) + assert "absolute path" not in str(edit_result.content) + assert (tmp_path / ".reme/daily/new.md").read_text(encoding="utf-8") == "new" + assert source.read_text(encoding="utf-8") == "after" + + @pytest.mark.asyncio async def test_agentscope_bash_uses_managed_proxy_without_changing_subprocess_environment(tmp_path): """AgentScope applies the managed proxy only to its command backend.""" diff --git a/tests/unit/test_crud_steps.py b/tests/unit/test_crud_steps.py index 1132a686..0eb4d98e 100644 --- a/tests/unit/test_crud_steps.py +++ b/tests/unit/test_crud_steps.py @@ -225,6 +225,29 @@ def test_list_respects_limit_and_non_recursive(): asyncio.run(run()) +def test_list_can_sort_by_most_recent_modification(): + """Filtering precedes the limit so unrelated generated files cannot hide recent notes.""" + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + older = _seed_md(root, "older.md", "old") + newer = _seed_md(root, "newer.md", "new") + generated = root / "newest.json" + generated.write_text("{}", encoding="utf-8") + os.utime(older, (100, 100)) + os.utime(newer, (200, 200)) + os.utime(generated, (300, 300)) + + files = crud_list.ListStep._walk_files( + root, + recursive=True, + limit=2, + sort_by="mtime", + extensions=frozenset({"md"}), + ) + + assert [path.name for path in files] == ["newer.md", "older.md"] + + # -- download ------------------------------------------------------------ # # DownloadStep lives in reme_cc (the local plugin overlay), not reme diff --git a/tests/unit/test_workspace_web_steps.py b/tests/unit/test_workspace_web_steps.py index 7071ce8e..ad592c2a 100644 --- a/tests/unit/test_workspace_web_steps.py +++ b/tests/unit/test_workspace_web_steps.py @@ -3,6 +3,7 @@ import asyncio import os from datetime import datetime +from zoneinfo import ZoneInfo from reme.components.agent_wrapper import BaseAgentWrapper from reme.components.application_context import ApplicationContext @@ -17,8 +18,8 @@ from reme.steps.file_io.save import SaveStep class _StreamingAgent(BaseAgentWrapper): """Minimal agent that records resume/tool arguments and emits rich chunks.""" - def __init__(self): - super().__init__() + def __init__(self, **kwargs): + super().__init__(**kwargs) self.reply_kwargs = {} async def reply(self, inputs, **kwargs) -> dict: @@ -28,6 +29,7 @@ class _StreamingAgent(BaseAgentWrapper): self.reply_kwargs = kwargs yield StreamChunk(chunk_type=ChunkEnum.TOOL_CALL, chunk="{}", tool_call_id="tool-1", tool_call_name="search") yield StreamChunk(chunk_type=ChunkEnum.CONTENT, chunk="hello", session_id="session-new") + yield StreamChunk(chunk_type=ChunkEnum.REPLY_END, chunk="", session_id="session-new") def test_save_step_preserves_complete_markdown(tmp_path): @@ -113,21 +115,47 @@ def test_save_step_rejects_external_change(tmp_path): asyncio.run(run()) -def test_chat_step_streams_rich_chunks_and_resumes_session(): +def test_chat_step_streams_rich_chunks_and_resumes_session(tmp_path): """Web chat keeps StreamChunk metadata and exposes only read-only tools.""" async def run(): - agent = _StreamingAgent() + app_context = ApplicationContext(workspace_dir=str(tmp_path), timezone="Asia/Shanghai") + agent = _StreamingAgent(app_context=app_context) queue = asyncio.Queue() context = RuntimeContext(stream_queue=queue, query="hello", session_id="session-old") - response = await ChatStep(agent_wrapper=agent)(context) + response = await ChatStep(agent_wrapper=agent, app_context=app_context)(context) - chunks = [await queue.get(), await queue.get()] + chunks = [await queue.get(), await queue.get(), await queue.get()] assert response.answer == "hello" assert chunks[0].tool_call_name == "search" assert chunks[1].session_id == "session-new" + assert chunks[2].metadata["answer"] == "hello" assert agent.reply_kwargs["resume"] == "session-old" - assert agent.reply_kwargs["job_tools"] == ["search", "list", "read", "stat", "traverse"] + assert agent.reply_kwargs["job_tools"] == ["search", "list", "stat", "traverse"] + system_prompt = agent.reply_kwargs["system_prompt"] + assert f"Current date: {datetime.now(ZoneInfo('Asia/Shanghai')).date().isoformat()}" in system_prompt + assert f"Current working directory: {tmp_path.resolve()}" in system_prompt + + asyncio.run(run()) + + +def test_chat_step_appends_environment_context_to_system_prompt_override(tmp_path): + """A caller prompt override retains request-time date and cwd context.""" + + async def run(): + app_context = ApplicationContext(workspace_dir=str(tmp_path), timezone="Asia/Shanghai") + agent = _StreamingAgent(app_context=app_context) + + await ChatStep(agent_wrapper=agent, app_context=app_context)( + query="hello", + system_prompt="Custom prompt.", + stream_queue=asyncio.Queue(), + ) + + system_prompt = agent.reply_kwargs["system_prompt"] + assert system_prompt.startswith("Custom prompt.\n\n") + assert "Current date:" in system_prompt + assert f"Current working directory: {tmp_path.resolve()}" in system_prompt asyncio.run(run()) From f31daf1949b8cc30ef3533d8618b07655c001104 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Wed, 5 Aug 2026 23:17:08 +0800 Subject: [PATCH 3/3] Revert "feat(backend): improve workspace support for web clients (#417)" (#419) This reverts commit b00eb0a9ea292924d57915fc42477f0697525b54. --- .../agent_wrapper/as_agent_wrapper.py | 86 +------------------ reme/config/default.yaml | 12 --- reme/schema/application_config.py | 15 +--- reme/steps/common/chat.py | 29 +------ reme/steps/file_io/list.py | 38 ++------ tests/unit/test_base_agent_wrapper.py | 36 -------- tests/unit/test_crud_steps.py | 23 ----- tests/unit/test_workspace_web_steps.py | 42 ++------- 8 files changed, 22 insertions(+), 259 deletions(-) diff --git a/reme/components/agent_wrapper/as_agent_wrapper.py b/reme/components/agent_wrapper/as_agent_wrapper.py index 7360543c..3b539bf5 100644 --- a/reme/components/agent_wrapper/as_agent_wrapper.py +++ b/reme/components/agent_wrapper/as_agent_wrapper.py @@ -123,86 +123,6 @@ class WorkspaceBackend(LocalBackend): return ExecResult(exit_code=process.returncode or 0, stdout=stdout, stderr=stderr) -class WorkspaceRead(Read): - """AgentScope Read accepting paths relative to the configured agent workspace.""" - - description = Read.description.replace( - "must be an absolute path, not a relative path", - "may be absolute or relative to the workspace", - ) - input_schema = { - **Read.input_schema, - "properties": { - **Read.input_schema["properties"], - "file_path": { - "type": "string", - "description": "Absolute path or path relative to the workspace.", - }, - }, - } - - async def call( - self, - file_path: str, - offset: int = 1, - limit: int = 2000, - _agent_state: AgentState | None = None, - ) -> ToolChunk: - """Resolve a workspace-relative path before delegating to AgentScope.""" - cwd = await self._backend.getcwd() # pylint: disable=protected-access - target = self._backend.abspath(file_path, cwd=cwd) # pylint: disable=protected-access - return await super().call(target, offset, limit, _agent_state) - - -class WorkspaceWrite(Write): - """AgentScope Write accepting paths relative to the configured agent workspace.""" - - input_schema = { - **Write.input_schema, - "properties": { - **Write.input_schema["properties"], - "file_path": { - "type": "string", - "description": "Absolute path or path relative to the workspace.", - }, - }, - } - - async def call(self, file_path: str, content: str, _agent_state: AgentState | None = None) -> ToolChunk: - """Resolve a workspace-relative path before delegating to AgentScope.""" - cwd = await self._backend.getcwd() # pylint: disable=protected-access - target = self._backend.abspath(file_path, cwd=cwd) # pylint: disable=protected-access - return await super().call(target, content, _agent_state) - - -class WorkspaceEdit(Edit): - """AgentScope Edit accepting paths relative to the configured agent workspace.""" - - input_schema = { - **Edit.input_schema, - "properties": { - **Edit.input_schema["properties"], - "file_path": { - "type": "string", - "description": "Absolute path or path relative to the workspace.", - }, - }, - } - - async def call( - self, - file_path: str, - old_string: str, - new_string: str, - replace_all: bool = False, - _agent_state: AgentState | None = None, - ) -> ToolChunk: - """Resolve a workspace-relative path before delegating to AgentScope.""" - cwd = await self._backend.getcwd() # pylint: disable=protected-access - target = self._backend.abspath(file_path, cwd=cwd) # pylint: disable=protected-access - return await super().call(target, old_string, new_string, replace_all, _agent_state) - - class BypassAnalysisBash(Bash): """Bash variant that delegates permission decisions to PermissionEngine. @@ -272,11 +192,11 @@ class AsAgentWrapper(BaseAgentWrapper): backend = WorkspaceBackend(cwd, self.bash_environment) factories = { "bash": lambda: BypassAnalysisBash(cwd=cwd, backend=backend), - "edit": lambda: WorkspaceEdit(backend=backend), + "edit": lambda: Edit(backend=backend), "glob": lambda: Glob(backend=backend), "grep": lambda: Grep(backend=backend), - "read": lambda: WorkspaceRead(backend=backend), - "write": lambda: WorkspaceWrite(backend=backend), + "read": lambda: Read(backend=backend), + "write": lambda: Write(backend=backend), } if names is False: selected_names = [] diff --git a/reme/config/default.yaml b/reme/config/default.yaml index 36036b65..30b29f1a 100644 --- a/reme/config/default.yaml +++ b/reme/config/default.yaml @@ -1,5 +1,3 @@ -workspace_dir: ~/.copaw/workspaces/default/ - service: backend: http @@ -505,16 +503,6 @@ jobs: type: boolean description: "recurse" default: false - sort_by: - type: string - enum: - - mtime - description: "optional ordering; mtime returns most recently modified files first" - extensions: - type: array - items: - type: string - description: "optional extension allowlist applied before sorting and limiting" limit: type: integer description: "max results" diff --git a/reme/schema/application_config.py b/reme/schema/application_config.py index ff1b8697..00a0b286 100644 --- a/reme/schema/application_config.py +++ b/reme/schema/application_config.py @@ -1,9 +1,8 @@ """Application configuration models.""" import os -from pathlib import Path -from pydantic import BaseModel, ConfigDict, Field, field_validator +from pydantic import BaseModel, ConfigDict, Field from ..enumeration import ComponentEnum @@ -33,11 +32,7 @@ class ApplicationConfig(BaseModel): default_factory=dict, description="Environment variables loaded once at startup and passed to agent subprocesses", ) - workspace_dir: str = Field( - default=".reme", - description="Workspace root directory for runtime files", - validate_default=True, - ) + workspace_dir: str = Field(default=".reme", description="Workspace root directory for runtime files") metadata_dir: str = Field(default="metadata", description="Subdirectory for ReMe persistent state") session_dir: str = Field(default="session", description="Subdirectory for persisted agent sessions") mem_session_dir: str = Field(default="mem_session", description="Subdirectory for persisted agent sessions") @@ -58,9 +53,3 @@ class ApplicationConfig(BaseModel): default_factory=dict, description="Component registry keyed by type then name", ) - - @field_validator("workspace_dir", mode="before") - @classmethod - def normalize_workspace_dir(cls, value) -> str: - """Expand home-relative paths once so every component sees the same absolute workspace.""" - return str(Path(value).expanduser().resolve(strict=False)) diff --git a/reme/steps/common/chat.py b/reme/steps/common/chat.py index bf8002de..56f18de1 100644 --- a/reme/steps/common/chat.py +++ b/reme/steps/common/chat.py @@ -1,8 +1,5 @@ """Workspace-aware streaming chat step for the ReMe web interface.""" -import datetime -import zoneinfo - from ..base_step import BaseStep from ...components import R from ...enumeration import ChunkEnum @@ -16,25 +13,7 @@ class ChatStep(BaseStep): Use the available ReMe tools when workspace facts are needed. Cite workspace-relative file paths when referring to notes. Never invent file contents, and do not claim to have changed files because this chat intentionally provides read-only tools. Reply in the user's language.""" - READ_ONLY_TOOLS = ["search", "list", "stat", "traverse"] - - def _system_prompt(self) -> str: - """Append request-time environment facts to the configured prompt.""" - timezone = self.app_context.app_config.timezone if self.app_context is not None else None - try: - current = datetime.datetime.now(zoneinfo.ZoneInfo(timezone)) if timezone else datetime.datetime.now() - except (zoneinfo.ZoneInfoNotFoundError, ValueError): - current = datetime.datetime.now() - - assert self.context is not None - base_prompt = str(self.context.get("system_prompt") or self.DEFAULT_SYSTEM_PROMPT).rstrip() - return ( - f"{base_prompt}\n\n" - "\n" - f"Current date: {current.date().isoformat()}\n" - f"Current working directory: {self.agent_wrapper.cwd.resolve(strict=False)}\n" - "" - ) + READ_ONLY_TOOLS = ["search", "list", "read", "stat", "traverse"] async def execute(self): assert self.context is not None @@ -47,7 +26,7 @@ this chat intentionally provides read-only tools. Reply in the user's language." raise RuntimeError("chat_step requires an agent_wrapper") wrapper_kwargs = { - "system_prompt": self._system_prompt(), + "system_prompt": self.context.get("system_prompt") or self.DEFAULT_SYSTEM_PROMPT, "job_tools": self.READ_ONLY_TOOLS, } if session_id := self.context.get("session_id"): @@ -55,10 +34,6 @@ this chat intentionally provides read-only tools. Reply in the user's language." parts: list[str] = [] async for chunk in self.agent_wrapper.reply_stream(query, **wrapper_kwargs): - if chunk.chunk_type == ChunkEnum.REPLY_END: - answer = "".join(parts).strip() - if answer: - chunk.metadata["answer"] = answer await self.context.add_stream_chunk(chunk) if chunk.chunk_type == ChunkEnum.CONTENT and isinstance(chunk.chunk, str): parts.append(chunk.chunk) diff --git a/reme/steps/file_io/list.py b/reme/steps/file_io/list.py index cee573ba..ebacc967 100644 --- a/reme/steps/file_io/list.py +++ b/reme/steps/file_io/list.py @@ -8,8 +8,6 @@ Parameters: path — dir to list under (relative to the workspace or absolute). Empty = workspace root. limit — cap on the number of returned items (default 100, must be > 0). recursive — descend into subdirectories. Default False = direct children only. - sort_by — optional ordering; ``mtime`` returns most recently modified files first. - extensions — optional extension allowlist applied before sorting and limiting. No frontmatter is read. Callers needing frontmatter-based filtering should iterate the result and call ``frontmatter_read`` per candidate. @@ -38,7 +36,7 @@ class ListStep(BaseStep): if meta: self.context.response.metadata.update(meta) - def _collect_params(self) -> tuple[str, bool, int, str, frozenset[str]]: + def _collect_params(self) -> tuple[str, bool, int]: """Read ``path`` / ``recursive`` / ``limit`` from context; coerce permissively.""" assert self.context is not None path = str(self.context.get("path") or "") @@ -49,37 +47,20 @@ class ListStep(BaseStep): limit = int(raw_limit) if raw_limit is not None else DEFAULT_LIMIT except (TypeError, ValueError): limit = DEFAULT_LIMIT - sort_by = str(self.context.get("sort_by") or "") - raw_extensions = self.context.get("extensions") or [] - if isinstance(raw_extensions, str): - raw_extensions = raw_extensions.split(",") - extensions = frozenset( - normalized for value in raw_extensions if (normalized := str(value).strip().lower().lstrip(".")) - ) - return path, recursive, limit if limit > 0 else DEFAULT_LIMIT, sort_by, extensions + return path, recursive, limit if limit > 0 else DEFAULT_LIMIT @staticmethod - def _walk_files( - target_dir: Path, - recursive: bool, - limit: int, - sort_by: str = "", - extensions: frozenset[str] = frozenset(), - ) -> list[Path]: - """Return up to ``limit`` regular files, scanning all entries only when sorting.""" + def _walk_files(target_dir: Path, recursive: bool, limit: int) -> list[Path]: + """Return up to ``limit`` regular files under ``target_dir``; short-circuits at the cap.""" entries: Iterable[Path] = target_dir.rglob("*") if recursive else target_dir.iterdir() files: list[Path] = [] for entry in entries: if not entry.is_file(): # skip dirs, sockets, broken links, etc. continue - if extensions and entry.suffix.lower().lstrip(".") not in extensions: - continue files.append(entry) - if sort_by != "mtime" and len(files) >= limit: + if len(files) >= limit: break - if sort_by == "mtime": - files.sort(key=lambda entry: (-entry.stat().st_mtime, entry.as_posix())) - return files[:limit] + return files @staticmethod def _format_relative(files: list[Path], workspace_dir: Path) -> list[str]: @@ -94,7 +75,7 @@ class ListStep(BaseStep): async def execute(self): assert self.context is not None - path, recursive, limit, sort_by, extensions = self._collect_params() + path, recursive, limit = self._collect_params() workspace_dir = Path(self.file_store.workspace_path or ".").resolve() target_dir, err = resolve_path(workspace_dir, path, allow_empty=True) if err or target_dir is None: @@ -108,10 +89,7 @@ class ListStep(BaseStep): self._fail(f"path {target_dir} is not a directory", path=str(target_dir)) return None - items = self._format_relative( - self._walk_files(target_dir, recursive, limit, sort_by, extensions), - workspace_dir, - ) + items = self._format_relative(self._walk_files(target_dir, recursive, limit), workspace_dir) self.context.response.success = True location = path or "." diff --git a/tests/unit/test_base_agent_wrapper.py b/tests/unit/test_base_agent_wrapper.py index 53fda2bb..17a1b1fa 100644 --- a/tests/unit/test_base_agent_wrapper.py +++ b/tests/unit/test_base_agent_wrapper.py @@ -139,42 +139,6 @@ async def test_agentscope_backend_passes_configured_environment_to_bash(tmp_path assert result.stdout == b"configured\n" -def test_workspace_dir_expands_user_home(monkeypatch, tmp_path): - """Home-relative workspace config is normalized before components consume it.""" - monkeypatch.setenv("HOME", str(tmp_path)) - - context = ApplicationContext(workspace_dir="~/.copaw/workspaces/default") - - assert context.app_config.workspace_dir == str(tmp_path / ".copaw/workspaces/default") - - -@pytest.mark.asyncio -async def test_agentscope_file_tools_accept_workspace_relative_paths(tmp_path): - """Built-in Read/Write/Edit share the workspace-relative path convention used by ReMe jobs.""" - wrapper = AsAgentWrapper( - app_context=ApplicationContext(workspace_dir=str(tmp_path)), - as_llm="", - ) - read, write, edit = wrapper._builtin_tools(["read", "write", "edit"]) # pylint: disable=protected-access - source = tmp_path / ".reme/daily/note.md" - source.parent.mkdir(parents=True) - source.write_text("before", encoding="utf-8") - - read_result = await read.call(file_path=".reme/daily/note.md") - write_result = await write.call(file_path=".reme/daily/new.md", content="new") - edit_result = await edit.call( - file_path=".reme/daily/note.md", - old_string="before", - new_string="after", - ) - - assert "absolute path" not in str(read_result.content) - assert "absolute path" not in str(write_result.content) - assert "absolute path" not in str(edit_result.content) - assert (tmp_path / ".reme/daily/new.md").read_text(encoding="utf-8") == "new" - assert source.read_text(encoding="utf-8") == "after" - - @pytest.mark.asyncio async def test_agentscope_bash_uses_managed_proxy_without_changing_subprocess_environment(tmp_path): """AgentScope applies the managed proxy only to its command backend.""" diff --git a/tests/unit/test_crud_steps.py b/tests/unit/test_crud_steps.py index 0eb4d98e..1132a686 100644 --- a/tests/unit/test_crud_steps.py +++ b/tests/unit/test_crud_steps.py @@ -225,29 +225,6 @@ def test_list_respects_limit_and_non_recursive(): asyncio.run(run()) -def test_list_can_sort_by_most_recent_modification(): - """Filtering precedes the limit so unrelated generated files cannot hide recent notes.""" - with tempfile.TemporaryDirectory() as tmp: - root = Path(tmp) - older = _seed_md(root, "older.md", "old") - newer = _seed_md(root, "newer.md", "new") - generated = root / "newest.json" - generated.write_text("{}", encoding="utf-8") - os.utime(older, (100, 100)) - os.utime(newer, (200, 200)) - os.utime(generated, (300, 300)) - - files = crud_list.ListStep._walk_files( - root, - recursive=True, - limit=2, - sort_by="mtime", - extensions=frozenset({"md"}), - ) - - assert [path.name for path in files] == ["newer.md", "older.md"] - - # -- download ------------------------------------------------------------ # # DownloadStep lives in reme_cc (the local plugin overlay), not reme diff --git a/tests/unit/test_workspace_web_steps.py b/tests/unit/test_workspace_web_steps.py index ad592c2a..7071ce8e 100644 --- a/tests/unit/test_workspace_web_steps.py +++ b/tests/unit/test_workspace_web_steps.py @@ -3,7 +3,6 @@ import asyncio import os from datetime import datetime -from zoneinfo import ZoneInfo from reme.components.agent_wrapper import BaseAgentWrapper from reme.components.application_context import ApplicationContext @@ -18,8 +17,8 @@ from reme.steps.file_io.save import SaveStep class _StreamingAgent(BaseAgentWrapper): """Minimal agent that records resume/tool arguments and emits rich chunks.""" - def __init__(self, **kwargs): - super().__init__(**kwargs) + def __init__(self): + super().__init__() self.reply_kwargs = {} async def reply(self, inputs, **kwargs) -> dict: @@ -29,7 +28,6 @@ class _StreamingAgent(BaseAgentWrapper): self.reply_kwargs = kwargs yield StreamChunk(chunk_type=ChunkEnum.TOOL_CALL, chunk="{}", tool_call_id="tool-1", tool_call_name="search") yield StreamChunk(chunk_type=ChunkEnum.CONTENT, chunk="hello", session_id="session-new") - yield StreamChunk(chunk_type=ChunkEnum.REPLY_END, chunk="", session_id="session-new") def test_save_step_preserves_complete_markdown(tmp_path): @@ -115,47 +113,21 @@ def test_save_step_rejects_external_change(tmp_path): asyncio.run(run()) -def test_chat_step_streams_rich_chunks_and_resumes_session(tmp_path): +def test_chat_step_streams_rich_chunks_and_resumes_session(): """Web chat keeps StreamChunk metadata and exposes only read-only tools.""" async def run(): - app_context = ApplicationContext(workspace_dir=str(tmp_path), timezone="Asia/Shanghai") - agent = _StreamingAgent(app_context=app_context) + agent = _StreamingAgent() queue = asyncio.Queue() context = RuntimeContext(stream_queue=queue, query="hello", session_id="session-old") - response = await ChatStep(agent_wrapper=agent, app_context=app_context)(context) + response = await ChatStep(agent_wrapper=agent)(context) - chunks = [await queue.get(), await queue.get(), await queue.get()] + chunks = [await queue.get(), await queue.get()] assert response.answer == "hello" assert chunks[0].tool_call_name == "search" assert chunks[1].session_id == "session-new" - assert chunks[2].metadata["answer"] == "hello" assert agent.reply_kwargs["resume"] == "session-old" - assert agent.reply_kwargs["job_tools"] == ["search", "list", "stat", "traverse"] - system_prompt = agent.reply_kwargs["system_prompt"] - assert f"Current date: {datetime.now(ZoneInfo('Asia/Shanghai')).date().isoformat()}" in system_prompt - assert f"Current working directory: {tmp_path.resolve()}" in system_prompt - - asyncio.run(run()) - - -def test_chat_step_appends_environment_context_to_system_prompt_override(tmp_path): - """A caller prompt override retains request-time date and cwd context.""" - - async def run(): - app_context = ApplicationContext(workspace_dir=str(tmp_path), timezone="Asia/Shanghai") - agent = _StreamingAgent(app_context=app_context) - - await ChatStep(agent_wrapper=agent, app_context=app_context)( - query="hello", - system_prompt="Custom prompt.", - stream_queue=asyncio.Queue(), - ) - - system_prompt = agent.reply_kwargs["system_prompt"] - assert system_prompt.startswith("Custom prompt.\n\n") - assert "Current date:" in system_prompt - assert f"Current working directory: {tmp_path.resolve()}" in system_prompt + assert agent.reply_kwargs["job_tools"] == ["search", "list", "read", "stat", "traverse"] asyncio.run(run())