Experimental/brain-of-reese
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
ducoterra
2026-08-22 01:27:02 -04:00
parent 2f738a7f19
commit 7e8d14702e
36 changed files with 2018 additions and 290 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/rag/chunker.py
+172 -17
View File
@@ -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)
app/rag/importer.py
+54 -16
View File
@@ -1,23 +1,30 @@
"""Knowledge-base importer (PLAN §5 / §9 / §11).
Walks ``*.md`` files (A9 exclusion list), diffs by sha256 against
``documents.content_hash`` and, for every new or changed file, runs the
two-phase upsert:
Walks the A9-format files (``md, markdown, txt, yaml, yml, json, py`` by
default — ``BOR_IMPORT_EXTENSIONS``; case-insensitive), diffs by sha256
against ``documents.content_hash`` and, for every new or changed file, runs
the two-phase upsert:
1. upsert the document row and replace its chunk rows (embeddings NULL)
2. embed the new chunks in batches and attach the vectors
3. commit — one transaction per file, so a failed embedding leaves the
database untouched and the file is simply retried on the next run
Scope (A9, revised 2026-08-21): any path containing a dot-prefixed
component (hidden dirs — vendored caches like ``.esphome/.espressif/**`` —
or hidden files) is skipped, plus the well-known exclusion list.
``prune=True`` deletes documents (of the imported sources only) whose files
no longer exist. Per-file logging uses the verbs
``added | updated | unchanged | pruned`` plus a summary line (PLAN §9).
no longer exist **or no longer match the format filter** — this is how
previously-imported junk (e.g. dot-dir READMEs) leaves the index. Per-file
logging uses the verbs ``added | updated | unchanged | pruned`` plus a
summary line with per-format counts (PLAN §9).
"""
from __future__ import annotations
import hashlib
import logging
from dataclasses import dataclass
from dataclasses import dataclass, field
from datetime import UTC, datetime
from pathlib import Path
from typing import Protocol
@@ -28,7 +35,7 @@ from sqlalchemy.orm import Session
from app.config import Settings
from app.db import SessionLocal
from app.models import Chunk, Document
from app.rag.chunker import chunk_markdown, extract_title
from app.rag.chunker import chunk_document, extract_title
from app.rag.llm import EmbeddingError
logger = logging.getLogger("app.importer")
@@ -60,11 +67,20 @@ class ImportSummary:
errors: int = 0
chunks: int = 0
embed_batches: int = 0
#: Files walked, keyed by lowercased extension (``md``, ``yaml``, …).
formats: dict[str, int] = field(default_factory=dict)
def format_counts(self) -> str:
"""``md:203,yaml:267,py:14`` — highest count first (PLAN §9)."""
if not self.formats:
return "none"
ordered = sorted(self.formats.items(), key=lambda kv: (-kv[1], kv[0]))
return ",".join(f"{ext}:{count}" for ext, count in ordered)
def log(self) -> None:
logger.info(
"import: summary files=%d added=%d updated=%d unchanged=%d pruned=%d "
"errors=%d chunks=%d embed_batches=%d",
"errors=%d chunks=%d embed_batches=%d formats=%s",
self.files,
self.added,
self.updated,
@@ -73,17 +89,32 @@ class ImportSummary:
self.errors,
self.chunks,
self.embed_batches,
self.format_counts(),
)
def iter_markdown_files(root: Path, excluded: frozenset[str] = EXCLUDED_DIRS) -> list[Path]:
"""All ``*.md`` files under *root* (sorted), skipping excluded dirs (A9)."""
def iter_importable_files(
root: Path,
extensions: frozenset[str],
excluded: frozenset[str] = EXCLUDED_DIRS,
) -> list[Path]:
"""All importable files under *root* (sorted), per the A9 scope rules.
*extensions* is a set of lowercased dotted suffixes (``{'.md', '.py'}``).
Skips: any path with a dot-prefixed component (hidden dirs/files —
vendored caches like ``.esphome/.espressif/**``) and the well-known
non-content directories in *excluded*.
"""
if not root.is_dir():
return []
files: list[Path] = []
for path in sorted(root.rglob("*.md")):
for path in sorted(root.rglob("*")):
if not path.is_file():
continue
rel = path.relative_to(root)
if any(part in excluded for part in rel.parts[:-1]):
if any(part.startswith(".") or part in excluded for part in rel.parts):
continue
if path.suffix.lower() not in extensions:
continue
files.append(path)
return files
@@ -97,7 +128,7 @@ async def import_sources(
limit: int | None = None,
session: Session | None = None,
) -> ImportSummary:
"""Import every ``*.md`` under *sources* (see module docstring).
"""Import every A9-format file under *sources* (see module docstring).
``session`` may be supplied (tests); a private one is opened and closed
otherwise. ``limit`` caps the number of files processed (debug only) and
@@ -120,12 +151,14 @@ async def import_sources(
break
source = root.name
source_names.add(source)
for path in iter_markdown_files(root):
for path in iter_importable_files(root, llm.settings.import_extension_set):
if limit is not None and summary.files >= limit:
break
rel = path.relative_to(root).as_posix()
seen.add((source, rel))
summary.files += 1
ext = path.suffix.lower().lstrip(".") or "unknown"
summary.formats[ext] = summary.formats.get(ext, 0) + 1
try:
await _index_file(
session, source=source, rel=rel, full_path=path, llm=llm,
@@ -172,7 +205,12 @@ async def _index_file(
return
verb = "updated" if doc is not None else "added"
title = extract_title(content, fallback=full_path.stem)
# A ``#`` line is a real heading in markdown but a comment in every
# other format — titles for those come from the file stem.
if full_path.suffix.lower() in (".md", ".markdown"):
title = extract_title(content, fallback=full_path.stem)
else:
title = full_path.stem
if doc is None:
doc = Document(
source=source,
@@ -200,7 +238,7 @@ async def _index_file(
# policy stays intact for the rest of the KB.
target = max(400, settings.chunk_target_chars)
while True:
chunks_text = chunk_markdown(content, target, settings.chunk_overlap_chars)
chunks_text = chunk_document(content, rel, target, settings.chunk_overlap_chars)
doc.chunks = [
Chunk(document_id=doc.id, position=i, content=c) for i, c in enumerate(chunks_text)
]
app/rag/retriever.py
+194 -23
View File
@@ -1,19 +1,31 @@
"""pgvector cosine retrieval → parent-document mapping (PLAN §3/§6, A7).
"""Hybrid retrieval: pgvector cosine ∪ Postgres FTS, RRF-fused (PLAN §6, A7).
Retrieval returns the *chunks* closest to the question embedding (top-K by
cosine distance). The product requirement is that the LLM receives the
**entire relevant document**, not just the chunk (LOCKED A7) — so
:meth:`select_documents` maps chunk hits back to their parent documents
(``chunks.document_id → documents``), dedupes, ranks by best chunk score,
and caps the combined context at ``BOR_MAX_CONTEXT_CHARS``.
* **Vector list** — top-N chunks by cosine distance (``embedding <=> $1``),
each carrying its cosine ``1 − distance`` (the honesty-gate input).
* **Lexical list** — top-N chunks matching an OR-``tsquery`` over the
question's tokens, ordered by ``ts_rank``. This is what finds
name-your-tool questions ("gitlab") that vector similarity buries.
* **Fusion** — Reciprocal Rank Fusion (``score = Σ 1/(k + rank)`` over the
lists a chunk appears in; chunks hit by both lists get both terms). The
fused score ranks; :meth:`select_documents` and :func:`weak_hit_titles`
keep working off ``score``.
The product requirement is unchanged (LOCKED A7): the LLM receives the
**entire relevant document**, not just the chunk — chunk hits map back to
their parents, dedupe, rank by best fused score, and the combined context
is capped at ``BOR_MAX_CONTEXT_CHARS``.
Deterministic tie-break for equal fused scores:
``(−fused, −cosine, document.path, chunk.position)``.
"""
from __future__ import annotations
import re
import uuid
from collections.abc import Sequence
from dataclasses import dataclass
from dataclasses import dataclass, replace
from sqlalchemy import select
from sqlalchemy import select, text
from sqlalchemy.orm import Session
from app.config import get_settings
@@ -22,50 +34,209 @@ from app.models import Chunk, Document
#: Marker appended when the context budget is exceeded (PLAN §6).
TRUNCATION_MARKER = "[…truncated…]"
#: Alphanumeric tokens of a question (``to_tsquery`` input, OR-joined).
_TOKEN_RE = re.compile(r"[a-z0-9]+")
#: One row of the lexical candidate query (all fields needed to build a
#: detached :class:`Document` plus the chunk fields and ``ts_rank``).
_LEXICAL_SQL = text(
"""
SELECT c.id AS chunk_id,
c.position AS position,
c.content AS content,
d.id AS doc_id,
d.source AS source,
d.path AS path,
d.full_path AS full_path,
d.title AS title,
d.content AS doc_content,
d.content_hash AS content_hash,
d.indexed_at AS indexed_at,
ts_rank(c.tsv, to_tsquery('english', :tsquery)) AS rank
FROM chunks c
JOIN documents d ON d.id = c.document_id
WHERE c.tsv @@ to_tsquery('english', :tsquery)
ORDER BY rank DESC, d.path ASC, c.position ASC
LIMIT :limit
"""
)
@dataclass
class RetrievedChunk:
"""One chunk hit: its cosine score plus the parent document row."""
"""One retrieval candidate: fused rank score + parent document row.
* ``score`` — RRF fused score (the ranking key for document selection
and weak-hit titles).
* ``cosine`` — vector similarity ``1 − distance`` (the honesty-gate
input; ``0.0`` for lexical-only hits that have no vector rank).
* ``fts_hit`` — the chunk matched the question's OR-tsquery.
"""
chunk_id: uuid.UUID
position: int
content: str
score: float # 1 − cosine_distance (higher is more similar)
score: float
document: Document
cosine: float = 0.0
fts_hit: bool = False
def retrieve(
db: Session, question_embedding: list[float], top_k: int | None = None
) -> list[RetrievedChunk]:
"""Top-*top_k* chunks by pgvector cosine distance (``<=>``).
def lexical_tsquery(question: str) -> str | None:
"""OR-joined token string for ``to_tsquery('english', …)``, or ``None``.
``score = 1 − distance``. Results are ordered by ascending distance, so
index 0 is the best hit. Chunks whose embedding is still NULL (two-phase
import in progress) are skipped.
Tokens are lowercased ``[a-z0-9]+`` runs, de-duplicated in order of
first appearance. Postgres does the lexing/stemming; a question whose
tokens are all stopwords lexes to an *empty* tsquery (which matches
nothing), so no special-casing is needed there. Pure-symbol questions
("???", "🔧") yield no tokens → ``None`` → no lexical query at all.
"""
seen: set[str] = set()
tokens: list[str] = []
for tok in _TOKEN_RE.findall(question.lower()):
if tok not in seen:
seen.add(tok)
tokens.append(tok)
return " | ".join(tokens) if tokens else None
def fuse(
vector: Sequence[RetrievedChunk],
lexical: Sequence[RetrievedChunk],
k: int,
) -> list[RetrievedChunk]:
"""Reciprocal Rank Fusion over the two ranked candidate lists.
``score(chunk) = Σ 1/(k + rank)`` — one term per list the chunk appears
in (ranks are 1-based; a chunk in both lists gets both terms). Returns
the union ordered by ``(-score, -cosine, document.path, position)``.
Lexical-only hits (no vector rank) enter with ``cosine=0.0`` and
``fts_hit=True``; vector chunks matched by the lexical list get
``fts_hit=True`` in place (the input objects are mutated — callers
should not reuse them afterwards).
"""
if k <= 0:
raise ValueError("rrf k must be > 0")
by_id: dict[uuid.UUID, RetrievedChunk] = {}
fused: dict[uuid.UUID, float] = {}
for rank, rc in enumerate(vector, start=1):
by_id[rc.chunk_id] = rc
fused[rc.chunk_id] = fused.get(rc.chunk_id, 0.0) + 1.0 / (k + rank)
for rank, rc in enumerate(lexical, start=1):
term = 1.0 / (k + rank)
if rc.chunk_id in by_id:
existing = by_id[rc.chunk_id]
by_id[rc.chunk_id] = replace(existing, fts_hit=True)
fused[rc.chunk_id] += term
else:
rc = replace(rc, fts_hit=True)
by_id[rc.chunk_id] = rc
fused[rc.chunk_id] = fused.get(rc.chunk_id, 0.0) + term
out = [replace(rc, score=fused[rc.chunk_id]) for rc in by_id.values()]
out.sort(key=lambda rc: (-rc.score, -rc.cosine, rc.document.path, rc.position))
return out
def _vector_candidates(
db: Session, question_embedding: list[float], limit: int
) -> list[RetrievedChunk]:
"""Top-*limit* chunks by pgvector cosine distance (``<=>``).
``cosine = 1 − distance``. Chunks whose embedding is still NULL
(two-phase import in progress) are skipped.
"""
k = top_k if top_k is not None else get_settings().top_k_chunks
distance = Chunk.embedding.cosine_distance(question_embedding)
rows = db.execute(
select(Chunk, distance.label("distance"), Document)
.join(Document, Chunk.document_id == Document.id)
.where(Chunk.embedding.is_not(None))
.order_by(distance)
.limit(k)
.limit(limit)
).all()
return [
RetrievedChunk(
chunk_id=chunk.id,
position=chunk.position,
content=chunk.content,
score=round(1.0 - float(dist), 6),
score=0.0, # fused score is filled in by :func:`fuse`
document=doc,
cosine=round(1.0 - float(dist), 6),
)
for chunk, dist, doc in rows
]
def _lexical_candidates(db: Session, question: str, limit: int) -> list[RetrievedChunk]:
"""Top-*limit* chunks matching the question's OR-tsquery (A7).
Ordered by ``ts_rank`` (with ``d.path, c.position`` as the
deterministic tie-break); an empty tsquery (stopword-only question)
simply matches nothing.
"""
tsquery = lexical_tsquery(question)
if tsquery is None:
return []
rows = db.execute(
_LEXICAL_SQL, {"tsquery": tsquery, "limit": limit}
).all()
out: list[RetrievedChunk] = []
for row in rows:
doc = Document(
id=row.doc_id,
source=row.source,
path=row.path,
full_path=row.full_path,
title=row.title,
content=row.doc_content,
content_hash=row.content_hash,
indexed_at=row.indexed_at,
)
out.append(
RetrievedChunk(
chunk_id=row.chunk_id,
position=row.position,
content=row.content,
score=0.0, # filled in by :func:`fuse`
document=doc,
cosine=0.0, # no vector rank — lexical-only hit
fts_hit=True,
)
)
return out
def retrieve(
db: Session,
question: str,
question_embedding: list[float],
vector_candidates: int | None = None,
lexical_candidates: int | None = None,
) -> list[RetrievedChunk]:
"""Hybrid retrieval (A7): vector top-N ∪ FTS top-N, RRF-fused.
Returns the fused candidate list in rank order (best first). Each
:class:`RetrievedChunk` carries the fused ``score`` (ranking), the
``cosine`` similarity (honesty gate) and the ``fts_hit`` flag.
"""
settings = get_settings()
v_n = (
settings.hybrid_vector_candidates if vector_candidates is None else vector_candidates
)
l_n = (
settings.hybrid_lexical_candidates if lexical_candidates is None else lexical_candidates
)
if v_n <= 0:
raise ValueError("vector_candidates must be >= 1")
if l_n <= 0:
raise ValueError("lexical_candidates must be >= 1")
vector = _vector_candidates(db, question_embedding, v_n)
lexical = _lexical_candidates(db, question, l_n)
return fuse(vector, lexical, settings.rrf_k)
def weak_hit_titles(chunks: Sequence[RetrievedChunk]) -> list[str]:
"""Distinct parent-document titles of *chunks*, best chunk score first.
"""Distinct parent-document titles of *chunks*, best fused score first.
Deflection mode (PLAN §6, A8) is built from these *titles only* — the
LOW prompt and the "Maybe try" chips never see document content.
@@ -85,7 +256,7 @@ def select_documents(
n: int | None = None,
max_chars: int | None = None,
) -> list[Document]:
"""Map chunk hits to distinct parent documents, ranked by best chunk score.
"""Map chunk hits to distinct parent documents, ranked by best fused score.
At most *n* documents are returned (default ``BOR_TOP_N_DOCS``). The
returned rows carry the full document content; if the combined content