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
        ),
    }