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,247 @@
|
||||
"""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:
|
||||
|
||||
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
|
||||
|
||||
``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).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime
|
||||
from pathlib import Path
|
||||
from typing import Protocol
|
||||
|
||||
from sqlalchemy import select
|
||||
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.llm import EmbeddingError
|
||||
|
||||
logger = logging.getLogger("app.importer")
|
||||
|
||||
#: Non-content directories never imported (PLAN anchor A9).
|
||||
EXCLUDED_DIRS: frozenset[str] = frozenset(
|
||||
{".venv", "node_modules", ".git", "__pycache__", ".pytest_cache", "dist", "build"}
|
||||
)
|
||||
|
||||
|
||||
class Embedder(Protocol):
|
||||
"""Everything the importer needs from the LLM client (duck-typed for tests)."""
|
||||
|
||||
settings: Settings
|
||||
embed_batches: int
|
||||
|
||||
async def embed(self, texts: list[str]) -> list[list[float]]: ...
|
||||
|
||||
|
||||
@dataclass
|
||||
class ImportSummary:
|
||||
"""Counts for one import run (also printed by the CLI)."""
|
||||
|
||||
files: int = 0
|
||||
added: int = 0
|
||||
updated: int = 0
|
||||
unchanged: int = 0
|
||||
pruned: int = 0
|
||||
errors: int = 0
|
||||
chunks: int = 0
|
||||
embed_batches: int = 0
|
||||
|
||||
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",
|
||||
self.files,
|
||||
self.added,
|
||||
self.updated,
|
||||
self.unchanged,
|
||||
self.pruned,
|
||||
self.errors,
|
||||
self.chunks,
|
||||
self.embed_batches,
|
||||
)
|
||||
|
||||
|
||||
def iter_markdown_files(root: Path, excluded: frozenset[str] = EXCLUDED_DIRS) -> list[Path]:
|
||||
"""All ``*.md`` files under *root* (sorted), skipping excluded dirs (A9)."""
|
||||
if not root.is_dir():
|
||||
return []
|
||||
files: list[Path] = []
|
||||
for path in sorted(root.rglob("*.md")):
|
||||
rel = path.relative_to(root)
|
||||
if any(part in excluded for part in rel.parts[:-1]):
|
||||
continue
|
||||
files.append(path)
|
||||
return files
|
||||
|
||||
|
||||
async def import_sources(
|
||||
sources: list[Path],
|
||||
llm: Embedder,
|
||||
*,
|
||||
prune: bool = False,
|
||||
limit: int | None = None,
|
||||
session: Session | None = None,
|
||||
) -> ImportSummary:
|
||||
"""Import every ``*.md`` 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
|
||||
disables pruning, since an incomplete walk must not drive deletions.
|
||||
"""
|
||||
if limit is not None and limit <= 0:
|
||||
raise ValueError("limit must be >= 1")
|
||||
summary = ImportSummary()
|
||||
owns_session = session is None
|
||||
if session is None:
|
||||
session = SessionLocal()
|
||||
seen: set[tuple[str, str]] = set()
|
||||
source_names: set[str] = set()
|
||||
try:
|
||||
for root in sources:
|
||||
if not root.is_dir():
|
||||
logger.warning("import: source dir not found, skipping: %s", root)
|
||||
continue
|
||||
if limit is not None and summary.files >= limit:
|
||||
break
|
||||
source = root.name
|
||||
source_names.add(source)
|
||||
for path in iter_markdown_files(root):
|
||||
if limit is not None and summary.files >= limit:
|
||||
break
|
||||
rel = path.relative_to(root).as_posix()
|
||||
seen.add((source, rel))
|
||||
summary.files += 1
|
||||
try:
|
||||
await _index_file(
|
||||
session, source=source, rel=rel, full_path=path, llm=llm,
|
||||
summary=summary,
|
||||
)
|
||||
except EmbeddingError as e:
|
||||
# A pathological file (e.g. content the embedding endpoint
|
||||
# refuses) must not abort the whole KB: roll back its
|
||||
# uncommitted rows, log loudly, and keep going. The next
|
||||
# run retries it.
|
||||
session.rollback()
|
||||
summary.errors += 1
|
||||
logger.error("import: error source=%s path=%s — %s", source, rel, e)
|
||||
if prune:
|
||||
if limit is not None:
|
||||
logger.warning("import: --prune ignored because --limit was given")
|
||||
else:
|
||||
summary.pruned = _prune(session, source_names, seen)
|
||||
summary.embed_batches = llm.embed_batches
|
||||
summary.log()
|
||||
return summary
|
||||
finally:
|
||||
if owns_session:
|
||||
session.close()
|
||||
|
||||
|
||||
async def _index_file(
|
||||
session: Session,
|
||||
*,
|
||||
source: str,
|
||||
rel: str,
|
||||
full_path: Path,
|
||||
llm: Embedder,
|
||||
summary: ImportSummary,
|
||||
) -> None:
|
||||
"""Upsert one file: doc row + chunk rows + embeddings, one transaction."""
|
||||
settings = llm.settings
|
||||
content = full_path.read_text(encoding="utf-8", errors="replace")
|
||||
digest = hashlib.sha256(content.encode("utf-8")).hexdigest()
|
||||
doc = session.scalar(select(Document).where(Document.source == source, Document.path == rel))
|
||||
if doc is not None and doc.content_hash == digest:
|
||||
summary.unchanged += 1
|
||||
logger.info("import: unchanged source=%s path=%s", source, rel)
|
||||
return
|
||||
|
||||
verb = "updated" if doc is not None else "added"
|
||||
title = extract_title(content, fallback=full_path.stem)
|
||||
if doc is None:
|
||||
doc = Document(
|
||||
source=source,
|
||||
path=rel,
|
||||
full_path=str(full_path),
|
||||
title=title,
|
||||
content=content,
|
||||
content_hash=digest,
|
||||
indexed_at=datetime.now(UTC),
|
||||
)
|
||||
session.add(doc)
|
||||
else:
|
||||
doc.full_path = str(full_path)
|
||||
doc.title = title
|
||||
doc.content = content
|
||||
doc.content_hash = digest
|
||||
doc.indexed_at = datetime.now(UTC)
|
||||
|
||||
session.flush() # guarantees doc.id even for brand-new rows
|
||||
|
||||
# Phase 1+2 — chunk, replace the chunk rows (embeddings NULL), embed,
|
||||
# and commit the whole file atomically. If the endpoint rejects a chunk
|
||||
# as over its input token cap (URL-dense paragraphs tokenize at ~1.1
|
||||
# chars/token), halve this file's chunk target and retry — the global
|
||||
# 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)
|
||||
doc.chunks = [
|
||||
Chunk(document_id=doc.id, position=i, content=c) for i, c in enumerate(chunks_text)
|
||||
]
|
||||
session.flush() # delete-orphan cascade drops the previous rows
|
||||
if not doc.chunks:
|
||||
break
|
||||
try:
|
||||
vectors = await llm.embed([c.content for c in doc.chunks])
|
||||
for row, vec in zip(doc.chunks, vectors, strict=True):
|
||||
row.embedding = vec
|
||||
break
|
||||
except EmbeddingError as e:
|
||||
if "token cap" not in str(e) or target <= 400:
|
||||
raise
|
||||
logger.info(
|
||||
"import: re-chunking at %d chars after endpoint token cap: %s",
|
||||
target // 2,
|
||||
rel,
|
||||
)
|
||||
target //= 2
|
||||
session.commit()
|
||||
|
||||
if verb == "added":
|
||||
summary.added += 1
|
||||
else:
|
||||
summary.updated += 1
|
||||
summary.chunks += len(chunks_text)
|
||||
logger.info("import: %s source=%s path=%s chunks=%d", verb, source, rel, len(chunks_text))
|
||||
|
||||
|
||||
def _prune(session: Session, source_names: set[str], seen: set[tuple[str, str]]) -> int:
|
||||
"""Delete documents of *source_names* whose file is no longer in *seen*."""
|
||||
if not source_names:
|
||||
return 0
|
||||
pruned = 0
|
||||
docs = session.scalars(select(Document).where(Document.source.in_(source_names))).all()
|
||||
for doc in docs:
|
||||
if (doc.source, doc.path) not in seen:
|
||||
session.delete(doc)
|
||||
pruned += 1
|
||||
logger.info("import: pruned source=%s path=%s", doc.source, doc.path)
|
||||
if pruned:
|
||||
session.commit()
|
||||
return pruned
|
||||
Reference in New Issue
Block a user