From 490f4662baed731954e6c0a9872b61992633c4e6 Mon Sep 17 00:00:00 2001 From: jp170na Date: Sun, 16 Aug 2026 22:25:33 +0200 Subject: [PATCH] adicia k utils --- scripts/rag_document_expansion.py | 1078 +++++++++++++++++++++++++++++ 1 file changed, 1078 insertions(+) create mode 100644 scripts/rag_document_expansion.py diff --git a/scripts/rag_document_expansion.py b/scripts/rag_document_expansion.py new file mode 100644 index 0000000..edf0045 --- /dev/null +++ b/scripts/rag_document_expansion.py @@ -0,0 +1,1078 @@ +from __future__ import annotations + +import json +import re +import sqlite3 +from pathlib import Path +from typing import Any + +from scripts.search_utils import ( + make_source_url, + normalize_for_compare, +) + + +YEAR_RE = re.compile( + r"\b(?:19|20)\d{2}\b" +) + + +SECTION_INTENTS: dict[ + str, + tuple[str, ...], +] = { + "diploma_thesis": ( + "diplomov", + "prac", + ), + "bachelor_thesis": ( + "bakalarsk", + "prac", + ), + "doctoral_thesis": ( + "dizertacn", + "prac", + ), + "diploma_project": ( + "diplomov", + "projekt", + ), + "team_project": ( + "timov", + "projekt", + ), +} + + +WORK_INTENTS = { + "diploma_thesis", + "bachelor_thesis", + "doctoral_thesis", +} + + +QUERY_STOPWORDS = { + "a", + "aj", + "aka", + "ake", + "aky", + "je", + "k", + "kde", + "ktora", + "ktore", + "ktory", + "ma", + "mal", + "mala", + "malo", + "na", + "o", + "osoba", + "osobe", + "osoby", + "pri", + "pre", + "sa", + "s", + "so", + "su", + "ta", + "ten", + "to", + "u", + "v", + "vo", + "z", + "za", +} + + +def normalized_tokens( + value: str, +) -> list[str]: + normalized = ( + normalize_for_compare( + value + ) + ) + + if not normalized: + return [] + + return normalized.split() + + +def has_prefix( + tokens: list[str], + prefix: str, +) -> bool: + return any( + token.startswith( + prefix + ) + for token in tokens + ) + + +def detect_section_intent( + query: str, +) -> str | None: + tokens = normalized_tokens( + query + ) + + if not tokens: + return None + + # Diplomový projekt musí byť pred + # diplomovou prácou, pretože oba + # začínajú tokenom "diplomov". + ordered_intents = ( + "diploma_project", + "team_project", + "diploma_thesis", + "bachelor_thesis", + "doctoral_thesis", + ) + + for intent in ordered_intents: + prefixes = ( + SECTION_INTENTS[ + intent + ] + ) + + if all( + has_prefix( + tokens, + prefix, + ) + for prefix in prefixes + ): + return intent + + return None + + +def extract_years( + value: str, +) -> list[int]: + return [ + int( + match + ) + for match in YEAR_RE.findall( + value + ) + ] + + +def parse_heading_paths( + value: Any, +) -> list[Any]: + if isinstance( + value, + list, + ): + return value + + if not value: + return [] + + try: + parsed = json.loads( + str( + value + ) + ) + + except json.JSONDecodeError: + return [] + + if not isinstance( + parsed, + list, + ): + return [] + + return parsed + + +def flatten_heading_paths( + heading_paths: list[Any], +) -> str: + parts: list[str] = [] + + for item in heading_paths: + if isinstance( + item, + str, + ): + value = ( + item.strip() + ) + + if value: + parts.append( + value + ) + + continue + + if isinstance( + item, + ( + list, + tuple, + ), + ): + for part in item: + value = str( + part + ).strip() + + if value: + parts.append( + value + ) + + return " ".join( + parts + ) + + +def section_matches_intent( + value: str, + intent: str, +) -> bool: + prefixes = ( + SECTION_INTENTS.get( + intent + ) + ) + + if not prefixes: + return False + + tokens = normalized_tokens( + value + ) + + return all( + has_prefix( + tokens, + prefix, + ) + for prefix in prefixes + ) + + +def is_primary_student_document( + document_path: str, +) -> bool: + parts = Path( + document_path + ).parts + + return ( + len( + parts + ) + == 5 + and parts[0] + == "pages" + and parts[1] + == "students" + and parts[-1] + == "README.md" + ) + + +def title_occurs_in_query( + title: str, + query: str, +) -> bool: + normalized_title = ( + normalize_for_compare( + title + ) + ) + + normalized_query = ( + normalize_for_compare( + query + ) + ) + + if ( + not normalized_title + or not normalized_query + ): + return False + + # Jednoslovné názvy sú príliš široké + # na bezpečné exact-document rozšírenie. + if len( + normalized_title.split() + ) < 2: + return False + + return ( + f" {normalized_title} " + in f" {normalized_query} " + ) + + +def sqlite_table_exists( + conn: sqlite3.Connection, + table_name: str, +) -> bool: + row = conn.execute( + """ + SELECT 1 + FROM sqlite_master + WHERE type = 'table' + AND name = ? + LIMIT 1 + """, + ( + table_name, + ), + ).fetchone() + + return row is not None + + +def find_exact_student_document( + conn: sqlite3.Connection, + query: str, + *, + published_only: bool, +) -> dict[str, Any] | None: + rows = conn.execute( + """ + SELECT + path, + title, + author, + published + FROM documents + WHERE ( + ? = 0 + OR published = 1 + ) + ORDER BY + LENGTH(title) DESC, + path ASC + """, + ( + 1 + if published_only + else 0, + ), + ).fetchall() + + candidates: list[ + dict[str, Any] + ] = [] + + for row in rows: + document_path = str( + row[ + "path" + ] + or "" + ).strip() + + title = str( + row[ + "title" + ] + or "" + ).strip() + + if ( + not document_path + or not title + ): + continue + + if not is_primary_student_document( + document_path + ): + continue + + if not title_occurs_in_query( + title, + query, + ): + continue + + candidates.append( + { + "document_path": ( + document_path + ), + "title": title, + "author": row[ + "author" + ], + "published": ( + row[ + "published" + ] + ), + } + ) + + if not candidates: + return None + + candidates.sort( + key=lambda item: ( + -len( + normalized_tokens( + str( + item[ + "title" + ] + ) + ) + ), + -len( + str( + item[ + "title" + ] + ) + ), + str( + item[ + "document_path" + ] + ), + ) + ) + + return candidates[ + 0 + ] + + +def query_content_tokens( + query: str, + title: str, +) -> list[str]: + query_tokens = ( + normalized_tokens( + query + ) + ) + + title_tokens = set( + normalized_tokens( + title + ) + ) + + result: list[str] = [] + + for token in query_tokens: + if token in title_tokens: + continue + + if token in QUERY_STOPWORDS: + continue + + if len(token) < 3: + continue + + result.append( + token + ) + + return result + + +def local_chunk_score( + *, + query: str, + title: str, + intent: str, + heading_paths: list[Any], + text: str, + chunk_index: int, +) -> tuple[ + int, + int, + int, +]: + heading_text = ( + flatten_heading_paths( + heading_paths + ) + ) + + heading_normalized = ( + normalize_for_compare( + heading_text + ) + ) + + text_normalized = ( + normalize_for_compare( + text + ) + ) + + score = 0 + + heading_matches = ( + section_matches_intent( + heading_text, + intent, + ) + ) + + text_matches = ( + section_matches_intent( + text, + intent, + ) + ) + + if heading_matches: + score += 1000 + + if text_matches: + score += 500 + + # Ak chunk vôbec nepatrí k požadovanému + # typu sekcie, nemá byť lokálnym víťazom. + if ( + not heading_matches + and not text_matches + ): + return ( + -1, + -1, + -chunk_index, + ) + + query_years = extract_years( + query + ) + + heading_years = extract_years( + heading_text + ) + + text_years = extract_years( + text + ) + + if query_years: + if any( + year in heading_years + for year in query_years + ): + score += 1500 + + elif any( + year in text_years + for year in query_years + ): + score += 750 + + else: + score -= 500 + + content_tokens = ( + query_content_tokens( + query, + title, + ) + ) + + for token in content_tokens: + if token in heading_normalized: + score += 30 + + if token in text_normalized: + score += 8 + + if intent in WORK_INTENTS: + if ( + "nazov diplomovej prace" + in text_normalized + or "nazov bakalarskej prace" + in text_normalized + or "nazov dizertacnej prace" + in text_normalized + ): + score += 120 + + elif "nazov" in text_normalized: + score += 80 + + if "tema" in text_normalized: + score += 50 + + if has_prefix( + normalized_tokens( + query + ), + "rok", + ): + if heading_years: + score += 100 + + # Ak rok nie je v dopyte a dokument má + # viac rovnakých typov prác, preferujeme + # najnovšiu sekciu. + latest_year = ( + max( + heading_years + ) + if heading_years + else -1 + ) + + return ( + score, + latest_year, + -chunk_index, + ) + + +def load_best_document_section_chunk( + conn: sqlite3.Connection, + *, + query: str, + document: dict[str, Any], + intent: str, + published_only: bool, +) -> dict[str, Any] | None: + document_path = str( + document[ + "document_path" + ] + ) + + rows = conn.execute( + """ + SELECT + chunk_id, + document_path, + title, + author, + published, + chunk_index, + heading_paths_json, + text + FROM chunks + WHERE document_path = ? + AND ( + ? = 0 + OR published = 1 + ) + ORDER BY + chunk_index ASC, + id ASC + """, + ( + document_path, + ( + 1 + if published_only + else 0 + ), + ), + ).fetchall() + + best_item: dict[ + str, + Any, + ] | None = None + + best_score: tuple[ + int, + int, + int, + ] | None = None + + for row in rows: + heading_paths = ( + parse_heading_paths( + row[ + "heading_paths_json" + ] + ) + ) + + text = str( + row[ + "text" + ] + or "" + ).strip() + + if not text: + continue + + chunk_index = int( + row[ + "chunk_index" + ] + ) + + score = local_chunk_score( + query=query, + title=str( + document[ + "title" + ] + ), + intent=intent, + heading_paths=( + heading_paths + ), + text=text, + chunk_index=( + chunk_index + ), + ) + + if score[ + 0 + ] < 0: + continue + + if ( + best_score is not None + and score <= best_score + ): + continue + + best_score = score + + best_item = { + "chunk_id": str( + row[ + "chunk_id" + ] + ), + "document_path": ( + document_path + ), + "title": ( + row[ + "title" + ] + or document[ + "title" + ] + ), + "author": ( + row[ + "author" + ] + or document.get( + "author" + ) + ), + "published": ( + bool( + row[ + "published" + ] + ) + if row[ + "published" + ] + is not None + else None + ), + "chunk_index": ( + chunk_index + ), + "heading_paths": ( + heading_paths + ), + "text": text, + "source_url": ( + make_source_url( + document_path + ) + ), + } + + return best_item + + +def exact_document_expansion_metadata( + *, + applied: bool, + document_path: str | None, + chunk_id: str | None, + chunk_index: int | None, + added_source: bool, +) -> dict[str, Any]: + return { + "strategy": ( + "exact_document_section" + ), + "applied": applied, + "document_path": ( + document_path + ), + "chunk_id": ( + chunk_id + ), + "chunk_index": ( + chunk_index + ), + "added_source": ( + added_source + ), + } + + +def expand_results_with_exact_document_section( + db_path: Path, + query: str, + results: list[ + dict[str, Any] + ], + *, + published_only: bool, + limit: int, +) -> list[dict[str, Any]]: + if not results: + base_results: list[ + dict[str, Any] + ] = [] + + else: + base_results = [ + dict( + result + ) + for result in results + ] + + intent = detect_section_intent( + query + ) + + if intent is None: + return base_results + + if not db_path.exists(): + return base_results + + with sqlite3.connect( + db_path, + timeout=5.0, + ) as conn: + conn.row_factory = ( + sqlite3.Row + ) + + conn.execute( + "PRAGMA query_only = ON" + ) + + # Niektoré unit testy používajú zámerne + # minimalistickú SQLite databázu iba + # s tabuľkou chunks. + # + # Exact-document expansion je voliteľná + # RAG-only nadstavba. Ak potrebná schéma + # nie je dostupná, musí zostať bezpečný + # pôvodný výsledok retrievalu. + if ( + not sqlite_table_exists( + conn, + "documents", + ) + or not sqlite_table_exists( + conn, + "chunks", + ) + ): + return base_results + + document = ( + find_exact_student_document( + conn, + query, + published_only=( + published_only + ), + ) + ) + + if document is None: + return base_results + + local_chunk = ( + load_best_document_section_chunk( + conn, + query=query, + document=document, + intent=intent, + published_only=( + published_only + ), + ) + ) + + if local_chunk is None: + return base_results + + document_path = str( + document[ + "document_path" + ] + ) + + existing_index: int | None = None + + for index, result in enumerate( + base_results + ): + if str( + result.get( + "document_path" + ) + or "" + ) == document_path: + existing_index = ( + index + ) + + break + + if existing_index is not None: + item = dict( + base_results[ + existing_index + ] + ) + + local_chunk_id = str( + local_chunk[ + "chunk_id" + ] + ) + + primary_chunk_id = str( + item.get( + "chunk_id" + ) + or "" + ) + + if ( + local_chunk_id + != primary_chunk_id + ): + item[ + "exact_document_text" + ] = local_chunk[ + "text" + ] + + item[ + "exact_document_heading_paths" + ] = local_chunk[ + "heading_paths" + ] + + item[ + "document_expansion" + ] = ( + exact_document_expansion_metadata( + applied=( + local_chunk_id + != primary_chunk_id + ), + document_path=( + document_path + ), + chunk_id=( + local_chunk_id + ), + chunk_index=int( + local_chunk[ + "chunk_index" + ] + ), + added_source=False, + ) + ) + + del base_results[ + existing_index + ] + + # Presne zhodný dokument má byť + # v RAG kontexte pred dokumentmi, + # kde je meno iba v author poli. + base_results.insert( + 0, + item, + ) + + return ( + base_results[ + :limit + ] + if limit > 0 + else base_results + ) + + synthetic = { + **local_chunk, + "match_strategy": ( + "exact_document_section" + ), + "fts_rank": None, + "vector_rank": None, + "vector_score": None, + "hybrid_score": None, + "document_expansion": ( + exact_document_expansion_metadata( + applied=True, + document_path=( + document_path + ), + chunk_id=str( + local_chunk[ + "chunk_id" + ] + ), + chunk_index=int( + local_chunk[ + "chunk_index" + ] + ), + added_source=True, + ) + ), + } + + expanded = [ + synthetic, + *base_results, + ] + + if limit > 0: + return expanded[ + :limit + ] + + return expanded