feat(rag): hybrid FTS+vector retrieval and multi-format ingestion — name-your-tool questions find the right document

2026-08-22 01:27:02 -04:00
parent 2f738a7f19
commit 7e8d14702e
36 changed files with 2018 additions and 290 deletions
@@ -1,11 +1,25 @@
-"""Unit tests: markdown-aware chunker (PLAN §5 policy)."""
+"""Unit tests: format-aware chunker (PLAN §5 policy, A9 formats).
+
+The markdown policy tests are the original contract (md output stays
+unchanged); the per-format tests cover the phase-09 dispatcher
+(yaml/yml, json, py, txt) and the 1200-char hard cap for every format.
+"""
 from __future__ import annotations

 from itertools import pairwise

 import pytest

-from app.rag.chunker import HARD_MAX_CHARS, chunk_markdown, extract_title
+from app.rag.chunker import (
+    HARD_MAX_CHARS,
+    chunk_document,
+    chunk_json,
+    chunk_markdown,
+    chunk_python,
+    chunk_text,
+    chunk_yaml,
+    extract_title,
+)

 ANCHOR = "## Big"
 ANCHOR_PREFIX = f"{ANCHOR}\n\n"
@@ -157,3 +171,213 @@ def test_extract_title_prefers_h1() -> None:
    assert extract_title("## not a title\n\nbody") == ""
    assert extract_title("## sub only", fallback="stem") == "stem"
    assert extract_title("", fallback="fallback") == "fallback"
