setup tests
treefmt / nix fmt (pull_request) Failing after 7s
pytest / pytest (pull_request) Failing after 12s
build_systems / build-brain (pull_request) Successful in 44s
build_systems / build-bob (pull_request) Successful in 46s
build_systems / build-leviathan (pull_request) Successful in 54s
build_systems / build-rhapsody-in-green (pull_request) Successful in 59s
build_systems / build-jeeves (pull_request) Successful in 2m29s
treefmt / nix fmt (pull_request) Failing after 7s
pytest / pytest (pull_request) Failing after 12s
build_systems / build-brain (pull_request) Successful in 44s
build_systems / build-bob (pull_request) Successful in 46s
build_systems / build-leviathan (pull_request) Successful in 54s
build_systems / build-rhapsody-in-green (pull_request) Successful in 59s
build_systems / build-jeeves (pull_request) Successful in 2m29s
This commit is contained in:
@@ -0,0 +1,380 @@
|
|||||||
|
"""Tests for EPUB search core helpers."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from dataclasses import replace
|
||||||
|
from datetime import UTC, datetime
|
||||||
|
from os import environ
|
||||||
|
from pathlib import Path
|
||||||
|
from threading import Event
|
||||||
|
from types import ModuleType
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from sqlalchemy import create_engine, select
|
||||||
|
from sqlalchemy.orm import sessionmaker
|
||||||
|
|
||||||
|
from python.ebook_search.answer import answer_query
|
||||||
|
from python.ebook_search.bm25_corpus import (
|
||||||
|
BM25Corpus,
|
||||||
|
BM25CorpusUnavailableError,
|
||||||
|
BM25Manifest,
|
||||||
|
ensure_bm25_corpus,
|
||||||
|
load_bm25_corpus,
|
||||||
|
)
|
||||||
|
from python.ebook_search.config import EbookSearchConfig, RerankConfig, load_config, normalize_embedding_model
|
||||||
|
from python.ebook_search.embeddings import MODEL_DIMENSIONS, ensure_embedding_models
|
||||||
|
from python.ebook_search.ingest import chunk_text, find_existing_source
|
||||||
|
from python.ebook_search.search import (
|
||||||
|
SearchResponse,
|
||||||
|
SearchResult,
|
||||||
|
bm25_candidates,
|
||||||
|
reciprocal_rank_fusion,
|
||||||
|
retrieval_query_from_text,
|
||||||
|
search_ebooks,
|
||||||
|
)
|
||||||
|
from python.ebook_search.timing import RuntimeStep
|
||||||
|
from python.orm.richie import EbookEmbeddingModel, EbookSource, RichieBase
|
||||||
|
|
||||||
|
|
||||||
|
def test_chunk_text_uses_overlap() -> None:
|
||||||
|
chunks = chunk_text(" ".join(str(index) for index in range(100)), chunk_tokens=20, overlap_tokens=5)
|
||||||
|
|
||||||
|
assert len(chunks) > 1
|
||||||
|
assert chunks[0].token_start == 0
|
||||||
|
assert chunks[1].token_start == 15
|
||||||
|
assert all(chunk.token_count <= 20 for chunk in chunks)
|
||||||
|
|
||||||
|
|
||||||
|
def test_reciprocal_rank_fusion_combines_vector_and_bm25_rankings() -> None:
|
||||||
|
vector_results = [
|
||||||
|
SearchResult(chunk_id=1, text="a", source_title="A", score=0.9, vector_score=0.9),
|
||||||
|
SearchResult(chunk_id=2, text="b", source_title="B", score=0.8, vector_score=0.8),
|
||||||
|
]
|
||||||
|
lexical_results = [
|
||||||
|
SearchResult(chunk_id=2, text="b", source_title="B", score=4.2, bm25_score=4.2),
|
||||||
|
SearchResult(chunk_id=3, text="c", source_title="C", score=2.1, bm25_score=2.1),
|
||||||
|
]
|
||||||
|
|
||||||
|
fused = reciprocal_rank_fusion(vector_results, lexical_results)
|
||||||
|
|
||||||
|
assert [result.chunk_id for result in fused] == [2, 1, 3]
|
||||||
|
assert fused[0].rank_source == "Hybrid"
|
||||||
|
assert fused[0].vector_score == 0.8
|
||||||
|
assert fused[0].bm25_score == 4.2
|
||||||
|
assert fused[0].fused_score == fused[0].score
|
||||||
|
|
||||||
|
|
||||||
|
def test_find_existing_source_matches_path_or_hash() -> None:
|
||||||
|
engine = create_engine("sqlite+pysqlite:///:memory:", future=True)
|
||||||
|
RichieBase.metadata.create_all(engine)
|
||||||
|
with sessionmaker(bind=engine, expire_on_commit=False, future=True)() as session:
|
||||||
|
source = EbookSource(
|
||||||
|
title="Book",
|
||||||
|
author=None,
|
||||||
|
language=None,
|
||||||
|
publisher=None,
|
||||||
|
identifier=None,
|
||||||
|
file_path="/old/book.epub",
|
||||||
|
file_sha256="a" * 64,
|
||||||
|
file_mtime=datetime.now(tz=UTC),
|
||||||
|
file_size=10,
|
||||||
|
)
|
||||||
|
session.add(source)
|
||||||
|
session.commit()
|
||||||
|
|
||||||
|
assert find_existing_source(session, Path("/old/book.epub"), "b" * 64) == source
|
||||||
|
assert find_existing_source(session, Path("/new/book.epub"), "a" * 64) == source
|
||||||
|
|
||||||
|
|
||||||
|
def test_reciprocal_rank_fusion_marks_hybrid_source() -> None:
|
||||||
|
vector_results = [SearchResult(chunk_id=1, text="a", source_title="A")]
|
||||||
|
lexical_results = [SearchResult(chunk_id=2, text="b", source_title="B")]
|
||||||
|
|
||||||
|
fused = reciprocal_rank_fusion(vector_results, lexical_results)
|
||||||
|
|
||||||
|
assert {result.rank_source for result in fused} == {"Hybrid"}
|
||||||
|
|
||||||
|
|
||||||
|
def test_search_response_sums_runtime_steps() -> None:
|
||||||
|
response = SearchResponse(
|
||||||
|
query="query",
|
||||||
|
results=[],
|
||||||
|
rank_label="Hybrid",
|
||||||
|
timings=(
|
||||||
|
RuntimeStep(name="A", duration_ms=1.25),
|
||||||
|
RuntimeStep(name="B", duration_ms=2.75),
|
||||||
|
RuntimeStep(name="Parallel detail", duration_ms=10.0, counts_toward_total=False),
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
assert response.total_runtime_ms == 4.0
|
||||||
|
|
||||||
|
|
||||||
|
def test_search_ebooks_runs_vector_and_bm25_in_parallel(monkeypatch) -> None:
|
||||||
|
engine = create_engine("sqlite+pysqlite:///:memory:", future=True)
|
||||||
|
vector_started = Event()
|
||||||
|
bm25_started = Event()
|
||||||
|
received_engines: list[object] = []
|
||||||
|
|
||||||
|
def fake_vector_candidates(received_engine, query, _config):
|
||||||
|
"""Return vector candidates after confirming BM25 has started."""
|
||||||
|
received_engines.append(received_engine)
|
||||||
|
assert query == "parallel"
|
||||||
|
vector_started.set()
|
||||||
|
assert bm25_started.wait(timeout=2)
|
||||||
|
return [SearchResult(chunk_id=1, text="vector", source_title="Vector", vector_score=0.9)]
|
||||||
|
|
||||||
|
def fake_bm25_candidates(query, _config):
|
||||||
|
"""Return BM25 candidates after confirming vector search has started."""
|
||||||
|
assert query == "parallel"
|
||||||
|
bm25_started.set()
|
||||||
|
assert vector_started.wait(timeout=2)
|
||||||
|
return [SearchResult(chunk_id=2, text="bm25", source_title="BM25", bm25_score=2.0)]
|
||||||
|
|
||||||
|
monkeypatch.setattr("python.ebook_search.search.vector_candidates", fake_vector_candidates)
|
||||||
|
monkeypatch.setattr("python.ebook_search.search.bm25_candidates", fake_bm25_candidates)
|
||||||
|
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
|
||||||
|
|
||||||
|
response = search_ebooks(engine, "parallel", config)
|
||||||
|
|
||||||
|
timings = {step.name: step for step in response.timings}
|
||||||
|
assert [result.chunk_id for result in response.results] == [1, 2]
|
||||||
|
assert timings["Embedding + vector search"].counts_toward_total is False
|
||||||
|
assert timings["BM25 search"].counts_toward_total is False
|
||||||
|
assert timings["Hybrid retrieval"].counts_toward_total is True
|
||||||
|
assert received_engines == [engine]
|
||||||
|
|
||||||
|
|
||||||
|
def test_retrieval_query_keeps_entity_and_series_terms() -> None:
|
||||||
|
assert retrieval_query_from_text("what does Damien Montgomery stand for in starship mage") == (
|
||||||
|
"damien montgomery stand starship mage"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_bm25_candidates_scores_whole_corpus(monkeypatch) -> None:
|
||||||
|
record = {
|
||||||
|
"chunk_id": 2,
|
||||||
|
"text": "high",
|
||||||
|
"source_title": "B",
|
||||||
|
"source_author": None,
|
||||||
|
"chapter_title": None,
|
||||||
|
"page_label": None,
|
||||||
|
"bm25_text": "high",
|
||||||
|
}
|
||||||
|
manifest = BM25Manifest(created_at=datetime.now(tz=UTC), db_updated_at=None, chunk_count=1)
|
||||||
|
corpus = BM25Corpus(retriever=object(), records=(record,), manifest=manifest)
|
||||||
|
captured: dict[str, object] = {}
|
||||||
|
|
||||||
|
def fake_score_bm25_corpus(query, saved_corpus, *, limit):
|
||||||
|
captured["query"] = query
|
||||||
|
captured["corpus"] = saved_corpus
|
||||||
|
captured["limit"] = limit
|
||||||
|
return [(record, 1.5)]
|
||||||
|
|
||||||
|
monkeypatch.setattr("python.ebook_search.search.load_bm25_corpus", lambda _config: corpus)
|
||||||
|
monkeypatch.setattr("python.ebook_search.search.score_bm25_corpus", fake_score_bm25_corpus)
|
||||||
|
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
|
||||||
|
|
||||||
|
results = bm25_candidates("high", config)
|
||||||
|
|
||||||
|
assert captured["query"] == "high"
|
||||||
|
assert captured["corpus"] == corpus
|
||||||
|
assert captured["limit"] == 120
|
||||||
|
assert [result.chunk_id for result in results] == [2]
|
||||||
|
assert [result.bm25_score for result in results] == [1.5]
|
||||||
|
|
||||||
|
|
||||||
|
def test_bm25_candidates_raises_when_corpus_is_unavailable(monkeypatch) -> None:
|
||||||
|
def fake_load_bm25_corpus(_config):
|
||||||
|
raise BM25CorpusUnavailableError
|
||||||
|
|
||||||
|
monkeypatch.setattr("python.ebook_search.search.load_bm25_corpus", fake_load_bm25_corpus)
|
||||||
|
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
|
||||||
|
|
||||||
|
with pytest.raises(BM25CorpusUnavailableError):
|
||||||
|
bm25_candidates("high", config)
|
||||||
|
|
||||||
|
|
||||||
|
def test_load_bm25_corpus_caches_disk_load(monkeypatch, tmp_path) -> None:
|
||||||
|
load_bm25_corpus.cache_clear()
|
||||||
|
manifest = BM25Manifest(created_at=datetime.now(tz=UTC), db_updated_at=None, chunk_count=1)
|
||||||
|
record = {
|
||||||
|
"chunk_id": 2,
|
||||||
|
"text": "cached",
|
||||||
|
"source_title": "B",
|
||||||
|
"source_author": None,
|
||||||
|
"chapter_title": None,
|
||||||
|
"page_label": None,
|
||||||
|
"bm25_text": "cached",
|
||||||
|
}
|
||||||
|
load_count = 0
|
||||||
|
|
||||||
|
class FakeRetriever:
|
||||||
|
"""Fake persisted BM25 retriever."""
|
||||||
|
|
||||||
|
corpus = (record,)
|
||||||
|
|
||||||
|
class FakeBM25:
|
||||||
|
"""Fake BM25 class with observable load count."""
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def load(index_path, *, load_corpus, mmap):
|
||||||
|
nonlocal load_count
|
||||||
|
load_count += 1
|
||||||
|
assert index_path == tmp_path
|
||||||
|
assert load_corpus is True
|
||||||
|
assert mmap is True
|
||||||
|
return FakeRetriever()
|
||||||
|
|
||||||
|
fake_bm25s = ModuleType("bm25s")
|
||||||
|
fake_bm25s.BM25 = FakeBM25
|
||||||
|
monkeypatch.setattr("python.ebook_search.bm25_corpus.read_bm25_manifest", lambda _path: manifest)
|
||||||
|
monkeypatch.setattr("python.ebook_search.bm25_corpus.bm25_index_exists", lambda _path, _manifest: True)
|
||||||
|
monkeypatch.setattr("python.ebook_search.bm25_corpus.bm25s", fake_bm25s)
|
||||||
|
config = EbookSearchConfig(rerank=RerankConfig(enabled=False), bm25_index_dir=str(tmp_path))
|
||||||
|
|
||||||
|
try:
|
||||||
|
first = load_bm25_corpus(config)
|
||||||
|
second = load_bm25_corpus(config)
|
||||||
|
finally:
|
||||||
|
load_bm25_corpus.cache_clear()
|
||||||
|
|
||||||
|
assert first is second
|
||||||
|
assert first is not None
|
||||||
|
assert first.records == (record,)
|
||||||
|
assert load_count == 1
|
||||||
|
|
||||||
|
|
||||||
|
def test_load_bm25_corpus_raises_when_index_is_missing(monkeypatch, tmp_path) -> None:
|
||||||
|
load_bm25_corpus.cache_clear()
|
||||||
|
monkeypatch.setattr("python.ebook_search.bm25_corpus.read_bm25_manifest", lambda _path: None)
|
||||||
|
monkeypatch.setattr("python.ebook_search.bm25_corpus.bm25_index_exists", lambda _path, _manifest: False)
|
||||||
|
config = EbookSearchConfig(rerank=RerankConfig(enabled=False), bm25_index_dir=str(tmp_path))
|
||||||
|
|
||||||
|
try:
|
||||||
|
with pytest.raises(BM25CorpusUnavailableError, match="BM25 corpus is not available"):
|
||||||
|
load_bm25_corpus(config)
|
||||||
|
finally:
|
||||||
|
load_bm25_corpus.cache_clear()
|
||||||
|
|
||||||
|
|
||||||
|
def test_ensure_bm25_corpus_refreshes_missing_index(monkeypatch) -> None:
|
||||||
|
refreshed: list[object] = []
|
||||||
|
db_updated_at = datetime.now(tz=UTC)
|
||||||
|
|
||||||
|
monkeypatch.setattr("python.ebook_search.bm25_corpus.read_bm25_manifest", lambda _path: None)
|
||||||
|
monkeypatch.setattr("python.ebook_search.bm25_corpus.bm25_index_exists", lambda _path, _manifest: False)
|
||||||
|
monkeypatch.setattr("python.ebook_search.bm25_corpus.corpus_last_updated_at", lambda _session: db_updated_at)
|
||||||
|
monkeypatch.setattr(
|
||||||
|
"python.ebook_search.bm25_corpus.refresh_bm25_corpus",
|
||||||
|
lambda session, config, *, db_updated_at: refreshed.append((session, config, db_updated_at)),
|
||||||
|
)
|
||||||
|
|
||||||
|
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
|
||||||
|
session = object()
|
||||||
|
|
||||||
|
ensure_bm25_corpus(session, config)
|
||||||
|
|
||||||
|
assert refreshed == [(session, config, db_updated_at)]
|
||||||
|
|
||||||
|
|
||||||
|
def test_ensure_bm25_corpus_refreshes_stale_index(monkeypatch) -> None:
|
||||||
|
refreshed: list[object] = []
|
||||||
|
created_at = datetime(2026, 1, 1, tzinfo=UTC)
|
||||||
|
db_updated_at = datetime(2026, 1, 2, tzinfo=UTC)
|
||||||
|
manifest = BM25Manifest(created_at=created_at, db_updated_at=created_at, chunk_count=10)
|
||||||
|
|
||||||
|
monkeypatch.setattr("python.ebook_search.bm25_corpus.read_bm25_manifest", lambda _path: manifest)
|
||||||
|
monkeypatch.setattr("python.ebook_search.bm25_corpus.bm25_index_exists", lambda _path, _manifest: True)
|
||||||
|
monkeypatch.setattr("python.ebook_search.bm25_corpus.corpus_last_updated_at", lambda _session: db_updated_at)
|
||||||
|
monkeypatch.setattr(
|
||||||
|
"python.ebook_search.bm25_corpus.refresh_bm25_corpus",
|
||||||
|
lambda session, config, *, db_updated_at: refreshed.append((session, config, db_updated_at)),
|
||||||
|
)
|
||||||
|
|
||||||
|
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
|
||||||
|
session = object()
|
||||||
|
|
||||||
|
ensure_bm25_corpus(session, config)
|
||||||
|
|
||||||
|
assert refreshed == [(session, config, db_updated_at)]
|
||||||
|
|
||||||
|
|
||||||
|
def test_supported_embedding_models_match_service_names() -> None:
|
||||||
|
assert MODEL_DIMENSIONS == {
|
||||||
|
"qwen3-embedding-0.6b": 1024,
|
||||||
|
"qwen3-embedding-4b": 2560,
|
||||||
|
"qwen3-embedding-8b": 4096,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def test_ensure_embedding_models_registers_service_names() -> None:
|
||||||
|
engine = create_engine("sqlite+pysqlite:///:memory:", future=True)
|
||||||
|
RichieBase.metadata.create_all(engine)
|
||||||
|
with sessionmaker(bind=engine, expire_on_commit=False, future=True)() as session:
|
||||||
|
ensure_embedding_models(session)
|
||||||
|
session.commit()
|
||||||
|
|
||||||
|
models = list(session.scalars(select(EbookEmbeddingModel).order_by(EbookEmbeddingModel.name)))
|
||||||
|
|
||||||
|
assert [(model.name, model.dimension) for model in models] == [
|
||||||
|
("qwen3-embedding-0.6b", 1024),
|
||||||
|
("qwen3-embedding-4b", 2560),
|
||||||
|
("qwen3-embedding-8b", 4096),
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
def test_embedding_model_aliases_normalize_to_provider_names() -> None:
|
||||||
|
assert normalize_embedding_model() == "qwen3-embedding-0.6b"
|
||||||
|
|
||||||
|
environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "qwen3-embedding-0.6b"
|
||||||
|
assert normalize_embedding_model() == "qwen3-embedding-0.6b"
|
||||||
|
|
||||||
|
environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "Qwen3-Embedding-0.6B"
|
||||||
|
assert normalize_embedding_model() == "qwen3-embedding-0.6b"
|
||||||
|
|
||||||
|
environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "Qwen/Qwen3-Embedding-4B"
|
||||||
|
|
||||||
|
assert normalize_embedding_model() == "qwen3-embedding-4b"
|
||||||
|
|
||||||
|
environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "qwen3-embedding:8b"
|
||||||
|
assert normalize_embedding_model() == "qwen3-embedding-8b"
|
||||||
|
|
||||||
|
environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "qwen3-embedding-8b"
|
||||||
|
assert normalize_embedding_model() == "qwen3-embedding-8b"
|
||||||
|
|
||||||
|
|
||||||
|
def test_answer_generation_is_enabled_by_default(monkeypatch) -> None:
|
||||||
|
monkeypatch.delenv("EBOOK_SEARCH_ANSWER_ENABLED", raising=False)
|
||||||
|
|
||||||
|
config = load_config()
|
||||||
|
|
||||||
|
assert config.answer_enabled is True
|
||||||
|
|
||||||
|
|
||||||
|
def test_chat_defaults_use_ollama_cloud(monkeypatch) -> None:
|
||||||
|
monkeypatch.delenv("EBOOK_SEARCH_VLLM_BASE_URL", raising=False)
|
||||||
|
monkeypatch.delenv("EBOOK_SEARCH_CHAT_MODEL", raising=False)
|
||||||
|
|
||||||
|
config = load_config()
|
||||||
|
|
||||||
|
assert config.vllm_base_url == "https://ollama.com/v1"
|
||||||
|
assert config.chat_model == "deepseek-v4-flash"
|
||||||
|
|
||||||
|
|
||||||
|
def test_chat_api_key_falls_back_to_ollama_api_key(monkeypatch) -> None:
|
||||||
|
monkeypatch.delenv("EBOOK_SEARCH_VLLM_API_KEY", raising=False)
|
||||||
|
monkeypatch.setenv("OLLAMA_API_KEY", "ollama-key")
|
||||||
|
|
||||||
|
config = load_config()
|
||||||
|
|
||||||
|
assert config.vllm_api_key == "ollama-key"
|
||||||
|
|
||||||
|
|
||||||
|
def test_answer_query_does_not_call_model_when_disabled() -> None:
|
||||||
|
config = replace(load_config(), answer_enabled=False)
|
||||||
|
result = SearchResult(chunk_id=1, text="source text", source_title="Book")
|
||||||
|
|
||||||
|
answer = answer_query("question", [result], config)
|
||||||
|
|
||||||
|
assert "Answer generation is disabled" in answer
|
||||||
@@ -0,0 +1,84 @@
|
|||||||
|
"""Tests for EPUB search HTTP model adapters."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from python.ebook_search.answer import answer_query
|
||||||
|
from python.ebook_search.config import EbookSearchConfig, RerankConfig
|
||||||
|
from python.ebook_search.embeddings import embed_texts
|
||||||
|
from python.ebook_search.search import SearchResult
|
||||||
|
|
||||||
|
|
||||||
|
def test_answer_query_uses_httpx_chat_completions(monkeypatch) -> None:
|
||||||
|
captured: dict[str, object] = {}
|
||||||
|
|
||||||
|
def fake_post(url: str, **kwargs: object) -> httpx.Response:
|
||||||
|
captured["url"] = url
|
||||||
|
captured["kwargs"] = kwargs
|
||||||
|
return httpx.Response(
|
||||||
|
200,
|
||||||
|
json={"choices": [{"message": {"content": "grounded answer"}}]},
|
||||||
|
request=httpx.Request("POST", url),
|
||||||
|
)
|
||||||
|
|
||||||
|
monkeypatch.setattr(httpx, "post", fake_post)
|
||||||
|
config = EbookSearchConfig(
|
||||||
|
rerank=RerankConfig(enabled=False),
|
||||||
|
vllm_base_url="https://ollama.com/v1",
|
||||||
|
vllm_api_key="secret",
|
||||||
|
chat_model="deepseek-v4-flash",
|
||||||
|
)
|
||||||
|
|
||||||
|
answer = answer_query("question", [SearchResult(chunk_id=1, text="source", source_title="Book")], config)
|
||||||
|
|
||||||
|
assert answer == "grounded answer"
|
||||||
|
assert captured["url"] == "https://ollama.com/v1/chat/completions"
|
||||||
|
kwargs = captured["kwargs"]
|
||||||
|
assert isinstance(kwargs, dict)
|
||||||
|
assert kwargs["headers"] == {"Authorization": "Bearer secret"}
|
||||||
|
payload = kwargs["json"]
|
||||||
|
assert isinstance(payload, dict)
|
||||||
|
assert payload["model"] == "deepseek-v4-flash"
|
||||||
|
|
||||||
|
|
||||||
|
def test_embed_texts_uses_httpx_embeddings(monkeypatch) -> None:
|
||||||
|
captured: dict[str, object] = {}
|
||||||
|
vector = [0.0] * 1024
|
||||||
|
|
||||||
|
def fake_post(url: str, **kwargs: object) -> httpx.Response:
|
||||||
|
captured["url"] = url
|
||||||
|
captured["kwargs"] = kwargs
|
||||||
|
return httpx.Response(
|
||||||
|
200,
|
||||||
|
json={"data": [{"embedding": vector}]},
|
||||||
|
request=httpx.Request("POST", url),
|
||||||
|
)
|
||||||
|
|
||||||
|
monkeypatch.setattr(httpx, "post", fake_post)
|
||||||
|
config = EbookSearchConfig(
|
||||||
|
rerank=RerankConfig(enabled=False),
|
||||||
|
embedding_base_url="http://bob:8000/v1",
|
||||||
|
embedding_model="qwen3-embedding-0.6b",
|
||||||
|
)
|
||||||
|
|
||||||
|
embeddings = embed_texts(["hello"], config)
|
||||||
|
|
||||||
|
assert embeddings == [vector]
|
||||||
|
assert captured["url"] == "http://bob:8000/v1/embeddings"
|
||||||
|
kwargs = captured["kwargs"]
|
||||||
|
assert isinstance(kwargs, dict)
|
||||||
|
assert kwargs["headers"] == {}
|
||||||
|
assert kwargs["json"] == {"model": "qwen3-embedding-0.6b", "input": ["hello"]}
|
||||||
|
|
||||||
|
|
||||||
|
def test_embed_texts_rejects_bad_response_shape(monkeypatch) -> None:
|
||||||
|
def fake_post(url: str, **_kwargs: object) -> httpx.Response:
|
||||||
|
return httpx.Response(200, json={"data": [{}]}, request=httpx.Request("POST", url))
|
||||||
|
|
||||||
|
monkeypatch.setattr(httpx, "post", fake_post)
|
||||||
|
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="Embedding request failed"):
|
||||||
|
embed_texts(["hello"], config)
|
||||||
@@ -0,0 +1,150 @@
|
|||||||
|
"""Tests for EPUB search reranking."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from python.ebook_search.config import EbookSearchConfig, RerankConfig, load_rerank_config
|
||||||
|
from python.ebook_search.rerank import rerank_chunks
|
||||||
|
from python.ebook_search.search import SearchResult, apply_rerank, skip_rerank
|
||||||
|
|
||||||
|
|
||||||
|
def candidates() -> list[SearchResult]:
|
||||||
|
return [
|
||||||
|
SearchResult(chunk_id=1, text="alpha", source_title="A", score=0.9),
|
||||||
|
SearchResult(chunk_id=2, text="beta", source_title="B", score=0.8),
|
||||||
|
SearchResult(chunk_id=3, text="gamma", source_title="C", score=0.7),
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
def rerank_response(payload: dict[str, object] | None = None, *, content: bytes | None = None) -> httpx.Response:
|
||||||
|
return httpx.Response(
|
||||||
|
200,
|
||||||
|
content=content,
|
||||||
|
json=payload,
|
||||||
|
request=httpx.Request("POST", "http://rerank.test/rerank"),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_config_defaults_keep_reranking_optional(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||||
|
monkeypatch.delenv("EBOOK_SEARCH_RERANK_ENABLED", raising=False)
|
||||||
|
monkeypatch.delenv("EBOOK_SEARCH_RERANK_BASE_URL", raising=False)
|
||||||
|
monkeypatch.delenv("EBOOK_SEARCH_RERANK_MODEL", raising=False)
|
||||||
|
monkeypatch.delenv("EBOOK_SEARCH_RERANK_CANDIDATES", raising=False)
|
||||||
|
monkeypatch.delenv("EBOOK_SEARCH_RERANK_TIMEOUT_SECONDS", raising=False)
|
||||||
|
|
||||||
|
config = load_rerank_config()
|
||||||
|
|
||||||
|
assert config.enabled is False
|
||||||
|
assert config.base_url == "http://192.168.90.25:8001"
|
||||||
|
assert config.model == "qwen3-reranker-06b"
|
||||||
|
assert config.candidates == 24
|
||||||
|
assert config.timeout_seconds == 30
|
||||||
|
|
||||||
|
|
||||||
|
def test_reranking_disabled_returns_original_fused_order() -> None:
|
||||||
|
config = EbookSearchConfig(rerank=RerankConfig(enabled=False), top_k=2)
|
||||||
|
|
||||||
|
response = skip_rerank("query", candidates(), config)
|
||||||
|
|
||||||
|
assert response.rank_label == "Hybrid"
|
||||||
|
assert [result.chunk_id for result in response.results] == [1, 2]
|
||||||
|
|
||||||
|
|
||||||
|
def test_reranking_enabled_reorders_candidates(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||||
|
def fake_post(_url: str, *, json: dict[str, object], timeout: float) -> httpx.Response:
|
||||||
|
assert timeout == 30
|
||||||
|
assert json == {
|
||||||
|
"model": "qwen3-reranker-06b",
|
||||||
|
"query": "query",
|
||||||
|
"documents": ["alpha", "beta", "gamma"],
|
||||||
|
}
|
||||||
|
return rerank_response(
|
||||||
|
{
|
||||||
|
"results": [
|
||||||
|
{"index": 0, "relevance_score": 0.1},
|
||||||
|
{"index": 1, "relevance_score": 0.9},
|
||||||
|
{"index": 2, "relevance_score": 0.4},
|
||||||
|
]
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
monkeypatch.setattr(httpx, "post", fake_post)
|
||||||
|
|
||||||
|
results = rerank_chunks("query", candidates(), RerankConfig())
|
||||||
|
|
||||||
|
assert [result.chunk_id for result in results] == [2, 1, 3]
|
||||||
|
assert [round(result.score, 3) for result in results] == [0.45, 0.1, 0.0]
|
||||||
|
assert [result.rerank_score for result in results] == [0.9, 0.1, 0.4]
|
||||||
|
|
||||||
|
|
||||||
|
def test_reranking_cannot_ignore_hybrid_score(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||||
|
candidates = [
|
||||||
|
SearchResult(chunk_id=1, text="strong hybrid", source_title="A", score=1.0),
|
||||||
|
SearchResult(chunk_id=2, text="weak hybrid", source_title="B", score=0.1),
|
||||||
|
]
|
||||||
|
|
||||||
|
def fake_post(_url: str, **_kwargs: object) -> httpx.Response:
|
||||||
|
return rerank_response(
|
||||||
|
{
|
||||||
|
"results": [
|
||||||
|
{"index": 0, "relevance_score": 0.7},
|
||||||
|
{"index": 1, "relevance_score": 1.0},
|
||||||
|
]
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
monkeypatch.setattr(httpx, "post", fake_post)
|
||||||
|
|
||||||
|
results = rerank_chunks("query", candidates, RerankConfig())
|
||||||
|
|
||||||
|
assert [result.chunk_id for result in results] == [1, 2]
|
||||||
|
assert results[0].score == 0.7
|
||||||
|
assert results[1].score == 0.0
|
||||||
|
assert results[1].rerank_score == 1.0
|
||||||
|
|
||||||
|
|
||||||
|
def test_vllm_rerank_timeout_raises(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||||
|
def fake_rerank_chunks(
|
||||||
|
_query: str,
|
||||||
|
_candidates: list[SearchResult],
|
||||||
|
_config: RerankConfig,
|
||||||
|
) -> list[SearchResult]:
|
||||||
|
message = "timeout"
|
||||||
|
raise httpx.TimeoutException(message)
|
||||||
|
|
||||||
|
monkeypatch.setattr("python.ebook_search.search.rerank_chunks", fake_rerank_chunks)
|
||||||
|
config = EbookSearchConfig(rerank=RerankConfig(enabled=True), top_k=2)
|
||||||
|
|
||||||
|
with pytest.raises(httpx.TimeoutException, match="timeout"):
|
||||||
|
apply_rerank("query", candidates(), config)
|
||||||
|
|
||||||
|
|
||||||
|
def test_malformed_vllm_rerank_json_does_not_crash_search(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||||
|
def fake_post(_url: str, **_kwargs: object) -> httpx.Response:
|
||||||
|
return rerank_response(content=b"not-json")
|
||||||
|
|
||||||
|
monkeypatch.setattr(httpx, "post", fake_post)
|
||||||
|
|
||||||
|
results = rerank_chunks("query", candidates()[:1], RerankConfig())
|
||||||
|
|
||||||
|
assert results[0].score == 0.0
|
||||||
|
|
||||||
|
|
||||||
|
def test_vllm_rerank_scores_are_clamped(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||||
|
def fake_post(_url: str, **_kwargs: object) -> httpx.Response:
|
||||||
|
return rerank_response(
|
||||||
|
{
|
||||||
|
"results": [
|
||||||
|
{"index": 0, "relevance_score": -1},
|
||||||
|
{"index": 1, "relevance_score": 2},
|
||||||
|
]
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
monkeypatch.setattr(httpx, "post", fake_post)
|
||||||
|
|
||||||
|
results = rerank_chunks("query", candidates()[:2], RerankConfig())
|
||||||
|
|
||||||
|
assert [result.rerank_score for result in results] == [0.0, 1.0]
|
||||||
@@ -0,0 +1,265 @@
|
|||||||
|
"""Tests for EPUB search HTMX routes."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from fastapi.testclient import TestClient
|
||||||
|
from sqlalchemy import create_engine
|
||||||
|
|
||||||
|
from python.ebook_search.api.main import create_app
|
||||||
|
from python.ebook_search.config import EbookSearchConfig, RerankConfig
|
||||||
|
from python.ebook_search.embeddings import EmbeddingModelStats
|
||||||
|
from python.ebook_search.search import SearchResponse, SearchResult
|
||||||
|
from python.ebook_search.timing import RuntimeStep
|
||||||
|
|
||||||
|
|
||||||
|
def patch_app_runtime(monkeypatch):
|
||||||
|
"""Patch app startup dependencies used by UI route tests."""
|
||||||
|
monkeypatch.setattr("python.ebook_search.api.main.get_postgres_engine", fake_get_postgres_engine)
|
||||||
|
monkeypatch.setattr("python.ebook_search.api.main.ensure_bm25_corpus", lambda _session, _config: None)
|
||||||
|
|
||||||
|
|
||||||
|
def fake_get_postgres_engine(**_kwargs):
|
||||||
|
"""Return an in-memory engine for route tests."""
|
||||||
|
return create_engine("sqlite+pysqlite:///:memory:", future=True)
|
||||||
|
|
||||||
|
|
||||||
|
def test_ui_form_passes_rerank_flag_to_search_handler(monkeypatch) -> None:
|
||||||
|
captured: dict[str, object] = {}
|
||||||
|
|
||||||
|
def fake_search_ebooks(_engine, query, config, *, rerank=False):
|
||||||
|
captured["query"] = query
|
||||||
|
captured["rerank"] = rerank
|
||||||
|
captured["config"] = config
|
||||||
|
return SearchResponse(query=query, results=[], rank_label="Hybrid + rerank")
|
||||||
|
|
||||||
|
monkeypatch.setattr("python.ebook_search.api.routes.search.search_ebooks", fake_search_ebooks)
|
||||||
|
monkeypatch.setattr(
|
||||||
|
"python.ebook_search.api.routes.search.answer_query",
|
||||||
|
lambda _query, _results, _config: "answer",
|
||||||
|
)
|
||||||
|
patch_app_runtime(monkeypatch)
|
||||||
|
app = create_app()
|
||||||
|
app.state.config = EbookSearchConfig(rerank=RerankConfig(enabled=False), top_k=12, answer_enabled=True)
|
||||||
|
|
||||||
|
with TestClient(app) as client:
|
||||||
|
response = client.post("/search", data={"query": "where is the quote?", "rerank": "true"})
|
||||||
|
|
||||||
|
assert response.status_code == 200
|
||||||
|
assert "Hybrid + rerank" in response.text
|
||||||
|
assert captured["query"] == "where is the quote?"
|
||||||
|
assert captured["rerank"] is True
|
||||||
|
|
||||||
|
|
||||||
|
def test_ui_search_failure_returns_visible_error(monkeypatch) -> None:
|
||||||
|
def fake_search_ebooks(_engine, _query, _config, *, rerank=False):
|
||||||
|
del rerank
|
||||||
|
msg = "search exploded"
|
||||||
|
raise RuntimeError(msg)
|
||||||
|
|
||||||
|
monkeypatch.setattr("python.ebook_search.api.routes.search.search_ebooks", fake_search_ebooks)
|
||||||
|
patch_app_runtime(monkeypatch)
|
||||||
|
app = create_app()
|
||||||
|
app.state.config = EbookSearchConfig(rerank=RerankConfig(enabled=False), top_k=12)
|
||||||
|
|
||||||
|
with TestClient(app) as client:
|
||||||
|
response = client.post("/search", data={"query": "where is the quote?"})
|
||||||
|
|
||||||
|
assert response.status_code == 500
|
||||||
|
assert "search exploded" in response.text
|
||||||
|
|
||||||
|
|
||||||
|
def test_ui_answer_failure_still_returns_sources(monkeypatch) -> None:
|
||||||
|
def fake_search_ebooks(_engine, query, _config, *, rerank=False):
|
||||||
|
del rerank
|
||||||
|
return SearchResponse(query=query, results=[], rank_label="Hybrid")
|
||||||
|
|
||||||
|
def fake_answer_query(_query, _results, _config):
|
||||||
|
msg = "answer exploded"
|
||||||
|
raise RuntimeError(msg)
|
||||||
|
|
||||||
|
monkeypatch.setattr("python.ebook_search.api.routes.search.search_ebooks", fake_search_ebooks)
|
||||||
|
monkeypatch.setattr("python.ebook_search.api.routes.search.answer_query", fake_answer_query)
|
||||||
|
patch_app_runtime(monkeypatch)
|
||||||
|
app = create_app()
|
||||||
|
app.state.config = EbookSearchConfig(rerank=RerankConfig(enabled=False), top_k=12, answer_enabled=True)
|
||||||
|
|
||||||
|
with TestClient(app) as client:
|
||||||
|
response = client.post("/search", data={"query": "where is the quote?"})
|
||||||
|
|
||||||
|
assert response.status_code == 200
|
||||||
|
assert "Answer generation failed" in response.text
|
||||||
|
|
||||||
|
|
||||||
|
def test_ui_skips_answer_when_disabled(monkeypatch) -> None:
|
||||||
|
called = False
|
||||||
|
|
||||||
|
def fake_search_ebooks(_engine, query, _config, *, rerank=False):
|
||||||
|
del rerank
|
||||||
|
return SearchResponse(query=query, results=[], rank_label="Hybrid")
|
||||||
|
|
||||||
|
def fake_answer_query(_query, _results, _config):
|
||||||
|
nonlocal called
|
||||||
|
called = True
|
||||||
|
return "answer"
|
||||||
|
|
||||||
|
monkeypatch.setattr("python.ebook_search.api.routes.search.search_ebooks", fake_search_ebooks)
|
||||||
|
monkeypatch.setattr("python.ebook_search.api.routes.search.answer_query", fake_answer_query)
|
||||||
|
patch_app_runtime(monkeypatch)
|
||||||
|
app = create_app()
|
||||||
|
app.state.config = EbookSearchConfig(rerank=RerankConfig(enabled=False), answer_enabled=False)
|
||||||
|
|
||||||
|
with TestClient(app) as client:
|
||||||
|
response = client.post("/search", data={"query": "where is the quote?"})
|
||||||
|
|
||||||
|
assert response.status_code == 200
|
||||||
|
assert called is False
|
||||||
|
assert "Answer generation is disabled" in response.text
|
||||||
|
|
||||||
|
|
||||||
|
def test_ui_shows_component_scores(monkeypatch) -> None:
|
||||||
|
def fake_search_ebooks(_engine, query, _config, *, rerank=False):
|
||||||
|
del rerank
|
||||||
|
return SearchResponse(
|
||||||
|
query=query,
|
||||||
|
rank_label="Hybrid + rerank",
|
||||||
|
results=[
|
||||||
|
SearchResult(
|
||||||
|
chunk_id=1,
|
||||||
|
text="source text",
|
||||||
|
source_title="Book",
|
||||||
|
score=0.9,
|
||||||
|
rerank_score=0.9,
|
||||||
|
vector_score=0.8,
|
||||||
|
bm25_score=2.5,
|
||||||
|
fused_score=0.03,
|
||||||
|
)
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
|
monkeypatch.setattr("python.ebook_search.api.routes.search.search_ebooks", fake_search_ebooks)
|
||||||
|
monkeypatch.setattr(
|
||||||
|
"python.ebook_search.api.routes.search.answer_query",
|
||||||
|
lambda _query, _results, _config: "answer",
|
||||||
|
)
|
||||||
|
patch_app_runtime(monkeypatch)
|
||||||
|
app = create_app()
|
||||||
|
app.state.config = EbookSearchConfig(rerank=RerankConfig(enabled=False), answer_enabled=True)
|
||||||
|
|
||||||
|
with TestClient(app) as client:
|
||||||
|
response = client.post("/search", data={"query": "where is the quote?"})
|
||||||
|
|
||||||
|
assert response.status_code == 200
|
||||||
|
assert "rerank" in response.text
|
||||||
|
assert "vector cosine" in response.text
|
||||||
|
assert "BM25" in response.text
|
||||||
|
assert "RRF" in response.text
|
||||||
|
|
||||||
|
|
||||||
|
def test_ui_shows_search_runtime_chart(monkeypatch) -> None:
|
||||||
|
def fake_search_ebooks(_engine, query, _config, *, rerank=False):
|
||||||
|
del rerank
|
||||||
|
return SearchResponse(
|
||||||
|
query=query,
|
||||||
|
rank_label="Hybrid",
|
||||||
|
results=[],
|
||||||
|
timings=(
|
||||||
|
RuntimeStep(name="Embedding + vector search", duration_ms=12.5),
|
||||||
|
RuntimeStep(name="BM25 search", duration_ms=4.0),
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
monkeypatch.setattr("python.ebook_search.api.routes.search.search_ebooks", fake_search_ebooks)
|
||||||
|
monkeypatch.setattr(
|
||||||
|
"python.ebook_search.api.routes.search.answer_query",
|
||||||
|
lambda _query, _results, _config: "answer",
|
||||||
|
)
|
||||||
|
patch_app_runtime(monkeypatch)
|
||||||
|
app = create_app()
|
||||||
|
app.state.config = EbookSearchConfig(rerank=RerankConfig(enabled=False), answer_enabled=True)
|
||||||
|
|
||||||
|
with TestClient(app) as client:
|
||||||
|
response = client.post("/search", data={"query": "where is the quote?"})
|
||||||
|
|
||||||
|
assert response.status_code == 200
|
||||||
|
assert "Runtime" in response.text
|
||||||
|
assert "Total" in response.text
|
||||||
|
assert "Embedding + vector search" in response.text
|
||||||
|
assert "BM25 search" in response.text
|
||||||
|
assert "Answer generation" in response.text
|
||||||
|
assert "ms left" in response.text
|
||||||
|
|
||||||
|
|
||||||
|
def test_ui_embed_all_batches_until_complete(monkeypatch) -> None:
|
||||||
|
counts = iter([32, 32, 5, 0])
|
||||||
|
batch_sizes: list[int] = []
|
||||||
|
|
||||||
|
def fake_embed_missing_chunks(_session, config):
|
||||||
|
batch_sizes.append(config.embedding_batch_size)
|
||||||
|
return next(counts)
|
||||||
|
|
||||||
|
monkeypatch.setattr("python.ebook_search.api.routes.admin.embed_missing_chunks", fake_embed_missing_chunks)
|
||||||
|
patch_app_runtime(monkeypatch)
|
||||||
|
app = create_app()
|
||||||
|
|
||||||
|
with TestClient(app) as client:
|
||||||
|
response = client.post("/admin/embed-all")
|
||||||
|
|
||||||
|
assert response.status_code == 200
|
||||||
|
assert "Embedded 69 chunks in 3 batches of 32" in response.text
|
||||||
|
assert batch_sizes == [32, 32, 32, 32]
|
||||||
|
|
||||||
|
|
||||||
|
def test_ui_scan_schedules_bm25_refresh_after_database_change(monkeypatch) -> None:
|
||||||
|
scheduled = False
|
||||||
|
|
||||||
|
def fake_ingest_configured_paths(_session, _config):
|
||||||
|
return 1
|
||||||
|
|
||||||
|
def fake_schedule_bm25_refresh(_app):
|
||||||
|
nonlocal scheduled
|
||||||
|
scheduled = True
|
||||||
|
|
||||||
|
monkeypatch.setattr("python.ebook_search.api.routes.admin.ingest_configured_paths", fake_ingest_configured_paths)
|
||||||
|
monkeypatch.setattr("python.ebook_search.api.routes.admin.schedule_bm25_refresh", fake_schedule_bm25_refresh)
|
||||||
|
patch_app_runtime(monkeypatch)
|
||||||
|
app = create_app()
|
||||||
|
|
||||||
|
with TestClient(app) as client:
|
||||||
|
response = client.post("/admin/scan")
|
||||||
|
|
||||||
|
assert response.status_code == 200
|
||||||
|
assert "Indexed 1 EPUBs" in response.text
|
||||||
|
assert scheduled is True
|
||||||
|
|
||||||
|
|
||||||
|
def test_admin_page_shows_embedding_counts_by_model(monkeypatch) -> None:
|
||||||
|
def fake_embedding_model_stats(_session):
|
||||||
|
return [
|
||||||
|
EmbeddingModelStats(
|
||||||
|
model_name="qwen3-embedding-0.6b",
|
||||||
|
dimension=1024,
|
||||||
|
embedded_chunks=40,
|
||||||
|
total_chunks=64,
|
||||||
|
),
|
||||||
|
EmbeddingModelStats(
|
||||||
|
model_name="qwen3-embedding-4b",
|
||||||
|
dimension=2560,
|
||||||
|
embedded_chunks=8,
|
||||||
|
total_chunks=64,
|
||||||
|
),
|
||||||
|
]
|
||||||
|
|
||||||
|
monkeypatch.setattr("python.ebook_search.api.routes.admin.embedding_model_stats", fake_embedding_model_stats)
|
||||||
|
patch_app_runtime(monkeypatch)
|
||||||
|
app = create_app()
|
||||||
|
|
||||||
|
with TestClient(app) as client:
|
||||||
|
response = client.get("/admin")
|
||||||
|
|
||||||
|
assert response.status_code == 200
|
||||||
|
assert "qwen3-embedding-0.6b" in response.text
|
||||||
|
assert "1024" in response.text
|
||||||
|
assert "40" in response.text
|
||||||
|
assert "24" in response.text
|
||||||
|
assert "qwen3-embedding-4b" in response.text
|
||||||
|
assert "2560" in response.text
|
||||||
Reference in New Issue
Block a user