pablan/backend/app/rag/similarity.py

"""Similarity: "what else is close to this text", with a threshold.

Deliberately NOT the hybrid path. RRF produces a fusion rank, not a
similarity, and its `vector_distance` is None for hits that surfaced only
through full text — a threshold needs a comparable number. What both paths DO
share is the permission filter: the same `allowed` CTE, so a suggestion can
never point at something the caller may not read.

Two callers, two calibrated limits: a refinement grounds on loosely related
material, while a duplicate check may only propose merging on a close match.
"""

import logging
import time
import uuid
from dataclasses import dataclass

from sqlalchemy import and_, select
from sqlalchemy.ext.asyncio import AsyncSession

from app.llm.client import embed
from app.metrics import metrics
from app.models import Chunk, Document, User
from app.rag.permissions import searchable_documents_filter
from app.rag.retrieval import CANDIDATES_PER_SOURCE, heading_path

logger = logging.getLogger("pablan.rag")

# One mechanic at two moments: during capture a loose limit is useful
# because a near-miss still makes good context, while at review time only a
# high-confidence match may propose merging into an existing document.
# CHANGING THE EMBEDDING MODEL MEANS RE-MEASURING ALL THREE constants;
# tests/evals/test_duplicate_eval.py prints the numbers to do it with.
#
# Measured against bge-m3 on the fixture corpus (2026-07-20): drafts that
# duplicate an existing document land at 0.116-0.274, genuinely new topics
# at 0.402-0.486. 0.35 sits in that gap. The upper end of the duplicate
# range comes from REAL capture drafts, which are compressed notes rather
# than full prose and therefore sit further from their source than a
# hand-written paraphrase does — calibrating on paraphrases alone gives a
# threshold that misses real duplicates (it did: 0.25 missed one at 0.256).
CAPTURE_CONTEXT_MAX_DISTANCE = 0.45
DUPLICATE_MAX_DISTANCE = 0.35


@dataclass
class SimilarChunk:
    chunk_id: uuid.UUID
    document_id: uuid.UUID
    title: str
    heading_path: str
    content: str
    distance: float


@dataclass
class SimilarDocument:
    document_id: uuid.UUID
    title: str
    distance: float  # the closest chunk of that document


async def similar_chunks(
    db: AsyncSession,
    text: str,
    *,
    user: User,
    top_k: int = 5,
    max_distance: float,
    exclude_builtin: bool = False,
    exclude_document_id: uuid.UUID | None = None,
) -> list[SimilarChunk]:
    """Pure vector neighbours of a text, permission-filtered like everything
    else — the same `allowed` CTE `search()` uses.

    Deliberately NOT the hybrid path: RRF produces a fusion rank, not a
    similarity, and its `vector_distance` is None for hits that surfaced
    only through full text. A threshold needs a comparable number.

    `max_distance` is keyword-only and has no default on purpose: every
    caller names one of the two calibrated constants, so "similar" means
    exactly two things in this product and both are written down.

    `exclude_builtin` drops Pablan's own help pages. They are answerable
    through query mode on purpose (the product documents itself), but a
    capture or duplicate check asks "what does the COMPANY already know" —
    proposing to extend a help page, or telling an author their topic is
    "already documented" because a help page mentions it, is wrong.
    """
    started = time.monotonic()
    vector = (await embed([text]))[0]

    allowed_filter = searchable_documents_filter(user)
    if exclude_builtin:
        allowed_filter = and_(allowed_filter, Document.is_builtin.is_(False))
    if exclude_document_id is not None:
        # A document must never ground on itself (the extend flow re-opens a
        # published document and would otherwise retrieve its own chunks).
        allowed_filter = and_(allowed_filter, Document.id != exclude_document_id)
    allowed = select(Document.id, Document.title).where(allowed_filter).cte("allowed")
    distance = Chunk.embedding.cosine_distance(vector)

    # Threshold in Python, after ORDER BY ... LIMIT: a distance predicate in
    # WHERE fights the HNSW index, ordering and limiting is what it serves.
    rows = (
        await db.execute(
            select(
                Chunk.id,
                Chunk.document_id,
                Chunk.content,
                Chunk.meta,
                allowed.c.title,
                distance.label("distance"),
            )
            .join(allowed, allowed.c.id == Chunk.document_id)
            .order_by(distance)
            .limit(top_k)
        )
    ).all()

    results = [
        SimilarChunk(
            chunk_id=row.id,
            document_id=row.document_id,
            title=row.title,
            heading_path=heading_path(row.meta),
            content=row.content,
            distance=float(row.distance),
        )
        for row in rows
        if float(row.distance) <= max_distance
    ]

    duration = time.monotonic() - started
    metrics.inc("similarity_searches_total")
    metrics.observe("similarity_seconds", duration)
    logger.info(
        "similarity",
        extra={
            "event": "similarity",
            "duration_ms": round(duration * 1000),
            "candidate_count": len(rows),
            "result_count": len(results),
            "top_k": top_k,
        },
    )
    return results


async def similar_documents(
    db: AsyncSession,
    text: str,
    *,
    user: User,
    top_k: int = 3,
    max_distance: float,
    exclude_builtin: bool = False,
) -> list[SimilarDocument]:
    """Documents near a text, ranked by their closest chunk.

    Overfetches chunks and groups them, so this is literally the same search
    as `similar_chunks` — one notion of "similar" in the product, not two
    implementations that drift apart.
    """
    chunks = await similar_chunks(
        db,
        text,
        user=user,
        top_k=CANDIDATES_PER_SOURCE,
        max_distance=max_distance,
        exclude_builtin=exclude_builtin,
    )
    best: dict[uuid.UUID, SimilarDocument] = {}
    for chunk in chunks:  # distance-ordered, so the first hit per document wins
        best.setdefault(
            chunk.document_id,
            SimilarDocument(
                document_id=chunk.document_id,
                title=chunk.title,
                distance=chunk.distance,
            ),
        )
    return list(best.values())[:top_k]