kb/engine/kb/search.py

"""Hybrid search — FTS5 + sqlite-vec with Reciprocal Rank Fusion."""

import json
import logging
import struct
import sqlite3

logger = logging.getLogger("kb.search")


def hybrid_search(
    conn: sqlite3.Connection,
    query: str,
    cfg,
    top: int = 10,
    tags: list[str] | None = None,
    doc_type: str | None = None,
    fts_only: bool = False,
    vec_only: bool = False,
    threshold: float | None = None,
    explain: bool = False,
    rerank: bool | None = None,
) -> dict:
    """Run hybrid search and return merged, enriched results.

    Args:
        conn: SQLite connection (with row_factory = sqlite3.Row).
        query: User search query string.
        cfg: Config object with ``model`` and ``device`` attributes.
        top: Maximum number of results to return.
        tags: Optional tag filter — documents must have *all* listed tags.
        doc_type: Optional document-type filter.
        fts_only: Only use FTS5 (skip vector search).
        vec_only: Only use vector search (skip FTS5).
        threshold: Optional minimum score; results below are dropped.
        explain: Attach a per-result score breakdown (arm scores, ranks,
            RRF contributions, rerank blend) under an ``explain`` key.
        rerank: Cross-encoder rerank the top candidates. None uses the
            engine default (cfg.rerank_enabled); True still requires a
            loaded reranker and degrades silently to plain retrieval
            otherwise. Never applied to fts_only / vec_only searches.

    Returns:
        Dict with keys: query, results, total_matches, returned, reranked.
    """
    from kb import reranker

    want_rerank = cfg.rerank_enabled if rerank is None else rerank
    do_rerank = (
        want_rerank
        and not fts_only
        and not vec_only
        and reranker.is_available()
    )

    candidate_count = top * 3
    if do_rerank:
        candidate_count = max(candidate_count, cfg.rerank_candidates)

    fts_results: dict[int, float] = {}
    vec_results: dict[int, float] = {}

    if not vec_only:
        fts_results = _fts_search(conn, query, candidate_count, tags, doc_type)

    if not fts_only:
        vec_results = _vector_search(conn, query, candidate_count, tags, doc_type)

    # --- merge ---------------------------------------------------------------
    if fts_only:
        merged = sorted(fts_results.items(), key=lambda x: x[1], reverse=True)
        details = _single_arm_details("fts", fts_results)
    elif vec_only:
        merged = sorted(vec_results.items(), key=lambda x: x[1], reverse=True)
        details = _single_arm_details("vec", vec_results)
    else:
        merged, details = _rrf_merge(fts_results, vec_results)

    # Apply threshold filter — use config default if not specified per-query
    effective_threshold = threshold if threshold is not None else cfg.search_threshold
    if effective_threshold > 0:
        merged = [(cid, score) for cid, score in merged if score >= effective_threshold]

    total_matches = len(merged)

    # --- rerank --------------------------------------------------------------
    # Blended scores are 0-1 normalised, a different scale from RRF scores;
    # the threshold above was applied to RRF scores and is NOT re-applied.
    reranked = False
    if do_rerank and merged:
        candidates = merged[: cfg.rerank_candidates]
        rr_scores = reranker.rerank_scores(
            query, _fetch_chunk_texts(conn, [cid for cid, _ in candidates])
        )
        merged = _blend_rerank(candidates, rr_scores, details)
        reranked = True

    merged = merged[:top]

    # --- enrich --------------------------------------------------------------
    results = _enrich(conn, merged, details if explain else None)

    return {
        "query": query,
        "results": results,
        "total_matches": total_matches,
        "returned": len(results),
        "reranked": reranked,
    }


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

def _sanitize_fts_query(query: str) -> str:
    """Escape a raw user query for safe use with FTS5 MATCH.

    Splits on whitespace, strips double quotes from each token, wraps each
    token in double quotes (making FTS5 treat all content as literals), and
    joins with spaces.  Returns empty string if no valid tokens remain.
    """
    tokens = []
    for token in query.split():
        token = token.replace('"', '')
        if token:
            tokens.append(f'"{token}"')
    return " ".join(tokens)


def _fts_search(
    conn: sqlite3.Connection,
    query: str,
    limit: int,
    tags: list[str] | None,
    doc_type: str | None,
) -> dict[int, float]:
    """FTS5 search on ``chunks_fts``.

    Returns:
        {chunk_id: bm25_score} where scores are positive (higher = better).
    """
    safe_query = _sanitize_fts_query(query)
    if not safe_query:
        return {}

    sql = "SELECT f.rowid AS chunk_id, bm25(chunks_fts) AS rank FROM chunks_fts f"
    joins: list[str] = []
    where: list[str] = ["chunks_fts MATCH ?"]
    params: list = [safe_query]

    if tags or doc_type:
        joins.append("JOIN chunks c ON f.rowid = c.id")
        joins.append("JOIN documents d ON c.document_id = d.id")

    if doc_type:
        where.append("d.doc_type = ?")
        params.append(doc_type)

    if tags:
        for i, tag in enumerate(tags):
            joins.append(f"JOIN document_tags dt{i} ON d.id = dt{i}.document_id")
            joins.append(f"JOIN tags t{i} ON dt{i}.tag_id = t{i}.id")
            where.append(f"t{i}.name = ?")
            params.append(tag.strip().lower())

    sql += " " + " ".join(joins)
    sql += " WHERE " + " AND ".join(where)
    sql += " ORDER BY rank LIMIT ?"
    params.append(limit)

    try:
        rows = conn.execute(sql, params).fetchall()
    except sqlite3.OperationalError:
        logger.warning("FTS5 query failed for input: %r", query)
        return {}

    # BM25 returns negative values (lower = better match); negate so
    # higher = better.
    return {row[0]: -row[1] for row in rows}