+
+
+# ---------------------------------------------------------------------------
+# Format dispatcher (chunk_document) — A9 multi-format ingestion
+# ---------------------------------------------------------------------------
+
+
+def test_dispatch_by_lowercased_suffix() -> None:
+    md = "# T\n\n## A\n\nbody\n"
+    assert chunk_document(md, "notes/Doc.MD") == chunk_markdown(md)
+    assert chunk_document(md, "notes/doc.MARKDOWN") == chunk_markdown(md)
+    assert chunk_document("p1\n\np2\n", "x.TXT") == chunk_text("p1\n\np2\n")
+    assert chunk_document("a: 1\n", "x.YAML") == chunk_yaml("a: 1\n")
+    assert chunk_document("a: 1\n", "x.Yml") == chunk_yaml("a: 1\n")
+    assert chunk_document('{"a": 1}', "x.Json") == chunk_json('{"a": 1}')
+    assert chunk_document("def f(): pass\n", "x.PY") == chunk_python("def f(): pass\n")
+
+
+def test_dispatch_unknown_suffix_falls_back_to_paragraphs() -> None:
+    assert chunk_document("hello\n\nworld", "data.csv") == ["hello\nworld"]
+
+
+def test_dispatch_ignores_directory_part_of_path() -> None:
+    assert chunk_document("def f(): pass\n", "a/b/c/script.py") == chunk_python("def f(): pass\n")
+
+
+# ---------------------------------------------------------------------------
+# yaml / yml
+# ---------------------------------------------------------------------------
+
+
+def test_yaml_splits_on_top_level_keys_and_keeps_key_anchors() -> None:
+    doc = (
+        "# leading comment\n"
+        "services:\n"
+        "  gitlab:\n"
+        "    image: gitlab/gitlab-ce\n"
+        "  prometheus:\n"
+        "    image: prom/prometheus\n"
+        "volumes:\n"
+        "  gitlab-data:\n"
+    )
+    chunks = chunk_yaml(doc)
+    joined = "\n".join(chunks)
+    for key in ("services:", "volumes:"):
+        assert key in joined
+    # Indented keys are NOT block starts — they stay inside their parent block.
+    assert not any(c.startswith("  gitlab:") for c in chunks)
+    # The leading comment stays with the first block (preamble).
+    assert chunks[0].startswith("# leading comment")
+    assert "gitlab/gitlab-ce" in joined and "prom/prometheus" in joined
+
+
+def test_yaml_document_separators_start_new_blocks() -> None:
+    a = "site_a: " + "a" * 500 + "\n"
+    b = "site_b: " + "b" * 500 + "\n"
+    chunks = chunk_yaml(a + "---\n" + b, target_chars=600, overlap_chars=0)
+    # Each site is long enough to force its own chunk; the separator must not
+    # glue them into one over-budget chunk.
+    assert len(chunks) >= 2
+    assert all(len(c) <= 600 for c in chunks)
+    assert not any("site_a" in c and "site_b" in c for c in chunks)
+
+
+def test_yaml_oversized_key_block_is_split_under_hard_cap() -> None:
+    doc = "big_list:\n" + ("  - " + "x" * 60 + "\n") * 60  # one ~3800-char block
+    chunks = chunk_yaml(doc)
+    assert len(chunks) >= 2
+    assert all(len(c) <= HARD_MAX_CHARS for c in chunks)
+    # Overlap re-prints (≤50 chars per split), so only a little content is
+    # re-stated — the bulk of the block must survive.
+    assert sum(len(c) for c in chunks) >= len(doc) - 300
+
+
+def test_yaml_empty_content() -> None:
+    assert chunk_yaml("") == []
+    assert chunk_yaml("\n\n  \n") == []
+
+
+# ---------------------------------------------------------------------------
+# json
+# ---------------------------------------------------------------------------
+
+
+def test_json_splits_on_top_level_keys_pretty_printed() -> None:
+    doc = '{"hosts": {"kafkabridge": "10.0.3.7"}, "count": 3}'
+    chunks = chunk_json(doc, target_chars=45, overlap_chars=0)  # force 1 chunk/block
+    assert len(chunks) == 2
+    first, second = chunks
+    assert '"hosts"' in first and "kafkabridge" in first
+    assert '"count"' in second
+    # Pretty-printed (indent=2), not the compact input form.
+    assert '"kafkabridge": "10.0.3.7"' in first
+    assert not any('{"hosts"' in c for c in chunks)
+
+
+def test_json_each_key_block_is_self_contained() -> None:
+    doc = '{"a": "x", "b": "y"}'
+    chunks = chunk_json(doc, target_chars=13, overlap_chars=0)  # force 1 chunk/block
+    assert [c for c in chunks if '"a"' in c] and [c for c in chunks if '"b"' in c]
+    assert not any('"a"' in c and '"b"' in c for c in chunks)
+
+
+def test_json_oversized_value_falls_under_hard_cap() -> None:
+    doc = '{"blob": "' + "z" * 4000 + '"}'
+    chunks = chunk_json(doc)
+    assert len(chunks) >= 2
+    assert all(len(c) <= HARD_MAX_CHARS for c in chunks)
+    assert "".join(chunks).count("z") >= 4000
+
+
+def test_json_top_level_list_is_one_pretty_block() -> None:
+    chunks = chunk_json("[1, 2, 3]")
+    assert chunks == ["[\n  1,\n  2,\n  3\n]"]
+
+
+def test_json_unparseable_falls_back_to_paragraph_packing() -> None:
+    doc = "{broken json\n\nsecond paragraph here\n"
+    assert chunk_json(doc) == chunk_text(doc)
+    assert chunk_json("not json at all") == chunk_text("not json at all")
+
+
+# ---------------------------------------------------------------------------
+# python
+# ---------------------------------------------------------------------------
+
+
+def test_python_splits_on_top_level_defs_and_classes() -> None:
+    doc = (
+        '"""Module doc."""\n'
+        "import asyncio\n"
+        "\n"
+        "CONST = 1\n"
+        "\n"
+        "def alpha():\n"
+        "    return 1\n"
+        "\n"
+        "class Beta:\n"
+        "    def run(self):\n"
+        "        return 2\n"
+    )
+    chunks = chunk_python(doc, target_chars=60, overlap_chars=0)  # force 1 chunk/block
+    assert len(chunks) == 3
+    assert chunks[0].startswith('"""Module doc."""')
+    assert "CONST = 1" in chunks[0]  # preamble ends at the first def/class
+    assert chunks[1].startswith("def alpha")
+    assert chunks[2].startswith("class Beta")
+    assert "def run" in chunks[2]  # nested def stays inside the class block
+
+
+def test_python_decorators_stay_with_their_definition() -> None:
+    doc = "@app.get('/x')\ndef handler():\n    return 'x'\n"
+    chunks = chunk_python(doc)
+    assert chunks[0].startswith("@app.get")
+
+
+def test_python_oversized_function_falls_back_to_line_packing() -> None:
+    doc = "def big():\n" + "\n".join(f"    val_{i:03d} = {i}  # padding" for i in range(80))
+    chunks = chunk_python(doc)
+    assert len(chunks) >= 2
+    assert all(len(c) <= HARD_MAX_CHARS for c in chunks)
+    assert "val_000" in chunks[0]
+    assert "val_079" in chunks[-1]
+    assert sum(len(c) for c in chunks) >= len(doc) - 100
+
+
+def test_python_unparseable_source_falls_back_to_paragraphs() -> None:
+    src = "def broken(:\n\nstill text\n"
+    assert chunk_python(src) == chunk_text(src)
+
+
+# ---------------------------------------------------------------------------
+# txt
+# ---------------------------------------------------------------------------
+
+
+def test_txt_paragraph_packing() -> None:
+    doc = "alpha\n\nbeta\n\ngamma\n"
+    chunks = chunk_text(doc)
+    assert chunks == ["alpha\nbeta\ngamma"]  # all three fit the target
+
+
+def test_txt_long_doc_packs_with_overlap() -> None:
+    doc = "\n\n".join(f"para {i} " + "l" * 300 for i in range(6))
+    chunks = chunk_text(doc, target_chars=800, overlap_chars=100)
+    assert len(chunks) >= 2
+    assert all(len(c) <= 800 for c in chunks)
+    assert all(f"para {i}" in "\n".join(chunks) for i in range(6))
+
+
+# ---------------------------------------------------------------------------
+# Hard cap across every format (aipi ~1024-token request cap)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    ("content", "path"),
+    [
+        ("# T\n\n" + "word " * 1200, "big.md"),
+        ("key: " + "v" * 5000 + "\n", "big.yaml"),
+        ('{"blob": "' + "z" * 5000 + '"}', "big.json"),
+        ("def f():\n" + "    x = 1\n" * 1000, "big.py"),
+        ("line of text\n\n" * 800, "big.txt"),
+    ],
+)
+def test_hard_cap_holds_for_every_format(content: str, path: str) -> None:
+    chunks = chunk_document(content, path)
+    assert chunks, "expected at least one chunk"
+    for c in chunks:
+        assert len(c) <= HARD_MAX_CHARS, f"{path}: {len(c)} chars"