Spaces:
Configuration error
Configuration error
File size: 6,480 Bytes
6e49b19 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 | #!/usr/bin/env python3
"""Stub test: markdown-heading section metadata in ingest_lib -- NO models.
Proves:
1. parse_page_sections derives a running h1 > h2 stack from '#' lines and
keeps per-page section paths (nested headings, empty-heading lines).
2. One kind='section' record per distinct heading, _summary_record pattern
(id/kind/section_path/text/page/page_end/sha/doc_name).
3. _section_for_page gives chunks the section active on their start page.
4. retrieve.py chunk mapping carries section_path through.
"""
from __future__ import annotations
import os
import sys
import tempfile
from pathlib import Path
REPO = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO))
os.environ["FACTORY_RAG_DATA"] = tempfile.mkdtemp(prefix="fr-section-")
import config # noqa: E402
import ingest_lib # noqa: E402
def main() -> int:
failures = []
# --- 1: running section stack over nested headings --------------------
pages_md = {
1: "# 4 Fuel System\n\nsome prose about pumps.\n",
2: "## 4.2 Fuel Filter\nfilter details.\n#tag-not-heading\n## \n## 4.3 Startup Procedure\nsteps here.\n",
3: "plain page with no headings.\n",
4: "# 5 Cooling\nnew chapter.\n",
}
by_page, recs = ingest_lib.parse_page_sections(pages_md)
if by_page.get(1) != "4 Fuel System":
failures.append(f"page1 section wrong: {by_page.get(1)!r}")
# h2 pops back under the running h1; empty '## ' line must be ignored
if by_page.get(2) != "4 Fuel System > 4.3 Startup Procedure":
failures.append(f"page2 section wrong: {by_page.get(2)!r}")
# heading-free page inherits the running section
if by_page.get(3) != "4 Fuel System > 4.3 Startup Procedure":
failures.append(f"page3 section wrong: {by_page.get(3)!r}")
if by_page.get(4) != "5 Cooling":
failures.append(f"page4 section wrong: {by_page.get(4)!r}")
# distinct headings only: 4 Fuel System, 4.2 Fuel Filter, 4.3 Startup
# Procedure, 5 Cooling (the empty '## ' and '#tag-not-heading' excluded)
paths = sorted(r["section_path"] for r in recs)
if paths != sorted([
"4 Fuel System", "4 Fuel System > 4.2 Fuel Filter",
"4 Fuel System > 4.3 Startup Procedure", "5 Cooling"]):
failures.append(f"section records wrong: {paths}")
# --- 2: record shape follows the _summary_record pattern --------------
r = next(x for x in recs if x["section_path"] == "4 Fuel System")
for key in ("id", "kind", "section_path", "text", "page", "page_end",
"sha", "doc_name"):
if key not in r:
failures.append(f"section record missing key {key}")
if r["kind"] != "section" or r["page"] != 1 or r["page_end"] != 1:
failures.append(f"section record kind/pages wrong: {r}")
if "4 Fuel System" not in r["text"] or "section heading" not in r["text"]:
failures.append(f"section record text wrong: {r['text']!r}")
# filled in by the caller, same slot usage as doc_summary records
filled = dict(r, sha="deadbeef", doc_name="TM_5-811-6.pdf")
if filled["sha"] != "deadbeef" or filled["doc_name"] != "TM_5-811-6.pdf":
failures.append("section record sha/doc_name fill-in broken")
# --- 3: chunk page -> active section ----------------------------------
if ingest_lib._section_for_page(1, by_page) != "4 Fuel System":
failures.append("_section_for_page p1 wrong")
if ingest_lib._section_for_page(3, by_page) != "4 Fuel System > 4.3 Startup Procedure":
failures.append("_section_for_page p3 wrong")
if ingest_lib._section_for_page(4, by_page) != "5 Cooling":
failures.append("_section_for_page p4 wrong")
# a chunk starting on a page BEFORE any heading: '' (nothing active yet)
if ingest_lib._section_for_page(0, by_page) != "":
failures.append("_section_for_page p0 must be '' before first heading")
if ingest_lib._section_for_page(2, {}) != "":
failures.append("_section_for_page with empty map must be ''")
# no headings at all -> no records, empty paths
by2, recs2 = ingest_lib.parse_page_sections({1: "just text\n", 2: "more\n"})
if recs2 or any(by2.values()):
failures.append("heading-free doc must yield no sections")
# --- 4: chunk mapping carries section_path ---------------------------
import retrieve # noqa: E402
import types
seam = types.ModuleType("sentence_transformers")
class FakeST:
def __init__(self, *a, **kw):
pass
def encode(self, texts, **kw):
import numpy as np
return np.zeros((len(texts), 4), dtype=np.float32)
seam.SentenceTransformer = FakeST
sys.modules["sentence_transformers"] = seam
fake_points = [
{"text": "t1", "page": 2, "page_end": 2, "sha": "s1", "kind": "text",
"doc_name": "d.pdf", "section_path": "4 Fuel System > 4.2 Fuel Filter",
"score": 0.9},
{"text": "t2", "page": 3, "sha": "s1", "kind": "text",
"doc_name": "d.pdf", "score": 0.4}, # no section_path key
]
retr = retrieve.Retriever.__new__(retrieve.Retriever)
retr.embedder = FakeST()
retr._reranker = "passthrough"
orig_rerank = retrieve.Retriever._rerank
retrieve.Retriever._rerank = lambda self, q, pts, top_n: pts[:top_n]
orig_hybrid = retrieve.Retriever._hybrid_query
retrieve.Retriever._hybrid_query = lambda self, dq, sq, top_k: list(fake_points)
orig_visual = retrieve.Retriever._visual_rrf
retrieve.Retriever._visual_rrf = lambda self, q, pts, top_k: pts
orig_gate = config.ENABLE_PRE_REFUSAL_GATE
config.ENABLE_PRE_REFUSAL_GATE = False
try:
res = retr.retrieve("q", top_n=2)
c0, c1 = res["chunks"][0], res["chunks"][1]
if c0.get("section_path") != "4 Fuel System > 4.2 Fuel Filter":
failures.append(f"mapping lost section_path: {c0}")
if c1.get("section_path", None) != "":
failures.append(f"mapping must default section_path to '': {c1}")
finally:
retrieve.Retriever._rerank = orig_rerank
retrieve.Retriever._hybrid_query = orig_hybrid
retrieve.Retriever._visual_rrf = orig_visual
config.ENABLE_PRE_REFUSAL_GATE = orig_gate
if failures:
for f in failures:
print(f"FAIL: {f}")
return 1
print("PASS test_section_metadata_stub")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|