def _vector_search(
    conn: sqlite3.Connection,
    query: str,
    limit: int,
    tags: list[str] | None,
    doc_type: str | None,
) -> dict[int, float]:
    """Embed *query* and search ``chunks_vec`` via sqlite-vec.

    Returns:
        {chunk_id: similarity} where similarity = 1 / (1 + distance).
    """
    from kb.embeddings import embed_texts

    query_embedding = embed_texts([query])[0]
    blob = struct.pack(f"{len(query_embedding)}f", *query_embedding)

    rows = conn.execute(
        """
        SELECT chunk_id, distance
        FROM chunks_vec
        WHERE embedding MATCH ?
        ORDER BY distance
        LIMIT ?
        """,
        (blob, limit),
    ).fetchall()

    results: dict[int, float] = {}
    for row in rows:
        chunk_id = row[0]
        distance = row[1]
        similarity = 1.0 / (1.0 + distance)

        # Post-hoc tag / doc_type filtering for vector results
        if tags or doc_type:
            if not _passes_filters(conn, chunk_id, tags, doc_type):
                continue

        results[chunk_id] = similarity

    return results


def _passes_filters(
    conn: sqlite3.Connection,
    chunk_id: int,
    tags: list[str] | None,
    doc_type: str | None,
) -> bool:
    """Return True if chunk passes tag and doc_type filters."""
    sql = """
        SELECT d.id FROM chunks c
        JOIN documents d ON c.document_id = d.id
        WHERE c.id = ?
    """
    params: list = [chunk_id]

    if doc_type:
        sql += " AND d.doc_type = ?"
        params.append(doc_type)

    doc_row = conn.execute(sql, params).fetchone()
    if not doc_row:
        return False

    if tags:
        doc_id = doc_row[0]
        placeholders = ",".join("?" * len(tags))
        normalised = [t.strip().lower() for t in tags]
        count = conn.execute(
            f"""
            SELECT COUNT(DISTINCT t.name) FROM document_tags dt
            JOIN tags t ON dt.tag_id = t.id
            WHERE dt.document_id = ? AND t.name IN ({placeholders})
            """,
            [doc_id, *normalised],
        ).fetchone()[0]
        if count < len(tags):
            return False

    return True


def _rrf_merge(
    fts_results: dict[int, float],
    vec_results: dict[int, float],
    k: int = 60,
) -> tuple[list[tuple[int, float]], dict[int, dict]]:
    """Reciprocal Rank Fusion over two scored result sets.

    Each set is ranked independently (highest score first, rank starts at 1).
    RRF score for a document = sum of 1/(k + rank) across sets it appears in,
    plus a top-rank bonus per arm: +0.05 for rank 1, +0.02 for ranks 2-3.
    The bonus preserves exact matches — a chunk at the top of either arm is
    nearly impossible to displace via mid-rank RRF accumulation alone.

    Score scale: base RRF maxes at 2/(k+1) ≈ 0.033 (k=60); with bonuses the
    ceiling is ≈ 0.133. Bonuses only raise scores, so the threshold filter
    (default 0.01) can never drop a result that base RRF would have kept.

    Returns:
        (scores, details) — scores is a sorted list of (chunk_id, rrf_score),
        highest first; details maps chunk_id to a per-arm score breakdown
        suitable for the ``explain`` response field.
    """
    fts_ranked = _rank_by_score(fts_results)
    vec_ranked = _rank_by_score(vec_results)

    all_ids = set(fts_ranked) | set(vec_ranked)
    scores: list[tuple[int, float]] = []
    details: dict[int, dict] = {}

    for chunk_id in all_ids:
        fts_rank = fts_ranked.get(chunk_id)
        vec_rank = vec_ranked.get(chunk_id)
        rrf_fts = 1.0 / (k + fts_rank) if fts_rank is not None else None
        rrf_vec = 1.0 / (k + vec_rank) if vec_rank is not None else None
        bonus = _top_rank_bonus(fts_rank) + _top_rank_bonus(vec_rank)
        rrf = (rrf_fts or 0.0) + (rrf_vec or 0.0) + bonus

        details[chunk_id] = {
            "fts_score": _round6(fts_results.get(chunk_id)),
            "fts_rank": fts_rank,
            "vec_score": _round6(vec_results.get(chunk_id)),
            "vec_rank": vec_rank,
            "rrf_fts": _round6(rrf_fts),
            "rrf_vec": _round6(rrf_vec),
            "bonus": bonus,
            "final_score": _round6(rrf),
        }
        scores.append((chunk_id, rrf))

    scores.sort(key=lambda x: x[1], reverse=True)
    return scores, details


