chatbot_v3/tests/unit/infrastructure/embedding/test_openai_compatible.py

"""`OpenAICompatibleEmbedder` against a mocked `/embeddings` endpoint.

Backs both `dense_nomic` and `dense_openai` (ADR-0001) -- `httpx.MockTransport`
stands in for the real self-hosted/OpenAI server so this stays a unit test
with no network dependency.
"""

import json
from collections.abc import Callable

import httpx
import pytest

from src.infrastructure.embedding.openai_compatible import (
    OpenAICompatibleEmbedder,
    is_ollama_base_url,
)

pytestmark = pytest.mark.unit


def _client(handler: Callable[[httpx.Request], httpx.Response]) -> httpx.AsyncClient:
    return httpx.AsyncClient(transport=httpx.MockTransport(handler), base_url="http://embedder")


@pytest.mark.asyncio
async def test_embed_batch_returns_vectors_in_input_order() -> None:
    def handler(request: httpx.Request) -> httpx.Response:
        # Respond out of order to prove the adapter re-sorts by `index`.
        return httpx.Response(
            200,
            json={
                "data": [
                    {"index": 1, "embedding": [0.2, 0.2]},
                    {"index": 0, "embedding": [0.1, 0.1]},
                ],
                "model": "test-model",
            },
        )

    embedder = OpenAICompatibleEmbedder(_client(handler), name="dense_nomic", model="test-model")
    vectors = await embedder.embed_batch(["first", "second"])

    assert vectors == [[0.1, 0.1], [0.2, 0.2]]


@pytest.mark.asyncio
async def test_embed_batch_sends_model_and_input() -> None:
    captured: dict[str, object] = {}

    def handler(request: httpx.Request) -> httpx.Response:
        captured.update(json.loads(request.read()))
        return httpx.Response(
            200, json={"data": [{"index": 0, "embedding": [0.0]}], "model": "test-model"}
        )

    embedder = OpenAICompatibleEmbedder(_client(handler), name="dense_openai", model="test-model")
    await embedder.embed_batch(["only text"])

    assert captured["model"] == "test-model"
    assert captured["input"] == ["only text"]
    assert "dimensions" not in captured


@pytest.mark.asyncio
async def test_embed_batch_sends_dimensions_when_configured() -> None:
    captured: dict[str, object] = {}

    def handler(request: httpx.Request) -> httpx.Response:
        captured.update(json.loads(request.read()))
        return httpx.Response(
            200, json={"data": [{"index": 0, "embedding": [0.0] * 256}], "model": "test-model"}
        )

    embedder = OpenAICompatibleEmbedder(
        _client(handler), name="dense_openai", model="test-model", dimensions=256
    )
    await embedder.embed_batch(["only text"])

    assert captured["dimensions"] == 256


@pytest.mark.asyncio
async def test_embed_batch_raises_on_non_2xx_response() -> None:
    def handler(request: httpx.Request) -> httpx.Response:
        return httpx.Response(500, json={"error": "boom"})

    embedder = OpenAICompatibleEmbedder(_client(handler), name="dense_nomic", model="test-model")

    with pytest.raises(httpx.HTTPStatusError):
        await embedder.embed_batch(["text"])


@pytest.mark.asyncio
async def test_embed_batch_sends_no_prefix_by_default() -> None:
    """emet parity: the benchmarked run used no task prefix (ADR-0004)."""
    captured: dict[str, object] = {}

    def handler(request: httpx.Request) -> httpx.Response:
        captured.update(json.loads(request.read()))
        return httpx.Response(200, json={"data": [{"index": 0, "embedding": [0.0]}], "model": "m"})

    embedder = OpenAICompatibleEmbedder(_client(handler), name="dense_nomic", model="m")
    await embedder.embed_batch(["سلام"])

    assert captured["input"] == ["سلام"]


@pytest.mark.asyncio
async def test_embed_batch_applies_document_prefix_when_configured() -> None:
    captured: dict[str, object] = {}

    def handler(request: httpx.Request) -> httpx.Response:
        captured.update(json.loads(request.read()))
        return httpx.Response(200, json={"data": [{"index": 0, "embedding": [0.0]}], "model": "m"})

    embedder = OpenAICompatibleEmbedder(
        _client(handler), name="dense_nomic", model="m", document_prefix="search_document: "
    )
    await embedder.embed_batch(["سلام"])

    assert captured["input"] == ["search_document: سلام"]


@pytest.mark.asyncio
async def test_embed_batch_sends_keep_alive_when_configured() -> None:
    """Keeps an Ollama-hosted model resident; a cold load outruns the
    ingestion timeout entirely.
    """
    captured: dict[str, object] = {}

    def handler(request: httpx.Request) -> httpx.Response:
        captured.update(json.loads(request.read()))
        return httpx.Response(200, json={"data": [{"index": 0, "embedding": [0.0]}], "model": "m"})

    embedder = OpenAICompatibleEmbedder(
        _client(handler), name="dense_nomic", model="m", keep_alive="30m"
    )
    await embedder.embed_batch(["text"])

    assert captured["keep_alive"] == "30m"


@pytest.mark.asyncio
async def test_embed_batch_omits_keep_alive_when_not_configured() -> None:
    """OpenAI would reject an unknown field, so it must not be sent there."""
    captured: dict[str, object] = {}

    def handler(request: httpx.Request) -> httpx.Response:
        captured.update(json.loads(request.read()))
        return httpx.Response(200, json={"data": [{"index": 0, "embedding": [0.0]}], "model": "m"})

    embedder = OpenAICompatibleEmbedder(_client(handler), name="dense_openai", model="m")
    await embedder.embed_batch(["text"])

    assert "keep_alive" not in captured


@pytest.mark.parametrize(
    ("base_url", "expected"),
    [
        ("http://192.168.10.10:11435/v1", True),
        ("http://127.0.0.1:11434/v1", True),
        ("http://ollama.internal/v1", True),
        ("https://api.openai.com/v1", False),
        ("http://127.0.0.1:8081/v1", False),
    ],
)
def test_is_ollama_base_url(base_url: str, expected: bool) -> None:
    assert is_ollama_base_url(base_url) is expected