feat(rag): hybrid FTS+vector retrieval and multi-format ingestion — name-your-tool questions find the right document
This commit is contained in:
+172
-17
@@ -1,25 +1,36 @@
|
||||
"""Markdown-aware chunker (PLAN §5 chunking policy).
|
||||
"""Format-aware chunker (PLAN §5 chunking policy).
|
||||
|
||||
Pure functions, no I/O — fully unit-testable.
|
||||
Pure functions, no I/O — fully unit-testable. Stdlib only.
|
||||
|
||||
Policy
|
||||
------
|
||||
* **Sections** are split on ATX headings of level ≥ 2 (``## ``/``### ``/…).
|
||||
* A section that fits in ``target_chars`` becomes a single chunk.
|
||||
* A longer section is sub-split at paragraph boundaries (blank lines outside
|
||||
code fences); each chunk after the first starts with the trailing
|
||||
``overlap_chars`` of the previous chunk so context survives the cut.
|
||||
* Every chunk keeps its nearest preceding heading line (the section anchor),
|
||||
so a retrieval hit is always readable in context.
|
||||
* **Code fences** (``` / ~~~) are atomic: a chunk boundary never falls
|
||||
inside one, and lines inside a fence are never mistaken for headings or
|
||||
paragraph breaks. One exception: a fence *larger than* :data:`HARD_MAX_CHARS`
|
||||
is split by line, because aipi's local embedding model rejects requests
|
||||
over ~1024 input tokens and a single 5000-char code block would blow
|
||||
past that on its own.
|
||||
:func:`chunk_document` dispatches on the file's lowercased suffix;
|
||||
per-format policies:
|
||||
|
||||
* **md / markdown** — sections are split on ATX headings of level ≥ 2
|
||||
(``## ``/``### ``/…); a section that fits in ``target_chars`` becomes a
|
||||
single chunk, a longer one is sub-split at paragraph boundaries (blank
|
||||
lines outside code fences) with ``overlap_chars`` carry-over, and every
|
||||
chunk keeps its nearest preceding heading line (the section anchor). Code
|
||||
fences are atomic (a boundary never falls inside one) except a fence
|
||||
larger than :data:`HARD_MAX_CHARS`, which is split by line.
|
||||
* **yaml / yml** — blocks start at ``---`` document separators and at
|
||||
top-level (indent-0) ``key:`` lines; every chunk keeps its key lines as
|
||||
anchors, so a hit is always readable in context.
|
||||
* **json** — pretty-printed (``json.dumps(obj, indent=2)``) and split on
|
||||
top-level keys (one ``{key: value}`` block per key); unparseable input
|
||||
falls back to paragraph packing.
|
||||
* **py** — split at top-level defs/classes via the stdlib ``ast`` (the
|
||||
module preamble — imports, constants — is its own block); an oversized
|
||||
definition falls back to line packing.
|
||||
* **txt** (and any unknown suffix) — paragraph packing.
|
||||
|
||||
Every format honors :data:`HARD_MAX_CHARS` (1200 — the aipi ~1024-token
|
||||
request cap) and the target/overlap settings; oversized blocks are split
|
||||
by line so no chunk can exceed the cap.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Sequence
|
||||
|
||||
@@ -203,3 +214,147 @@ def chunk_markdown(
|
||||
for start, end in _section_ranges(lines, flags):
|
||||
chunks.extend(_chunk_section(lines[start:end], flags[start:end], target, overlap))
|
||||
return [c for c in chunks if c.strip()]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Non-markdown formats (A9, revised 2026-08-21): yaml/yml, json, py, txt
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
#: YAML document separator (column 0).
|
||||
_YAML_DOC_SEP_RE = re.compile(r"^-{3,}\s*$")
|
||||
#: Top-level YAML key (column 0, no leading whitespace) — the block anchor.
|
||||
_YAML_KEY_RE = re.compile(r"^[A-Za-z0-9_.\-]+\s*:")
|
||||
|
||||
|
||||
def _yaml_blocks(lines: Sequence[str]) -> list[str]:
|
||||
"""Group YAML lines into blocks: ``---`` separators and indent-0
|
||||
``key:`` lines each start a new block (the key line stays the anchor)."""
|
||||
blocks: list[str] = []
|
||||
cur: list[str] = []
|
||||
for line in lines:
|
||||
if cur and (_YAML_DOC_SEP_RE.match(line) or _YAML_KEY_RE.match(line)):
|
||||
blocks.append("\n".join(cur))
|
||||
cur = []
|
||||
cur.append(line)
|
||||
if cur:
|
||||
blocks.append("\n".join(cur))
|
||||
return [b for b in blocks if b.strip()]
|
||||
|
||||
|
||||
def chunk_yaml(content: str, target_chars: int = 2000, overlap_chars: int = 200) -> list[str]:
|
||||
"""Split YAML on document separators + top-level keys (see module docstring)."""
|
||||
target, overlap = _normalize_target_overlap(target_chars, overlap_chars)
|
||||
return _pack_blocks(_yaml_blocks(content.splitlines()), target, overlap)
|
||||
|
||||
|
||||
def _json_blocks(content: str) -> list[str] | None:
|
||||
"""Pretty-printed per-top-level-key blocks, or ``None`` if unparseable."""
|
||||
try:
|
||||
obj = json.loads(content)
|
||||
except (ValueError, TypeError):
|
||||
return None
|
||||
if isinstance(obj, dict):
|
||||
return [json.dumps({k: v}, indent=2) for k, v in obj.items()]
|
||||
# Top-level list/scalar: nothing to key on — one pretty-printed block.
|
||||
return [json.dumps(obj, indent=2)]
|
||||
|
||||
|
||||
def chunk_json(content: str, target_chars: int = 2000, overlap_chars: int = 200) -> list[str]:
|
||||
"""Pretty-print JSON and split on top-level keys (unparseable → paragraphs)."""
|
||||
target, overlap = _normalize_target_overlap(target_chars, overlap_chars)
|
||||
blocks = _json_blocks(content)
|
||||
if blocks is None:
|
||||
return chunk_text(content, target, overlap)
|
||||
return _pack_blocks(blocks, target, overlap)
|
||||
|
||||
|
||||
def _python_blocks(content: str) -> list[str] | None:
|
||||
"""Line blocks: module preamble, then one per top-level def/class.
|
||||
|
||||
Returns ``None`` when the source does not parse (→ line/paragraph
|
||||
packing fallback).
|
||||
"""
|
||||
try:
|
||||
tree = ast.parse(content)
|
||||
except (SyntaxError, ValueError):
|
||||
return None
|
||||
lines = content.splitlines()
|
||||
tops = [
|
||||
node
|
||||
for node in tree.body
|
||||
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef))
|
||||
]
|
||||
tops.sort(key=lambda n: n.lineno)
|
||||
ranges: list[tuple[int, int]] = []
|
||||
for node in tops:
|
||||
start = node.lineno - 1
|
||||
for dec in node.decorator_list:
|
||||
start = min(start, dec.lineno - 1)
|
||||
end = node.end_lineno or node.lineno # end_lineno is None on odd parses
|
||||
ranges.append((start, end)) # 0-based start, 1-based end
|
||||
blocks: list[str] = []
|
||||
cursor = 0
|
||||
for start, end in ranges:
|
||||
if start > cursor:
|
||||
blocks.append("\n".join(lines[cursor:start]))
|
||||
blocks.append("\n".join(lines[start:end]))
|
||||
cursor = end
|
||||
if cursor < len(lines):
|
||||
blocks.append("\n".join(lines[cursor:]))
|
||||
return [b for b in blocks if b.strip()]
|
||||
|
||||
|
||||
def chunk_python(content: str, target_chars: int = 2000, overlap_chars: int = 200) -> list[str]:
|
||||
"""Split Python on top-level defs/classes (stdlib ``ast``; see module doc)."""
|
||||
target, overlap = _normalize_target_overlap(target_chars, overlap_chars)
|
||||
blocks = _python_blocks(content)
|
||||
if blocks is None:
|
||||
return chunk_text(content, target, overlap)
|
||||
return _pack_blocks(blocks, target, overlap)
|
||||
|
||||
|
||||
def chunk_text(content: str, target_chars: int = 2000, overlap_chars: int = 200) -> list[str]:
|
||||
"""Plain-text paragraph packing (blank lines separate paragraphs)."""
|
||||
target, overlap = _normalize_target_overlap(target_chars, overlap_chars)
|
||||
lines = content.splitlines()
|
||||
return _pack_blocks(_paragraph_blocks(lines, [False] * len(lines)), target, overlap)
|
||||
|
||||
|
||||
def _normalize_target_overlap(target_chars: int, overlap_chars: int) -> tuple[int, int]:
|
||||
"""Validate + clamp the size policy (shared by every format)."""
|
||||
if target_chars <= 0:
|
||||
raise ValueError("target_chars must be > 0")
|
||||
if overlap_chars < 0:
|
||||
raise ValueError("overlap_chars must be >= 0")
|
||||
# The endpoint's token cap is absolute — a larger target is unsafe.
|
||||
target = min(target_chars, HARD_MAX_CHARS)
|
||||
return target, min(overlap_chars, target - 1)
|
||||
|
||||
|
||||
#: suffix → chunker (A9, revised: md, markdown, txt, yaml, yml, json, py).
|
||||
_FORMAT_CHUNKERS = {
|
||||
".md": chunk_markdown,
|
||||
".markdown": chunk_markdown,
|
||||
".txt": chunk_text,
|
||||
".yaml": chunk_yaml,
|
||||
".yml": chunk_yaml,
|
||||
".json": chunk_json,
|
||||
".py": chunk_python,
|
||||
}
|
||||
|
||||
|
||||
def chunk_document(
|
||||
content: str,
|
||||
path: str,
|
||||
target_chars: int = 2000,
|
||||
overlap_chars: int = 200,
|
||||
) -> list[str]:
|
||||
"""Chunk *content* according to *path*'s lowercased suffix.
|
||||
|
||||
Unknown suffixes fall back to plain-text paragraph packing (the
|
||||
importer only passes A9-format files, so this is belt-and-braces).
|
||||
"""
|
||||
name = path.rsplit("/", 1)[-1]
|
||||
suffix = "." + name.rsplit(".", 1)[-1].lower() if "." in name else ""
|
||||
chunker = _FORMAT_CHUNKERS.get(suffix, chunk_text)
|
||||
return chunker(content, target_chars, overlap_chars)
|
||||
|
||||
Reference in New Issue
Block a user