118 lines
4.0 KiB
Python
118 lines
4.0 KiB
Python
"""pgvector cosine retrieval → parent-document mapping (PLAN §3/§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``.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import uuid
|
||
from collections.abc import Sequence
|
||
from dataclasses import dataclass
|
||
|
||
from sqlalchemy import select
|
||
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…]"
|
||
|
||
|
||
@dataclass
|
||
class RetrievedChunk:
|
||
"""One chunk hit: its cosine score plus the parent document row."""
|
||
|
||
chunk_id: uuid.UUID
|
||
position: int
|
||
content: str
|
||
score: float # 1 − cosine_distance (higher is more similar)
|
||
document: Document
|
||
|
||
|
||
def retrieve(
|
||
db: Session, question_embedding: list[float], top_k: int | None = None
|
||
) -> list[RetrievedChunk]:
|
||
"""Top-*top_k* chunks by pgvector cosine distance (``<=>``).
|
||
|
||
``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.
|
||
"""
|
||
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)
|
||
).all()
|
||
return [
|
||
RetrievedChunk(
|
||
chunk_id=chunk.id,
|
||
position=chunk.position,
|
||
content=chunk.content,
|
||
score=round(1.0 - float(dist), 6),
|
||
document=doc,
|
||
)
|
||
for chunk, dist, doc in rows
|
||
]
|
||
|
||
|
||
def weak_hit_titles(chunks: Sequence[RetrievedChunk]) -> list[str]:
|
||
"""Distinct parent-document titles of *chunks*, best chunk 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 chunk 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
|