289 lines
10 KiB
Python
289 lines
10 KiB
Python
"""Hybrid retrieval: pgvector cosine ∪ Postgres FTS, RRF-fused (PLAN §6, A7).
|
||
|
||
* **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, replace
|
||
|
||
from sqlalchemy import select, text
|
||
from sqlalchemy.orm import Session
|
||
|
||
from app.config import get_settings
|
||
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 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
|
||
document: Document
|
||
cosine: float = 0.0
|
||
fts_hit: bool = False
|
||
|
||
|
||
def lexical_tsquery(question: str) -> str | None:
|
||
"""OR-joined token string for ``to_tsquery('english', …)``, or ``None``.
|
||
|
||
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.
|
||
"""
|
||
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(limit)
|
||
).all()
|
||
return [
|
||
RetrievedChunk(
|
||
chunk_id=chunk.id,
|
||
position=chunk.position,
|
||
content=chunk.content,
|
||
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 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.
|
||
"""
|
||
titles: list[str] = []
|
||
seen: set[uuid.UUID] = set()
|
||
for rc in sorted(chunks, key=lambda c: c.score, reverse=True):
|
||
if rc.document.id in seen:
|
||
continue
|
||
seen.add(rc.document.id)
|
||
titles.append(rc.document.title)
|
||
return titles
|
||
|
||
|
||
def select_documents(
|
||
chunks: Sequence[RetrievedChunk],
|
||
n: int | None = None,
|
||
max_chars: int | None = None,
|
||
) -> list[Document]:
|
||
"""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
|
||
would exceed *max_chars* (default ``BOR_MAX_CONTEXT_CHARS``), the
|
||
lowest-ranked overflowing document is truncated in place with the
|
||
``[…truncated…]`` marker so the assembled context never exceeds the
|
||
budget (PLAN §6).
|
||
"""
|
||
top_n = n if n is not None else get_settings().top_n_docs
|
||
budget = max_chars if max_chars is not None else get_settings().max_context_chars
|
||
|
||
docs: list[Document] = []
|
||
seen: set[uuid.UUID] = set()
|
||
for rc in sorted(chunks, key=lambda c: c.score, reverse=True):
|
||
if rc.document.id in seen:
|
||
continue
|
||
seen.add(rc.document.id)
|
||
docs.append(rc.document)
|
||
docs = docs[:top_n]
|
||
|
||
remaining = budget
|
||
for doc in docs:
|
||
if len(doc.content) <= remaining:
|
||
remaining -= len(doc.content)
|
||
else:
|
||
keep = max(0, remaining - len(TRUNCATION_MARKER))
|
||
doc.content = doc.content[:keep] + TRUNCATION_MARKER
|
||
remaining = 0
|
||
return docs
|