from __future__ import annotations from tools.publication import DEFAULT_PUBLISHED_ONLY import sqlite3 from pathlib import Path from typing import Any from tools.index_state import rag_snapshot from tools.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 ), }