feat(rag): hybrid FTS+vector retrieval and multi-format ingestion — name-your-tool questions find the right document
This commit is contained in:
+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