diff --git a/reme/components/file_chunker/markdown_file_chunker.py b/reme/components/file_chunker/markdown_file_chunker.py index dd93e0e9..9847f669 100644 --- a/reme/components/file_chunker/markdown_file_chunker.py +++ b/reme/components/file_chunker/markdown_file_chunker.py @@ -115,12 +115,16 @@ def __init__( encoding: str = "utf-8", chunk_chars: int = 10000, embed_toc: bool = True, + include_frontmatter_in_metadata: bool = False, + include_frontmatter_keys_in_metadata: list[str] | None = None, **kwargs, ): super().__init__(**kwargs) self.encoding = encoding self.chunk_chars = max(100, chunk_chars) self.embed_toc = embed_toc + self.include_frontmatter_in_metadata = include_frontmatter_in_metadata + self.include_frontmatter_keys_in_metadata = list(include_frontmatter_keys_in_metadata or []) async def chunk(self, path: str | Path) -> tuple[FileNode, list[FileChunk]]: from mistletoe.markdown_renderer import MarkdownRenderer @@ -135,6 +139,13 @@ async def chunk(self, path: str | Path) -> tuple[FileNode, list[FileChunk]]: with MarkdownRenderer() as renderer: tree = self._build_tree(Document(content), renderer, line_offset=line_offset) chunks = self._chunk_node(tree, "", "", rel_path, renderer) + if self.include_frontmatter_in_metadata: + chunk_metadata = self._chunk_metadata( + front_matter, + allow_keys=self.include_frontmatter_keys_in_metadata or None, + ) + for chunk in chunks: + chunk.metadata = chunk_metadata.copy() links = WikilinkHandler.extract_links(content, rel_path) if content else [] @@ -172,6 +183,20 @@ def _parse_front_matter(text: str) -> tuple[FileFrontMatter, str, int]: return front_matter, "".join(lines[close_idx + 1 :]), close_idx + 1 + @staticmethod + def _chunk_metadata( + front_matter: FileFrontMatter, + allow_keys: list[str] | None = None, + ) -> dict: + """Expose frontmatter fields on chunks so search filters can match them. + + ``allow_keys`` (when non-empty) restricts the output to that subset of + field names; absent or empty means "all fields with non-empty values". + """ + dumped = front_matter.model_dump(mode="json") + filtered = dumped if not allow_keys else {k: dumped.get(k) for k in allow_keys if k in dumped} + return {key: value for key, value in filtered.items() if value not in (None, "")} + def _build_tree(self, doc: Any, renderer, line_offset: int = 0) -> MdNode: """Heading-level stack folds mistletoe's flat children into nested sections; non-headings attach as ``body`` to the current section diff --git a/reme/config/default.yaml b/reme/config/default.yaml index 37a09938..12928842 100644 --- a/reme/config/default.yaml +++ b/reme/config/default.yaml @@ -666,12 +666,14 @@ components: markdown: backend: markdown supported_extensions: [ "md" ] + include_frontmatter_in_metadata: false + include_frontmatter_keys_in_metadata: [] # empty = all non-empty frontmatter keys json: backend: json supported_extensions: [ "json" ] jsonl: backend: jsonl - supported_extensions: [ "jsonl" ] + supported_extensions: [ "jsonl" ] # noqa: keep #314 chunker scope intact after #325 default: backend: default supported_extensions: ["txt","log"] diff --git a/tests/unit/test_config_parser.py b/tests/unit/test_config_parser.py index 6a6dcf85..effe99b3 100644 --- a/tests/unit/test_config_parser.py +++ b/tests/unit/test_config_parser.py @@ -30,6 +30,18 @@ def test_default_config_registers_daily_write_job(): assert job["parameters"]["required"] == ["name", "description", "session_id", "content"] +def test_default_config_keeps_frontmatter_chunk_metadata_opt_in(): + """Markdown frontmatter-to-chunk metadata is disabled by default for compatibility.""" + cfg = _load_config("default.yaml") + + markdown = cfg["components"]["file_chunker"]["markdown"] + assert markdown["include_frontmatter_in_metadata"] is False + # Allow-list defaults to empty; combined with the False above, chunk metadata stays empty. + assert markdown["include_frontmatter_keys_in_metadata"] == [] or markdown.get( + "include_frontmatter_keys_in_metadata", + ) in (None, []) + + def test_parse_args_rejects_non_key_value_extra_argument(): """Extra CLI arguments must use key=value syntax.""" with pytest.raises(ValueError, match="expected key=value"): diff --git a/tests/unit/test_markdown_file_chunker.py b/tests/unit/test_markdown_file_chunker.py index 3848594a..7db22fd4 100644 --- a/tests/unit/test_markdown_file_chunker.py +++ b/tests/unit/test_markdown_file_chunker.py @@ -72,6 +72,100 @@ async def run(): asyncio.run(run()) +def test_parse_frontmatter_metadata_is_opt_in(): + """Chunk metadata preserves the old empty default unless explicitly enabled.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + body = ( + "---\n" + "name: locomo-event\n" + "description: Jon lost his job\n" + "conversation_date: 2023-01-19\n" + "---\n" + "Jon said he lost his job today.\n" + ) + path = _write_md(tmp, "daily/2023-01-19/locomo-event.md", body) + chunker = MarkdownFileChunker(chunk_chars=500) + _, chunks = await chunker.chunk(path) + + assert len(chunks) == 1 + assert chunks[0].metadata == {} + + chunker = MarkdownFileChunker(chunk_chars=500, include_frontmatter_in_metadata=True) + _, chunks = await chunker.chunk(path) + + assert len(chunks) == 1 + assert chunks[0].metadata == { + "name": "locomo-event", + "description": "Jon lost his job", + "conversation_date": "2023-01-19", + } + print("✓ test_parse_frontmatter_metadata_is_opt_in passed") + + asyncio.run(run()) + + +def test_parse_frontmatter_metadata_keys_allowlist(): + """``include_frontmatter_keys_in_metadata`` restricts copied keys to an allow-list.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + body = ( + "---\n" + "name: locomo-event\n" + "description: Jon lost his job\n" + "conversation_date: 2023-01-19\n" + "---\n" + "Jon said he lost his job today.\n" + ) + path = _write_md(tmp, "daily/2023-01-19/locomo-event.md", body) + + # Allow-list restricted to a single key. + chunker = MarkdownFileChunker( + chunk_chars=500, + include_frontmatter_in_metadata=True, + include_frontmatter_keys_in_metadata=["conversation_date"], + ) + _, chunks = await chunker.chunk(path) + assert len(chunks) == 1 + assert chunks[0].metadata == {"conversation_date": "2023-01-19"} + + # Allow-list with a key not present in frontmatter is a no-op for that key. + chunker = MarkdownFileChunker( + chunk_chars=500, + include_frontmatter_in_metadata=True, + include_frontmatter_keys_in_metadata=["conversation_date", "absent"], + ) + _, chunks = await chunker.chunk(path) + assert chunks[0].metadata == {"conversation_date": "2023-01-19"} + + # Empty allow-list (not None) keeps the legacy "all non-empty keys" behavior. + chunker = MarkdownFileChunker( + chunk_chars=500, + include_frontmatter_in_metadata=True, + include_frontmatter_keys_in_metadata=[], + ) + _, chunks = await chunker.chunk(path) + assert chunks[0].metadata == { + "name": "locomo-event", + "description": "Jon lost his job", + "conversation_date": "2023-01-19", + } + + # Allow-list is ignored when the master toggle is off (back-compat default). + chunker = MarkdownFileChunker( + chunk_chars=500, + include_frontmatter_in_metadata=False, + include_frontmatter_keys_in_metadata=["conversation_date"], + ) + _, chunks = await chunker.chunk(path) + assert chunks[0].metadata == {} + print("✓ test_parse_frontmatter_metadata_keys_allowlist passed") + + asyncio.run(run()) + + def test_parse_small_body_one_chunk(): """A body shorter than chunk_chars produces exactly one chunk that contains the body.""" @@ -275,6 +369,7 @@ async def run(): print("\n=== MarkdownFileChunker tests ===") test_parse_empty_file() test_parse_frontmatter_only() + test_parse_frontmatter_metadata_is_opt_in() test_parse_small_body_one_chunk() test_parse_oversized_body_splits() test_parse_chunk_ids_match_node_chunk_ids() diff --git a/tests/unit/test_search_step.py b/tests/unit/test_search_step.py index 7d347255..5d974781 100644 --- a/tests/unit/test_search_step.py +++ b/tests/unit/test_search_step.py @@ -165,6 +165,24 @@ async def run(): asyncio.run(run()) +def test_search_step_passes_metadata_filter_to_store(): + """Search filters can target chunk metadata such as conversation_date.""" + + async def run(): + hit = _chunk("hit", "daily/2023-01-19/event.md", "historical hit", "keyword", 3.0) + store = FakeSearchStore(keyword_results=[hit]) + step = SearchStep(file_store=store, expand_links=False) + search_filter = {"metadata": {"conversation_date": "2023-01-19"}} + + resp = await step(RuntimeContext(query="Jon job", limit=5, search_filter=search_filter)) + + assert resp.success is True + assert resp.metadata["results"][0]["id"] == "hit" + assert all(call[3] == search_filter for call in store.calls) + + asyncio.run(run()) + + def test_search_step_tool_context_seen_chunks_expire_after_ttl(): """Seen chunk ids under a tool_context_id are reusable after the configured TTL."""