feat(rag): index markdown KB — chunker, embed client, delta importer, Sources page
Phase 02 (story: import documents):
- fence-aware markdown chunker (heading sections, 200-char overlap,
heading anchor on every chunk, 1200-char hard cap, fence blocks
kept atomic and split under the cap)
- LLMClient over aipi (LiteLLM) reusing the openai client's httpx
transport to send a clean {model, input} payload — the openai SDK
injects encoding_format, which aipi's openai_like group rejects;
token-budget batching + halving retry for the endpoint's
~1024-token per-request input cap
- two-phase per-file upsert importer: sha256 delta (unchanged skip),
atomic commit, A9 exclusion walk, per-source prune, per-file error
tolerance (rollback + log + continue, non-zero CLI exit), adaptive
re-chunk at half target for URL-dense files the endpoint rejects
- scripts/import_docs CLI (repeatable --source, --prune, --limit,
defaults ~/Homelab + ~/Deployments)
- GET /api/docs with per-doc chunk counts; Sources page wired to the
real endpoint (stat cards, full-width a11y table, designed empty
state, DOM-built rows — no innerHTML)
- tests: 63 passed (chunker/llm/importer units, docs API + importer
integration), story E2E 3/3 (real endpoints, in-thread import);
app/ coverage 98%
- real KB imported: 672 docs / 8969 chunks in ~3m, idempotent
re-run (672 unchanged, 0 batches)
- harness: .agent/validate.sh now gates through uv (pytest +
coverage >90% + ruff + pyright) instead of system python3
This commit is contained in:
@@ -0,0 +1,279 @@
|
||||
"""Unit tests: importer directory walk + sha256 delta logic.
|
||||
|
||||
The walk tests are pure filesystem (``tmp_path``); the delta tests run
|
||||
against the local compose Postgres (preferred — a real vector table),
|
||||
skipping with clear instructions when the stack is not up.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import func, select
|
||||
|
||||
from app.models import Chunk, Document
|
||||
from app.rag.importer import (
|
||||
EXCLUDED_DIRS,
|
||||
import_sources,
|
||||
iter_markdown_files,
|
||||
)
|
||||
from app.rag.llm import EmbeddingError
|
||||
from tests.fakes import FakeEmbedder
|
||||
|
||||
|
||||
class _PoisonEmbedder(FakeEmbedder):
|
||||
"""Fails (like a real endpoint) on any text containing 'poison'."""
|
||||
|
||||
async def embed(self, texts: list[str]) -> list[list[float]]:
|
||||
if any("poison" in t for t in texts):
|
||||
raise EmbeddingError("embeddings endpoint refused the input (simulated)")
|
||||
return await super().embed(texts)
|
||||
|
||||
|
||||
class _CapEmbedder(FakeEmbedder):
|
||||
"""Simulates the endpoint's ~1024-token input cap at ~1.1 chars/token:
|
||||
any single text over 1000 chars is rejected (URL-dense worst case)."""
|
||||
|
||||
async def embed(self, texts: list[str]) -> list[list[float]]:
|
||||
if any(len(t) > 1000 for t in texts):
|
||||
raise EmbeddingError(
|
||||
"a single 1100-char chunk exceeded the endpoint's per-request "
|
||||
"input token cap — lower BOR_CHUNK_TARGET_CHARS and re-import"
|
||||
)
|
||||
return await super().embed(texts)
|
||||
|
||||
|
||||
def _cleanup_source(db, source: str) -> None:
|
||||
for doc in db.scalars(select(Document).where(Document.source == source)).all():
|
||||
db.delete(doc)
|
||||
db.commit()
|
||||
|
||||
|
||||
def test_iter_markdown_files_excludes_noncontent_dirs(tmp_path: Path) -> None:
|
||||
root = tmp_path / "proj"
|
||||
for d in (
|
||||
"notes/sub",
|
||||
".venv/lib",
|
||||
"node_modules/x",
|
||||
".git",
|
||||
"__pycache__",
|
||||
".pytest_cache",
|
||||
"dist",
|
||||
"build",
|
||||
):
|
||||
(root / d).mkdir(parents=True)
|
||||
files = {
|
||||
"README.md": "readme",
|
||||
"notes/sub/deep.md": "deep",
|
||||
".venv/lib/junk.md": "junk",
|
||||
"node_modules/x/j.md": "j",
|
||||
".git/c.md": "g",
|
||||
"__pycache__/c.md": "p",
|
||||
".pytest_cache/c.md": "pc",
|
||||
"dist/d.md": "d",
|
||||
"build/b.md": "b",
|
||||
}
|
||||
for rel, text in files.items():
|
||||
(root / rel).write_text(text)
|
||||
(root / "notes" / "not-md.txt").write_text("skip me")
|
||||
|
||||
found = {p.relative_to(root).as_posix() for p in iter_markdown_files(root)}
|
||||
assert found == {"README.md", "notes/sub/deep.md"}
|
||||
|
||||
|
||||
def test_iter_markdown_files_missing_dir_yields_nothing(tmp_path: Path) -> None:
|
||||
assert iter_markdown_files(tmp_path / "definitely-missing") == []
|
||||
|
||||
|
||||
def test_excluded_dirs_match_plan_anchor_a9() -> None:
|
||||
assert {
|
||||
".venv", "node_modules", ".git", "__pycache__", ".pytest_cache", "dist", "build"
|
||||
} == EXCLUDED_DIRS
|
||||
|
||||
|
||||
def test_added_then_unchanged_then_updated_then_pruned(db, tmp_path: Path) -> None:
|
||||
root = tmp_path / "src"
|
||||
root.mkdir()
|
||||
(root / "a.md").write_text("# A\n\nalpha\n\n## Sub\n\nmore alpha\n")
|
||||
(root / "b.md").write_text("# B\n\nbeta\n")
|
||||
llm = FakeEmbedder()
|
||||
try:
|
||||
s1 = asyncio.run(import_sources([root], llm, session=db))
|
||||
assert (s1.files, s1.added, s1.unchanged, s1.updated, s1.pruned) == (2, 2, 0, 0, 0)
|
||||
# a.md has two sections (2 chunks), b.md one (1 chunk).
|
||||
assert s1.chunks == 3
|
||||
# Embeddings are stored with the configured dimension.
|
||||
n = db.scalar(
|
||||
select(func.count())
|
||||
.select_from(Chunk)
|
||||
.join(Document, Document.id == Chunk.document_id)
|
||||
.where(Document.source == root.name)
|
||||
)
|
||||
assert n == 3
|
||||
for c in db.scalars(
|
||||
select(Chunk)
|
||||
.join(Document, Document.id == Chunk.document_id)
|
||||
.where(Document.source == root.name)
|
||||
).all():
|
||||
assert c.embedding is not None and len(c.embedding) == 768
|
||||
|
||||
s2 = asyncio.run(import_sources([root], llm, session=db))
|
||||
assert s2.added == 0 and s2.unchanged == 2
|
||||
|
||||
(root / "a.md").write_text("# A\n\nalpha CHANGED\n")
|
||||
s3 = asyncio.run(import_sources([root], llm, session=db))
|
||||
assert s3.updated == 1 and s3.unchanged == 1
|
||||
doc = db.scalar(
|
||||
select(Document).where(Document.source == root.name, Document.path == "a.md")
|
||||
)
|
||||
assert doc is not None and "CHANGED" in doc.content
|
||||
|
||||
(root / "a.md").unlink()
|
||||
s4 = asyncio.run(import_sources([root], llm, session=db, prune=True))
|
||||
assert s4.pruned == 1
|
||||
assert db.scalar(
|
||||
select(Document).where(Document.source == root.name, Document.path == "a.md")
|
||||
) is None
|
||||
# Chunks of the pruned document are gone (FK cascade).
|
||||
n_after = db.scalar(
|
||||
select(func.count())
|
||||
.select_from(Chunk)
|
||||
.join(Document, Document.id == Chunk.document_id)
|
||||
.where(Document.source == root.name)
|
||||
)
|
||||
assert n_after == 1
|
||||
finally:
|
||||
_cleanup_source(db, root.name)
|
||||
|
||||
|
||||
def test_embedding_failure_is_logged_and_import_continues(db, tmp_path: Path) -> None:
|
||||
"""A file the embedding endpoint refuses must not abort the whole KB:
|
||||
its rows are rolled back, the error is counted, and other files import."""
|
||||
root = tmp_path / "mixed"
|
||||
root.mkdir()
|
||||
(root / "bad.md").write_text("# Bad\n\npoison content that the endpoint refuses\n")
|
||||
(root / "good.md").write_text("# Good\n\nperfectly fine content\n")
|
||||
try:
|
||||
summary = asyncio.run(import_sources([root], _PoisonEmbedder(), session=db))
|
||||
assert summary.files == 2
|
||||
assert summary.errors == 1
|
||||
assert summary.added == 1 # only good.md
|
||||
# bad.md left no row and no orphan chunks behind (rolled back).
|
||||
assert db.scalar(
|
||||
select(Document).where(Document.source == root.name, Document.path == "bad.md")
|
||||
) is None
|
||||
assert db.scalar(
|
||||
select(func.count())
|
||||
.select_from(Chunk)
|
||||
.join(Document, Document.id == Chunk.document_id)
|
||||
.where(Document.source == root.name, Document.path == "bad.md")
|
||||
) == 0
|
||||
assert db.scalar(
|
||||
select(Document).where(Document.source == root.name, Document.path == "good.md")
|
||||
) is not None
|
||||
finally:
|
||||
_cleanup_source(db, root.name)
|
||||
|
||||
|
||||
def test_oversized_chunk_triggers_adaptive_rechunk(db, tmp_path: Path) -> None:
|
||||
"""A URL-dense paragraph the endpoint rejects must be re-chunked smaller
|
||||
for that file only — the import still succeeds."""
|
||||
root = tmp_path / "dense"
|
||||
root.mkdir()
|
||||
# One ~1165-char paragraph: under the 1200-char hard cap, over the
|
||||
# simulated token cap. The retry at 600 chars must split it.
|
||||
para = "see https://example.com/" + "a" * 1100
|
||||
(root / "dense.md").write_text(f"# D\n\n{para}\n")
|
||||
try:
|
||||
summary = asyncio.run(import_sources([root], _CapEmbedder(), session=db))
|
||||
assert summary.errors == 0
|
||||
assert summary.added == 1
|
||||
doc = db.scalar(
|
||||
select(Document).where(Document.source == root.name, Document.path == "dense.md")
|
||||
)
|
||||
assert doc is not None
|
||||
assert len(doc.chunks) >= 2 # re-chunked smaller than the hard cap
|
||||
assert all(len(c.content) <= 1000 for c in doc.chunks)
|
||||
assert all(c.embedding is not None for c in doc.chunks)
|
||||
# The content survives the split.
|
||||
assert "".join(c.content for c in doc.chunks).count("a" * 500) >= 1
|
||||
finally:
|
||||
_cleanup_source(db, root.name)
|
||||
|
||||
|
||||
def test_missing_source_dir_is_skipped(db, tmp_path: Path) -> None:
|
||||
llm = FakeEmbedder()
|
||||
summary = asyncio.run(import_sources([tmp_path / "missing"], llm, session=db))
|
||||
assert summary.files == 0 and summary.added == 0
|
||||
|
||||
|
||||
def test_limit_caps_files_and_disables_prune(db, tmp_path: Path) -> None:
|
||||
root = tmp_path / "limited"
|
||||
root.mkdir()
|
||||
for name in ("a.md", "b.md", "c.md"):
|
||||
(root / name).write_text(f"# {name}\n\nbody {name}\n")
|
||||
llm = FakeEmbedder()
|
||||
try:
|
||||
summary = asyncio.run(import_sources([root], llm, limit=2, session=db, prune=True))
|
||||
assert summary.files == 2 and summary.added == 2
|
||||
# c.md was never walked, so it must NOT be pruned (prune disabled
|
||||
# under --limit) — and nothing else disappears either.
|
||||
assert summary.pruned == 0
|
||||
assert db.scalar(
|
||||
select(func.count()).select_from(Document).where(Document.source == root.name)
|
||||
) == 2
|
||||
finally:
|
||||
_cleanup_source(db, root.name)
|
||||
|
||||
|
||||
def test_limit_must_be_positive(db, tmp_path: Path) -> None:
|
||||
with pytest.raises(ValueError):
|
||||
asyncio.run(import_sources([tmp_path], FakeEmbedder(), limit=0, session=db))
|
||||
|
||||
|
||||
def test_prune_is_scoped_to_the_given_sources(db, tmp_path: Path) -> None:
|
||||
src_x = tmp_path / "SourceX"
|
||||
src_y = tmp_path / "SourceY"
|
||||
src_x.mkdir()
|
||||
src_y.mkdir()
|
||||
(src_x / "x.md").write_text("# X\n\nx body\n")
|
||||
(src_y / "y.md").write_text("# Y\n\ny body\n")
|
||||
llm = FakeEmbedder()
|
||||
try:
|
||||
asyncio.run(import_sources([src_x, src_y], llm, session=db))
|
||||
# Re-import ONLY source Y (y.md removed) with prune: source X's doc
|
||||
# must survive — prune never touches sources not passed to this run.
|
||||
(src_y / "y.md").unlink()
|
||||
summary = asyncio.run(import_sources([src_y], llm, session=db, prune=True))
|
||||
assert summary.pruned == 1
|
||||
assert db.scalar(
|
||||
select(Document).where(Document.source == "SourceX", Document.path == "x.md")
|
||||
) is not None
|
||||
finally:
|
||||
_cleanup_source(db, "SourceX")
|
||||
_cleanup_source(db, "SourceY")
|
||||
|
||||
|
||||
def test_chunk_positions_and_titles(db, tmp_path: Path) -> None:
|
||||
root = tmp_path / "titled"
|
||||
root.mkdir()
|
||||
(root / "multi.md").write_text("# Real Title\n\n## One\n\na\n\n## Two\n\nb\n")
|
||||
(root / "noh1.md").write_text("## Only heading\n\nbody\n")
|
||||
llm = FakeEmbedder()
|
||||
try:
|
||||
asyncio.run(import_sources([root], llm, session=db))
|
||||
titles = {
|
||||
d.path: d.title
|
||||
for d in db.scalars(select(Document).where(Document.source == root.name)).all()
|
||||
}
|
||||
assert titles["multi.md"] == "Real Title" # H1 wins
|
||||
assert titles["noh1.md"] == "noh1" # …else the file stem
|
||||
doc = db.scalar(
|
||||
select(Document).where(Document.source == root.name, Document.path == "multi.md")
|
||||
)
|
||||
assert doc is not None
|
||||
positions = sorted(c.position for c in doc.chunks)
|
||||
assert positions == list(range(len(doc.chunks))) and len(doc.chunks) >= 2
|
||||
finally:
|
||||
_cleanup_source(db, root.name)
|
||||
Reference in New Issue
Block a user