from __future__ import annotations import json import re import sqlite3 from pathlib import Path from typing import Any from scripts.rag_query_evidence import ( build_focus_excerpt, evidence_score, parse_heading_paths_json, query_content_tokens, ) from scripts.search_utils import make_source_url, normalize_for_compare WORD_RE = re.compile(r"[^\W_]+", re.UNICODE) SUPPLEMENTAL_STOPWORDS = { "autor", "autora", "dokument", "dokumenty", "informacna", "informacnej", "projektova", "projektovej", "stranka", "stranku", "student", "studenta", "tema", "teme", "temou", } CANONICAL_TOKEN_ALIASES = { "multilingualna": "multilingual", "multilingualny": "multilingual", "multilingual": "multilingual", "viacjazycna": "multilingual", "viacjazycneho": "multilingual", "viacjazycny": "multilingual", "extrakcia": "extract", "extrakcie": "extract", "extrahovanie": "extract", "extrahovania": "extract", "extraction": "extract", "extract": "extract", "trojic": "triplet", "trojice": "triplet", "trojicami": "triplet", "triplet": "triplet", "triplets": "triplet", "medicinskych": "medical", "medicinsky": "medical", "medicinske": "medical", "lekarskych": "medical", "lekarsky": "medical", "medical": "medical", "dat": "data", "data": "data", "strojovy": "machine", "strojoveho": "machine", "machine": "machine", "preklad": "translation", "prekladu": "translation", "translation": "translation", "automaticky": "automatic", "automatickeho": "automatic", "automatic": "automatic", "pomenovane": "named", "pomenovanych": "named", "named": "named", "entity": "entity", "entities": "entity", "entit": "entity", "anotacia": "annotation", "anotacii": "annotation", "annotation": "annotation", "znalostny": "knowledge", "znalostne": "knowledge", "knowledge": "knowledge", "graf": "graph", "grafy": "graph", "graph": "graph", "graphy": "graph", "otazok": "question", "otazky": "question", "question": "question", "questions": "question", "odpovedi": "answer", "odpovede": "answer", "answer": "answer", "answers": "answer", "system": "system", "systemom": "system", "systemu": "system", "integrovat": "integrate", "integracia": "integrate", "integrate": "integrate", } MIN_SINGLE_GLOBAL_COVERAGE = 0.50 MIN_MULTI_GLOBAL_COVERAGE = 0.45 MIN_GLOBAL_MATCHES = 3 MAX_FAMILY_EVIDENCE = 3 def sqlite_table_exists(conn: sqlite3.Connection, table_name: str) -> bool: row = conn.execute( "SELECT 1 FROM sqlite_master " "WHERE type IN ('table', 'view') AND name = ? LIMIT 1", (table_name,), ).fetchone() return row is not None def sqlite_table_columns(conn: sqlite3.Connection, table_name: str) -> set[str]: if not sqlite_table_exists(conn, table_name): return set() rows = conn.execute(f"PRAGMA table_info({table_name})").fetchall() return { str(row[1]) for row in rows if len(row) > 1 } def normalized_words(value: str) -> list[str]: return WORD_RE.findall(normalize_for_compare(value)) def canonical_token(token: str) -> str: return CANONICAL_TOKEN_ALIASES.get(token, token) def canonical_tokens(value: str) -> list[str]: result: list[str] = [] for token in normalized_words(value): if token in SUPPLEMENTAL_STOPWORDS: continue result.append(canonical_token(token)) return list(dict.fromkeys(result)) def canonical_query_tokens(query: str) -> list[str]: base = [ token for token in query_content_tokens(query) if token not in SUPPLEMENTAL_STOPWORDS ] return list( dict.fromkeys( canonical_token(token) for token in base ) ) def canonical_coverage(query_tokens: list[str], value: str) -> tuple[int, float]: if not query_tokens: return 0, 0.0 value_tokens = set(canonical_tokens(value)) matches = sum( 1 for token in query_tokens if token in value_tokens ) return matches, matches / len(query_tokens) def canonical_student_main_path(document_path: str) -> str: parts = Path(document_path).parts if ( len(parts) >= 5 and parts[0] == "pages" and parts[1] == "students" ): return "/".join( [ parts[0], parts[1], parts[2], parts[3], "README.md", ] ) return document_path def same_student_family(left: str, right: str) -> bool: return ( canonical_student_main_path(left) == canonical_student_main_path(right) ) def is_multi_document_query(query: str) -> bool: normalized = normalize_for_compare(query) return any( marker in normalized for marker in ( "viacero dokument", "dokumenty", "ktore dokumenty", "najdi dokumenty", ) ) def _load_documents(conn: sqlite3.Connection) -> list[dict[str, Any]]: columns = sqlite_table_columns(conn, "documents") if not columns: return [] path_column = ( "path" if "path" in columns else "document_path" if "document_path" in columns else None ) if path_column is None: return [] selected = [ f"{path_column} AS document_path", ] for optional in ("title", "author", "published"): if optional in columns: selected.append(optional) query = ( "SELECT " + ", ".join(selected) + " FROM documents" ) return [ dict(row) for row in conn.execute(query).fetchall() ] def _load_chunks( conn: sqlite3.Connection, *, published_only: bool, ) -> list[dict[str, Any]]: columns = sqlite_table_columns(conn, "chunks") required = { "chunk_id", "document_path", "chunk_index", "heading_paths_json", "text", } if not required <= columns: return [] selected = [ "chunk_id", "document_path", "chunk_index", "heading_paths_json", "text", ] for optional in ("title", "author", "published"): if optional in columns: selected.append(optional) where = "" if published_only and "published" in columns: where = " WHERE published = 1" order = " ORDER BY document_path ASC, chunk_index ASC" if "id" in columns: order += ", id ASC" query = ( "SELECT " + ", ".join(selected) + " FROM chunks" + where + order ) rows = [ dict(row) for row in conn.execute(query).fetchall() ] for row in rows: row["heading_paths"] = parse_heading_paths_json( row.get("heading_paths_json") ) return rows def _document_metadata_by_path( documents: list[dict[str, Any]], ) -> dict[str, dict[str, Any]]: return { str(item.get("document_path") or ""): item for item in documents if str(item.get("document_path") or "").strip() } def _best_main_metadata( canonical_path: str, metadata_by_path: dict[str, dict[str, Any]], fallback_chunk: dict[str, Any], ) -> dict[str, Any]: direct = metadata_by_path.get(canonical_path) if direct is not None: return direct evidence_path = str( fallback_chunk.get("document_path") or "" ) evidence_meta = metadata_by_path.get(evidence_path) if evidence_meta is not None: return evidence_meta return { "document_path": canonical_path, "title": fallback_chunk.get("title"), "author": fallback_chunk.get("author"), "published": fallback_chunk.get("published"), } def _source_from_chunk( canonical_path: str, chunk: dict[str, Any], metadata_by_path: dict[str, dict[str, Any]], *, reason: str, score: float, ) -> dict[str, Any]: meta = _best_main_metadata( canonical_path, metadata_by_path, chunk, ) return { "chunk_id": chunk.get("chunk_id"), "document_path": canonical_path, "title": meta.get("title") or chunk.get("title"), "author": meta.get("author") or chunk.get("author"), "published": ( meta.get("published") if meta.get("published") is not None else chunk.get("published") ), "chunk_index": int(chunk.get("chunk_index") or 0), "heading_paths": chunk.get("heading_paths") or [], "text": str(chunk.get("text") or ""), "source_url": make_source_url(canonical_path), "match_strategy": "rag_supplemental", "fts_rank": None, "vector_rank": None, "vector_score": None, "hybrid_score": None, "supplemental_expansion": { "strategy": "global_or_entity_anchor", "applied": True, "reason": reason, "score": round(score, 6), "evidence_document_path": chunk.get("document_path"), "evidence_chunk_id": chunk.get("chunk_id"), }, } def _title_anchor_candidates( query: str, documents: list[dict[str, Any]], chunks: list[dict[str, Any]], ) -> list[tuple[float, str, dict[str, Any], str]]: normalized_query = normalize_for_compare(query) chunks_by_path: dict[str, list[dict[str, Any]]] = {} for chunk in chunks: chunks_by_path.setdefault( str(chunk.get("document_path") or ""), [], ).append(chunk) candidates: list[tuple[float, str, dict[str, Any], str]] = [] for document in documents: title = str(document.get("title") or "").strip() path = str(document.get("document_path") or "").strip() if not title or not path: continue normalized_title = normalize_for_compare(title) title_tokens = normalized_words(title) if len(title_tokens) < 2: continue if normalized_title not in normalized_query: continue family_path = canonical_student_main_path(path) family_chunks = [ chunk for chunk in chunks if canonical_student_main_path( str(chunk.get("document_path") or "") ) == family_path ] path_chunks = chunks_by_path.get(path, []) pool = family_chunks or path_chunks if not pool: continue best = max( pool, key=lambda chunk: ( evidence_score( query, text=str(chunk.get("text") or ""), heading_paths=chunk.get("heading_paths"), ), -int(chunk.get("chunk_index") or 0), ), ) score = 500.0 + 10.0 * len(title_tokens) candidates.append( (score, family_path, best, "document_title_in_query") ) return candidates def _global_chunk_candidates( query: str, chunks: list[dict[str, Any]], *, multi_document: bool, ) -> list[tuple[float, str, dict[str, Any], str]]: query_tokens = canonical_query_tokens(query) if not query_tokens: return [] threshold = ( MIN_MULTI_GLOBAL_COVERAGE if multi_document else MIN_SINGLE_GLOBAL_COVERAGE ) normalized_query = normalize_for_compare(query) explicitly_entity_or_student = any( marker in normalized_query for marker in ( "student", "osob", "kto", "autor", ) ) best_by_source: dict[ str, tuple[float, str, dict[str, Any], str], ] = {} for chunk in chunks: text = str(chunk.get("text") or "") headings = json.dumps( chunk.get("heading_paths") or [], ensure_ascii=False, ) combined = f"{headings}\n{text}" matches, coverage = canonical_coverage( query_tokens, combined, ) raw_path = str(chunk.get("document_path") or "") topic_page = raw_path.startswith("pages/topics/") required_coverage = threshold required_matches = ( 2 if multi_document else MIN_GLOBAL_MATCHES ) # Generic factual questions without a named person often belong to a # canonical project/topic page. Allow such pages to enter with two # strong concept matches, then rank them with an explicit topic-page # bonus. This handles short follow-up-like questions without changing # the locked global retrieval scorer. if topic_page and not explicitly_entity_or_student: required_coverage = min(required_coverage, 0.33) required_matches = min(required_matches, 2) if ( coverage < required_coverage or matches < required_matches ): continue lexical = evidence_score( query, text=text, heading_paths=chunk.get("heading_paths"), ) score = ( coverage * 160.0 + matches * 12.0 + lexical ) if topic_page and not explicitly_entity_or_student: score += 25.0 topic_parts = Path(raw_path).parts title_matches, _ = canonical_coverage( query_tokens, str(chunk.get("title") or ""), ) if ( title_matches >= 1 and len(topic_parts) == 4 and topic_parts[0] == "pages" and topic_parts[1] == "topics" and topic_parts[3] == "README.md" ): score += 45.0 if ( len(topic_parts) == 4 and topic_parts[0] == "pages" and topic_parts[1] == "topics" and topic_parts[3] == "README.md" ): score += 20.0 canonical_path = canonical_student_main_path( str(chunk.get("document_path") or "") ) candidate = ( score, canonical_path, chunk, "global_query_match", ) previous = best_by_source.get(canonical_path) if previous is None or candidate[0] > previous[0]: best_by_source[canonical_path] = candidate return sorted( best_by_source.values(), key=lambda item: (-item[0], item[1]), ) def _family_evidence_blocks( query: str, canonical_path: str, chunks: list[dict[str, Any]], *, max_blocks: int = MAX_FAMILY_EVIDENCE, ) -> list[dict[str, Any]]: family_chunks = [ chunk for chunk in chunks if same_student_family( str(chunk.get("document_path") or ""), canonical_path, ) ] scored: list[tuple[float, dict[str, Any]]] = [] for chunk in family_chunks: score = evidence_score( query, text=str(chunk.get("text") or ""), heading_paths=chunk.get("heading_paths"), ) focus = build_focus_excerpt( query, str(chunk.get("text") or ""), ) if score < 20.0 or not focus: continue item = dict(chunk) item["score"] = score item["focus_text"] = focus scored.append((score, item)) scored.sort( key=lambda pair: ( -pair[0], str(pair[1].get("document_path") or ""), int(pair[1].get("chunk_index") or 0), ) ) blocks: list[dict[str, Any]] = [] seen: set[str] = set() for _, chunk in scored: focus = str(chunk.get("focus_text") or "").strip() key = normalize_for_compare(focus) if key in seen: continue seen.add(key) blocks.append( { "document_path": chunk.get("document_path"), "chunk_id": chunk.get("chunk_id"), "chunk_index": int(chunk.get("chunk_index") or 0), "heading_paths": chunk.get("heading_paths") or [], "score": round(float(chunk.get("score") or 0.0), 6), "focus_text": focus, "text": str(chunk.get("text") or "").strip(), } ) if len(blocks) >= max_blocks: break return blocks def expand_results_with_supplemental_sources( db_path: Path, query: str, results: list[dict[str, Any]], *, published_only: bool = False, limit: int = 5, ) -> list[dict[str, Any]]: base_results = [dict(item) for item in results] if ( not query.strip() or not db_path.exists() or limit <= 0 ): return base_results with sqlite3.connect(db_path, timeout=5.0) as conn: conn.row_factory = sqlite3.Row conn.execute("PRAGMA query_only = ON") if not sqlite_table_exists(conn, "chunks"): return base_results documents = _load_documents(conn) chunks = _load_chunks( conn, published_only=published_only, ) if not chunks: return base_results metadata_by_path = _document_metadata_by_path(documents) # Attach strongest evidence from child documents belonging to the same student. for item in base_results: path = str(item.get("document_path") or "").strip() if not path: continue canonical_path = canonical_student_main_path(path) if canonical_path == path and "/students/" not in f"/{path}": continue blocks = _family_evidence_blocks( query, canonical_path, chunks, ) if blocks: item["family_evidence_blocks"] = blocks item["family_expansion"] = { "strategy": "student_family_evidence", "applied": True, "canonical_document_path": canonical_path, "evidence_count": len(blocks), } title_candidates = _title_anchor_candidates( query, documents, chunks, ) has_title_anchor = bool(title_candidates) multi_document = is_multi_document_query(query) global_candidates: list[ tuple[float, str, dict[str, Any], str] ] = [] # If the query already names an exact document/person, anchor that document. # Otherwise use a conservative global scan. Multi-document questions always # receive the scan because one base retrieval result is not sufficient. if not has_title_anchor or multi_document: global_candidates = _global_chunk_candidates( query, chunks, multi_document=multi_document, ) all_candidates = sorted( title_candidates + global_candidates, key=lambda item: (-item[0], item[1]), ) existing_by_path: dict[str, dict[str, Any]] = { canonical_student_main_path( str(item.get("document_path") or "") ): item for item in base_results if str(item.get("document_path") or "").strip() } synthetic: list[dict[str, Any]] = [] seen_candidate_paths: set[str] = set() for score, canonical_path, chunk, reason in all_candidates: if canonical_path in seen_candidate_paths: continue seen_candidate_paths.add(canonical_path) existing = existing_by_path.get(canonical_path) if existing is not None: existing["supplemental_expansion"] = { "strategy": "global_or_entity_anchor", "applied": True, "reason": reason, "score": round(score, 6), "evidence_document_path": chunk.get("document_path"), "evidence_chunk_id": chunk.get("chunk_id"), } continue synthetic.append( _source_from_chunk( canonical_path, chunk, metadata_by_path, reason=reason, score=score, ) ) if len(synthetic) >= limit: break # Exact title/person anchors are intentionally promoted. For broad # multi-document queries, supplemental sources are merged by score ahead of # weaker base-only results, but the final output still respects the tool limit. if title_candidates: promoted_paths = { canonical_path for _, canonical_path, _, _ in title_candidates } promoted_existing = [ item for item in base_results if canonical_student_main_path( str(item.get("document_path") or "") ) in promoted_paths ] remainder = [ item for item in base_results if item not in promoted_existing ] combined = promoted_existing + synthetic + remainder elif multi_document: combined = synthetic + base_results else: combined = base_results + synthetic deduped: list[dict[str, Any]] = [] seen_paths: set[str] = set() for item in combined: path = canonical_student_main_path( str(item.get("document_path") or "") ) dedupe_key = path or str(item.get("source_url") or "") if dedupe_key in seen_paths: continue seen_paths.add(dedupe_key) deduped.append(item) if len(deduped) >= limit: break return deduped