qmd_py.search.hybrid¶

async def _rerank_safe_text(llm_client: LlmClient, text: str, model: str) -> str:
    """Truncates `text` to fit the router's per-pair rerank token budget,
    measuring with the model's own tokenizer rather than guessing another
    fixed chars/token ratio - that guess is exactly what broke the first
    time (see module-level comment above)."""
    if len(text) <= _RERANK_SAFE_CHAR_SKIP:
        return text
    tokens = await llm_client.tokenize(text, model)
    if len(tokens) <= _RERANK_TOKEN_BUDGET:
        return text
    ratio = len(text) / len(tokens)
    return text[: max(1, int(_RERANK_TOKEN_BUDGET * ratio))]

def parse_structured_query(query: str) -> tuple[list[ExpandedQuery], str | None] | None:
    """Multi-line `lex:`/`vec:`/`hyde:` (+ optional `intent:`) query
    syntax - port of the TS reference's `parseStructuredQuery`, minus its
    strict error-throwing for malformed input: an unprefixed line among
    several, or no typed lines at all, just falls through to `None`
    (treated as an ordinary single query, auto-expanded) rather than
    raising - a softer degrade that seemed more appropriate for a CLI
    argument than a hard parse error.
    """
    lines = [line.strip() for line in query.split("\n") if line.strip()]
    if len(lines) <= 1:
        return None

    typed: list[ExpandedQuery] = []
    intent: str | None = None
    for line in lines:
        type_match = _TYPED_LINE_RE.match(line)
        if type_match:
            typed.append(ExpandedQuery(type_match.group(1).lower(), type_match.group(2).strip()))
            continue
        intent_match = _INTENT_LINE_RE.match(line)
        if intent_match:
            intent = intent_match.group(1).strip()
            continue
        return None

    return (typed, intent) if typed else None

def validate_typed_queries(queries: list[ExpandedQuery]) -> str | None:
    """First problem found among explicitly typed sub-queries, or None.

    Deliberately applied only to sub-queries the *caller* spelled out -
    the `lex:`/`vec:`/`hyde:` document syntax and the MCP `query` tool's
    `searches` - never to `expand_query()`'s LLM-generated variants: a
    stray `-term` in those is the model's doing, not a user mistake worth
    failing the search over.
    """
    for q in queries:
        error = (
            validate_lex_query(q.query)
            if q.type == "lex"
            else validate_semantic_query(q.query)
        )
        if error is not None:
            return f"{q.type}: {error}"
    return None

async def expand_query(
    llm_client: LlmClient, query: str, model: str, intent: str | None = None
) -> list[ExpandedQuery]:
    """Typed query variants (lex/vec/hyde) for RRF fusion - port of the TS
    reference's `expandQuery()`.

    Uses JSON-schema-constrained chat completion rather than the TS
    reference's local node-llama-cpp GBNF-grammar-constrained generation.
    llama.cpp's server *does* accept a raw `grammar` field over HTTP, but
    that specific line-grammar (`type ": " content "\\n"`) proved
    unreliable in manual testing against this project's own
    qwen2.5-3b-instruct-q4_k_m preset: the model satisfied the grammar's
    syntax while ignoring its semantics (emitting dozens of "lex: " lines,
    never "vec"/"hyde"). JSON-schema output was reliable in the same
    testing and is the more portable mechanism for a pure-HTTP client
    generally, so that's what qmd-py builds on instead.
    """
    user_content = f"Expand this search query: {query}"
    if intent:
        user_content += f"\nQuery intent: {intent}"

    try:
        data = await llm_client.chat_json(
            messages=[
                {"role": "system", "content": _EXPAND_SYSTEM_PROMPT},
                {"role": "user", "content": user_content},
            ],
            model=model,
            json_schema=_EXPAND_JSON_SCHEMA,
        )
    except (httpx.HTTPError, IndexError, KeyError, ValueError, TypeError) as exc:
        # IndexError included because chat_json subscripts choices[0]: a
        # router answering {"choices": []} is exactly the misbehavior this
        # fallback exists for, and it used to escape and fail the search.
        logger.warning(
            "query expansion failed (%s: %s), falling back to unexpanded lex+vec search",
            type(exc).__name__,
            exc,
        )
        return [ExpandedQuery("lex", query), ExpandedQuery("vec", query)]

    expanded = []
    for type_ in ("lex", "vec", "hyde"):
        text = str(data.get(type_, "")).strip()
        if text and text != query:
            expanded.append(ExpandedQuery(type_, text))
    return expanded or [ExpandedQuery("lex", query), ExpandedQuery("vec", query)]

def reciprocal_rank_fusion(
    result_lists: list[list[RankedResult]], weights: list[float] | None = None, k: int = 60
) -> list[RankedResult]:
    """Fuse several ranked lists into one by reciprocal rank.

    Each list contributes `weight / (k + rank + 1)` per document, summed
    across lists, so appearing in several lists beats ranking highly in
    one. Scores from the underlying engines are ignored entirely - only
    positions matter, which is what makes lexical and vector results
    comparable at all.

    Args:
        result_lists: Ranked lists, best first. Empty lists are harmless.
        weights: Per-list multipliers, positional. Lists beyond the end
            default to 1.0, so a short list is not an error.
        k: Damping constant. Larger flattens the difference between
            ranks; the default of 60 is the value from the original RRF
            paper.

    Returns:
        One entry per distinct `file`, ordered by descending fused score.
        Documents ranked first in any list get a small bonus, and ones in
        the top three a smaller one, to break ties toward strong single
        signals.
    """
    weights = weights or []
    entries: dict[str, _RrfEntry] = {}

    for list_idx, result_list in enumerate(result_lists):
        weight = weights[list_idx] if list_idx < len(weights) else 1.0
        for rank, result in enumerate(result_list):
            contribution = weight / (k + rank + 1)
            entry = entries.get(result.file)
            if entry is not None:
                entry.rrf_score += contribution
                entry.top_rank = min(entry.top_rank, rank)
            else:
                entries[result.file] = _RrfEntry(result, contribution, rank)

    fused: list[RankedResult] = []
    for entry in entries.values():
        score = entry.rrf_score
        if entry.top_rank == 0:
            score += 0.05
        elif entry.top_rank <= 2:
            score += 0.02
        fused.append(replace(entry.result, score=score))

    fused.sort(key=lambda r: r.score, reverse=True)
    return fused

def _hybrid_rrf_weights(query_types: list[str]) -> list[float]:
    """Original-query retrieval paths (the primary evidence) get 2x
    weight; expansion-derived (lex/vec/hyde) lists stay at 1x regardless
    of insertion order."""
    return [2.0 if t == "original" else 1.0 for t in query_types]

async def hybrid_query(
    session: AsyncSession,
    user: CurrentUser,
    query: str,
    llm_client: LlmClient,
    models: ModelConfig,
    options: QueryOptions | None = None,
) -> list[HybridQueryResult]:
    """BM25 + vector + query expansion + RRF + chunked reranking - port of
    the TS reference's `hybridQuery`. See `QueryOptions` for the tunables.
    """
    with log_duration(logger, "query") as timing:
        results = await _hybrid_query_impl(
            session, user, query, llm_client, models, options, timing
        )
        # Counts and durations only - never the query text or any document
        # content (see log.py's privacy note); the query itself is DEBUG.
        timing["results"] = len(results)
        logger.debug("query text: %r (intent: %r)", query, (options or QueryOptions()).intent)
    return results