From 06fb46fa48c1706f398d649dadc26814eb82c62a Mon Sep 17 00:00:00 2001 From: xyf2020 <75460675+xyf2020@users.noreply.github.com> Date: Thu, 10 Sep 2026 18:01:09 +0800 Subject: [PATCH] feat(tags): add configurable tag indexing and filtered hybrid search (#530) * Add optional tag generation and normalization to auto memory * Add tag index components and clean up temporary JSONL files * Preserve tag index state when reconciliation fails * Refactor and streamline application implementation * Fix pylint C1803 warnings in tag normalization tests * Document optional tag index configuration * Make tag index failures non-blocking and disable auto-memory tags * Add configurable tag indexing and tag listing * Add tag-filtered hybrid search with exact candidate ranking * Remove obsolete generated files * Rename tag index key to tag_key and reject reserved fields * Extract automatic tagging into a dedicated step * Restrict frontmatter updates to authorized keys * Refine tag filtering and automatic memory tagging * Require underscore-separated tags in auto-tag prompts - Forbid spaces in tags and require underscores (e.g. sam_altman) in both English and Chinese auto_tag prompts, with English examples switched to English entities (OpenAI, gold) - Drop prompt-string assertions superseded by the new tagging rule - Merge construction/runtime tag_key validation tests into one parametrized case * Make max_tags_per_file configurable in auto-tag step * Consolidate tag index tests * Fix search test fixture lint warnings * Align tag contracts and index health behavior * Fall back when tag index is unavailable --------- Co-authored-by: jinli.yl --- reme/components/file_store/base_file_store.py | 18 ++ .../file_store/faiss_local_file_store.py | 7 +- .../components/file_store/local_file_store.py | 87 ++++--- .../file_store/zvec_local_file_store.py | 2 +- .../keyword_index/base_keyword_index.py | 26 +- reme/components/keyword_index/bm25_index.py | 80 +++++- reme/components/tag_index/base_tag_index.py | 55 +++- reme/components/tag_index/local_tag_index.py | 86 ++++++- reme/config/default.yaml | 57 ++++- reme/steps/evolve/__init__.py | 2 + reme/steps/evolve/auto_memory.py | 87 ++----- reme/steps/evolve/auto_memory.yaml | 14 -- reme/steps/evolve/auto_tag.py | 235 ++++++++++++++++++ reme/steps/evolve/auto_tag.yaml | 45 ++++ reme/steps/evolve/dream/utils.py | 3 +- reme/steps/file_io/frontmatter_update.py | 18 ++ reme/steps/index/__init__.py | 2 + reme/steps/index/list_tags.py | 22 ++ reme/steps/index/search.py | 78 ++++++ reme/steps/index/update_changes.py | 8 +- tests/unit/test_auto_tag.py | 212 ++++++++++++++++ tests/unit/test_background_steps.py | 72 +----- tests/unit/test_evolve_utils.py | 35 ++- tests/unit/test_faiss_index_maintenance.py | 1 + tests/unit/test_file_store_consistency.py | 37 +++ tests/unit/test_frontmatter_steps.py | 53 ++++ tests/unit/test_injected_job_kwargs.py | 59 +---- tests/unit/test_keyword_index.py | 26 ++ tests/unit/test_search_step.py | 125 +++++++++- tests/unit/test_tag_index.py | 219 ++++++++++++++-- tests/unit/test_zvec_file_store.py | 1 + 31 files changed, 1465 insertions(+), 307 deletions(-) create mode 100644 reme/steps/evolve/auto_tag.py create mode 100644 reme/steps/evolve/auto_tag.yaml create mode 100644 reme/steps/index/list_tags.py create mode 100644 tests/unit/test_auto_tag.py diff --git a/reme/components/file_store/base_file_store.py b/reme/components/file_store/base_file_store.py index 91bcaa73..414be787 100644 --- a/reme/components/file_store/base_file_store.py +++ b/reme/components/file_store/base_file_store.py @@ -6,6 +6,7 @@ from contextlib import asynccontextmanager from functools import wraps from ..base_component import BaseComponent +from ..tag_index import BaseTagIndex from ...enumeration import ComponentEnum, LinkScopeEnum from ...schema import FileChunk, FileLink, FileNode @@ -23,9 +24,26 @@ class BaseFileStore(BaseComponent): def __init__(self, **kwargs): super().__init__(**kwargs) + self.tag_index: BaseTagIndex | None = None self._maintenance_lock = asyncio.Lock() self._maintenance_lock_owner = None + @property + def tag_index_enabled(self) -> bool: + """Whether this file-store backend has a tag index configured.""" + return self.tag_index is not None + + def require_tag_index(self) -> BaseTagIndex: + """Return the configured tag index or fail with one consistent error.""" + if self.tag_index is None: + raise RuntimeError("tag index is not configured") + return self.tag_index + + @property + def embedding_dimensions(self) -> int: + """Vector dimensions used for memory estimates; zero when unavailable.""" + return 0 + @asynccontextmanager async def _maintenance_guard(self): """Serialize maintenance and mutations, allowing nested backend overrides.""" diff --git a/reme/components/file_store/faiss_local_file_store.py b/reme/components/file_store/faiss_local_file_store.py index 12684e81..02613eab 100644 --- a/reme/components/file_store/faiss_local_file_store.py +++ b/reme/components/file_store/faiss_local_file_store.py @@ -648,12 +648,7 @@ class FaissLocalFileStore(LocalFileStore): async def vector_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]: index_empty = self._faiss_index is None or self._faiss_index.ntotal == 0 embedding_unavailable = self.embedding_store is None or self._embedding_rebuild_pending - if ( - embedding_unavailable - or not query - or limit <= 0 - or (index_empty and getattr(self.embedding_store, "is_healthy", True)) - ): + if embedding_unavailable or not query or limit <= 0 or (index_empty and self.embedding_store.is_healthy): return [] query_embedding = await self._get_query_embedding(query) diff --git a/reme/components/file_store/local_file_store.py b/reme/components/file_store/local_file_store.py index 6075aaee..bccd8fbf 100644 --- a/reme/components/file_store/local_file_store.py +++ b/reme/components/file_store/local_file_store.py @@ -81,6 +81,12 @@ class LocalFileStore(BaseFileStore): self._tag_indexed_file_count = 0 self._closing = False + @property + def embedding_dimensions(self) -> int: + if self.embedding_store is None: + return 0 + return max(0, int(getattr(self.embedding_store, "dimensions", 0))) + # -- lifecycle ------------------------------------------------------------ async def _start(self) -> None: @@ -136,7 +142,7 @@ class LocalFileStore(BaseFileStore): self.embedding_store is None or self._embedding_rebuild_pending or was_healthy - or not getattr(self.embedding_store, "is_healthy", True) + or not self.embedding_store.is_healthy ): return self.logger.info(f"{self.name}: embedding provider recovered; scheduling missing-vector backfill") @@ -196,7 +202,7 @@ class LocalFileStore(BaseFileStore): return None embedding_generation = self._embedding_space_generation - was_healthy = bool(getattr(embedding_store, "is_healthy", True)) + was_healthy = embedding_store.is_healthy try: query_embedding = await embedding_store.get_embedding(query) except Exception as e: @@ -248,7 +254,7 @@ class LocalFileStore(BaseFileStore): tag_sync_started_at = time.monotonic() tag_synced = await self._rebuild_tag_index("startup synchronization") self.logger.info( - f"{self.name}: tag index sync complete: enabled={self.tag_index is not None}, " + f"{self.name}: tag index sync complete: enabled={self.tag_index_enabled}, " f"healthy={tag_synced}, " f"elapsed={time.monotonic() - tag_sync_started_at:.3f}s", ) @@ -502,7 +508,7 @@ class LocalFileStore(BaseFileStore): return total = len(missing) - batch_size = max(1, int(getattr(self.embedding_store, "max_batch_size", 10))) + batch_size = max(1, int(self.embedding_store.max_batch_size)) self.logger.info(f"{self.name}: embedding backfill started: total={total}, batch_size={batch_size}") try: if not skip_health_check: @@ -636,13 +642,14 @@ class LocalFileStore(BaseFileStore): async def _rebuild_tag_index(self, reason: str) -> bool: """Best-effort rebuild that never lets an optional tag index block the file store.""" - if self.tag_index is None: + if not self.tag_index_enabled: self._tag_index_rebuild_required = False self._tag_indexed_file_count = 0 return True + tag_index = self.require_tag_index() try: nodes = await self.file_graph.get_nodes() - await self.tag_index.rebuild(nodes) + await tag_index.rebuild(nodes) except Exception: self._tag_index_rebuild_required = True self.logger.exception( @@ -651,20 +658,20 @@ class LocalFileStore(BaseFileStore): await self._quarantine_tag_index(reason) return False self._tag_index_rebuild_required = False - self.tag_index.set_healthy(True) - self._tag_indexed_file_count = self.tag_index.n_files + tag_index.set_healthy(True) + self._tag_indexed_file_count = tag_index.n_files return True async def _quarantine_tag_index(self, reason: str) -> None: """Fail closed after a reconciliation error without propagating cleanup failures.""" - assert self.tag_index is not None - self.tag_index.set_healthy(False) + tag_index = self.require_tag_index() + tag_index.set_healthy(False) try: - await self.tag_index.clear() + await tag_index.clear() except Exception: self.logger.exception(f"{self.name}: failed to clear unhealthy tag index during {reason}") finally: - self.tag_index.set_healthy(False) + tag_index.set_healthy(False) async def _reindex_tag(self) -> dict: """Synchronously rebuild the optional tag index from the authoritative file graph.""" @@ -674,33 +681,35 @@ class LocalFileStore(BaseFileStore): async def _upsert_tag_nodes(self, nodes: list[FileNode]) -> None: """Update tags without allowing optional-index failures to block other indexes.""" - if self.tag_index is None: + if not self.tag_index_enabled: return + tag_index = self.require_tag_index() if self._tag_index_rebuild_required: await self._rebuild_tag_index("retry before incremental update") return try: - await self.tag_index.upsert_nodes(nodes) - self.tag_index.set_healthy(True) + await tag_index.upsert_nodes(nodes) + tag_index.set_healthy(True) except Exception: self._tag_index_rebuild_required = True - self.tag_index.set_healthy(False) + tag_index.set_healthy(False) self.logger.exception(f"{self.name}: incremental tag index update failed; rebuilding from graph") await self._rebuild_tag_index("incremental update recovery") async def _delete_tag_paths(self, paths: list[str]) -> None: """Delete tag paths without allowing optional-index failures to block other indexes.""" - if self.tag_index is None: + if not self.tag_index_enabled: return + tag_index = self.require_tag_index() if self._tag_index_rebuild_required: await self._rebuild_tag_index("retry before incremental delete") return try: - await self.tag_index.delete(paths) - self.tag_index.set_healthy(True) + await tag_index.delete(paths) + tag_index.set_healthy(True) except Exception: self._tag_index_rebuild_required = True - self.tag_index.set_healthy(False) + tag_index.set_healthy(False) self.logger.exception(f"{self.name}: incremental tag index delete failed; rebuilding from graph") await self._rebuild_tag_index("incremental delete recovery") @@ -831,7 +840,7 @@ class LocalFileStore(BaseFileStore): return embedding_store = self.embedding_store embedding_generation = self._embedding_space_generation - was_healthy = bool(getattr(embedding_store, "is_healthy", True)) + was_healthy = embedding_store.is_healthy try: await embedding_store.get_node_embeddings(chunks) except Exception as e: @@ -903,15 +912,16 @@ class LocalFileStore(BaseFileStore): if self.keyword_index: await self.keyword_index.clear() await self.file_graph.clear() - if self.tag_index is not None: + if self.tag_index_enabled: + tag_index = self.require_tag_index() try: - await self.tag_index.clear() + await tag_index.clear() self._tag_index_rebuild_required = False self._tag_indexed_file_count = 0 - self.tag_index.set_healthy(True) + tag_index.set_healthy(True) except Exception: self._tag_index_rebuild_required = True - self.tag_index.set_healthy(False) + tag_index.set_healthy(False) self.logger.exception(f"{self.name}: tag index clear failed; rebuilding from empty graph") await self._rebuild_tag_index("clear recovery") self._mutation_generation += 1 @@ -967,17 +977,32 @@ class LocalFileStore(BaseFileStore): ] async def keyword_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]: - if not self.keyword_index: + if not self.keyword_index or limit <= 0: return [] query = query.strip() if not query: return [] - retrieve_limit = limit if search_filter: - retrieve_limit = max(limit, getattr(self.keyword_index, "n_docs", limit)) - doc_id_score_dict = await self.keyword_index.retrieve(query, limit=retrieve_limit) + eligible_ids = { + chunk.id for chunk in self.file_chunks.values() if self._matches_search_filter(chunk, search_filter) + } + if not eligible_ids: + return [] + try: + eligible_ids.intersection_update(self.keyword_index.document_ids) + if not eligible_ids: + return [] + doc_id_score_dict = await self.keyword_index.retrieve_filtered(query, limit, eligible_ids) + except NotImplementedError: + # Compatibility path for third-party indexes implementing only retrieve(). + doc_id_score_dict = await self.keyword_index.retrieve( + query, + limit=max(limit, getattr(self.keyword_index, "n_docs", limit)), + ) + else: + doc_id_score_dict = await self.keyword_index.retrieve(query, limit=limit) results = [] for doc_id, score in doc_id_score_dict.items(): chunk = self.file_chunks.get(doc_id) @@ -1033,10 +1058,12 @@ class LocalFileStore(BaseFileStore): return True exact_paths = set() + has_exact_path_filter = False for key in ("path", "paths"): if key in search_filter: + has_exact_path_filter = True exact_paths.update(cls._as_filter_values(search_filter[key])) - if exact_paths and chunk.path not in exact_paths: + if has_exact_path_filter and chunk.path not in exact_paths: return False prefixes = [] diff --git a/reme/components/file_store/zvec_local_file_store.py b/reme/components/file_store/zvec_local_file_store.py index afad4b4d..58a0300a 100644 --- a/reme/components/file_store/zvec_local_file_store.py +++ b/reme/components/file_store/zvec_local_file_store.py @@ -441,7 +441,7 @@ class ZvecLocalFileStore(LocalFileStore): if self.embedding_store is None or self._embedding_rebuild_pending or not query or limit <= 0: return [] index_empty = self._collection is None or not self._indexed_ids - if index_empty and getattr(self.embedding_store, "is_healthy", True): + if index_empty and self.embedding_store.is_healthy: return [] query_embedding = await self._get_query_embedding(query) diff --git a/reme/components/keyword_index/base_keyword_index.py b/reme/components/keyword_index/base_keyword_index.py index feac9e50..186c65b3 100644 --- a/reme/components/keyword_index/base_keyword_index.py +++ b/reme/components/keyword_index/base_keyword_index.py @@ -1,7 +1,7 @@ """Abstract base class for keyword indexes (BM25 and other lexical backends).""" from abc import abstractmethod -from collections.abc import Set +from collections.abc import Collection, Set from ..base_component import BaseComponent from ..tokenizer import BaseTokenizer @@ -54,6 +54,30 @@ class BaseKeywordIndex(BaseComponent): async def retrieve(self, query: str, limit: int = 3) -> dict[str, float]: """Return top-`limit` doc_id → score for the given query.""" + async def score_documents(self, query: str, document_ids: Collection[str]) -> dict[str, float]: + """Score selected documents using the index's normal corpus statistics. + + This compatibility implementation uses the existing public retrieval + contract. Backends can override it to avoid ranking the whole corpus. + """ + allowed = set(document_ids) + if not allowed: + return {} + results = await self.retrieve(query, limit=len(self.document_ids)) + return {doc_id: score for doc_id, score in results.items() if doc_id in allowed} + + async def retrieve_filtered( + self, + query: str, + limit: int, + document_ids: Collection[str], + ) -> dict[str, float]: + """Return top results restricted to document IDs, preserving old APIs.""" + if limit <= 0: + return {} + scores = await self.score_documents(query, document_ids) + return dict(list(scores.items())[:limit]) + @abstractmethod async def clear(self) -> None: """Wipe in-memory state and remove any persisted artifacts.""" diff --git a/reme/components/keyword_index/bm25_index.py b/reme/components/keyword_index/bm25_index.py index 0169726f..4bd4c14d 100644 --- a/reme/components/keyword_index/bm25_index.py +++ b/reme/components/keyword_index/bm25_index.py @@ -22,7 +22,7 @@ import math import pickle import re from collections import Counter -from collections.abc import KeysView +from collections.abc import Collection, KeysView from pathlib import Path from uuid import uuid4 @@ -280,32 +280,65 @@ class BM25Index(BaseKeywordIndex): self._remove_doc(doc_id) self._idf_cache = {} - def _score_query(self, query_ids: list[int], n_docs: int) -> np.ndarray: - """Compute BM25 scores across all docs; deleted docs zeroed out.""" + def _score_query(self, query_ids: list[int], candidate_idxs: np.ndarray | None = None) -> np.ndarray: + """Compute BM25 scores globally or for selected live document indexes.""" + n_docs = self.n_docs + if n_docs == 0: + size = self._doc_lens.size if candidate_idxs is None else candidate_idxs.size + return np.zeros(size, dtype=np.float32) + avg_len = self.total_len / n_docs k1, b = self.k1, self.b denom_base = k1 * (1.0 - b) denom_norm = k1 * b / avg_len if avg_len > 0 else 0.0 - scores = np.zeros(self._doc_lens.size, dtype=np.float32) + size = self._doc_lens.size if candidate_idxs is None else candidate_idxs.size + scores = np.zeros(size, dtype=np.float32) for tid in query_ids: - doc_idxs = self._posting_doc_idxs.get(tid) - if doc_idxs is None or doc_idxs.size == 0: + posting_idxs = self._posting_doc_idxs.get(tid) + if posting_idxs is None or posting_idxs.size == 0: continue idf = self._get_idf(tid, n_docs) if idf == 0.0: continue - tfs = self._posting_tfs[tid].astype(np.float32) - d_lens = self._doc_lens[doc_idxs].astype(np.float32) + + if candidate_idxs is None: + posting_positions = slice(None) + score_positions = posting_idxs + else: + _common, posting_positions, score_positions = np.intersect1d( + posting_idxs, + candidate_idxs, + assume_unique=True, + return_indices=True, + ) + if score_positions.size == 0: + continue + + doc_idxs = posting_idxs[posting_positions] + tfs = self._posting_tfs[tid][posting_positions].astype(np.float32) + doc_lens = self._doc_lens[doc_idxs].astype(np.float32) # Each doc_idx appears at most once per posting list (Counter dedups # within a doc, and updates allocate fresh idxs), so fancy-index # accumulation is safe here. - scores[doc_idxs] += idf * tfs * (k1 + 1.0) / (tfs + denom_base + denom_norm * d_lens) + scores[score_positions] += idf * tfs * (k1 + 1.0) / (tfs + denom_base + denom_norm * doc_lens) - if self._deleted.any(): + if candidate_idxs is None and self._deleted.any(): scores[self._deleted] = 0.0 return scores + def _score_query_documents( + self, + query_ids: list[int], + document_ids: Collection[str], + ) -> tuple[np.ndarray, np.ndarray]: + """Score live selected documents while retaining global BM25 statistics.""" + candidate_idxs = np.array( + sorted({self._doc_id_to_idx[doc_id] for doc_id in document_ids if doc_id in self._doc_id_to_idx}), + dtype=np.int32, + ) + return candidate_idxs, self._score_query(query_ids, candidate_idxs) + async def retrieve(self, query: str, limit: int = 3) -> dict[str, float]: """BM25 retrieval; returns {doc_id: score} sorted by score descending.""" n_docs = self.n_docs @@ -315,10 +348,35 @@ class BM25Index(BaseKeywordIndex): if not query_ids: return {} - scores = self._score_query(query_ids, n_docs) + scores = self._score_query(query_ids) top_idxs = self._top_k(scores, limit) return {self._doc_ids[int(i)]: float(scores[int(i)]) for i in top_idxs} + async def score_documents(self, query: str, document_ids: Collection[str]) -> dict[str, float]: + """Return every positive-scoring selected document in descending order.""" + query_ids = self._encode_query(query) + if not query_ids: + return {} + candidate_idxs, scores = self._score_query_documents(query_ids, document_ids) + ranked = self._top_k(scores, scores.size) + return {self._doc_ids[int(candidate_idxs[i])]: float(scores[i]) for i in ranked} + + async def retrieve_filtered( + self, + query: str, + limit: int, + document_ids: Collection[str], + ) -> dict[str, float]: + """Return exact BM25 top-k within selected documents.""" + if limit <= 0: + return {} + query_ids = self._encode_query(query) + if not query_ids: + return {} + candidate_idxs, scores = self._score_query_documents(query_ids, document_ids) + ranked = self._top_k(scores, limit) + return {self._doc_ids[int(candidate_idxs[i])]: float(scores[i]) for i in ranked} + # -- Persistence ---------------------------------------------------------- def _snapshot(self) -> dict: diff --git a/reme/components/tag_index/base_tag_index.py b/reme/components/tag_index/base_tag_index.py index 01a9a9d8..95882945 100644 --- a/reme/components/tag_index/base_tag_index.py +++ b/reme/components/tag_index/base_tag_index.py @@ -1,21 +1,59 @@ """Abstract interface for file-level tag indexes derived from graph nodes.""" from abc import abstractmethod +from typing import ClassVar, Literal, TypedDict from ..base_component import BaseComponent from ...enumeration import ComponentEnum from ...schema import FileNode +TagOrderBy = Literal["tag", "file_count"] +TagOrder = Literal["asc", "desc"] +TagListItem = tuple[str, int] + + +class TagListResult(TypedDict): + """One page of active tags and their indexed-file counts.""" + + total_tags: int + total_pages: int + page: int + range: tuple[int, int] + items: list[TagListItem] + class BaseTagIndex(BaseComponent): """A rebuildable index of normalized ``FileNode`` frontmatter tags.""" component_type = ComponentEnum.TAG_INDEX + reserved_tag_keys: ClassVar[frozenset[str]] = frozenset() - def __init__(self, **kwargs): + def __init__(self, tag_key: object = "memory_tags", **kwargs): super().__init__(**kwargs) + self._tag_key = self._validate_tag_key(tag_key) self.is_healthy = True + @property + def tag_key(self) -> str: + """Frontmatter field from which this index derives tags.""" + return self._tag_key + + @tag_key.setter + def tag_key(self, value: object) -> None: + tag_key = self._validate_tag_key(value) + if tag_key != self._tag_key: + self._tag_key = tag_key + self.set_healthy(False) + + def _validate_tag_key(self, value: object) -> str: + """Validate and normalize the configured frontmatter tag field.""" + if not isinstance(value, str) or not value.strip(): + raise ValueError("tag_key must be a non-empty string") + tag_key = value.strip() + if tag_key in self.reserved_tag_keys: + raise ValueError(f"tag_key must not be a reserved frontmatter key: {tag_key!r}") + return tag_key + def set_healthy(self, healthy: bool) -> None: """Mark whether lookups can safely use the current derived state.""" self.is_healthy = healthy @@ -29,6 +67,10 @@ class BaseTagIndex(BaseComponent): def normalize_tags(self, value: object) -> list[str]: """Return canonical tags according to this index's configured limits.""" + @abstractmethod + def normalize_query_tags(self, value: object) -> list[str]: + """Return canonical lookup tags without per-file count limits.""" + @abstractmethod async def rebuild(self, nodes: list[FileNode]) -> None: """Replace the complete index with relationships derived from ``nodes``.""" @@ -49,6 +91,17 @@ class BaseTagIndex(BaseComponent): async def tags_for_path(self, path: str) -> list[str]: """Return normalized tags for one workspace-relative path.""" + @abstractmethod + async def list_tags( + self, + *, + page: int = 1, + order_by: TagOrderBy = "tag", + order: TagOrder | None = None, + page_size: int = 100, + ) -> TagListResult: + """Return a page of active tags with counts and a 1-based result range.""" + @abstractmethod async def clear(self) -> None: """Clear memory and persisted state.""" diff --git a/reme/components/tag_index/local_tag_index.py b/reme/components/tag_index/local_tag_index.py index c3e97780..e2f4e209 100644 --- a/reme/components/tag_index/local_tag_index.py +++ b/reme/components/tag_index/local_tag_index.py @@ -3,17 +3,31 @@ import asyncio from pathlib import PurePosixPath -from .base_tag_index import BaseTagIndex +from .base_tag_index import BaseTagIndex, TagListItem, TagListResult, TagOrder, TagOrderBy from ..component_registry import R -from ...schema import FileNode +from ...schema import FileFrontMatter, FileNode @R.register("local") class LocalTagIndex(BaseTagIndex): """Maintain bidirectional path/tag relationships without separate source I/O.""" - def __init__(self, max_tags_per_file: int = 8, max_tag_length: int = 64, **kwargs): - super().__init__(**kwargs) + reserved_tag_keys = frozenset(FileFrontMatter.model_fields) | { + "kind", + "session_id", + "source_conversation", + "source_resource", + "status", + } + + def __init__( + self, + tag_key: object = "memory_tags", + max_tags_per_file: int = 3, + max_tag_length: int = 64, + **kwargs, + ): + super().__init__(tag_key=tag_key, **kwargs) self.max_tags_per_file = self._positive_int("max_tags_per_file", max_tags_per_file) self.max_tag_length = self._positive_int("max_tag_length", max_tag_length) self.path_to_tags: dict[str, tuple[str, ...]] = {} @@ -39,8 +53,8 @@ class LocalTagIndex(BaseTagIndex): for item in value: if isinstance(item, bool) or not isinstance(item, (str, int)): continue - raw = str(item).strip() - if not raw or len(raw) > self.max_tag_length or any(char.isspace() for char in raw): + raw = "_".join(str(item).split()) + if not raw or len(raw) > self.max_tag_length: continue if not any(char.isalnum() for char in raw): continue @@ -57,6 +71,10 @@ class LocalTagIndex(BaseTagIndex): """Normalize frontmatter tags according to the per-file count limit.""" return self._normalize_tags(value, limit=self.max_tags_per_file) + def normalize_query_tags(self, value: object) -> list[str]: + """Normalize lookup tags without truncating the query expression.""" + return self._normalize_tags(value, limit=None) + @staticmethod def _validate_path(path: str) -> str: if not isinstance(path, str) or not path or "\\" in path: @@ -70,7 +88,8 @@ class LocalTagIndex(BaseTagIndex): prepared: list[tuple[str, tuple[str, ...]]] = [] for node in nodes: path = self._validate_path(node.path) - tags = self.normalize_tags(node.front_matter.model_dump().get("tags")) + frontmatter = node.front_matter.model_extra or {} + tags = self.normalize_tags(frontmatter.get(self.tag_key)) prepared.append((path, tuple(tags))) return prepared @@ -127,7 +146,7 @@ class LocalTagIndex(BaseTagIndex): # ``max_tags_per_file`` constrains indexed documents, not lookup # expressions. Truncating here would silently weaken AND queries and # omit valid matches from OR queries. - normalized = self._normalize_tags(tags, limit=None) + normalized = self.normalize_query_tags(tags) if not normalized: return [] async with self._maintenance_lock: @@ -142,6 +161,57 @@ class LocalTagIndex(BaseTagIndex): async with self._maintenance_lock: return list(self.path_to_tags.get(path, ())) + async def list_tags( + self, + *, + page: int = 1, + order_by: TagOrderBy = "tag", + order: TagOrder | None = None, + page_size: int = 100, + ) -> TagListResult: + """Return a deterministic page of active tags, counts, and its 1-based range.""" + page = self._positive_int("page", page) + page_size = self._positive_int("page_size", page_size) + if page_size > 1000: + raise ValueError("page_size must be less than or equal to 1000") + + if not isinstance(order_by, str): + raise ValueError("order_by must be one of ['file_count', 'tag']") + order_by = order_by.lower() + if order_by not in {"tag", "file_count"}: + raise ValueError("order_by must be one of ['file_count', 'tag']") + if order is None: + order = "asc" if order_by == "tag" else "desc" + if not isinstance(order, str): + raise ValueError("order must be one of ['asc', 'desc']") + order = order.lower() + if order not in {"asc", "desc"}: + raise ValueError("order must be one of ['asc', 'desc']") + + async with self._maintenance_lock: + if not self.is_healthy: + raise RuntimeError("tag index is unavailable") + items: list[TagListItem] = [(tag, len(paths)) for tag, paths in self.tag_to_paths.items()] + + if order_by == "tag": + items.sort(key=lambda item: item[0], reverse=order == "desc") + else: + direction = 1 if order == "asc" else -1 + items.sort(key=lambda item: (direction * item[1], item[0])) + + total_tags = len(items) + total_pages = (total_tags + page_size - 1) // page_size + start = (page - 1) * page_size + page_items = items[start : start + page_size] + item_range = (start + 1, start + len(page_items)) if page_items else (0, 0) + return { + "total_tags": total_tags, + "total_pages": total_pages, + "page": page, + "range": item_range, + "items": page_items, + } + async def clear(self) -> None: async with self._maintenance_lock: self.path_to_tags = {} diff --git a/reme/config/default.yaml b/reme/config/default.yaml index 4b26a2f2..9f62ccb8 100644 --- a/reme/config/default.yaml +++ b/reme/config/default.yaml @@ -188,7 +188,8 @@ jobs: - messages steps: - backend: auto_memory_step - enable_tags: false + - backend: auto_tag_step + max_tags_per_file: 3 auto_memory_cc: backend: base @@ -206,6 +207,8 @@ jobs: - session_id steps: - backend: auto_memory_cc_step + - backend: auto_tag_step + max_tags_per_file: 3 auto_resource: backend: base @@ -388,6 +391,37 @@ jobs: steps: - backend: graph_snapshot_step + list_tags: + backend: base + description: >- + List active tags with file counts. Returns {total_tags, total_pages, page, range, items}; range is the inclusive + 1-based item range and each item is [tag, file_count]. A page past the end returns empty items and range [0, 0]. + parameters: + type: object + properties: + page: + type: integer + minimum: 1 + default: 1 + description: "1-based page; values past the end return an empty page" + order_by: + type: string + enum: [tag, file_count] + default: tag + description: "sort by tag or file count" + order: + type: string + enum: [asc, desc] + description: "defaults to asc for tag, desc for file_count" + page_size: + type: integer + minimum: 1 + maximum: 1000 + default: 100 + description: "items per page" + steps: + - backend: list_tags_step + reindex: backend: base description: "rebuild BM25, embedding, and/or tag indexes without rescanning workspace files" @@ -418,6 +452,12 @@ jobs: type: number description: "min fused score" default: 0.0 + tags: + type: array + items: + type: string + description: "optional tags filter; a file matches when it contains at least one valid tag" + default: [] start_date: type: string description: "optional inclusive start date filter (YYYY-MM-DD); results earlier than this date are excluded" @@ -515,7 +555,7 @@ jobs: frontmatter_update: backend: base - description: "Merge key-values into a file's frontmatter." + description: "Merge key-value pairs into a Markdown file's frontmatter while preserving its body and unrelated fields." parameters: type: object properties: @@ -896,11 +936,12 @@ components: backend: bm25 tokenizer: default -# tag_index: -# default: -# backend: local -# max_tags_per_file: 8 -# max_tag_length: 64 + tag_index: + default: + backend: local + tag_key: memory_tags + max_tags_per_file: 3 + max_tag_length: 64 file_store: default: @@ -911,4 +952,4 @@ components: embedding_store: "" keyword_index: default file_graph: default -# tag_index: default + tag_index: default diff --git a/reme/steps/evolve/__init__.py b/reme/steps/evolve/__init__.py index c32b3b9b..fe9e6dbe 100644 --- a/reme/steps/evolve/__init__.py +++ b/reme/steps/evolve/__init__.py @@ -3,6 +3,7 @@ from ._evolve import now, passthrough_response from .auto_image_resource import AutoImageResourceStep from .auto_memory import AutoMemoryStep +from .auto_tag import AutoTagStep from .auto_memory_cc import AutoMemoryCCStep from .auto_resource import AutoResourceStep from .auto_text_resource import AutoTextResourceStep @@ -22,6 +23,7 @@ __all__ = [ "AutoImageResourceStep", "passthrough_response", "AutoMemoryStep", + "AutoTagStep", "AutoMemoryCCStep", "AutoResourceStep", "AutoTextResourceStep", diff --git a/reme/steps/evolve/auto_memory.py b/reme/steps/evolve/auto_memory.py index 9961371f..8cf29a1f 100644 --- a/reme/steps/evolve/auto_memory.py +++ b/reme/steps/evolve/auto_memory.py @@ -17,9 +17,6 @@ from ...components import R _SESSION_ID_KEY = "session_id" _SOURCE_CONVERSATION_KEY = "source_conversation" -_TAGS_KEY = "tags" -_MAX_TAGS = 8 -_MAX_TAG_LENGTH = 64 _MESSAGE_TIME_ALIASES = ("time_created", "timestamp", "createdAt", "timeCreated", "created_time") @@ -62,39 +59,6 @@ def _normalize_msg_timestamp(item: dict) -> dict: return item -def _normalize_tags(value) -> list[str]: - """Return up to eight unique, retrieval-friendly frontmatter tags. - - Tags are stored uniformly as strings. Technical tokens such as ``GPT-5``, - ``C++``, ``C#``, and ``.NET`` are valid; whitespace-delimited phrases and - punctuation-only values are not. - """ - if not isinstance(value, list): - return [] - - tags: list[str] = [] - seen: set[str] = set() - for item in value: - if isinstance(item, bool) or not isinstance(item, (str, int)): - continue - tag = str(item).strip() - if not tag or len(tag) > _MAX_TAG_LENGTH: - continue - if any(char.isspace() for char in tag): - continue - if not any(char.isalnum() for char in tag): - continue - - dedupe_key = tag.casefold() - if dedupe_key in seen: - continue - seen.add(dedupe_key) - tags.append(tag) - if len(tags) >= _MAX_TAGS: - break - return tags - - @R.register("auto_memory_step") class AutoMemoryStep(BaseStep): """Record conversation facts into a daily note via an Agent.""" @@ -119,10 +83,6 @@ class AutoMemoryStep(BaseStep): def _daily_note_path(self, day: str, name: str) -> str: return f"{self.config_value('daily_dir')}/{day}/{name}.md" - def _tags_enabled(self) -> bool: - """Whether this step should generate and normalize frontmatter tags.""" - return bool(self.kwargs.get("enable_tags", False)) - def _frontmatter(self, path: str) -> dict: post = frontmatter.loads((self.file_store.workspace_path / path).read_text(encoding="utf-8")) return dict(post.metadata or {}) @@ -158,14 +118,12 @@ class AutoMemoryStep(BaseStep): notes = list_response.metadata.get("notes") or [] return self._find_session_note(notes, session_id) - async def _ensure_memory_frontmatter(self, path: str, session_id: str) -> None: - current = self._frontmatter(path) + async def _ensure_session_frontmatter(self, path: str, session_id: str) -> None: metadata = { _SESSION_ID_KEY: session_id, _SOURCE_CONVERSATION_KEY: self._session_link(session_id), } - if self._tags_enabled(): - metadata[_TAGS_KEY] = _normalize_tags(current.get(_TAGS_KEY)) + current = self._frontmatter(path) if all(current.get(key) == value for key, value in metadata.items()): return response = await self.run_job( @@ -318,6 +276,7 @@ class AutoMemoryStep(BaseStep): # pylint: disable=too-many-return-statements async def execute(self): assert self.context is not None + self.context["changes"] = [] raw_messages = self.context.get("messages") or [] session_id: str = self.context.get("session_id", "") memory_hint: str = self.context.get("memory_hint", "") @@ -381,7 +340,6 @@ class AutoMemoryStep(BaseStep): template_key = "user_message_create" if created else "user_message_update" user_message = self.prompt_format( template_key, - enable_tags=self._tags_enabled(), today=day, note=memory_hint or "(none)", note_path=note_path, @@ -399,7 +357,7 @@ class AutoMemoryStep(BaseStep): reply_kwargs["injected_job_kwargs"] = {"_allowed_paths": [note_path]} result = await self.agent_wrapper.reply( user_message, - system_prompt=self.prompt_format("system_prompt", enable_tags=self._tags_enabled()), + system_prompt=self.prompt_format("system_prompt"), job_tools=self.create_tools if created else self.update_tools, **reply_kwargs, ) @@ -425,27 +383,28 @@ class AutoMemoryStep(BaseStep): self.logger.info(f"[{self.name}] done without note session_id={session_id!r} modified=False") return note_path = str(note["path"]) - try: - if not created or self._tags_enabled(): - await self._ensure_memory_frontmatter(note_path, session_id) - if not created: + else: + try: + await self._ensure_session_frontmatter(note_path, session_id) note_path = await self._rename_from_frontmatter_name(note_path, day) - except RuntimeError as exc: - self.context.response.success = False - self.context.response.answer = str(exc) - self.context.response.metadata.update( - { - "date": day, - "path": note_path, - "created": created, - "modified": self._note_modified(before_note_path, before_note_bytes, note_path), - "n_messages": len(messages), - }, - ) - self.logger.info(f"[{self.name}] post-write failed path={note_path} answer={str(exc)!r}") - return + except RuntimeError as exc: + self.context.response.success = False + self.context.response.answer = str(exc) + self.context.response.metadata.update( + { + "date": day, + "path": note_path, + "created": created, + "modified": self._note_modified(before_note_path, before_note_bytes, note_path), + "n_messages": len(messages), + }, + ) + self.logger.info(f"[{self.name}] post-update failed path={note_path} answer={str(exc)!r}") + return modified = self._note_modified(before_note_path, before_note_bytes, note_path) + if modified: + self.context["changes"] = [{"change": "added" if created else "modified", "path": note_path}] daily_dir = self.config_value("daily_dir") self.logger.info(f"[{self.name}] refresh index start date={day} daily_dir={daily_dir}") index_payload = await refresh_day_index(self.file_store, day, daily_dir) diff --git a/reme/steps/evolve/auto_memory.yaml b/reme/steps/evolve/auto_memory.yaml index 76df7878..0767a502 100644 --- a/reme/steps/evolve/auto_memory.yaml +++ b/reme/steps/evolve/auto_memory.yaml @@ -20,7 +20,6 @@ system_prompt: | - `name` = a concise, stable topic/event filename stem, such as `cold-remedies` or `project-kickoff-decision`. Do not include today's date or the daily directory date; the outer daily path already records the date. For existing notes, update it when a better filename is clearly warranted. - `description` = a thorough summary; vague descriptions like "notes" / "misc" are unacceptable. - [enable_tags] - `tags` = 0–8 unique retrieval keywords for the complete note. Always include the field, using `[]` when no useful tags exist. Each tag must be one string with no whitespace, may contain technical punctuation (for example `GPT-5`, `C++`, `C#`, or `.NET`), must contain at least one letter or digit, and must be at most 64 characters. Store numeric tags as strings, for example `"100"`. - **Never set `status`** — it is a field reserved for downstream processing. system_prompt_zh: | 你是自动记忆系统。你的职责是将最近对话中的核心信息记录到日记记忆中。思考人类会从这段对话中自然地记住什么——不是所有内容,而是真正重要的信息。 @@ -44,7 +43,6 @@ system_prompt_zh: | - `name` = 简洁、稳定的主题/事件文件名 stem,例如 `cold-remedies` 或 `project-kickoff-decision`。不要包含今天日期或日记目录日期;外层日记路径已经记录日期。对已有笔记,如果明显有更好的文件名,就更新它。 - `description` = 详细总结;模糊的描述如 "notes" / "misc" 不可接受。 - [enable_tags] - `tags` = 针对完整笔记的 0–8 个不重复检索关键词。该字段必须始终存在;没有合适关键词时使用 `[]`。每个 tag 必须是一个不含空白字符的字符串,可以包含技术名称中的标点(例如 `GPT-5`、`C++`、`C#` 或 `.NET`),必须至少包含一个字母或数字,且不得超过 64 个字符。数字 tag 也保存为字符串,例如 `"100"`。 - **永远不要设置 `status`**——它是下游处理保留的字段。 user_message_create: | @@ -70,13 +68,11 @@ user_message_create: | Create the note in one shot: `daily_write name= description= session_id={session_id} date={today} content=` - [enable_tags] Also pass `metadata={{"tags": [, ...]}}` in the same `daily_write` call. - Generate `name` as a concise, stable topic/event filename stem for this memory. Prefer a reusable topic or event summary, optionally in kebab-case. - Do not include today's date or the daily directory date in `name`; the note already lives under today's daily path. - `name` must be a valid single filename component: no slash, backslash, leading/trailing whitespace, or characters like `< > : " | ? *`. - `description` must be a thorough summary of the body — specific enough that the description alone conveys all key information. - [enable_tags] - Generate 0–8 `tags` from the complete note. Always pass `metadata={{"tags": [...]}}`, including an empty list when there are no useful tags. Follow the system prompt's tag format exactly. ## Step 3 — Summary @@ -108,13 +104,11 @@ user_message_create_zh: | 一次性创建笔记: `daily_write name= description= session_id={session_id} date={today} content=<正文>` - [enable_tags] 在同一次 `daily_write` 调用中另外传入 `metadata={{"tags": [, ...]}}`。 - 由你生成 `name`,作为这条记忆简洁、稳定的主题/事件文件名 stem。优先使用可复用的主题或事件总结,可以采用 kebab-case。 - `name` 不要包含今天日期或日记目录日期;笔记已经位于当天日记路径下。 - `name` 必须是合法的单个文件名组件:不能包含 slash、反斜杠、首尾空白,或 `< > : " | ? *` 等字符。 - `description` 必须是正文的详尽总结——具体到仅凭 description 就能传达全部核心信息。 - [enable_tags] - 根据完整笔记生成 0–8 个 `tags`。必须始终传入 `metadata={{"tags": [...]}}`,没有合适关键词时也要传入空数组。严格遵守 system prompt 的 tag 格式。 ## 步骤 3 — 总结 @@ -161,22 +155,18 @@ user_message_update: | Execution: 1. Use `edit path={note_path} old= new=` for each section that needs updating. You may call `edit` multiple times. 2. After body changes, refresh frontmatter with `frontmatter_update path={note_path} metadata={{"name": "", "description": ""}}`. - [enable_tags] Also regenerate 0–8 tags from the complete merged note and include `"tags": [, ...]` in the same metadata object. Always include `tags`, using `[]` when none are useful. - Keep the existing `name` only when it is already the best concise topic/event filename stem. The system will rename the file after your final response. - Do not add today's date or the daily directory date to `name`. 3. If `edit` fails repeatedly (e.g., cannot find the original text due to formatting mismatch), fall back to `write path={note_path} name= description= content=` for a complete rewrite. - [enable_tags] In this fallback, also pass `metadata={{"tags": [, ...]}}`. ## Step 3b — Full Write (Empty File Fallback) The file exists but its body is empty. Write the full content in one shot: `write path={note_path} name= description= content=` - [enable_tags] Also pass `metadata={{"tags": [, ...]}}` in this call. - Use a concise, stable topic/event `name`; filename changes are applied after your final response. - Do not include today's date or the daily directory date in `name`. - `description` must be a thorough summary of the body — specific enough that the description alone conveys all key information. - [enable_tags] - Generate 0–8 tags from the complete body and always include `metadata={{"tags": [...]}}`, using `[]` when none are useful. ## Step 4 — Summary @@ -224,22 +214,18 @@ user_message_update_zh: | 执行: 1. 对需要更新的每个部分使用 `edit path={note_path} old=<原文片段> new=<替换片段>`。可以多次调用 `edit`。 2. 正文变更后,刷新 frontmatter:`frontmatter_update path={note_path} metadata={{"name": "<更新后的文件名 stem>", "description": "<更新后的总结>"}}`。 - [enable_tags] 同时根据合并后的完整笔记重新生成 0–8 个 tags,并在同一个 metadata 对象中包含 `"tags": [, ...]`。必须始终包含 `tags`,没有合适关键词时使用 `[]`。 - 只有当前 `name` 已经是最合适的简洁主题/事件文件名 stem 时才保留。系统会在你最终回复后负责重命名文件。 - 不要在 `name` 中加入今天日期或日记目录日期。 3. 如果 `edit` 多次失败(如因格式不匹配找不到原文),退回 `write path={note_path} name= description= content=<完整正文>` 全量重写。 - [enable_tags] 在该 fallback 中还要传入 `metadata={{"tags": [, ...]}}`。 ## 步骤 3b — 全量写入(空文件 fallback) 文件存在但正文为空。一次性写入完整内容: `write path={note_path} name= description= content=<正文>` - [enable_tags] 在该调用中另外传入 `metadata={{"tags": [, ...]}}`。 - 使用简洁、稳定的主题/事件 `name`;文件名变化会在你最终回复后应用。 - `name` 不要包含今天日期或日记目录日期。 - `description` 必须是正文的详尽总结——具体到仅凭 description 就能传达全部核心信息。 - [enable_tags] - 根据完整正文生成 0–8 个 tags,并始终包含 `metadata={{"tags": [...]}}`;没有合适关键词时使用 `[]`。 ## 步骤 4 — 总结 diff --git a/reme/steps/evolve/auto_tag.py b/reme/steps/evolve/auto_tag.py new file mode 100644 index 00000000..19fe0dd5 --- /dev/null +++ b/reme/steps/evolve/auto_tag.py @@ -0,0 +1,235 @@ +"""Generate entity-oriented memory tags for added or modified Markdown files.""" + +from dataclasses import dataclass +from pathlib import Path +from typing import Literal + +import frontmatter + +from ._evolve import agent_reply_result_text +from ..base_step import BaseStep +from ..file_io import parse_daily_date, refresh_day_index +from ..file_io._path import display_path, resolve_path +from ..index import normalize_posix_path +from ...components import R + +_DEFAULT_MAX_MEMORY_TAGS = 3 +_DEFAULT_MAX_MEMORY_TAG_LENGTH = 64 +_SUPPORTED_CHANGES = {"added", "modified"} + + +@dataclass(frozen=True) +class _TagTarget: + change: Literal["added", "modified"] + path: str + + +def _positive_int(value: object, *, name: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError(f"{name} must be a positive integer") + return value + + +def normalize_memory_tags( + value: object, + *, + max_tags_per_file: int = _DEFAULT_MAX_MEMORY_TAGS, + max_tag_length: int = _DEFAULT_MAX_MEMORY_TAG_LENGTH, +) -> list[str]: + """Normalize human-readable entity labels for frontmatter storage.""" + max_tags_per_file = _positive_int(max_tags_per_file, name="max_tags_per_file") + max_tag_length = _positive_int(max_tag_length, name="max_tag_length") + if not isinstance(value, list): + return [] + + tags: list[str] = [] + seen: set[str] = set() + for item in value: + if not isinstance(item, str): + continue + tag = "_".join(item.split()) + if not tag or len(tag) > max_tag_length or not any(char.isalnum() for char in tag): + continue + canonical = tag.casefold() + if canonical in seen: + continue + seen.add(canonical) + tags.append(tag) + if len(tags) >= max_tags_per_file: + break + return tags + + +@R.register("auto_tag_step") +class AutoTagStep(BaseStep): + """Update memory tags for Markdown files described by the common ``changes`` contract.""" + + def __init__(self, max_tags_per_file: int = _DEFAULT_MAX_MEMORY_TAGS, **kwargs): + super().__init__(**kwargs) + self.max_tags_per_file = _positive_int(max_tags_per_file, name="max_tags_per_file") + self.tools = ["read", "list_tags", "frontmatter_read", "frontmatter_update"] + + @staticmethod + def _index_limit(tag_index, name: str, fallback: int | None) -> int | None: + """Read one positive integer index limit, falling back when unavailable or invalid.""" + try: + value = getattr(tag_index, name) + except (AttributeError, TypeError, ValueError): + return fallback + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + return fallback + return value + + def _targets(self) -> tuple[list[_TagTarget], list[dict[str, str]]]: + """Validate, normalize, and de-duplicate added/modified Markdown changes.""" + assert self.context is not None + raw_changes = self.context.get("changes") or [] + if not isinstance(raw_changes, list): + raise ValueError("AutoTagStep requires changes: list[dict]") + + workspace = Path(self.file_store.workspace_path or ".").resolve() + targets: dict[str, _TagTarget] = {} + ignored: list[dict[str, str]] = [] + for item in raw_changes: + if not isinstance(item, dict): + ignored.append({"path": "", "reason": "change must be an object"}) + continue + + change = str(item.get("change") or "").strip().lower() + raw_path = str(item.get("path") or "").strip() + if change not in _SUPPORTED_CHANGES: + ignored.append({"path": raw_path, "reason": f"unsupported change: {change or 'missing'}"}) + continue + + target, error = resolve_path(workspace, raw_path) + if error or target is None: + ignored.append({"path": raw_path, "reason": error or "invalid path"}) + continue + if not target.is_file(): + ignored.append({"path": raw_path, "reason": "not a file"}) + continue + if target.suffix.lower() != ".md": + ignored.append({"path": raw_path, "reason": "not a Markdown file"}) + continue + + path = normalize_posix_path(display_path(workspace, target)) + previous = targets.get(path) + was_added = previous is not None and previous.change == "added" + normalized_change: Literal["added", "modified"] = "added" if change == "added" or was_added else "modified" + targets[path] = _TagTarget(change=normalized_change, path=path) + return list(targets.values()), ignored + + async def _process_target(self, target: _TagTarget, tag_key: str, max_tag_length: int) -> str: + result = await self.agent_wrapper.reply( + self.prompt_format("user_message", path=target.path, change=target.change, tag_key=tag_key), + system_prompt=self.prompt_format( + "system_prompt", + tag_key=tag_key, + max_tags_per_file=self.max_tags_per_file, + ), + job_tools=self.tools, + injected_job_kwargs={ + "_allowed_paths": [target.path], + "_allowed_frontmatter_keys": [tag_key], + }, + ) + + path = Path(self.file_store.workspace_path or ".") / target.path + metadata = dict(frontmatter.loads(path.read_text(encoding="utf-8")).metadata or {}) + normalized = normalize_memory_tags( + metadata.get(tag_key), + max_tags_per_file=self.max_tags_per_file, + max_tag_length=max_tag_length, + ) + if metadata.get(tag_key) != normalized: + response = await self.run_job( + "frontmatter_update", + path=target.path, + metadata={tag_key: normalized}, + _allowed_paths=[target.path], + _allowed_frontmatter_keys=[tag_key], + ) + if not response.success: + raise RuntimeError(str(response.answer)) + return agent_reply_result_text(result) + + def _daily_date(self, path: str) -> str | None: + daily_dir = normalize_posix_path(str(self.config_value("daily_dir"))).strip("/") + prefix = f"{daily_dir}/" + if not path.startswith(prefix): + return None + parts = path[len(prefix) :].split("/") + return parse_daily_date(parts[0]) if len(parts) == 2 else None + + async def execute(self): + assert self.context is not None + initial_success = self.context.response.success + initial_answer = self.context.response.answer + try: + targets, ignored = self._targets() + except ValueError as exc: + self.context.response.success = False + self.context.response.answer = str(exc) + return self.context.response + + results: list[dict] = [] + indexes: list[dict] = [] + if targets and not self.file_store.tag_index_enabled: + self.context.response.success = False + if initial_success: + self.context.response.answer = "Error: tag index is not configured" + return self.context.response + + tag_index = self.file_store.require_tag_index() if targets else None + if tag_index is not None and not tag_index.is_healthy: + self.context.response.success = False + if initial_success: + self.context.response.answer = "Error: tag index unavailable" + return self.context.response + max_tag_length = self._index_limit(tag_index, "max_tag_length", _DEFAULT_MAX_MEMORY_TAG_LENGTH) + index_max_tags = self._index_limit(tag_index, "max_tags_per_file", None) + if index_max_tags is not None and self.max_tags_per_file > index_max_tags: + self.context.response.success = False + if initial_success: + self.context.response.answer = ( + f"Error: auto_tag max_tags_per_file ({self.max_tags_per_file}) exceeds " + f"tag index limit ({index_max_tags})" + ) + return self.context.response + + dates: set[str] = set() + for target in targets: + if day := self._daily_date(target.path): + dates.add(day) + try: + summary = await self._process_target(target, tag_index.tag_key, max_tag_length) + results.append( + {"change": target.change, "path": target.path, "success": True, "summary": summary}, + ) + except Exception as exc: # pylint: disable=broad-exception-caught + results.append( + {"change": target.change, "path": target.path, "success": False, "error": str(exc)}, + ) + self.logger.warning(f"[{self.name}] failed path={target.path}: {exc}") + + for day in sorted(dates): + indexes.append(await refresh_day_index(self.file_store, day, self.config_value("daily_dir"))) + + failed = sum(not item["success"] for item in results) + succeeded = len(results) - failed + self.context.response.success = initial_success and failed == 0 + if initial_success and failed: + self.context.response.answer = f"Tagged {succeeded} file(s); {failed} failed" + elif initial_success and not initial_answer and succeeded: + self.context.response.answer = f"Tagged {succeeded} file(s)" + else: + self.context.response.answer = initial_answer + self.context.response.metadata["auto_tag"] = { + "processed": len(results), + "succeeded": succeeded, + "failed": failed, + "ignored": ignored, + "results": results, + "indexes": indexes, + } + return self.context.response diff --git a/reme/steps/evolve/auto_tag.yaml b/reme/steps/evolve/auto_tag.yaml new file mode 100644 index 00000000..0b3f5e45 --- /dev/null +++ b/reme/steps/evolve/auto_tag.yaml @@ -0,0 +1,45 @@ +system_prompt: | + You maintain entity-oriented memory tags for one Markdown memory file at a time. + + The `{tag_key}` field answers: "Who or what real-world entity is this memory about?" + + Required workflow: + 1. Use `read` and `frontmatter_read` to inspect the complete target document and its current frontmatter. + 2. Use `list_tags` to inspect tags already used in the workspace. Check additional pages when useful. + 3. Prefer one primary entity that best answers the question. Use multiple tags only when the memory is genuinely + about multiple independent central entities, and never use more than {max_tags_per_file}. Prefer an existing tag + whenever it represents the same entity. Tags must not contain spaces; replace spaces with underscores, for example + `sam_altman`. Examples include a person (`sam_altman`), an organization (`OpenAI`), or an asset/entity (`gold`). + 4. Use `frontmatter_update` to write the complete `{tag_key}` list, including `[]` when no entity is central. + + Choose entities, not broad topics, activities, attributes, dates, or generic keywords. Keep the canonical real-world + name and use one concise string per entity. Store multiple entities as separate YAML list items, for example + `[OpenAI, gold]`; never combine them into one comma-delimited string. Do not create aliases or near-duplicate + spellings for an existing tag. Update only `{tag_key}`; never change the document body or another frontmatter field. +system_prompt_zh: | + 你负责为每一份 Markdown 记忆维护以实体为中心的记忆标签。 + + `{tag_key}` 字段回答的问题是:“这份记忆是关于谁或什么现实实体的?” + + 必须遵循以下流程: + 1. 使用 `read` 和 `frontmatter_read` 阅读目标文档的完整内容及当前 frontmatter。 + 2. 使用 `list_tags` 查看工作区已经使用的标签;必要时继续查看后续分页。 + 3. 默认只选择一个最能回答上述问题的主实体。只有当记忆确实同时围绕多个相互独立的核心实体时,才选择 + 多个标签,并且绝不能超过 {max_tags_per_file} 个。同一实体已有标签时必须优先复用。标签中不能包含空格; + 如有空格,必须替换为下划线,例如将 `Sam Altman` 写成 `sam_altman`。实体可以是人物(如 `sam_altman`)、 + 组织或公司(如 `宁德时代`),也可以是资产等实体(如 `黄金`)。 + 4. 使用 `frontmatter_update` 写入完整的 `{tag_key}` 列表;没有核心实体时也要写入 `[]`。 + + 只选择实体,不要使用宽泛主题、行为、属性、日期或普通关键词。每个实体使用一个简洁、规范的现实名称。 + 多个实体必须保存为独立的 YAML 列表项,例如 `[宁德时代, 黄金]`;不要把它们拼成一个逗号分隔的字符串。 + 不要为已有实体创造别名或近似拼写。只能更新 `{tag_key}`,绝不能修改正文或其他 frontmatter 字段。 +user_message: | + Change: {change} + Target path: {path} + + Follow the required workflow and update this document's `{tag_key}` field. +user_message_zh: | + 变更类型:{change} + 目标路径:{path} + + 按照规定流程更新这份文档的 `{tag_key}` 字段。 diff --git a/reme/steps/evolve/dream/utils.py b/reme/steps/evolve/dream/utils.py index 7701d005..2f41d27d 100644 --- a/reme/steps/evolve/dream/utils.py +++ b/reme/steps/evolve/dream/utils.py @@ -31,8 +31,7 @@ def store_state(step: BaseStep, state: DreamState) -> None: def workspace_dir(step: BaseStep) -> Path: """Get workspace directory.""" - vr = getattr(step.file_store, "workspace_path", None) - return Path(vr).resolve() if vr else Path.cwd().resolve() + return step.file_store.workspace_path.resolve() def daily_dir(step: BaseStep) -> str: diff --git a/reme/steps/file_io/frontmatter_update.py b/reme/steps/file_io/frontmatter_update.py index cc4b8f2d..e8801c13 100644 --- a/reme/steps/file_io/frontmatter_update.py +++ b/reme/steps/file_io/frontmatter_update.py @@ -30,6 +30,10 @@ class FrontmatterUpdateStep(BaseStep): injected by the server into the RuntimeContext; without it, restricting read/edit/write alone would still leave frontmatter of arbitrary workspace Markdown files mutable. + + ``_allowed_frontmatter_keys`` optionally limits updates to an injected + list of top-level keys. Omitting it (or setting it to ``None``) preserves + the unrestricted historical behavior. """ async def execute(self): @@ -39,6 +43,18 @@ class FrontmatterUpdateStep(BaseStep): metadata = self.context.get("metadata") or {} assert isinstance(metadata, dict), "metadata must be a dict" + allowed_keys = self.context.get("_allowed_frontmatter_keys") + key_error: str | None = None + if allowed_keys is not None: + if not isinstance(allowed_keys, list) or any(not isinstance(key, str) for key in allowed_keys): + key_error = "_allowed_frontmatter_keys must be a list of strings" + else: + allowed = set(allowed_keys) + denied = [key for key in metadata if not isinstance(key, str) or key not in allowed] + if denied: + names = ", ".join(sorted(repr(key) for key in denied)) + key_error = f"frontmatter key(s) not allowed: {names}" + workspace_dir = Path(self.file_store.workspace_path or ".").resolve() target, err = resolve_path(workspace_dir, path) if err or target is None: @@ -57,6 +73,8 @@ class FrontmatterUpdateStep(BaseStep): payload = {"path": path, "error": "not markdown"} elif not metadata: payload = {"path": path, "error": "no fields to update"} + elif key_error: + payload = {"path": path, "error": key_error} else: post = frontmatter.loads(target.read_text(encoding="utf-8")) post.metadata.update(metadata) diff --git a/reme/steps/index/__init__.py b/reme/steps/index/__init__.py index b5dbb59e..2780b251 100644 --- a/reme/steps/index/__init__.py +++ b/reme/steps/index/__init__.py @@ -7,6 +7,7 @@ from .clear_store import ClearStoreStep from .draft import AddDraftStep, ReadAllDraftStep from .graph_snapshot import GraphSnapshotStep from .log_changes import LogChangesStep +from .list_tags import ListTagsStep from .node_search import NodeSearchStep from .init_changes import InitChangesStep from .optimize_index import OptimizeIndexStep @@ -35,6 +36,7 @@ __all__ = [ "GraphSnapshotStep", "InitChangesStep", "LogChangesStep", + "ListTagsStep", "NodeSearchStep", "normalize_posix_path", "ReadAllDraftStep", diff --git a/reme/steps/index/list_tags.py b/reme/steps/index/list_tags.py new file mode 100644 index 00000000..aaa172d9 --- /dev/null +++ b/reme/steps/index/list_tags.py @@ -0,0 +1,22 @@ +"""Paginated listing of active file tags.""" + +from ..base_step import BaseStep +from ...components import R + + +@R.register("list_tags_step") +class ListTagsStep(BaseStep): + """Return one compact page of active tags and their indexed file counts.""" + + async def execute(self): + assert self.context is not None + tag_index = self.file_store.require_tag_index() + + result = await tag_index.list_tags( + page=self.context.get("page", 1), + order_by=self.context.get("order_by", "tag"), + order=self.context.get("order"), + page_size=self.context.get("page_size", 100), + ) + self.context.response.answer = result + return self.context.response diff --git a/reme/steps/index/search.py b/reme/steps/index/search.py index 9fee48a5..a7bc42a9 100644 --- a/reme/steps/index/search.py +++ b/reme/steps/index/search.py @@ -27,6 +27,13 @@ def _default_limit() -> int: return _DEFAULT_LIMIT +def _filter_values(value: object) -> set: + """Normalize a scalar or collection-valued search filter.""" + if isinstance(value, (list, tuple, set, frozenset)): + return set(value) + return {value} + + @R.register("search_step") class SearchStep(BaseStep): """Hybrid search: run vector + keyword in parallel, fuse via RRF, filter, truncate.""" @@ -134,6 +141,66 @@ class SearchStep(BaseStep): "ttl_seconds": ttl, } + async def _resolve_tag_filter( + self, + raw_tags: object, + search_filter: dict, + ) -> tuple[dict, dict | None, str | None]: + """Merge tag-derived paths into the ordinary file-store filter.""" + if not raw_tags: + return search_filter, None, None + + if not self.file_store.tag_index_enabled: + self.logger.warning( + f"[{self.name}] tags requested but tag_index_unavailable", + ) + return ( + search_filter, + {"requested": True, "applied": False, "reason": "tag_index_unavailable"}, + None, + ) + tag_index = self.file_store.require_tag_index() + if not tag_index.is_healthy: + self.logger.warning(f"[{self.name}] tags requested but tag_index_unavailable") + return ( + search_filter, + {"requested": True, "applied": False, "reason": "tag_index_unavailable"}, + None, + ) + + normalized_tags = tag_index.normalize_query_tags(raw_tags) + if not normalized_tags: + return search_filter, None, "Error: tags contained no valid values" + + allowed_paths = set(await tag_index.paths_for_tags(normalized_tags, match_all=False)) + matched_path_count = len(allowed_paths) + if not tag_index.is_healthy: + self.logger.warning( + f"[{self.name}] tag index became unhealthy during lookup", + ) + return ( + search_filter, + {"requested": True, "applied": False, "reason": "tag_index_unavailable"}, + None, + ) + + exact_paths = set() + has_exact_path_filter = False + for key in ("path", "paths"): + if key in search_filter: + has_exact_path_filter = True + exact_paths.update(_filter_values(search_filter.pop(key))) + if has_exact_path_filter: + allowed_paths.intersection_update(exact_paths) + search_filter["paths"] = sorted(allowed_paths) + metadata = { + "requested": True, + "applied": True, + "tags": normalized_tags, + "matched_paths": matched_path_count, + } + return search_filter, metadata, None + async def execute(self): assert self.context is not None query: str = (self.context.get("query", "") or "").strip() @@ -170,6 +237,7 @@ class SearchStep(BaseStep): candidates = min(_MAX_CANDIDATES, max(1, int(limit * candidate_multiplier))) search_filter: dict = dict(self.context.get("search_filter", {}) or {}) + raw_tags = self.context.get("tags", []) or [] # Promote top-level date parameters into search_filter for file_store. for date_key in ("start_date", "end_date"): @@ -208,6 +276,14 @@ class SearchStep(BaseStep): if strict_date_filter: search_filter["strict_date_filter"] = True + search_filter, tag_filter_metadata, tag_error = await self._resolve_tag_filter(raw_tags, search_filter) + if tag_error is not None: + self.context.response.success = False + self.context.response.answer = tag_error + if tag_filter_metadata is not None: + self.context.response.metadata["tag_filter"] = tag_filter_metadata + return self.context.response + text_weight = 1.0 - vector_weight use_vector = vector_weight > 0.0 use_keyword = text_weight > 0.0 @@ -272,6 +348,8 @@ class SearchStep(BaseStep): "returned": len(fused), "hybrid": hybrid, } + if tag_filter_metadata is not None: + self.context.response.metadata["tag_filter"] = tag_filter_metadata if dedup is not None: self.context.response.metadata["dedup"] = dedup return self.context.response diff --git a/reme/steps/index/update_changes.py b/reme/steps/index/update_changes.py index 72c7f6c9..6ce461f0 100644 --- a/reme/steps/index/update_changes.py +++ b/reme/steps/index/update_changes.py @@ -351,13 +351,7 @@ class UpdateIndexStep(ChangeApplyStep): return self._estimate_index_memory(size_bytes, len(item[1])) def _estimate_index_memory(self, size_bytes: int, chunk_count: int) -> int: - embedding_bytes = 0 - embedding_store = getattr(self.file_store, "embedding_store", None) - if embedding_store is not None: - try: - embedding_bytes = max(0, int(embedding_store.dimensions)) * self.float16_bytes - except (AttributeError, TypeError, ValueError): - embedding_bytes = 0 + embedding_bytes = self.file_store.embedding_dimensions * self.float16_bytes expanded_content = int(size_bytes * self.batch_memory_expansion_factor) per_chunk = self.chunk_memory_overhead_bytes + embedding_bytes return expanded_content + self.file_memory_overhead_bytes + max(0, chunk_count) * per_chunk diff --git a/tests/unit/test_auto_tag.py b/tests/unit/test_auto_tag.py new file mode 100644 index 00000000..2255a816 --- /dev/null +++ b/tests/unit/test_auto_tag.py @@ -0,0 +1,212 @@ +"""Focused tests for the standalone automatic tagging Step.""" + +# pylint: disable=missing-function-docstring + +from pathlib import Path + +import frontmatter +import pytest + +from reme.components.agent_wrapper import BaseAgentWrapper +from reme.components.file_store import LocalFileStore +from reme.components.runtime_context import RuntimeContext +from reme.components.tag_index import LocalTagIndex +from reme.schema import Response +from reme.steps.evolve.auto_tag import AutoTagStep, normalize_memory_tags + + +class _TaggingWrapper(BaseAgentWrapper): + def __init__( + self, + workspace: Path, + *, + fail_name: str = "", + tag_key: str = "memory_tags", + tags: list[object] | None = None, + ) -> None: + super().__init__(name="tagger") + self.workspace = workspace + self.fail_name = fail_name + self.tag_key = tag_key + self.tags = ["宁德时代", "黄金"] if tags is None else tags + self.calls: list[tuple[str, dict]] = [] + + async def reply(self, inputs, **kwargs) -> dict: + text = str(inputs) + self.calls.append((text, kwargs)) + path = kwargs["injected_job_kwargs"]["_allowed_paths"][0] + if Path(path).name == self.fail_name: + raise RuntimeError("tagging failed") + target = self.workspace / path + post = frontmatter.loads(target.read_text(encoding="utf-8")) + post.metadata[self.tag_key] = self.tags + target.write_text(frontmatter.dumps(post), encoding="utf-8") + return {"result": f"tagged {path}", "last_message": {}} + + +def _write_note(path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("---\nname: note\ndescription: useful note\n---\nbody\n", encoding="utf-8") + + +@pytest.mark.asyncio +async def test_auto_tag_handles_noop_and_rejects_invalid_preconditions(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + wrapper = _TaggingWrapper(tmp_path) + unindexed_store = LocalFileStore(name="store", embedding_store="", tag_index="") + + context = RuntimeContext(changes=[]) + context.response.answer = "Skipped: no messages" + response = await AutoTagStep(file_store=unindexed_store, agent_wrapper=wrapper)(context) + assert response.success is True + assert response.answer == "Skipped: no messages" + assert response.metadata["auto_tag"]["processed"] == 0 + + response = await AutoTagStep(file_store=unindexed_store, agent_wrapper=wrapper)( + RuntimeContext(changes="daily/note.md"), + ) + assert response.success is False + assert response.answer == "AutoTagStep requires changes: list[dict]" + + note = tmp_path / "daily/2026-09-09/note.md" + _write_note(note) + before = note.read_bytes() + change = RuntimeContext(changes=[{"change": "added", "path": "daily/2026-09-09/note.md"}]) + response = await AutoTagStep(file_store=unindexed_store, agent_wrapper=wrapper)(change) + + assert response.success is False + assert response.answer == "Error: tag index is not configured" + assert not wrapper.calls + assert note.read_bytes() == before + + indexed_store = LocalFileStore(name="store", embedding_store="", tag_index="") + indexed_store.tag_index = LocalTagIndex(max_tags_per_file=2) + response = await AutoTagStep( + file_store=indexed_store, + agent_wrapper=wrapper, + max_tags_per_file=3, + )(RuntimeContext(changes=[{"change": "added", "path": "daily/2026-09-09/note.md"}])) + assert response.success is False + assert response.answer == "Error: auto_tag max_tags_per_file (3) exceeds tag index limit (2)" + assert not wrapper.calls + assert note.read_bytes() == before + + indexed_store.tag_index = LocalTagIndex() + indexed_store.tag_index.set_healthy(False) + response = await AutoTagStep(file_store=indexed_store, agent_wrapper=wrapper)( + RuntimeContext(changes=[{"change": "added", "path": "daily/2026-09-09/note.md"}]), + ) + assert response.success is False + assert response.answer == "Error: tag index unavailable" + assert not wrapper.calls + assert note.read_bytes() == before + + +@pytest.mark.asyncio +async def test_auto_tag_filters_paths_and_continues_after_one_file_fails(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + first = tmp_path / "daily/2026-09-09/first.md" + failed = tmp_path / "daily/2026-09-09/failed.md" + _write_note(first) + _write_note(failed) + (tmp_path / "daily/2026-09-09/notes").mkdir() + (tmp_path / "daily/2026-09-09/plain.txt").write_text("text", encoding="utf-8") + + store = LocalFileStore(name="store", embedding_store="", tag_index="") + store.tag_index = LocalTagIndex() + wrapper = _TaggingWrapper(tmp_path, fail_name="failed.md") + step = AutoTagStep(file_store=store, agent_wrapper=wrapper) + context = RuntimeContext( + changes=[ + {"change": "modified", "path": "daily/2026-09-09/failed.md"}, + {"change": "added", "path": "daily/2026-09-09/notes"}, + {"change": "added", "path": "daily/2026-09-09/plain.txt"}, + {"change": "modified", "path": "daily/2026-09-09/first.md"}, + {"change": "added", "path": "daily/2026-09-09/first.md"}, + {"change": "deleted", "path": "daily/2026-09-09/deleted.md"}, + ], + ) + + response = await step(context) + + assert response.success is False + assert [call[1]["injected_job_kwargs"] for call in wrapper.calls] == [ + { + "_allowed_paths": ["daily/2026-09-09/failed.md"], + "_allowed_frontmatter_keys": ["memory_tags"], + }, + { + "_allowed_paths": ["daily/2026-09-09/first.md"], + "_allowed_frontmatter_keys": ["memory_tags"], + }, + ] + assert all( + call[1]["job_tools"] == ["read", "list_tags", "frontmatter_read", "frontmatter_update"] + for call in wrapper.calls + ) + assert frontmatter.loads(first.read_text(encoding="utf-8")).metadata["memory_tags"] == ["宁德时代", "黄金"] + assert "memory_tags" not in frontmatter.loads(failed.read_text(encoding="utf-8")).metadata + assert response.metadata["auto_tag"]["processed"] == 2 + assert response.metadata["auto_tag"]["succeeded"] == 1 + assert response.metadata["auto_tag"]["failed"] == 1 + assert response.metadata["auto_tag"]["ignored"] == [ + {"path": "daily/2026-09-09/notes", "reason": "not a file"}, + {"path": "daily/2026-09-09/plain.txt", "reason": "not a Markdown file"}, + {"path": "daily/2026-09-09/deleted.md", "reason": "unsupported change: deleted"}, + ] + assert response.metadata["auto_tag"]["results"] == [ + { + "change": "modified", + "path": "daily/2026-09-09/failed.md", + "success": False, + "error": "tagging failed", + }, + { + "change": "added", + "path": "daily/2026-09-09/first.md", + "success": True, + "summary": "tagged daily/2026-09-09/first.md", + }, + ] + assert "memory_tags: ['宁德时代', '黄金']" in (tmp_path / "daily/2026-09-09.md").read_text(encoding="utf-8") + + +@pytest.mark.asyncio +async def test_auto_tag_uses_configured_key_and_normalizes_agent_output(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + note = tmp_path / "memory/note.md" + _write_note(note) + store = LocalFileStore(name="store", embedding_store="", tag_index="") + store.tag_index = LocalTagIndex(tag_key="keywords", max_tag_length=8) + wrapper = _TaggingWrapper( + tmp_path, + tag_key="keywords", + tags=["OpenAI", "openai", "Sam Altman", "++", 100, "宁德时代", "黄金"], + ) + step = AutoTagStep(file_store=store, agent_wrapper=wrapper, max_tags_per_file=2) + + async def update_frontmatter(name, /, **kwargs): + assert name == "frontmatter_update" + assert kwargs["_allowed_frontmatter_keys"] == ["keywords"] + post = frontmatter.loads(note.read_text(encoding="utf-8")) + post.metadata.update(kwargs["metadata"]) + note.write_text(frontmatter.dumps(post), encoding="utf-8") + return Response(answer="updated") + + monkeypatch.setattr(step, "run_job", update_frontmatter) + + context = RuntimeContext(changes=[{"change": "modified", "path": "memory/note.md"}]) + context.response.answer = "Created memory/note.md" + response = await step(context) + + assert response.success is True + assert response.answer == "Created memory/note.md" + assert frontmatter.loads(note.read_text(encoding="utf-8")).metadata["keywords"] == [ + "OpenAI", + "宁德时代", + ] + assert normalize_memory_tags( + ["one", "two", "three"], + max_tags_per_file=2, + max_tag_length=3, + ) == ["one", "two"] diff --git a/tests/unit/test_background_steps.py b/tests/unit/test_background_steps.py index 6de14302..9cdbfe22 100644 --- a/tests/unit/test_background_steps.py +++ b/tests/unit/test_background_steps.py @@ -837,6 +837,7 @@ def test_update_catalog_yields_to_event_loop_while_building_batch(): class _CountingEmbeddingStore: dimensions = 2 max_batch_size = 10 + is_healthy = True def __init__(self): self.calls = 0 @@ -1997,7 +1998,7 @@ def test_auto_memory_reports_modified_for_create_and_false_for_skip(): with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): cwd = Path.cwd() app_ctx = _make_app_context(cwd) - fs = LocalFileStore(name="test_store", embedding_store="") + fs = LocalFileStore(name="test_store", embedding_store="", tag_index="default") wrapper = _FakeAgentWrapper() await fs.start() _install_file_jobs(app_ctx, fs) @@ -2012,28 +2013,29 @@ def test_auto_memory_reports_modified_for_create_and_false_for_skip(): app_context=app_ctx, file_store=fs, agent_wrapper=wrapper, - enable_tags=True, ) - resp = await step( - RuntimeContext( - messages=[{"name": "user", "role": "user", "content": "remember project detail"}], - session_id="s1", - ), + context = RuntimeContext( + messages=[{"name": "user", "role": "user", "content": "remember project detail"}], + session_id="s1", ) + resp = await step(context) resp = resp or step.context.response assert resp.success is True assert resp.metadata["created"] is True assert resp.metadata["modified"] is True - assert "tags: []" in (cwd / "daily" / today / "memory.md").read_text(encoding="utf-8") + assert context["changes"] == [{"change": "added", "path": f"daily/{today}/memory.md"}] + assert "tags:" not in (cwd / "daily" / today / "memory.md").read_text(encoding="utf-8") wrapper.on_reply = None - resp = await step(RuntimeContext(messages=[], session_id="s2")) + context = RuntimeContext(messages=[], session_id="s2") + resp = await step(context) resp = resp or step.context.response assert resp.success is True assert resp.metadata["modified"] is False assert resp.metadata["n_messages"] == 0 + assert context["changes"] == [] finally: await fs.close() print("✓ test_auto_memory_reports_modified_for_create_and_false_for_skip passed") @@ -2041,58 +2043,6 @@ def test_auto_memory_reports_modified_for_create_and_false_for_skip(): asyncio.run(run()) -def test_auto_memory_normalizes_tags_after_existing_note_update(): - """Existing notes receive a refreshed, normalized tags field capped at eight entries.""" - - async def run(): - with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): - cwd = Path.cwd() - app_ctx = _make_app_context(cwd) - fs = LocalFileStore(name="test_store", embedding_store="") - wrapper = _FakeAgentWrapper() - await fs.start() - _install_file_jobs(app_ctx, fs) - try: - today = datetime.datetime.now().strftime("%Y-%m-%d") - note_path = cwd / "daily" / today / "memory.md" - write_file( - note_path, - "---\nname: memory\nsession_id: s1\n" - "source_conversation: '[[session/dialog/s1.jsonl]]'\n" - "tags: [old]\n---\nold body\n", - ) - - wrapper.on_reply = lambda *_: write_file( - note_path, - "---\nname: memory\nsession_id: s1\n" - "source_conversation: '[[session/dialog/s1.jsonl]]'\n" - "tags: [GPT-5, C++, C#, .NET, 100, 'memory system', '++', ReMe, reme, tag7, tag8, tag9]\n" - "---\nupdated body\n", - ) - - step = AutoMemoryStep( - app_context=app_ctx, - file_store=fs, - agent_wrapper=wrapper, - enable_tags=True, - ) - resp = await step( - RuntimeContext( - messages=[{"name": "user", "role": "user", "content": "updated project detail"}], - session_id="s1", - ), - ) - resp = resp or step.context.response - - assert resp.success is True - metadata = step._frontmatter(f"daily/{today}/memory.md") - assert metadata["tags"] == ["GPT-5", "C++", "C#", ".NET", "100", "ReMe", "tag7", "tag8"] - finally: - await fs.close() - - asyncio.run(run()) - - def test_auto_memory_uses_message_day_for_historical_create(): """AutoMemoryStep creates historical daily notes from message timestamps.""" diff --git a/tests/unit/test_evolve_utils.py b/tests/unit/test_evolve_utils.py index b8e3312c..c29cd0f9 100644 --- a/tests/unit/test_evolve_utils.py +++ b/tests/unit/test_evolve_utils.py @@ -9,7 +9,8 @@ from agentscope.message import Msg import pytest from reme.steps.evolve._evolve import agent_reply_result_text, format_history -from reme.steps.evolve.auto_memory import AutoMemoryStep, _normalize_tags, _sanitize_msg_for_save +from reme.steps.evolve.auto_memory import AutoMemoryStep, _sanitize_msg_for_save +from reme.steps.evolve.auto_tag import normalize_memory_tags def test_agent_reply_result_text_uses_last_text_block(): @@ -75,30 +76,24 @@ def test_sanitize_msg_for_save_drops_tool_results_and_base64_data(): assert sanitized.content[1].name == "memory_search" -def test_auto_memory_normalizes_frontmatter_tags(): - """Tags keep technical punctuation, reject phrases, de-duplicate, and stop at eight.""" - assert _normalize_tags( +def test_auto_tag_normalizes_frontmatter_tags(): + """Memory tags preserve entity names, de-duplicate, and stop at three.""" + assert normalize_memory_tags( [ - "GPT-5", - "C++", - "C#", - ".NET", - "100", - 100, - "memory system", + "OpenAI", + "openai", + "Sam Altman", "++", - "ReMe", - "reme", - "tag7", - "tag8", - "tag9", + 100, + "宁德时代", + "黄金", ], - ) == ["GPT-5", "C++", "C#", ".NET", "100", "ReMe", "tag7", "tag8"] + ) == ["OpenAI", "Sam_Altman", "宁德时代"] # pylint: disable=use-implicit-booleaness-not-comparison - assert _normalize_tags(None) == [] - assert _normalize_tags("GPT-5") == [] + assert normalize_memory_tags(None) == [] + assert normalize_memory_tags("OpenAI") == [] # pylint: enable=use-implicit-booleaness-not-comparison - assert _normalize_tags(["x" * 65, True, {}, "valid"]) == ["valid"] + assert normalize_memory_tags(["x" * 65, True, {}, "宁德时代"]) == ["宁德时代"] def test_auto_memory_accepts_message_timestamp_aliases(): diff --git a/tests/unit/test_faiss_index_maintenance.py b/tests/unit/test_faiss_index_maintenance.py index 15960f18..ca2a490c 100644 --- a/tests/unit/test_faiss_index_maintenance.py +++ b/tests/unit/test_faiss_index_maintenance.py @@ -36,6 +36,7 @@ class FakeEmbeddingStore: dimensions = 2 max_batch_size = 10 + is_healthy = True def _embed(self, text: str) -> np.ndarray: if "beta" in text or "fresh" in text: diff --git a/tests/unit/test_file_store_consistency.py b/tests/unit/test_file_store_consistency.py index a2693ed7..b7750017 100644 --- a/tests/unit/test_file_store_consistency.py +++ b/tests/unit/test_file_store_consistency.py @@ -49,6 +49,7 @@ class FakeEmbeddingStore: dimensions = 2 max_batch_size = 10 + is_healthy = True def _embed(self, text: str) -> np.ndarray: if "beta" in text or "fresh" in text: @@ -1519,6 +1520,42 @@ def test_search_filter_applies_to_vector_and_keyword_results(store_factory): run(go()) +def test_empty_exact_path_filter_matches_nothing(): + """An explicitly empty path domain must not widen into an unfiltered search.""" + candidate = chunk("a", "daily/a.md", "fresh topic") + + assert LocalFileStore._matches_search_filter(candidate, {"paths": []}) is False + + +@pytest.mark.parametrize("store_factory", [_new_local_store, _new_zvec_store]) +def test_path_filter_uses_one_domain_for_vector_and_keyword(store_factory): + """Vector and BM25 branches apply the same ordinary path and metadata filters.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = store_factory(name="t_filtered_domain") + await store.start() + store.embedding_store = FakeEmbeddingStore() + _ensure_zvec_collection(store) + await store.upsert( + [ + (node("daily/a.md"), [chunk("a", "daily/a.md", "fresh topic", kind="daily")]), + (node("daily/b.md"), [chunk("b", "daily/b.md", "fresh topic", kind="other")]), + (node("resource/c.md"), [chunk("c", "resource/c.md", "fresh topic", kind="daily")]), + ], + ) + + search_filter = { + "paths": ["daily/a.md", "daily/b.md"], + "metadata": {"kind": "daily"}, + } + assert [item.id for item in await store.vector_search("fresh", 5, search_filter)] == ["a"] + assert [item.id for item in await store.keyword_search("fresh", 5, search_filter)] == ["a"] + await store.close() + + run(go()) + + def test_faiss_rebuilds_stale_sidecar_and_updates_same_id_text(): """FAISS sidecar rebuilds when persisted rows no longer match chunks.""" diff --git a/tests/unit/test_frontmatter_steps.py b/tests/unit/test_frontmatter_steps.py index 1ea9a6ef..51a49f34 100644 --- a/tests/unit/test_frontmatter_steps.py +++ b/tests/unit/test_frontmatter_steps.py @@ -14,6 +14,7 @@ import os import tempfile from pathlib import Path +import frontmatter import pytest from reme.components.file_store import LocalFileStore @@ -89,6 +90,58 @@ async def test_update_no_suffix_autoappends_md(): await store.close() +@pytest.mark.asyncio +async def test_update_honors_injected_frontmatter_key_scope(): + """Injected key scope rejects the entire update when any key is outside it.""" + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + note = _seed(Path(tmp), NOTE, BODY) + store = await _make_store() + + before = note.read_bytes() + resp = await _run( + FrontmatterUpdateStep, + store, + path=NOTE, + metadata={"tags": ["new"], "name": "renamed"}, + _allowed_frontmatter_keys=["tags"], + ) + assert resp.success is False + assert resp.metadata["error"] == "frontmatter key(s) not allowed: 'name'" + assert note.read_bytes() == before + + resp = await _run( + FrontmatterUpdateStep, + store, + path=NOTE, + metadata={"tags": ["new"]}, + _allowed_frontmatter_keys=["tags"], + ) + assert resp.success is True + assert frontmatter.loads(note.read_text(encoding="utf-8")).metadata == {"name": "n", "tags": ["new"]} + await store.close() + + +@pytest.mark.asyncio +async def test_update_empty_frontmatter_key_scope_denies_all_updates(): + """An empty injected list is restrictive, unlike an omitted constraint.""" + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + note = _seed(Path(tmp), NOTE, BODY) + store = await _make_store() + before = note.read_bytes() + + resp = await _run( + FrontmatterUpdateStep, + store, + path=NOTE, + metadata={"tags": ["new"]}, + _allowed_frontmatter_keys=[], + ) + + assert resp.success is False + assert note.read_bytes() == before + await store.close() + + @pytest.mark.asyncio async def test_delete_no_suffix_autoappends_md(): """frontmatter_delete on a suffix-less path resolves to the ``.md`` file.""" diff --git a/tests/unit/test_injected_job_kwargs.py b/tests/unit/test_injected_job_kwargs.py index 678ae0e7..f4e32500 100644 --- a/tests/unit/test_injected_job_kwargs.py +++ b/tests/unit/test_injected_job_kwargs.py @@ -14,7 +14,6 @@ from reme.components.agent_wrapper import AsAgentWrapper, BaseAgentWrapper, CcAg from reme.components.file_store import LocalFileStore from reme.schema import Response from reme.steps.evolve.auto_memory import AutoMemoryStep -from reme.steps.evolve.auto_memory_cc import AutoMemoryCCStep class _Job: @@ -237,51 +236,6 @@ def test_auto_memory_keeps_original_tool_names(): assert step.update_tools == ["read", "edit", "frontmatter_update", "write"] -def test_auto_memory_tags_are_disabled_by_default_and_enabled_through_kwargs(): - assert AutoMemoryStep()._tags_enabled() is False - assert AutoMemoryCCStep()._tags_enabled() is False - assert AutoMemoryStep(enable_tags=True)._tags_enabled() is True - - -def test_auto_memory_tags_prompt_follows_step_kwarg(): - step = AutoMemoryStep() - prompt_kwargs = { - "today": "2026-09-01", - "note": "(none)", - "session_id": "s1", - "history": "user: remember GPT-5", - } - - disabled_system = step.prompt_format("system_prompt", enable_tags=step._tags_enabled()) - disabled_create = step.prompt_format("user_message_create", enable_tags=step._tags_enabled(), **prompt_kwargs) - disabled_update = step.prompt_format( - "user_message_update", - enable_tags=step._tags_enabled(), - note_path="daily/2026-09-01/memory.md", - **prompt_kwargs, - ) - assert "`tags`" not in disabled_system - assert 'metadata={"tags"' not in disabled_create - assert '"tags"' not in disabled_update - - enabled_step = AutoMemoryStep(enable_tags=True) - enabled_system = enabled_step.prompt_format("system_prompt", enable_tags=enabled_step._tags_enabled()) - enabled_create = enabled_step.prompt_format( - "user_message_create", - enable_tags=enabled_step._tags_enabled(), - **prompt_kwargs, - ) - enabled_update = enabled_step.prompt_format( - "user_message_update", - enable_tags=enabled_step._tags_enabled(), - note_path="daily/2026-09-01/memory.md", - **prompt_kwargs, - ) - assert "`tags`" in enabled_system - assert 'metadata={"tags"' in enabled_create - assert '"tags"' in enabled_update - - def test_auto_memory_create_prompts_match_upstream_date_arguments(): """Auto-memory prompts keep the upstream model-supplied date argument.""" from pathlib import Path @@ -289,8 +243,8 @@ def test_auto_memory_create_prompts_match_upstream_date_arguments(): prompt_file = Path("reme/steps/evolve/auto_memory.yaml") evolve_prompt = prompt_file.read_text(encoding="utf-8") assert "date={today}" in evolve_prompt or "`date`: {today}" in evolve_prompt or "`date`:{today}" in evolve_prompt - assert '"tags": [, ...]' in evolve_prompt - assert '"tags": []' in evolve_prompt or "using `[]`" in evolve_prompt + assert "enable_tags" not in evolve_prompt + assert "tags_key" not in evolve_prompt def test_configs_define_original_jobs_without_daily_variants(): @@ -305,7 +259,14 @@ def test_configs_define_original_jobs_without_daily_variants(): assert name not in jobs, f"{config_name} unexpectedly defines {name}" default = resolve_app_config(config="default", log_config=False) - assert default["jobs"]["auto_memory"]["steps"][0]["enable_tags"] is False + assert default["jobs"]["auto_memory"]["steps"] == [ + {"backend": "auto_memory_step"}, + {"backend": "auto_tag_step", "max_tags_per_file": 3}, + ] + assert default["jobs"]["auto_memory_cc"]["steps"] == [ + {"backend": "auto_memory_cc_step"}, + {"backend": "auto_tag_step", "max_tags_per_file": 3}, + ] if __name__ == "__main__": diff --git a/tests/unit/test_keyword_index.py b/tests/unit/test_keyword_index.py index 061c97b7..9d135284 100644 --- a/tests/unit/test_keyword_index.py +++ b/tests/unit/test_keyword_index.py @@ -351,6 +351,32 @@ def test_retrieve_score_ordering_by_tf(): run(go()) +def test_filtered_scoring_preserves_global_bm25_scores(): + """Selected-document scoring filters results without redefining the corpus.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + bm25 = await create_bm25() + await bm25.add_docs( + { + "high": "python python python", + "mid": "python python other", + "low": "python alpha beta", + "unrelated": "java only", + }, + ) + + global_scores = await bm25.retrieve("python", limit=4) + selected = await bm25.score_documents("python", {"mid", "low", "unrelated", "missing"}) + filtered = await bm25.retrieve_filtered("python", 1, {"mid", "low", "unrelated"}) + + assert selected == {"mid": global_scores["mid"], "low": global_scores["low"]} + assert filtered == {"mid": global_scores["mid"]} + await bm25.close() + + run(go()) + + def test_retrieve_idf_favours_rare_terms(): """In a query of {common, rare}, the doc containing the rare term wins.""" diff --git a/tests/unit/test_search_step.py b/tests/unit/test_search_step.py index 65824d55..bf9165b0 100644 --- a/tests/unit/test_search_step.py +++ b/tests/unit/test_search_step.py @@ -1,5 +1,7 @@ """Unit tests for workspace search Steps without embedding or LLM dependencies.""" +# pylint: disable=protected-access + import asyncio import importlib.util from pathlib import Path @@ -7,6 +9,8 @@ from pathlib import Path from agentscope.message import Msg from reme.components.file_store import BaseFileStore +from reme.components.file_store.local_file_store import LocalFileStore +from reme.components.tag_index import LocalTagIndex from reme.components import ApplicationContext from reme.components.runtime_context import RuntimeContext from reme.enumeration import LinkScopeEnum @@ -42,15 +46,22 @@ class FakeSearchStore(BaseFileStore): self.calls: list[tuple[str, str, int, dict]] = [] async def upsert(self, files: list[tuple[FileNode, list[FileChunk]]]) -> None: - raise NotImplementedError + """Ignore writes in this read-only search fixture.""" + del files async def delete(self, path: str | list[str]) -> None: - raise NotImplementedError + """Ignore deletes in this read-only search fixture.""" + del path async def clear(self) -> None: - raise NotImplementedError + """Clear the fixture's in-memory results and recorded calls.""" + self.vector_results.clear() + self.keyword_results.clear() + self.calls.clear() async def get_nodes(self, paths: list[str] | None = None) -> list[FileNode]: + """Return no graph nodes for search-only tests.""" + del paths return [] async def get_outlinks( @@ -58,6 +69,8 @@ class FakeSearchStore(BaseFileStore): path: str, scope: LinkScopeEnum = LinkScopeEnum.REAL, ) -> list[FileLink]: + """Return no outgoing links for search-only tests.""" + del path, scope return [] async def get_inlinks( @@ -65,17 +78,42 @@ class FakeSearchStore(BaseFileStore): path: str, scope: LinkScopeEnum = LinkScopeEnum.REAL, ) -> list[FileLink]: + """Return no incoming links for search-only tests.""" + del path, scope return [] async def vector_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]: + """Return the configured vector results.""" self.calls.append(("vector", query, limit, search_filter)) return self.vector_results[:limit] async def keyword_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]: + """Return the configured keyword results.""" self.calls.append(("keyword", query, limit, search_filter)) return self.keyword_results[:limit] +class TaggedFakeSearchStore(FakeSearchStore): + """Fake store with a tag index and ordinary file-store filtering.""" + + def __init__(self, *, chunks: list[FileChunk]): + super().__init__(vector_results=chunks, keyword_results=chunks) + self.file_chunks = {chunk.id: chunk for chunk in chunks} + self.tag_index = LocalTagIndex() + + async def vector_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]: + self.calls.append(("vector", query, limit, search_filter)) + return [chunk for chunk in self.vector_results if LocalFileStore._matches_search_filter(chunk, search_filter)][ + :limit + ] + + async def keyword_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]: + self.calls.append(("keyword", query, limit, search_filter)) + return [chunk for chunk in self.keyword_results if LocalFileStore._matches_search_filter(chunk, search_filter)][ + :limit + ] + + def _chunk( chunk_id: str, path: str, @@ -898,6 +936,87 @@ def test_search_step_empty_query_fails_before_store_calls(): asyncio.run(run()) +def test_search_step_falls_back_for_unavailable_index_and_rejects_invalid_tags(): + """Unavailable optional indexes preserve search, while malformed tags still fail.""" + + async def run(): + hit = _chunk("hit", "daily/a.md", "text", "keyword", 3.0) + missing = FakeSearchStore(keyword_results=[hit]) + unhealthy = TaggedFakeSearchStore(chunks=[hit]) + unhealthy.tag_index.set_healthy(False) + + for store in (missing, unhealthy): + resp = await SearchStep(file_store=store, expand_links=False)( + RuntimeContext(query="hello", limit=5, tags=["python"]), + ) + assert resp.success is True + assert [result["id"] for result in resp.metadata["results"]] == ["hit"] + assert {call[0] for call in store.calls} == {"vector", "keyword"} + assert all(call[3] == {} for call in store.calls) + assert resp.metadata["tag_filter"] == { + "requested": True, + "applied": False, + "reason": "tag_index_unavailable", + } + + invalid = TaggedFakeSearchStore(chunks=[]) + resp = await SearchStep(file_store=invalid, expand_links=False)( + RuntimeContext(query="hello", limit=5, tags=["!", "++", "x" * 65]), + ) + assert resp.success is False + assert resp.answer == "Error: tags contained no valid values" + assert not invalid.calls + + asyncio.run(run()) + + +def test_search_step_combines_tags_with_existing_chunk_filters(): + """Tags use OR internally and are ANDed with existing path/date behavior.""" + + async def run(): + matching = _chunk("a", "daily/2024-03-01/a.md", "match", "keyword", 3.0) + old = _chunk("b", "daily/2023-03-01/b.md", "old", "keyword", 2.0) + wrong_prefix = _chunk("c", "resource/2024-03-01/c.md", "resource", "keyword", 1.0) + store = TaggedFakeSearchStore(chunks=[matching, old, wrong_prefix]) + await store.tag_index.rebuild( + [ + FileNode(path=matching.path, st_mtime=1.0, front_matter={"memory_tags": ["Python"]}), + FileNode(path=old.path, st_mtime=1.0, front_matter={"memory_tags": ["ReMe"]}), + FileNode(path=wrong_prefix.path, st_mtime=1.0, front_matter={"memory_tags": ["python"]}), + ], + ) + step = SearchStep(file_store=store, expand_links=False) + + resp = await step( + RuntimeContext( + query="hello", + limit=5, + tags=[" PYTHON ", "REME"], + start_date="2024-01-01", + search_filter={"path_prefix": "daily/"}, + ), + ) + + assert resp.success is True + assert [result["id"] for result in resp.metadata["results"]] == ["a"] + assert resp.metadata["tag_filter"] == { + "requested": True, + "applied": True, + "tags": ["python", "reme"], + "matched_paths": 3, + } + assert {call[0] for call in store.calls} == {"vector", "keyword"} + assert all(set(call[3]["paths"]) == {matching.path, old.path, wrong_prefix.path} for call in store.calls) + store.calls.clear() + resp = await step(RuntimeContext(query="hello", limit=5, tags=["missing"])) + assert resp.success is True + assert resp.metadata["results"] == [] + assert resp.metadata["tag_filter"]["matched_paths"] == 0 + assert all(call[3]["paths"] == [] for call in store.calls) + + asyncio.run(run()) + + def test_search_step_start_end_date_promoted_into_search_filter(): """start_date and end_date from context are promoted into search_filter passed to store.""" diff --git a/tests/unit/test_tag_index.py b/tests/unit/test_tag_index.py index 407bd02b..d4189493 100644 --- a/tests/unit/test_tag_index.py +++ b/tests/unit/test_tag_index.py @@ -12,10 +12,11 @@ from reme.components.file_store import LocalFileStore from reme.components.tag_index import LocalTagIndex from reme.config import resolve_app_config from reme.schema import FileChunk, FileFrontMatter, FileNode +from reme.steps.index.list_tags import ListTagsStep -def _node(path: str, tags: object = None) -> FileNode: - metadata = {} if tags is None else {"tags": tags} +def _node(path: str, tags: object = None, *, key: str = "memory_tags") -> FileNode: + metadata = {} if tags is None else {key: tags} return FileNode(path=path, st_mtime=1.0, front_matter=FileFrontMatter(**metadata)) @@ -29,12 +30,12 @@ def test_tag_normalization_and_bidirectional_mutations() -> None: async def run() -> None: index = LocalTagIndex(max_tags_per_file=3) await index.start() - await index.upsert_nodes([_node("daily/a.md", ["Python", "PYTHON", "C++", ".NET", "ignored"])]) + await index.upsert_nodes([_node("daily/a.md", ["Python Tag", "PYTHON TAG", "C++", ".NET", "ignored"])]) - assert await index.tags_for_path("daily/a.md") == ["python", "c++", ".net"] - assert await index.paths_for_tags(["PYTHON"]) == ["daily/a.md"] + assert await index.tags_for_path("daily/a.md") == ["python_tag", "c++", ".net"] + assert await index.paths_for_tags(["PYTHON TAG"]) == ["daily/a.md"] assert index.tag_to_paths == { - "python": {"daily/a.md"}, + "python_tag": {"daily/a.md"}, "c++": {"daily/a.md"}, ".net": {"daily/a.md"}, } @@ -85,7 +86,7 @@ def test_queries_are_not_truncated_by_per_file_tag_limit() -> None: """Apply the count limit to indexed files without dropping lookup conditions.""" async def run() -> None: - index = LocalTagIndex(max_tags_per_file=2) + index = LocalTagIndex(max_tags_per_file=2, max_tag_length=8) await index.rebuild( [ _node("daily/a.md", ["a", "b"]), @@ -98,18 +99,181 @@ def test_queries_are_not_truncated_by_per_file_tag_limit() -> None: "daily/a.md", "daily/c.md", ] + assert index.normalize_query_tags([" A ", "B", "c c", "a", "too-long-tag", " "]) == ["a", "b", "c_c"] asyncio.run(run()) +def test_list_tags_paginates_and_applies_default_sort_orders() -> None: + """List only active tags with compact counts and deterministic pagination.""" + + async def run() -> None: + index = LocalTagIndex() + await index.rebuild( + [ + _node("daily/a.md", ["beta", "alpha"]), + _node("daily/b.md", ["gamma", "beta"]), + _node("daily/c.md", ["delta"]), + ], + ) + + assert await index.list_tags(page_size=2) == { + "total_tags": 4, + "total_pages": 2, + "page": 1, + "range": (1, 2), + "items": [("alpha", 1), ("beta", 2)], + } + assert await index.list_tags(page=2, page_size=2) == { + "total_tags": 4, + "total_pages": 2, + "page": 2, + "range": (3, 4), + "items": [("delta", 1), ("gamma", 1)], + } + assert (await index.list_tags(order_by="file_count"))["items"] == [ + ("beta", 2), + ("alpha", 1), + ("delta", 1), + ("gamma", 1), + ] + assert (await index.list_tags(order_by="file_count", order="asc"))["items"] == [ + ("alpha", 1), + ("delta", 1), + ("gamma", 1), + ("beta", 2), + ] + assert (await index.list_tags(order="desc"))["items"] == [ + ("gamma", 1), + ("delta", 1), + ("beta", 2), + ("alpha", 1), + ] + + await index.delete(["daily/c.md"]) + result = await index.list_tags(page=3, page_size=2) + assert result == { + "total_tags": 3, + "total_pages": 2, + "page": 3, + "range": (0, 0), + "items": [], + } + + await index.clear() + assert await index.list_tags(page=9) == { + "total_tags": 0, + "total_pages": 0, + "page": 9, + "range": (0, 0), + "items": [], + } + + invalid_cases = [ + ({"page": 0}, "page must be a positive integer"), + ({"page_size": True}, "page_size must be a positive integer"), + ({"page_size": 1001}, "page_size must be less than or equal to 1000"), + ({"order_by": "unknown"}, "order_by must be one of"), + ({"order": "sideways"}, "order must be one of"), + ] + for kwargs, message in invalid_cases: + with pytest.raises(ValueError, match=message): + await index.list_tags(**kwargs) + + index.set_healthy(False) + with pytest.raises(RuntimeError, match="tag index is unavailable"): + await index.list_tags() + + store = LocalFileStore(name="test", embedding_store="", tag_index="") + store.tag_index = LocalTagIndex() + await store.tag_index.rebuild([_node("daily/a.md", ["ReMe"])]) + + response = await ListTagsStep(file_store=store)(order_by="file_count") + + assert response.answer == { + "total_tags": 1, + "total_pages": 1, + "page": 1, + "range": (1, 1), + "items": [("reme", 1)], + } + + asyncio.run(run()) + + config = resolve_app_config(config="default", log_config=False) + job = config["jobs"]["list_tags"] + assert job["steps"] == [{"backend": "list_tags_step"}] + assert job["parameters"]["properties"]["page_size"]["default"] == 100 + assert "[tag, file_count]" in job["description"] + assert "range" in job["description"] + assert "empty page" in job["parameters"]["properties"]["page"]["description"] + + +def test_configured_frontmatter_key_contract() -> None: + """Validate, apply, and invalidate changes to the configured source key.""" + + async def run() -> None: + index = LocalTagIndex(tag_key="keywords") + await index.rebuild( + [ + _node("daily/a.md", ["ignored"]), + _node("daily/b.md", ["Python"], key="keywords"), + ], + ) + + assert index.tag_key == "keywords" + assert await index.paths_for_tags(["python"]) == ["daily/b.md"] + assert await index.paths_for_tags(["ignored"]) == [] + + asyncio.run(run()) + invalid_cases = [ + ("", "tag_key must be a non-empty string", False), + (" ", "tag_key must be a non-empty string", False), + (None, "tag_key must be a non-empty string", False), + (123, "tag_key must be a non-empty string", False), + ("name", "tag_key must not be a reserved frontmatter key", False), + ("description", "tag_key must not be a reserved frontmatter key", False), + ("kind", "tag_key must not be a reserved frontmatter key", False), + ("session_id", "tag_key must not be a reserved frontmatter key", False), + ("source_conversation", "tag_key must not be a reserved frontmatter key", False), + ("source_resource", "tag_key must not be a reserved frontmatter key", False), + ("status", "tag_key must not be a reserved frontmatter key", False), + ("", "tag_key must be a non-empty string", True), + ("name", "tag_key must not be a reserved frontmatter key", True), + ("description", "tag_key must not be a reserved frontmatter key", True), + ] + for tag_key, message, at_runtime in invalid_cases: + index = LocalTagIndex() + with pytest.raises(ValueError, match=message): + if at_runtime: + index.tag_key = tag_key + else: + LocalTagIndex(tag_key=tag_key) + if at_runtime: + assert index.tag_key == "memory_tags" + + assert LocalTagIndex(tag_key="model_config").tag_key == "model_config" + index = LocalTagIndex() + index.tag_key = "keywords" + assert index.tag_key == "keywords" + assert not index.is_healthy + + def test_file_store_updates_tag_index_from_file_nodes(monkeypatch, tmp_path: Path) -> None: """Keep daily and digest tags aligned through file-store mutations.""" + disabled_store = LocalFileStore(name="test", embedding_store="", tag_index="") + assert disabled_store.tag_index_enabled is False + with pytest.raises(RuntimeError, match="tag index is not configured"): + disabled_store.require_tag_index() + async def run() -> None: monkeypatch.chdir(tmp_path) store = LocalFileStore(name="test", embedding_store="", tag_index="default") await store.start() assert isinstance(store.tag_index, LocalTagIndex) + assert store.tag_index_enabled is True + assert store.require_tag_index() is store.tag_index await store.upsert( [ @@ -142,7 +306,7 @@ def test_tag_failures_do_not_block_other_indexes_and_retry_rebuild(monkeypatch, monkeypatch.chdir(tmp_path) store = LocalFileStore(name="test", embedding_store="", tag_index="default") await store.start() - assert store.tag_index is not None + assert store.tag_index_enabled original_rebuild = store.tag_index.rebuild async def fail_incremental(_nodes) -> None: @@ -182,7 +346,7 @@ def test_failed_tag_reconciliation_makes_queries_fail_closed(monkeypatch, tmp_pa monkeypatch.chdir(tmp_path) store = LocalFileStore(name="test", embedding_store="", tag_index="default") await store.start() - assert store.tag_index is not None + assert store.tag_index_enabled await store.upsert([(_node("daily/a.md", ["old"]), [])]) async def fail(_items) -> None: @@ -209,7 +373,7 @@ def test_tag_rebuild_graph_read_failure_does_not_block_upsert(monkeypatch, tmp_p monkeypatch.chdir(tmp_path) store = LocalFileStore(name="test", embedding_store="", tag_index="default") await store.start() - assert store.tag_index is not None + assert store.tag_index_enabled assert store.file_graph is not None store._tag_index_rebuild_required = True original_get_nodes = store.file_graph.get_nodes @@ -239,7 +403,7 @@ def test_explicit_reindex_restores_tag_index(monkeypatch, tmp_path: Path) -> Non monkeypatch.chdir(tmp_path) store = LocalFileStore(name="test", embedding_store="", tag_index="default") await store.start() - assert store.tag_index is not None + assert store.tag_index_enabled await store.upsert( [ (_node("daily/a.md", ["ReMe"]), []), @@ -256,6 +420,13 @@ def test_explicit_reindex_restores_tag_index(monkeypatch, tmp_path: Path) -> Non result = await store.reindex("all") assert result["tag"] == {"indexed": 1, "scope": "tag"} assert await store.tag_index.paths_for_tags(["reme"]) == ["daily/a.md"] + await store.file_graph.upsert_nodes([_node("daily/a.md", ["new"], key="keywords")]) + + store.tag_index.tag_key = "keywords" + await store.reindex("tag") + + assert await store.tag_index.paths_for_tags(["old"]) == [] + assert await store.tag_index.paths_for_tags(["new"]) == ["daily/a.md"] await store.close() asyncio.run(run()) @@ -268,7 +439,7 @@ def test_tag_delete_failures_do_not_block_core_deletion(monkeypatch, tmp_path: P monkeypatch.chdir(tmp_path) store = LocalFileStore(name="test", embedding_store="", tag_index="default") await store.start() - assert store.tag_index is not None + assert store.tag_index_enabled chunk = _chunk("chunk-a", "daily/a.md", "alpha memory") await store.upsert([(_node("daily/a.md", ["alpha"]), [chunk])]) @@ -298,7 +469,7 @@ def test_existing_markdown_chunker_supplies_frontmatter_tags(monkeypatch, tmp_pa monkeypatch.chdir(tmp_path) note = tmp_path / "daily" / "a.md" note.parent.mkdir() - note.write_text("---\ntags: [Python, ReMe]\n---\nbody\n", encoding="utf-8") + note.write_text("---\nmemory_tags: [Python, ReMe]\n---\nbody\n", encoding="utf-8") node, chunks = await MarkdownFileChunker().chunk(note) store = LocalFileStore(name="test", embedding_store="", tag_index="default") @@ -331,16 +502,22 @@ def test_file_store_rebuilds_non_persistent_tag_index_from_graph(monkeypatch, tm asyncio.run(run()) -def test_default_config_documents_optional_tag_index_without_enabling_it() -> None: - """Document tag indexing in the default config without enabling another index or watcher.""" +def test_default_config_enables_tag_index_with_explicit_key() -> None: + """Keep auto-memory generation and file-store indexing on the same configured key.""" config = resolve_app_config(config="default", log_config=False) assert config["jobs"]["index_update_loop"]["watch_dirs"] == ["daily_dir", "digest_dir"] assert "tag_index_loop" not in config["jobs"] - assert "tag_index" not in config["components"] - assert "tag_index" not in config["components"]["file_store"]["default"] - - default_yaml = Path("reme/config/default.yaml").read_text(encoding="utf-8") - assert "# tag_index:" in default_yaml - assert "# tag_index: default" in default_yaml + assert config["components"]["tag_index"]["default"]["tag_key"] == "memory_tags" + assert config["components"]["tag_index"]["default"]["max_tags_per_file"] == 3 + assert config["components"]["file_store"]["default"]["tag_index"] == "default" + assert config["jobs"]["search"]["parameters"]["properties"]["tags"]["default"] == [] + assert config["jobs"]["auto_memory"]["steps"] == [ + {"backend": "auto_memory_step"}, + {"backend": "auto_tag_step", "max_tags_per_file": 3}, + ] + assert config["jobs"]["auto_memory_cc"]["steps"] == [ + {"backend": "auto_memory_cc_step"}, + {"backend": "auto_tag_step", "max_tags_per_file": 3}, + ] diff --git a/tests/unit/test_zvec_file_store.py b/tests/unit/test_zvec_file_store.py index ba1566f4..037486ac 100644 --- a/tests/unit/test_zvec_file_store.py +++ b/tests/unit/test_zvec_file_store.py @@ -38,6 +38,7 @@ class FakeEmbeddingStore: dimensions = 2 max_batch_size = 10 + is_healthy = True def _embed(self, text: str) -> np.ndarray: if "beta" in text or "fresh" in text: