dp-zp-agent/scripts/rag_utils.py

from __future__ import annotations

from pathlib import Path
from typing import Any

from scripts.search_utils import search_database


RAG_INSTRUCTIONS = [
    (
        "Odpovedaj výhradne podľa informácií vrátených "
        "v poli context a sources zo ZP Wiki."
    ),
    (
        "Nepoužívaj vlastnú pamäť modelu, domnienky ani "
        "informácie, ktoré sa v poskytnutých zdrojoch nenachádzajú."
    ),
    (
        "Dôsledne rozlišuj názov dokumentu, autora dokumentu, "
        "osobu, o ktorej dokument pojednáva, rok začiatku štúdia "
        "a rok záverečnej práce."
    ),
    (
        "Rok začiatku štúdia nie je automaticky rokom "
        "záverečnej práce."
    ),
    (
        "Názov študentskej stránky nie je automaticky názvom "
        "záverečnej práce."
    ),
    (
        "Autor dokumentu nemusí byť osoba, o ktorej dokument "
        "pojednáva."
    ),
    (
        "Ak zdroje obsahujú viac možných interpretácií, "
        "stručne vysvetli rozdiel a nevytváraj jednoznačný "
        "záver bez opory v zdrojoch."
    ),
    (
        "Ak odpoveď nemožno zo zdrojov spoľahlivo určiť, povedz: "
        "V dostupných dokumentoch ZP Wiki sa túto informáciu "
        "nepodarilo spoľahlivo nájsť."
    ),
    (
        "Odpovedaj stručne, prirodzene a vetne po slovensky. "
        "Pri jednoduchej otázke zvyčajne stačí jedna alebo dve vety."
    ),
    (
        "Nepoužívaj odrážky, tabuľky, tučné písmo ani iné "
        "Markdown formátovanie pri jednoduchej faktickej odpovedi."
    ),
    (
        "Odpoveď formuluj prirodzenou vetou. Napríklad: "
        "'Ján Holp vypracovával diplomovú prácu v roku 2021.'"
    ),
    (
        "Dodržuj prirodzené medzery medzi slovami a číslami. "
        "Píš napríklad 'v roku 2021' a 'bol 2016'. "
        "Nikdy nepíš 'v roku2021', 'roku2021', 'bol2016' "
        "ani podobne spojené výrazy."
    ),
    (
        "Interné označenia zdrojov S1, S2, S3 a podobne slúžia "
        "iba na rozlíšenie vstupných zdrojov. "
        "Nevypisuj ich v konečnej odpovedi."
    ),
    (
        "V konečnej odpovedi nevypisuj interné retrieval údaje, "
        "ako sú fts_rank, vector_rank, vector_score, "
        "hybrid_score alebo match_strategy."
    ),
    (
        "Na konci odpovede uveď iba source_url zdrojov, "
        "z ktorých odpoveď skutočne vychádza."
    ),
    (
        "Pri jednom použitom zdroji po hlavnej odpovedi "
        "uveď samostatný riadok vo formáte "
        "'Zdroj: <source_url>'."
    ),
    (
        "Pri viacerých použitých zdrojoch napíš 'Zdroje:' "
        "a každý source_url uveď na samostatnom riadku."
    ),
    (
        "Medzi hlavnou odpoveďou a riadkom so zdrojom "
        "ponechaj prázdny riadok."
    ),
]


ANSWER_FORMAT = {
    "language": "slovak",
    "style": "stručný, prirodzený a vetný text",
    "internal_source_ids_visible": False,
    "source_section": True,
    "template": (
        "<stručná odpoveď v jednej alebo dvoch vetách>\n\n"
        "Zdroj: <source_url>"
    ),
    "single_source_template": (
        "<stručná odpoveď v jednej alebo dvoch vetách>\n\n"
        "Zdroj: <source_url>"
    ),
    "multiple_sources_template": (
        "<stručná odpoveď v prirodzených vetách>\n\n"
        "Zdroje:\n"
        "<source_url_1>\n"
        "<source_url_2>"
    ),
}


def build_source(
    result: dict[str, Any],
    number: int,
) -> dict[str, Any]:
    source_id = f"S{number}"

    return {
        "source_id": source_id,
        "title": result.get("title"),
        "author": result.get("author"),
        "document_path": result.get("document_path"),
        "source_url": result.get("source_url"),
        "published": result.get("published"),
        "section": result.get(
            "heading_paths",
            [],
        ),
        "text": result.get(
            "text",
            "",
        ),
        "retrieval": {
            "match_strategy": result.get(
                "match_strategy"
            ),
            "fts_rank": result.get(
                "fts_rank"
            ),
            "vector_rank": result.get(
                "vector_rank"
            ),
            "vector_score": result.get(
                "vector_score"
            ),
            "hybrid_score": result.get(
                "hybrid_score"
            ),
        },
    }


def build_context_text(
    sources: list[dict[str, Any]],
) -> str:
    if not sources:
        return (
            "V dostupných dokumentoch ZP Wiki "
            "sa k dotazu nenašli relevantné zdroje."
        )

    blocks: list[str] = []

    for source in sources:
        source_id = source[
            "source_id"
        ]

        title = (
            source.get("title")
            or "Neuvedené"
        )

        author = (
            source.get("author")
            or "Neuvedený"
        )

        source_url = (
            source.get("source_url")
            or "Neuvedené"
        )

        text = (
            source.get("text")
            or ""
        )

        block = (
            f"ZDROJ {source_id}\n"
            f"Názov dokumentu: {title}\n"
            f"Autor dokumentu: {author}\n"
            f"Source URL: {source_url}\n"
            "\n"
            f"{text}"
        )

        blocks.append(
            block
        )

    return (
        "\n\n"
        "--------------------"
        "\n\n"
    ).join(
        blocks
    )


def build_rag_context(
    db_path: Path,
    query: str,
    *,
    limit: int = 5,
    published_only: bool = False,
    max_per_document: int = 1,
) -> dict[str, Any]:
    response = search_database(
        db_path,
        query,
        limit,
        published_only=published_only,
        max_per_document=max_per_document,
    )

    results = response[
        "results"
    ]

    sources = [
        build_source(
            result,
            index,
        )
        for index, result in enumerate(
            results,
            start=1,
        )
    ]

    context = build_context_text(
        sources
    )

    return {
        "query": query,
        "engine": response[
            "engine"
        ],
        "strategies": response[
            "strategies"
        ],
        "source_count": len(
            sources
        ),
        "instructions": (
            RAG_INSTRUCTIONS
        ),
        "answer_format": (
            ANSWER_FORMAT
        ),
        "context": context,
        "sources": sources,
    }