dp-zp-agent/test/test_rag.py

from __future__ import annotations

from pathlib import Path
from typing import Any

import pytest
from fastapi.testclient import TestClient

import app.main as main_module
import scripts.rag_utils as rag_utils
from scripts.rag_utils import (
    ANSWER_FORMAT,
    RAG_INSTRUCTIONS,
    build_context_text,
    build_rag_context,
    build_source,
)


SEARCH_API_KEY = "a" * 64

@pytest.fixture
def client(
    security_environment,
) -> TestClient:
    return TestClient(
        main_module.app
    )

def sample_result() -> dict[str, Any]:
    return {
        "chunk_id": (
            "pages/students/2016/"
            "jan_holp/README.md::chunk-0"
        ),
        "document_path": (
            "pages/students/2016/"
            "jan_holp/README.md"
        ),
        "title": "Ján Holp",
        "author": "Daniel Hladek",
        "published": True,
        "heading_paths": [
            [
                "Ján Holp",
                "Diplomová práca 2021",
            ],
        ],
        "text": (
            "Dokument: Ján Holp\n"
            "Sekcia: Diplomová práca 2021\n\n"
            "Rok začiatku štúdia: 2016\n"
            "Názov diplomovej práce: "
            "Systém získavania informácií "
            "v slovenskom jazyku."
        ),
        "source_url": (
            "https://zp.kemt.fei.tuke.sk/"
            "students/2016/jan_holp"
        ),
        "match_strategy": "any_term",
        "fts_rank": 11,
        "vector_rank": 1,
        "vector_score": 0.863072,
        "hybrid_score": 0.02811129,
    }


def test_build_source() -> None:
    result = sample_result()

    source = build_source(
        result,
        1,
    )

    assert source["source_id"] == "S1"

    assert source["title"] == "Ján Holp"
    assert source["author"] == "Daniel Hladek"

    assert source["source_url"] == (
        "https://zp.kemt.fei.tuke.sk/"
        "students/2016/jan_holp"
    )

    assert source["published"] is True

    assert source["retrieval"] == {
        "match_strategy": "any_term",
        "fts_rank": 11,
        "vector_rank": 1,
        "vector_score": 0.863072,
        "hybrid_score": 0.02811129,
    }

    # Interná identifikácia zdroja nemá byť
    # používateľská citation hodnota.
    assert "citation" not in source


def test_build_context_text() -> None:
    source = build_source(
        sample_result(),
        1,
    )

    context = build_context_text(
        [source]
    )

    assert "ZDROJ S1" in context
    assert "Názov dokumentu: Ján Holp" in context

    assert (
        "Autor dokumentu: Daniel Hladek"
        in context
    )

    assert (
        "Sekcia: Diplomová práca 2021"
        in context
    )

    assert (
        "Rok začiatku štúdia: 2016"
        in context
    )

    assert (
        "https://zp.kemt.fei.tuke.sk/"
        "students/2016/jan_holp"
        in context
    )


def test_build_context_text_empty() -> None:
    context = build_context_text([])

    assert (
        "nenašli relevantné zdroje"
        in context
    )


def test_rag_instructions_require_grounding() -> None:
    instructions = " ".join(
        RAG_INSTRUCTIONS
    )

    assert (
        "výhradne podľa informácií"
        in instructions
    )

    assert (
        "v roku 2021"
        in instructions
    )

    assert (
        "roku2021"
        in instructions
    )

    assert (
        "Nevypisuj ich v konečnej odpovedi"
        in instructions
    )

    assert "source_url" in instructions


def test_answer_format() -> None:
    assert (
        ANSWER_FORMAT[
            "internal_source_ids_visible"
        ]
        is False
    )

    assert (
        ANSWER_FORMAT[
            "source_section"
        ]
        is True
    )

    assert (
        "<source_url>"
        in ANSWER_FORMAT["template"]
    )


