Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 25 additions & 0 deletions reme/components/file_chunker/markdown_file_chunker.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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 []

Expand Down Expand Up @@ -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
Expand Down
4 changes: 3 additions & 1 deletion reme/config/default.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand Down
12 changes: 12 additions & 0 deletions tests/unit/test_config_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"):
Expand Down
95 changes: 95 additions & 0 deletions tests/unit/test_markdown_file_chunker.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""

Expand Down Expand Up @@ -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()
Expand Down
18 changes: 18 additions & 0 deletions tests/unit/test_search_step.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""

Expand Down
Loading