dp-zp-agent/tools/rag/search_utils.py

from __future__ import annotations

from tools.retrieval.publication import DEFAULT_PUBLISHED_ONLY

import sqlite3
from pathlib import Path
from typing import Any
from tools.sync.index_state import rag_snapshot
from tools.retrieval.retrieval_contract import result_fields, source_references

# Tieto importy sú zámerne verejné.
# Zachovávajú spätnú kompatibilitu pre testy
# a evaluačné skripty, ktoré ich importujú
# z tools.rag.search_utils.
from tools.rag.search_core import (
    ANY_TERM_RRF_WEIGHT,
    BM25_SQL,
    BM25_WEIGHTS,
    DEFAULT_CANDIDATE_MULTIPLIER,
    FTS_RRF_WEIGHT,
    MIN_CANDIDATES,
    MIN_STEM_PREFIX_LENGTH,
    RRF_K,
    STRATEGY_PRIORITY,
    VECTOR_RRF_WEIGHT,
    WORD_RE,
    add_fts_metadata,
    add_vector_metadata,
    build_match_queries,
    database_bool,
    diversify_results,
    embedding_index_info,
    exact_match_bonus,
    fuse_hybrid_results,
    load_labels,
    load_result_labels,
    make_source_url,
    normalize_for_compare,
    parse_heading_paths,
    query_contains_title,
    query_tokens,
    quote_fts_token,
    run_fts_query,
    run_vector_query,
    tokens_match_for_title,
    verify_search_schema,
)


@rag_snapshot
def search_database(
    db_file: Path,
    query: str,
    limit: int = 10,
    published_only: bool = DEFAULT_PUBLISHED_ONLY,
    max_per_document: int = 3,
) -> dict[str, Any]:
    if limit < 1 or max_per_document < 0:
        raise ValueError("Limit musí byť kladný a max_per_document nezáporné.")
    if not db_file.exists():
        raise FileNotFoundError("Databáza neexistuje: " f"{db_file}")

    clean_query = query.strip()

    if not clean_query:
        return {
            **result_fields(status="no_answer", sources=[], reason="empty_query"),
            "engine": ("hybrid_fts5_embeddings"),
            "strategies": [],
            "results": [],
        }

    match_queries = build_match_queries(clean_query)

    candidate_limit = max(
        MIN_CANDIDATES,
        (limit * DEFAULT_CANDIDATE_MULTIPLIER),
    )

    with sqlite3.connect(
        db_file,
        timeout=5.0,
    ) as conn:
        conn.row_factory = sqlite3.Row

        conn.execute("PRAGMA query_only = ON")

        verify_search_schema(conn)

        # -------------------------
        # FTS5 kandidáti
        # -------------------------

        fts_candidates: list[dict[str, Any]] = []
        fts_capped = False

        used_strategies: list[str] = []

        for (
            strategy,
            match_query,
        ) in match_queries:
            rows = run_fts_query(
                conn,
                match_query,
                candidate_limit + 1,
                published_only,
            )

            if not rows:
                continue

            fts_capped = len(rows) > candidate_limit
            rows = rows[:candidate_limit]
            for row in rows:
                row["strategy"] = strategy

            fts_candidates = rows

            used_strategies = [strategy]

            # Používa sa prvá stratégia,
            # ktorá vráti výsledky.
            break

        fts_results = add_fts_metadata(
            conn,
            clean_query,
            fts_candidates,
        )

        # -------------------------
        # Embedding kandidáti
        # -------------------------

        vector_candidates = run_vector_query(
            conn,
            clean_query,
            candidate_limit + 1,
            published_only,
        )

        vector_capped = len(vector_candidates) > candidate_limit
        vector_candidates = vector_candidates[:candidate_limit]

        vector_results = add_vector_metadata(
            conn,
            vector_candidates,
        )

        # -------------------------
        # Presné FTS výsledky
        # -------------------------
        #
        # Pri all_terms alebo prefix_terms
        # embeddingy iba preradia chunky,
        # ktoré už našiel FTS.
        #
        # Tým sa zabráni pridávaniu
        # sémanticky podobných, ale
        # lexikálne nesúvisiacich výsledkov.

        if used_strategies and used_strategies[0] in {
            "all_terms",
            "prefix_terms",
        }:
            fts_chunk_ids = {item["chunk_id"] for item in fts_results}

            vector_results = [item for item in vector_results if item["chunk_id"] in fts_chunk_ids]

        # -------------------------
        # Hybrid fusion
        # -------------------------

        hybrid_results = fuse_hybrid_results(
            fts_results,
            vector_results,
        )

    # -------------------------
    # Document diversification
    # -------------------------

    final_results = diversify_results(
        hybrid_results,
        limit,
        max_per_document,
    )

    reasons = []
    if fts_candidates and fts_capped:
        reasons.append("candidate_limit")
    if vector_capped and (
        not used_strategies or used_strategies[0] not in {"all_terms", "prefix_terms"}
    ):
        reasons.append("candidate_limit")
    eligible = diversify_results(hybrid_results, max(1, len(hybrid_results)), max_per_document)
    if len(eligible) < len(hybrid_results):
        reasons.append("per_document_limit")
    if len(eligible) > limit:
        reasons.append("result_limit")

    return {
        **result_fields(
            status="ok" if final_results else "no_answer",
            sources=source_references(final_results),
            reason=None if final_results else "no_matching_sources",
            truncation_reasons=reasons,
        ),
        "engine": ("hybrid_fts5_embeddings"),
        "strategies": (used_strategies),
        "results": (final_results),
    }