feat(rag): hybrid FTS+vector retrieval and multi-format ingestion — name-your-tool questions find the right document
This commit is contained in:
+39
-26
@@ -1,18 +1,21 @@
|
||||
"""POST /api/chat — a RAG chat turn streamed over SSE (PLAN §3/§4).
|
||||
|
||||
Flow (LOCKED A7/A15): embed the question → pgvector cosine top-K chunks →
|
||||
the **honesty gate** (A8: best score < ``BOR_RELEVANCE_THRESHOLD`` ⇒
|
||||
deflection) → locked persona prompt (PLAN §6) → ``turbo`` streamed as
|
||||
``delta`` events → final ``done`` event (``deflected``, ``sources``,
|
||||
``suggestions``) + ``query_log`` row + the per-turn log line (PLAN §9).
|
||||
Flow (LOCKED A7/A15): embed the question → hybrid retrieval (cosine
|
||||
top-N ∪ Postgres FTS top-N, RRF-fused) → the **honesty gate** → locked
|
||||
persona prompt (PLAN §6) → ``turbo`` streamed as ``delta`` events → final
|
||||
``done`` event (``deflected``, ``sources``, ``suggestions``) +
|
||||
``query_log`` row + the per-turn log line (PLAN §9).
|
||||
Mid-stream failures become a structured ``error`` event; a pre-stream DB
|
||||
outage is a plain 503 JSON.
|
||||
|
||||
Honesty gate: a weak retrieval (score strictly below the threshold — or
|
||||
an empty KB) flips the turn to deflection mode: the LOW prompt carries
|
||||
weak-hit *titles only* (never document content) plus deterministic
|
||||
"Maybe try" chips, and the ``done`` event / ``query_log`` row record
|
||||
``deflected=true`` with the weak score.
|
||||
Honesty gate (A8, revised 2026-08-21): LOW — deflection — only when the
|
||||
best cosine is strictly below ``BOR_RELEVANCE_THRESHOLD`` **and** no
|
||||
candidate chunk FTS-matches the question (``fts_hits == 0``). A
|
||||
name-your-tool question with weak vector overlap but a lexical hit still
|
||||
gets a grounded answer. Deflection mode carries weak-hit *titles only*
|
||||
(never document content) plus deterministic "Maybe try" chips, and the
|
||||
``done`` event / ``query_log`` row record ``deflected=true``, the weak
|
||||
score and the ``fts_hits`` count.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -62,7 +65,8 @@ def sse_event(payload: dict[str, Any]) -> str:
|
||||
class TurnPlan:
|
||||
"""What one chat turn sends to the LLM and reports on ``done``."""
|
||||
|
||||
top_score: float
|
||||
top_score: float # best cosine across candidates (query_log.top_score)
|
||||
fts_hits: int # lexical (OR-tsquery) candidates matched
|
||||
deflected: bool
|
||||
system_prompt: str
|
||||
docs: list[Document] # cited sources (weak hits when deflected)
|
||||
@@ -70,25 +74,32 @@ class TurnPlan:
|
||||
|
||||
|
||||
def plan_turn(chunks: Sequence[RetrievedChunk], settings: Settings) -> TurnPlan:
|
||||
"""Apply the honesty gate (A8) and assemble prompt + context for a turn.
|
||||
"""Apply the honesty gate (A8, revised) and assemble prompt + context.
|
||||
|
||||
* ``top_score >= threshold`` → grounded: HIGH prompt with the full
|
||||
top-N documents, no suggestions. A score exactly at the threshold
|
||||
is an answer — the gate is strict (``score < threshold``).
|
||||
* ``top_score < threshold`` (or no hits at all) → deflected: LOW
|
||||
prompt (``DEFLECT_MODE``) with weak-hit titles only — never document
|
||||
content — plus deterministic alternative-question chips derived
|
||||
from those titles.
|
||||
* **HIGH (grounded)** when ``best_cosine >= threshold`` **or**
|
||||
``fts_hits > 0``: HIGH prompt with the full top-N documents, no
|
||||
suggestions. A cosine exactly at the threshold is an answer — the
|
||||
gate is strict (``< threshold``).
|
||||
* **LOW (deflected)** only when ``best_cosine < threshold`` **and**
|
||||
``fts_hits == 0`` (or no hits at all): LOW prompt (``DEFLECT_MODE``)
|
||||
with weak-hit titles only — never document content — plus
|
||||
deterministic alternative-question chips derived from those titles.
|
||||
|
||||
``top_score`` (stored in ``query_log``) is the best cosine, so the
|
||||
gate input is always a pure vector-similarity number; the lexical
|
||||
signal is recorded separately as ``fts_hits``.
|
||||
"""
|
||||
top_score = chunks[0].score if chunks else 0.0
|
||||
if top_score >= settings.relevance_threshold:
|
||||
best_cosine = max((c.cosine for c in chunks), default=0.0)
|
||||
fts_hits = sum(1 for c in chunks if c.fts_hit)
|
||||
if best_cosine >= settings.relevance_threshold or fts_hits > 0:
|
||||
docs = select_documents(
|
||||
chunks, n=settings.top_n_docs, max_chars=settings.max_context_chars
|
||||
)
|
||||
return TurnPlan(top_score, False, build_high_prompt(docs), docs, [])
|
||||
return TurnPlan(best_cosine, fts_hits, False, build_high_prompt(docs), docs, [])
|
||||
titles = weak_hit_titles(chunks)
|
||||
return TurnPlan(
|
||||
top_score,
|
||||
best_cosine,
|
||||
fts_hits,
|
||||
True,
|
||||
build_deflect_prompt(titles),
|
||||
select_documents(chunks, n=settings.top_n_docs, max_chars=settings.max_context_chars),
|
||||
@@ -143,7 +154,7 @@ async def chat(
|
||||
# HIGH (grounded) or LOW (deflected) prompt + context.
|
||||
settings = get_settings()
|
||||
try:
|
||||
chunks = retrieve(db, question_vec)
|
||||
chunks = retrieve(db, request.message, question_vec)
|
||||
plan = plan_turn(chunks, settings)
|
||||
except Exception: # noqa: BLE001 — DB failure mid-turn
|
||||
logger.exception(
|
||||
@@ -188,6 +199,7 @@ async def chat(
|
||||
QueryLog(
|
||||
question=request.message,
|
||||
top_score=plan.top_score,
|
||||
fts_hits=plan.fts_hits,
|
||||
chunk_hits=len(chunks),
|
||||
deflected=plan.deflected,
|
||||
sources=", ".join(source_paths),
|
||||
@@ -199,11 +211,12 @@ async def chat(
|
||||
logger.exception("chat: failed to write query_log question=%r", request.message)
|
||||
|
||||
logger.info(
|
||||
"question=%r embed_ms=%d top_score=%.3f threshold=%.2f deflected=%s "
|
||||
"sources=%r total_ms=%d",
|
||||
"question=%r embed_ms=%d top_score=%.3f fts_hits=%d threshold=%.2f "
|
||||
"deflected=%s sources=%r total_ms=%d",
|
||||
request.message,
|
||||
embed_ms,
|
||||
plan.top_score,
|
||||
plan.fts_hits,
|
||||
settings.relevance_threshold,
|
||||
plan.deflected,
|
||||
source_paths,
|
||||
|
||||
+62
-2
@@ -8,8 +8,15 @@ from __future__ import annotations
|
||||
import os
|
||||
from functools import lru_cache
|
||||
|
||||
from pydantic import field_validator
|
||||
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
|
||||
#: The A9 import formats (PLAN anchor A9, revised 2026-08-21).
|
||||
#: ``BOR_IMPORT_EXTENSIONS`` may narrow — but never widen — this set.
|
||||
_ALLOWED_IMPORT_EXTENSIONS: frozenset[str] = frozenset(
|
||||
{"md", "markdown", "txt", "yaml", "yml", "json", "py"}
|
||||
)
|
||||
|
||||
|
||||
class Settings(BaseSettings):
|
||||
model_config = SettingsConfigDict(
|
||||
@@ -37,14 +44,58 @@ class Settings(BaseSettings):
|
||||
|
||||
# --- RAG tuning ---
|
||||
embedding_dim: int = 768 # verified against aipi /v1 (embed model)
|
||||
top_k_chunks: int = 4
|
||||
top_n_docs: int = 2
|
||||
relevance_threshold: float = 0.30
|
||||
# Honesty gate (A8, re-tuned 2026-08-21): the ``embed`` model's cosine
|
||||
# scores compress into 0.41–0.84 on the real corpus, so the old 0.30
|
||||
# default never discriminated. LOW only fires when best cosine < this
|
||||
# AND no candidate chunk matches the question lexically (see A8).
|
||||
relevance_threshold: float = 0.62
|
||||
max_context_chars: int = 24_000
|
||||
chunk_target_chars: int = 2_000
|
||||
chunk_overlap_chars: int = 200
|
||||
embed_batch_size: int = 16
|
||||
|
||||
# --- Hybrid retrieval (A7, revised 2026-08-21) ---
|
||||
# cosine top-N ∪ Postgres FTS top-N, fused with Reciprocal Rank Fusion
|
||||
# (score = Σ 1/(rrf_k + rank) over the lists a chunk appears in).
|
||||
#
|
||||
# The vector window is deliberately wider than the lexical one: a
|
||||
# name-your-tool question's best *lexical* chunk (e.g. the "Install"
|
||||
# section of gitlab.md) can sit far down the vector ranking because the
|
||||
# question embeds close to generic templates. A 100-wide window is what
|
||||
# lets such chunks double-hit (one RRF term per list) and outrank a
|
||||
# template that owns vector rank 1 — measured 2026-08-22 against the
|
||||
# live 2774-chunk KB for "How did I install gitlab?" (gitlab.md:1 at
|
||||
# vrank 100 / lrank 3 → fused 0.0221 vs the template's 0.0164).
|
||||
hybrid_vector_candidates: int = 100
|
||||
hybrid_lexical_candidates: int = 30
|
||||
rrf_k: int = 60
|
||||
|
||||
# --- Import scope (A9, revised 2026-08-21) ---
|
||||
# Comma-separated list of lowercased file extensions (no dot) imported
|
||||
# by ``scripts/import_docs.py``. Hidden (dot) path components are always
|
||||
# skipped, plus the importer's exclusion list.
|
||||
# Stored as a raw CSV string (env-native — no JSON) and parsed on demand
|
||||
# via :py:meth:`import_extension_set`. ``mode="after"`` validation runs
|
||||
# against the raw string so a typo fails loudly at startup.
|
||||
import_extensions: str = "md,markdown,txt,yaml,yml,json,py"
|
||||
|
||||
@field_validator("import_extensions")
|
||||
@classmethod
|
||||
def _import_extensions_known(cls, v: str) -> str:
|
||||
"""Reject unknown/empty formats loudly instead of silently importing
|
||||
nothing (a typo like ``md,jsonn`` would otherwise walk zero files)."""
|
||||
exts = {part.strip().lstrip(".").lower() for part in v.split(",") if part.strip()}
|
||||
if not exts:
|
||||
raise ValueError("import_extensions must name at least one format")
|
||||
unknown = exts - _ALLOWED_IMPORT_EXTENSIONS
|
||||
if unknown:
|
||||
raise ValueError(
|
||||
f"unknown import extension(s): {', '.join(sorted(unknown))} — "
|
||||
f"allowed: {', '.join(sorted(_ALLOWED_IMPORT_EXTENSIONS))}"
|
||||
)
|
||||
return v
|
||||
|
||||
# Suggested questions (onboarding + empty state).
|
||||
suggestions: list[str] = [
|
||||
"How is my Kubernetes cluster set up?",
|
||||
@@ -53,6 +104,15 @@ class Settings(BaseSettings):
|
||||
"What's currently running in the homelab?",
|
||||
]
|
||||
|
||||
@property
|
||||
def import_extension_set(self) -> frozenset[str]:
|
||||
"""Lowercased, dotted extension set (``.md``) for path filtering."""
|
||||
return frozenset(
|
||||
f".{part.strip().lstrip('.').lower()}"
|
||||
for part in self.import_extensions.split(",")
|
||||
if part.strip()
|
||||
)
|
||||
|
||||
@property
|
||||
def effective_api_key(self) -> str:
|
||||
"""API key for aipi: explicit setting, then $AIPI_KEY, then a placeholder."""
|
||||
|
||||
+4
-1
@@ -2,7 +2,7 @@
|
||||
|
||||
Data model — see ``.agent/PLAN.md`` §Data Model:
|
||||
|
||||
* ``documents`` — one row per ``*.md`` file (full content, path, sha256 hash).
|
||||
* ``documents`` — one row per imported A9 file (full content, path, sha256 hash).
|
||||
* ``chunks`` — retrieval units; each chunk points at its parent document
|
||||
via ``document_id``. This is how an embedding maps back to
|
||||
a document path (the "feed the whole document" requirement).
|
||||
@@ -74,6 +74,9 @@ class QueryLog(Base):
|
||||
id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
||||
question: Mapped[str] = mapped_column(Text)
|
||||
top_score: Mapped[float] = mapped_column(Float, default=0.0) # best cosine similarity
|
||||
#: Lexical (FTS) candidates matched — the OR-tsquery hit count (A8). NULL
|
||||
#: for pre-hybrid rows (migration 0002).
|
||||
fts_hits: Mapped[int | None] = mapped_column(Integer)
|
||||
chunk_hits: Mapped[int] = mapped_column(Integer, default=0)
|
||||
deflected: Mapped[bool] = mapped_column(Boolean, default=False) # True = honest "no idea"
|
||||
sources: Mapped[str] = mapped_column(Text, default="") # comma-joined source paths
|
||||
|
||||
+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)
|
||||
|
||||
+54
-16
@@ -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)
|
||||
]
|
||||
|
||||
+194
-23
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user