shira-hermes/api/services/kb/search.py

"""Hybrid search: vector + full-text, fused with Reciprocal Rank Fusion,
and finally reranked with a Voyage cross-encoder for precision.

`search()` supports an `expand=True` mode that asks the LLM for query
variants, retrieves candidates for each, merges them, and reranks the
union against the original query. Costs one extra LLM call + one Voyage
embed per variant, but recovers documents that don't share the user's
exact phrasing — e.g. ספר הליקויים when the user types "מהי תקנה 37".
"""

from __future__ import annotations

import asyncio
import logging
import os
from typing import Literal

from openai import AsyncOpenAI

from api.services.kb import voyage
from api.services.kb.db import get_pool

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

_RRF_K = 60  # standard RRF constant
_CANDIDATES_PER_SIDE = 30
# Number of fused RRF candidates fed to the reranker. Higher = better
# recall before the final cut; capped at ~30 to keep latency reasonable.
_RERANK_POOL = 30
# How many alternative phrasings to ask the LLM to generate when expand=True.
_EXPANSION_VARIANTS = 3
_EXPANSION_TIMEOUT_S = 12.0


async def _retrieve_rrf(
    query: str,
    kind: Literal["law", "regulation", "circular", "caselaw", "tool", "any"],
    topic_id: int | None = None,
    label_id: int | None = None,
) -> list[dict]:
    """Run vector + lexical retrieval, fuse with RRF. No rerank, no top_k cut.

    Each candidate dict carries `chunk_id` (kb_chunk.id) so callers can
    deduplicate when merging results from multiple queries.

    topic_id, when provided, restricts retrieval to a single domain
    (kb_source.topic_id). When None, all topics are searched — kept for
    transitional callers; new code should always pass a resolved id.

    label_id, when provided, further restricts retrieval to sources
    tagged with that label (kb_source_label).
    """
    [vec] = await voyage.embed([query], input_type="query")

    filters: list[str] = []
    params: list = [vec, query, _CANDIDATES_PER_SIDE]
    if kind != "any":
        params.append(kind)
        filters.append(f"AND s.kind = ${len(params)}")
    if topic_id is not None:
        params.append(topic_id)
        filters.append(f"AND s.topic_id = ${len(params)}")
    if label_id is not None:
        params.append(label_id)
        filters.append(
            f"AND EXISTS (SELECT 1 FROM kb_source_label sl "
            f"WHERE sl.source_id = s.id AND sl.label_id = ${len(params)})"
        )
    extra = "\n        ".join(filters)

    sql = f"""
    WITH vec AS (
      SELECT c.id, ROW_NUMBER() OVER (ORDER BY c.embedding <=> $1::vector) AS rank
      FROM kb_chunk c
      JOIN kb_source s ON s.id = c.source_id
      WHERE c.embedding IS NOT NULL
        AND s.superseded_by IS NULL
        {extra}
      ORDER BY c.embedding <=> $1::vector
      LIMIT $3
    ),
    lex AS (
      SELECT c.id, ROW_NUMBER() OVER (ORDER BY ts_rank_cd(c.content_tsv, q) DESC) AS rank
      FROM kb_chunk c
      JOIN kb_source s ON s.id = c.source_id,
           websearch_to_tsquery('simple', $2) AS q
      WHERE c.content_tsv @@ q
        AND s.superseded_by IS NULL
        {extra}
      ORDER BY ts_rank_cd(c.content_tsv, q) DESC
      LIMIT $3
    ),
    fused AS (
      SELECT id, SUM(score) AS score FROM (
        SELECT id, 1.0 / ({_RRF_K} + rank) AS score FROM vec
        UNION ALL
        SELECT id, 1.0 / ({_RRF_K} + rank) AS score FROM lex
      ) u GROUP BY id
    )
    SELECT
      c.id AS chunk_id, c.chunk_index,
      s.id AS source_id, s.kind, s.title, s.identifier, s.source_url,
      s.published_at, s.effective_at, s.original_path,
      c.heading_path, c.section_ref, c.content, c.page_number,
      f.score
    FROM fused f
    JOIN kb_chunk c ON c.id = f.id
    JOIN kb_source s ON s.id = c.source_id
    ORDER BY f.score DESC
    LIMIT {_RERANK_POOL}
    """

    pool = await get_pool()
    async with pool.acquire() as conn:
        rows = await conn.fetch(sql, *params)
    return [dict(r) for r in rows]


