1079 lines
18 KiB
Python
1079 lines
18 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
|
|
import app.main as main_module
|
|
import app.routes as routes
|
|
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
|
|
),
|
|
}
|
|
|
|
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(
|
|
routes,
|
|
"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
|
|
)
|
|
|
|
payload = response.json()
|
|
|
|
assert (
|
|
payload["query"]
|
|
== "Ján Holp"
|
|
)
|
|
|
|
assert (
|
|
payload["engine"]
|
|
== "hybrid_fts5_embeddings"
|
|
)
|
|
|
|
assert (
|
|
payload["source_count"]
|
|
== 1
|
|
)
|
|
|
|
assert (
|
|
payload["sources"][0][
|
|
"source_url"
|
|
]
|
|
== (
|
|
"https://example.test/"
|
|
"jan_holp"
|
|
)
|
|
)
|
|
|
|
|
|
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(
|
|
routes,
|
|
"build_rag_context",
|
|
fake_build_rag_context,
|
|
)
|
|
|
|
response = client.post(
|
|
"/rag",
|
|
headers={
|
|
"Authorization": (
|
|
f"Bearer "
|
|
f"{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_with_wrong_api_key(
|
|
client: TestClient,
|
|
) -> None:
|
|
response = client.post(
|
|
"/rag",
|
|
headers={
|
|
"X-API-Key": (
|
|
"x" * 64
|
|
),
|
|
},
|
|
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_rag_endpoint_whitespace_query(
|
|
client: TestClient,
|
|
) -> None:
|
|
response = client.post(
|
|
"/rag",
|
|
headers={
|
|
"X-API-Key": (
|
|
SEARCH_API_KEY
|
|
),
|
|
},
|
|
json={
|
|
"query": " ",
|
|
},
|
|
)
|
|
|
|
assert (
|
|
response.status_code
|
|
== 422
|
|
)
|
|
|
|
|
|
def test_rag_endpoint_rejects_invalid_limit(
|
|
client: TestClient,
|
|
) -> None:
|
|
response = client.post(
|
|
"/rag",
|
|
headers={
|
|
"X-API-Key": (
|
|
SEARCH_API_KEY
|
|
),
|
|
},
|
|
json={
|
|
"query": "test",
|
|
"limit": 100,
|
|
},
|
|
)
|
|
|
|
assert (
|
|
response.status_code
|
|
== 422
|
|
)
|
|
|
|
|
|
def test_rag_endpoint_rejects_unknown_field(
|
|
client: TestClient,
|
|
) -> None:
|
|
response = client.post(
|
|
"/rag",
|
|
headers={
|
|
"X-API-Key": (
|
|
SEARCH_API_KEY
|
|
),
|
|
},
|
|
json={
|
|
"query": "test",
|
|
"unknown_field": True,
|
|
},
|
|
)
|
|
|
|
assert (
|
|
response.status_code
|
|
== 422
|
|
)
|
|
|
|
|
|
def test_rag_endpoint_returns_400_for_value_error(
|
|
client: TestClient,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
def invalid_rag(
|
|
*args,
|
|
**kwargs,
|
|
):
|
|
raise ValueError(
|
|
"Neplatný dotaz"
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
routes,
|
|
"build_rag_context",
|
|
invalid_rag,
|
|
)
|
|
|
|
response = client.post(
|
|
"/rag",
|
|
headers={
|
|
"X-API-Key": (
|
|
SEARCH_API_KEY
|
|
),
|
|
},
|
|
json={
|
|
"query": "test",
|
|
},
|
|
)
|
|
|
|
assert (
|
|
response.status_code
|
|
== 400
|
|
)
|
|
|
|
|
|
def test_rag_endpoint_returns_503_when_database_is_missing(
|
|
client: TestClient,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
def missing_database(
|
|
*args,
|
|
**kwargs,
|
|
):
|
|
raise FileNotFoundError(
|
|
"/private/path/zp_index.sqlite"
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
routes,
|
|
"build_rag_context",
|
|
missing_database,
|
|
)
|
|
|
|
response = client.post(
|
|
"/rag",
|
|
headers={
|
|
"X-API-Key": (
|
|
SEARCH_API_KEY
|
|
),
|
|
},
|
|
json={
|
|
"query": "test",
|
|
},
|
|
)
|
|
|
|
assert (
|
|
response.status_code
|
|
== 503
|
|
)
|
|
|
|
assert response.json()[
|
|
"detail"
|
|
] == (
|
|
routes.RETRIEVAL_UNAVAILABLE_DETAIL
|
|
)
|
|
|
|
assert (
|
|
"/private/path"
|
|
not in response.text
|
|
)
|
|
|
|
|
|
def test_rag_endpoint_returns_503_for_runtime_error(
|
|
client: TestClient,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
def unavailable_rag(
|
|
*args,
|
|
**kwargs,
|
|
):
|
|
raise RuntimeError(
|
|
"embedding model failed"
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
routes,
|
|
"build_rag_context",
|
|
unavailable_rag,
|
|
)
|
|
|
|
response = client.post(
|
|
"/rag",
|
|
headers={
|
|
"X-API-Key": (
|
|
SEARCH_API_KEY
|
|
),
|
|
},
|
|
json={
|
|
"query": "test",
|
|
},
|
|
)
|
|
|
|
assert (
|
|
response.status_code
|
|
== 503
|
|
)
|
|
|
|
assert response.json()[
|
|
"detail"
|
|
] == (
|
|
routes.RETRIEVAL_UNAVAILABLE_DETAIL
|
|
)
|
|
|
|
assert (
|
|
"embedding model failed"
|
|
not in response.text
|
|
)
|
|
|
|
|
|
def test_rag_endpoint_returns_generic_500_for_unexpected_error(
|
|
client: TestClient,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
def broken_rag(
|
|
*args,
|
|
**kwargs,
|
|
):
|
|
raise TypeError(
|
|
"private internal error"
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
routes,
|
|
"build_rag_context",
|
|
broken_rag,
|
|
)
|
|
|
|
response = client.post(
|
|
"/rag",
|
|
headers={
|
|
"X-API-Key": (
|
|
SEARCH_API_KEY
|
|
),
|
|
},
|
|
json={
|
|
"query": "test",
|
|
},
|
|
)
|
|
|
|
assert (
|
|
response.status_code
|
|
== 500
|
|
)
|
|
|
|
assert response.json()[
|
|
"detail"
|
|
] == (
|
|
routes.INTERNAL_ERROR_DETAIL
|
|
)
|
|
|
|
assert (
|
|
"private internal error"
|
|
not in response.text
|
|
)
|
|
|
|
|
|
def test_openapi_exposes_rag_only(
|
|
client: TestClient,
|
|
) -> None:
|
|
response = client.get(
|
|
"/openapi.json"
|
|
)
|
|
|
|
assert (
|
|
response.status_code
|
|
== 200
|
|
)
|
|
|
|
schema = response.json()
|
|
|
|
paths = schema[
|
|
"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
|
|
)
|
|
|
|
operation = paths[
|
|
"/rag"
|
|
][
|
|
"post"
|
|
]
|
|
|
|
assert operation[
|
|
"operationId"
|
|
] == (
|
|
"retrieve_zpwiki_context"
|
|
)
|
|
|
|
assert (
|
|
"requestBody"
|
|
in operation
|
|
)
|
|
|
|
assert (
|
|
"responses"
|
|
in operation
|
|
)
|
|
|
|
assert (
|
|
"200"
|
|
in operation[
|
|
"responses"
|
|
]
|
|
)
|
|
|
|
assert (
|
|
"401"
|
|
in operation[
|
|
"responses"
|
|
]
|
|
)
|
|
|
|
assert (
|
|
"422"
|
|
in operation[
|
|
"responses"
|
|
]
|
|
)
|
|
|
|
assert (
|
|
"503"
|
|
in operation[
|
|
"responses"
|
|
]
|
|
)
|
|
|
|
|
|
def test_rag_endpoint_live_end_to_end(
|
|
security_environment,
|
|
) -> None:
|
|
if (
|
|
os.getenv(
|
|
"RUN_LIVE_RAG_E2E",
|
|
"",
|
|
)
|
|
!= "1"
|
|
):
|
|
pytest.skip(
|
|
"Live RAG E2E test je vypnutý. "
|
|
"Spusti s RUN_LIVE_RAG_E2E=1."
|
|
)
|
|
|
|
if not routes.DB_FILE.exists():
|
|
pytest.fail(
|
|
"Live RAG E2E vyžaduje "
|
|
"existujúci SQLite index."
|
|
)
|
|
|
|
with TestClient(
|
|
main_module.app
|
|
) as live_client:
|
|
response = live_client.post(
|
|
"/rag",
|
|
headers={
|
|
"X-API-Key": (
|
|
SEARCH_API_KEY
|
|
),
|
|
},
|
|
json={
|
|
"query": (
|
|
"V akom roku robil "
|
|
"Ján Holp diplomovú prácu?"
|
|
),
|
|
"limit": 5,
|
|
},
|
|
)
|
|
|
|
assert (
|
|
response.status_code
|
|
== 200
|
|
)
|
|
|
|
payload = response.json()
|
|
|
|
assert (
|
|
payload["engine"]
|
|
== "hybrid_fts5_embeddings"
|
|
)
|
|
|
|
assert (
|
|
payload["source_count"]
|
|
>= 1
|
|
)
|
|
|
|
assert (
|
|
"2021"
|
|
in payload[
|
|
"context"
|
|
]
|
|
)
|
|
|
|
assert any(
|
|
source[
|
|
"source_url"
|
|
].endswith(
|
|
"/students/2016/jan_holp"
|
|
)
|
|
for source in payload[
|
|
"sources"
|
|
]
|
|
)
|