145 lines
5.2 KiB
Python
145 lines
5.2 KiB
Python
"""Tests for EPUB search HTTP model adapters."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
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.llm_interface import check_chat_endpoint, check_embedding_endpoint
|
|
from python.ebook_search.search import SearchResult
|
|
|
|
if TYPE_CHECKING:
|
|
from pytest_mock import MockerFixture
|
|
|
|
|
|
def make_async_client(mocker: MockerFixture, fake_post) -> httpx.AsyncClient:
|
|
"""Build a mock async client whose post call is served by fake_post."""
|
|
client = mocker.MagicMock(spec=httpx.AsyncClient)
|
|
client.post = mocker.AsyncMock(side_effect=fake_post)
|
|
return client
|
|
|
|
|
|
async def test_model_endpoint_checks_share_http_probe(mocker: MockerFixture) -> None:
|
|
client = mocker.MagicMock(spec=httpx.AsyncClient)
|
|
response = mocker.MagicMock(spec=httpx.Response)
|
|
client.get = mocker.AsyncMock(return_value=response)
|
|
config = EbookSearchConfig(
|
|
rerank=RerankConfig(enabled=False),
|
|
embedding_base_url="https://embedding.example/v1/",
|
|
vllm_base_url="https://chat.example/v1/",
|
|
vllm_api_key="secret",
|
|
)
|
|
|
|
assert await check_embedding_endpoint(client, config, timeout_seconds=2.0)
|
|
assert await check_chat_endpoint(client, config, timeout_seconds=3.0)
|
|
assert client.get.await_args_list == [
|
|
mocker.call("https://embedding.example/v1/models", headers={}, timeout=2.0),
|
|
mocker.call(
|
|
"https://chat.example/v1/models",
|
|
headers={"Authorization": "Bearer secret"},
|
|
timeout=3.0,
|
|
),
|
|
]
|
|
assert response.raise_for_status.call_count == 2
|
|
|
|
|
|
async def test_model_endpoint_checks_report_http_failures(mocker: MockerFixture) -> None:
|
|
client = mocker.MagicMock(spec=httpx.AsyncClient)
|
|
client.get = mocker.AsyncMock(
|
|
side_effect=[
|
|
httpx.ConnectError("embedding offline"),
|
|
httpx.ConnectError("chat offline"),
|
|
]
|
|
)
|
|
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
|
|
|
|
assert not await check_embedding_endpoint(client, config)
|
|
assert not await check_chat_endpoint(client, config)
|
|
|
|
|
|
async def test_answer_query_uses_httpx_chat_completions(mocker: MockerFixture) -> 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),
|
|
)
|
|
|
|
client = make_async_client(mocker, fake_post)
|
|
config = EbookSearchConfig(
|
|
rerank=RerankConfig(enabled=False),
|
|
vllm_base_url="https://ollama.com/v1",
|
|
vllm_api_key="secret",
|
|
chat_model="deepseek-v4-flash",
|
|
)
|
|
|
|
results = [SearchResult(chunk_id=1, text="source", source_title="Book")]
|
|
answer = await answer_query(client, "question", results, 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"
|
|
assert payload["messages"] == [
|
|
{
|
|
"role": "system",
|
|
"content": "Answer only from the provided context. Cite sources with bracketed numbers like [1]. "
|
|
"If the context is insufficient, say so.",
|
|
},
|
|
{"role": "user", "content": "Question:\nquestion\n\nContext:\n[1] Book\nsource"},
|
|
]
|
|
|
|
|
|
async def test_embed_texts_uses_httpx_embeddings(mocker: MockerFixture) -> 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),
|
|
)
|
|
|
|
client = make_async_client(mocker, fake_post)
|
|
config = EbookSearchConfig(
|
|
rerank=RerankConfig(enabled=False),
|
|
embedding_base_url="http://bob:8000/v1",
|
|
embedding_model="qwen3-embedding-0.6b",
|
|
)
|
|
|
|
embeddings = await embed_texts(client, ["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"]}
|
|
|
|
|
|
async def test_embed_texts_rejects_bad_response_shape(mocker: MockerFixture) -> None:
|
|
def fake_post(url: str, **_kwargs: object) -> httpx.Response:
|
|
return httpx.Response(200, json={"data": [{}]}, request=httpx.Request("POST", url))
|
|
|
|
client = make_async_client(mocker, fake_post)
|
|
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
|
|
|
|
with pytest.raises(RuntimeError, match="Embedding request failed"):
|
|
await embed_texts(client, ["hello"], config)
|