def _top_rank_bonus(rank: int | None) -> float:
    """Bonus for appearing at the top of one arm's ranking."""
    if rank == 1:
        return 0.05
    if rank in (2, 3):
        return 0.02
    return 0.0


def _single_arm_details(arm: str, results: dict[int, float]) -> dict[int, dict]:
    """Explain details for fts_only / vec_only searches (raw arm scores)."""
    ranked = _rank_by_score(results)
    return {
        chunk_id: {
            f"{arm}_score": _round6(score),
            f"{arm}_rank": ranked[chunk_id],
            "final_score": _round6(score),
        }
        for chunk_id, score in results.items()
    }


def _round6(value: float | None) -> float | None:
    return round(value, 6) if value is not None else None


def _fetch_chunk_texts(conn: sqlite3.Connection, chunk_ids: list[int]) -> list[str]:
    """Fetch chunk texts in the same order as *chunk_ids*."""
    placeholders = ",".join("?" * len(chunk_ids))
    rows = conn.execute(
        f"SELECT id, text FROM chunks WHERE id IN ({placeholders})", chunk_ids
    ).fetchall()
    by_id = {row[0]: row[1] for row in rows}
    return [by_id.get(cid, "") for cid in chunk_ids]


def _blend_rerank(
    candidates: list[tuple[int, float]],
    rr_scores: list[float],
    details: dict[int, dict],
) -> list[tuple[int, float]]:
    """Blend retrieval and cross-encoder scores, position-aware.

    Retrieval scores are min-max normalised within the candidate set; rerank
    scores are already 0-1. The retrieval weight depends on pre-rerank rank —
    75% for ranks 1-3, 60% for 4-10, 40% for 11+ — so the reranker can rescue
    mid-ranked semantic matches without destroying top exact-match hits.

    *candidates* must be in retrieval order; *rr_scores* aligned with it.
    Mutates *details* with the blend breakdown. Returns (chunk_id, blended)
    sorted highest first.
    """
    retrieval = [score for _, score in candidates]
    lo, hi = min(retrieval), max(retrieval)
    span = hi - lo

    blended: list[tuple[int, float]] = []
    for i, ((chunk_id, score), rr) in enumerate(zip(candidates, rr_scores)):
        rank = i + 1
        norm = (score - lo) / span if span > 0 else 1.0
        if rank <= 3:
            weight = 0.75
        elif rank <= 10:
            weight = 0.60
        else:
            weight = 0.40
        final = weight * norm + (1.0 - weight) * rr

        if chunk_id in details:
            details[chunk_id].update({
                "pre_rerank_rank": rank,
                "retrieval_norm": _round6(norm),
                "rerank_score": _round6(rr),
                "blend_weight": weight,
                "final_score": _round6(final),
            })
        blended.append((chunk_id, final))

    blended.sort(key=lambda x: x[1], reverse=True)
    return blended


def _rank_by_score(results: dict[int, float]) -> dict[int, int]:
    """Return {id: 1-based rank} sorted by score descending."""
    ordered = sorted(results, key=results.get, reverse=True)
    return {cid: rank for rank, cid in enumerate(ordered, start=1)}


def _enrich(
    conn: sqlite3.Connection,
    merged: list[tuple[int, float]],
    details: dict[int, dict] | None = None,
) -> list[dict]:
    """Fetch chunk text, document metadata, chunk metadata, and tags.

    When *details* is given, each result gains an ``explain`` key with its
    score breakdown.
    """
    results: list[dict] = []

    for chunk_id, score in merged:
        row = conn.execute(
            """
            SELECT c.id, c.text, c.chunk_index, c.metadata AS chunk_meta,
                   d.id AS doc_id, d.title, d.doc_type, d.source_path,
                   d.created_at
            FROM chunks c
            JOIN documents d ON c.document_id = d.id
            WHERE c.id = ?
            """,
            (chunk_id,),
        ).fetchone()

        if row is None:
            continue

        chunk_meta = json.loads(row[3]) if row[3] else {}

        tag_rows = conn.execute(
            """
            SELECT t.name, t.description FROM tags t
            JOIN document_tags dt ON t.id = dt.tag_id
            WHERE dt.document_id = ?
            ORDER BY t.name
            """,
            (row[4],),  # doc_id
        ).fetchall()

        result = {
            "chunk_id": row[0],
            "document_id": row[4],
            "score": round(score, 6),
            "text": row[1],
            "chunk_index": row[2],
            "chunk_metadata": chunk_meta,
            "title": row[5],
            "doc_type": row[6],
            "source_path": row[7],
            "created_at": row[8],
            "tags": [t[0] for t in tag_rows],
            "tag_contexts": {t[0]: t[1] for t in tag_rows if t[1]},
        }
        if details is not None and row[0] in details:
            result["explain"] = details[row[0]]
        results.append(result)

    return results