_EXPANSION_PROMPT = (
    "אתה עוזר חיפוש בבסיס ידע משפטי בנושא {topic_name}. "
    "המשתמש שואל שאלה. החזר עד {n} ניסוחים חלופיים שיעזרו לאחזר "
    "את אותו מידע ממסמכים משפטיים בעברית — תרגומים מסגנון שאלתי "
    "לסגנון מסמך, מילים נרדפות (למשל 'תקנה 37' → 'בדיקה מחדש', "
    "'נכות' → 'נכה'/'דרגת נכות'/'מוגבלות'), והוספת מונחים רפואיים/משפטיים "
    "שלא מופיעים במפורש בשאלה. החזר ניסוח אחד לשורה, ללא מספור או הקדמה. "
    "אל תחזור על הניסוח המקורי."
)


async def _expand_query(query: str, topic_name: str | None = None) -> list[str]:
    """Ask the LLM for alternative phrasings. Empty list on failure (caller
    just runs the original query)."""
    base = (os.environ.get("AI_GATEWAY_URL") or "http://localhost:3000").rstrip("/")
    api_key = os.environ.get("AI_GATEWAY_API_KEY", "")
    if not api_key:
        return []
    client = AsyncOpenAI(base_url=f"{base}/v1", api_key=api_key)
    prompt = _EXPANSION_PROMPT.format(
        n=_EXPANSION_VARIANTS,
        topic_name=topic_name or "כללי",
    )
    try:
        resp = await asyncio.wait_for(
            client.chat.completions.create(
                model=os.environ.get("CLAUDE_MODEL", "sonnet"),
                messages=[
                    {"role": "system", "content": prompt},
                    {"role": "user", "content": query},
                ],
                max_tokens=256,
                temperature=0.3,
            ),
            timeout=_EXPANSION_TIMEOUT_S,
        )
    except (asyncio.TimeoutError, Exception) as e:
        logger.warning("[kb.search] query expansion failed: %s", e)
        return []
    text = (resp.choices[0].message.content or "").strip()
    variants: list[str] = []
    for line in text.splitlines():
        v = line.strip(" -•·*\t").strip()
        if not v or v == query or v in variants:
            continue
        variants.append(v)
        if len(variants) >= _EXPANSION_VARIANTS:
            break
    return variants


async def _rerank_or_truncate(
    query: str, candidates: list[dict], top_k: int,
) -> list[dict]:
    """Cross-encoder rerank against `query`; fall back to RRF order on error."""
    rerank_disabled = os.environ.get("KB_RERANK_DISABLED", "").lower() in ("1", "true")
    if rerank_disabled or len(candidates) <= top_k:
        return candidates[:top_k]
    texts = [c["content"] for c in candidates]
    try:
        ranked = await voyage.rerank(query=query, documents=texts, top_k=top_k)
    except voyage.VoyageError as e:
        logger.warning("[kb.search] rerank failed, falling back to RRF: %s", e)
        return candidates[:top_k]
    out: list[dict] = []
    for r in ranked:
        i = r["index"]
        if 0 <= i < len(candidates):
            out.append({**candidates[i], "rerank_score": r["relevance_score"]})
    logger.info("[kb.search] rerank kept top %d / %d", len(out), len(candidates))
    return out


async def search(
    query: str,
    kind: Literal["law", "regulation", "circular", "caselaw", "tool", "any"] = "any",
    top_k: int = 8,
    expand: bool = False,
    topic_id: int | None = None,
    topic_name: str | None = None,
    label_id: int | None = None,
) -> list[dict]:
    query = (query or "").strip()
    if not query:
        return []

    if not expand:
        candidates = await _retrieve_rrf(query, kind, topic_id=topic_id, label_id=label_id)
        logger.info(
            "[kb.search] query=%r kind=%s topic_id=%s label_id=%s RRF_candidates=%d",
            query[:80], kind, topic_id, label_id, len(candidates),
        )
        if not candidates:
            return []
        return await _rerank_or_truncate(query, candidates, top_k)

    # Multi-query expansion: variants run in parallel with the original.
    variants = await _expand_query(query, topic_name=topic_name)
    queries = [query] + variants
    logger.info("[kb.search] expand=true topic_id=%s label_id=%s variants=%d",
                topic_id, label_id, len(variants))

    results = await asyncio.gather(
        *[_retrieve_rrf(q, kind, topic_id=topic_id, label_id=label_id) for q in queries],
        return_exceptions=True,
    )
    merged: dict[int, dict] = {}
    for r in results:
        if isinstance(r, Exception):
            logger.warning("[kb.search] sub-retrieval failed: %s", r)
            continue
        for c in r:
            cid = c["chunk_id"]
            if cid not in merged or c["score"] > merged[cid]["score"]:
                merged[cid] = c

    candidates = sorted(merged.values(), key=lambda x: -x["score"])[:_RERANK_POOL]
    logger.info(
        "[kb.search] merged candidates=%d (from %d queries)",
        len(candidates), len(queries),
    )
    if not candidates:
        return []
    return await _rerank_or_truncate(query, candidates, top_k)