def test_build_rag_context(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    captured: dict[str, Any] = {}

    def fake_search_database(
        db_path: Path,
        query: str,
        limit: int,
        *,
        published_only: bool,
        max_per_document: int,
    ) -> dict[str, Any]:
        captured["db_path"] = db_path
        captured["query"] = query
        captured["limit"] = limit
        captured["published_only"] = (
            published_only
        )
        captured["max_per_document"] = (
            max_per_document
        )

        return {
            "engine": (
                "hybrid_fts5_embeddings"
            ),
            "strategies": [
                "any_term"
            ],
            "results": [
                sample_result()
            ],
        }

    monkeypatch.setattr(
        rag_utils,
        "search_database",
        fake_search_database,
    )

    db_path = Path(
        "/tmp/test.sqlite"
    )

    response = build_rag_context(
        db_path,
        (
            "V akom roku robil Ján Holp "
            "diplomovú prácu?"
        ),
        limit=5,
        published_only=True,
        max_per_document=1,
    )

    assert captured == {
        "db_path": db_path,
        "query": (
            "V akom roku robil Ján Holp "
            "diplomovú prácu?"
        ),
        "limit": 5,
        "published_only": True,
        "max_per_document": 1,
    }

    assert response["engine"] == (
        "hybrid_fts5_embeddings"
    )

    assert response["strategies"] == [
        "any_term"
    ]

    assert response["source_count"] == 1

    assert (
        response["sources"][0]["title"]
        == "Ján Holp"
    )

    assert (
        "Diplomová práca 2021"
        in response["context"]
    )

    assert (
        response["answer_format"][
            "internal_source_ids_visible"
        ]
        is False
    )


def test_build_rag_context_without_results(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    def fake_search_database(
        db_path: Path,
        query: str,
        limit: int,
        *,
        published_only: bool,
        max_per_document: int,
    ) -> dict[str, Any]:
        return {
            "engine": (
                "hybrid_fts5_embeddings"
            ),
            "strategies": [],
            "results": [],
        }

    monkeypatch.setattr(
        rag_utils,
        "search_database",
        fake_search_database,
    )

    response = build_rag_context(
        Path("/tmp/test.sqlite"),
        "neexistujúca téma",
    )

    assert response["source_count"] == 0
    assert response["sources"] == []

    assert (
        "nenašli relevantné zdroje"
        in response["context"]
    )


def test_rag_endpoint(
    client: TestClient,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    expected = {
        "query": "Ján Holp",
        "engine": (
            "hybrid_fts5_embeddings"
        ),
        "strategies": [
            "all_terms"
        ],
        "source_count": 1,
        "instructions": (
            RAG_INSTRUCTIONS
        ),
        "answer_format": (
            ANSWER_FORMAT
        ),
        "context": (
            "ZDROJ S1\n"
            "Názov dokumentu: Ján Holp"
        ),
        "sources": [
            {
                "source_id": "S1",
                "title": "Ján Holp",
                "source_url": (
                    "https://example.test/"
                    "jan_holp"
                ),
            },
        ],
    }

    def fake_build_rag_context(
        db_path: Path,
        query: str,
        *,
        limit: int,
        published_only: bool,
        max_per_document: int,
    ) -> dict[str, Any]:
        assert query == "Ján Holp"
        assert limit == 5
        assert published_only is False
        assert max_per_document == 1

        return expected

    monkeypatch.setattr(
        main_module,
        "build_rag_context",
        fake_build_rag_context,
    )

    response = client.post(
        "/rag",
        headers={
            "X-API-Key": (
                SEARCH_API_KEY
            ),
        },
        json={
            "query": "Ján Holp",
        },
    )

    assert response.status_code == 200
    assert response.json() == expected


def test_rag_endpoint_with_bearer(
    client: TestClient,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    def fake_build_rag_context(
        db_path: Path,
        query: str,
        *,
        limit: int,
        published_only: bool,
        max_per_document: int,
    ) -> dict[str, Any]:
        return {
            "query": query,
            "engine": (
                "hybrid_fts5_embeddings"
            ),
            "strategies": [],
            "source_count": 0,
            "instructions": (
                RAG_INSTRUCTIONS
            ),
            "answer_format": (
                ANSWER_FORMAT
            ),
            "context": "bez výsledkov",
            "sources": [],
        }

    monkeypatch.setattr(
        main_module,
        "build_rag_context",
        fake_build_rag_context,
    )

    response = client.post(
        "/rag",
        headers={
            "Authorization": (
                f"Bearer {SEARCH_API_KEY}"
            ),
        },
        json={
            "query": "test",
        },
    )

    assert response.status_code == 200


def test_rag_endpoint_without_api_key(
    client: TestClient,
) -> None:
    response = client.post(
        "/rag",
        json={
            "query": "Ján Holp",
        },
    )

    assert response.status_code == 401


def test_rag_endpoint_empty_query(
    client: TestClient,
) -> None:
    response = client.post(
        "/rag",
        headers={
            "X-API-Key": (
                SEARCH_API_KEY
            ),
        },
        json={
            "query": "",
        },
    )

    assert response.status_code == 422


def test_openapi_exposes_rag_only(
    client: TestClient,
) -> None:
    response = client.get(
        "/openapi.json"
    )

    assert response.status_code == 200

    paths = response.json()[
        "paths"
    ]

    assert "/rag" in paths

    assert "/search" not in paths
    assert "/sync" not in paths
    assert "/health" not in paths
    assert "/webhook/gitea" not in paths