dp-zp-agent/test/test_search.py

from __future__ import annotations

import json
import sqlite3
from pathlib import Path

import pytest

import scripts.build_sqlite_index as indexer
from scripts.search_utils import search_database


def write_json(path: Path, value) -> None:
    path.write_text(
        json.dumps(
            value,
            ensure_ascii=False,
        ),
        encoding="utf-8",
    )


def build_search_database(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> Path:
    documents_file = tmp_path / "documents.json"
    chunks_file = tmp_path / "chunks.json"
    db_file = tmp_path / "search.sqlite"

    documents = [
        {
            "path": "pages/students/jan_ptak/README.md",
            "title": "Ján Pták",
            "author": "Daniel Hladek",
            "published": True,
            "content_length": 100,
            "metadata": {},
        },
        {
            "path": "pages/students/jan_holp/README.md",
            "title": "Ján Holp",
            "author": "Daniel Hladek",
            "published": True,
            "content_length": 100,
            "metadata": {},
        },
        {
            "path": "pages/topics/translation/README.md",
            "title": "Strojový preklad",
            "author": "Daniel Hladek",
            "published": True,
            "content_length": 200,
            "metadata": {},
        },
        {
            "path": "pages/topics/open/README.md",
            "title": "Otvorené projekty",
            "author": "Daniel Hladek",
            "published": True,
            "content_length": 100,
            "metadata": {},
        },
        {
            "path": "pages/topics/deep/README.md",
            "title": "Neurónové siete",
            "author": "Daniel Hladek",
            "published": False,
            "content_length": 100,
            "metadata": {},
        },
    ]

    def chunk(
        chunk_id: str,
        path: str,
        title: str,
        text: str,
        published: bool,
        index: int = 0,
        tags: list[str] | None = None,
        categories: list[str] | None = None,
    ) -> dict:
        return {
            "chunk_id": chunk_id,
            "document_path": path,
            "title": title,
            "author": "Daniel Hladek",
            "published": published,
            "chunk_index": index,
            "heading_paths": [],
            "text": text,
            "text_length": len(text),
            "token_count": 30,
            "content_hash": chunk_id,
            "tags": tags or [],
            "categories": categories or [],
        }

    chunks = [
        chunk(
            "jan-ptak::0",
            "pages/students/jan_ptak/README.md",
            "Ján Pták",
            "Dokument: Ján Pták. Agent pre manažment záverečných prác.",
            True,
            tags=["rag", "nlp"],
            categories=["dp2027"],
        ),
        chunk(
            "jan-holp::0",
            "pages/students/jan_holp/README.md",
            "Ján Holp",
            "Dokument: Ján Holp. Získavanie informácií a PageRank.",
            True,
            tags=["ir"],
        ),
        chunk(
            "translation::0",
            "pages/topics/translation/README.md",
            "Strojový preklad",
            "Dokument: Strojový preklad. Štatistický strojový preklad.",
            True,
            index=0,
            tags=["translation"],
            categories=["project"],
        ),
        chunk(
            "translation::1",
            "pages/topics/translation/README.md",
            "Strojový preklad",
            "Dokument: Strojový preklad. Neurónový preklad viet.",
            True,
            index=1,
            tags=["translation"],
            categories=["project"],
        ),
        chunk(
            "open::0",
            "pages/topics/open/README.md",
            "Otvorené projekty",
            "Téma pre strojový preklad slovenského jazyka.",
            True,
            categories=["info"],
        ),
        chunk(
            "deep::0",
            "pages/topics/deep/README.md",
            "Neurónové siete",
            (
                "Modely hlbokého učenia a trénovanie "
                "hlbokých neurónových sietí."
            ),
            False,
            tags=["nn"],
        ),
    ]

    write_json(
        documents_file,
        documents,
    )

    write_json(
        chunks_file,
        chunks,
    )

    monkeypatch.setattr(
        indexer,
        "DOCUMENTS_FILE",
        documents_file,
    )

    monkeypatch.setattr(
        indexer,
        "CHUNKS_FILE",
        chunks_file,
    )

    monkeypatch.setattr(
        indexer,
        "DB_FILE",
        db_file,
    )

    indexer.build_database()

    return db_file


def test_person_query_does_not_add_any_term_noise(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    db_file = build_search_database(
        tmp_path,
        monkeypatch,
    )

    response = search_database(
        db_file,
        "jan ptak",
        limit=5,
    )

    assert response["strategies"] == [
        "all_terms"
    ]

    assert [
        item["title"]
        for item in response["results"]
    ] == [
        "Ján Pták"
    ]


def test_topic_query_returns_best_topic_first(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    db_file = build_search_database(
        tmp_path,
        monkeypatch,
    )

    response = search_database(
        db_file,
        "strojovy preklad",
        limit=5,
    )

    assert response["strategies"] == [
        "all_terms"
    ]

    assert (
        response["results"][0]["title"]
        == "Strojový preklad"
    )


def test_prefix_fallback_handles_slovak_inflection(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    db_file = build_search_database(
        tmp_path,
        monkeypatch,
    )

    response = search_database(
        db_file,
        "hlboke ucenie",
        limit=5,
    )

    assert response["strategies"] == [
        "prefix_terms"
    ]

    assert (
        response["results"][0]["title"]
        == "Neurónové siete"
    )


def test_any_term_is_used_only_when_stricter_queries_find_nothing(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    db_file = build_search_database(
        tmp_path,
        monkeypatch,
    )

    response = search_database(
        db_file,
        "nezmysel preklad",
        limit=5,
    )

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

    assert response["results"]

    assert any(
        item.get("match_strategy")
        == "any_term"
        and "preklad"
        in item["text"].casefold()
        for item in response["results"]
    )


def test_published_only_excludes_unpublished_chunks(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    db_file = build_search_database(
        tmp_path,
        monkeypatch,
    )

    all_results = search_database(
        db_file,
        "hlboke ucenie",
        limit=5,
    )

    public_results = search_database(
        db_file,
        "hlboke ucenie",
        limit=5,
        published_only=True,
    )

    assert all_results["results"]

    assert any(
        item["published"] is False
        for item in all_results["results"]
    )

    assert all(
        item["published"] is True
        for item in public_results["results"]
    )


def test_results_are_diversified_by_document(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    db_file = build_search_database(
        tmp_path,
        monkeypatch,
    )

    response = search_database(
        db_file,
        "preklad",
        limit=5,
        max_per_document=1,
    )

    paths = [
        item["document_path"]
        for item in response["results"]
    ]

    assert len(paths) == len(
        set(paths)
    )


def test_public_result_format_has_no_internal_id_and_uses_boolean(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    db_file = build_search_database(
        tmp_path,
        monkeypatch,
    )

    result = search_database(
        db_file,
        "jan ptak",
        limit=1,
    )["results"][0]

    assert "id" not in result
    assert result["published"] is True
    assert result["chunk_id"] == "jan-ptak::0"

    assert result["source_url"].endswith(
        "students/jan_ptak"
    )


def test_empty_query_returns_empty_result(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    db_file = build_search_database(
        tmp_path,
        monkeypatch,
    )

    response = search_database(
        db_file,
        "   ",
        limit=5,
    )

    assert response == {
        "engine": "hybrid_fts5_embeddings",
        "strategies": [],
        "results": [],
    }


def test_missing_database_is_reported(
    tmp_path: Path,
) -> None:
    with pytest.raises(
        FileNotFoundError
    ):
        search_database(
            tmp_path / "missing.sqlite",
            "rag",
        )


def test_missing_fts_schema_is_reported(
    tmp_path: Path,
) -> None:
    db_file = (
        tmp_path / "broken.sqlite"
    )

    with sqlite3.connect(
        db_file
    ) as conn:
        conn.execute(
            """
            CREATE TABLE chunks(
                id INTEGER PRIMARY KEY
            )
            """
        )

    with pytest.raises(
        RuntimeError,
        match="FTS5 index",
    ):
        search_database(
            db_file,
            "rag",
        )