tests/test_embed_prefix.py

from __future__ import annotations

import httpx

from app.config import Settings
from backends.gemma import GemmaEmbed
from backends.openai_compat import apply_embed_prefix, EmbedDimensionError
import pytest


def test_prefix_query_and_passage() -> None:
    assert apply_embed_prefix("milk", "query", enabled=True) == "query: milk"
    assert apply_embed_prefix("milk 2%", "passage", enabled=True) == "passage: milk 2%"
    assert apply_embed_prefix("query: already", "query", enabled=True) == "query: already"


def test_embed_request_body_has_prefix_and_input_type(settings: Settings) -> None:
    recorded: list[tuple[str, dict]] = []

    def handler(request: httpx.Request) -> httpx.Response:
        recorded.append((request.url.path, request.read().decode()))
        return httpx.Response(
            200,
            json={
                "data": [
                    {"index": 0, "embedding": [0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5]}
                ]
            },
        )

    client = httpx.Client(transport=httpx.MockTransport(handler), base_url="http://llm.test/v1")
    embed = GemmaEmbed(settings, client=client)
    vecs = embed.embed(["milk"], input_type="query")
    assert len(vecs[0]) == 8
    path, body = recorded[0]
    assert path.endswith("/embeddings")
    assert "query: milk" in body
    assert '"input_type": "query"' in body or '"input_type":"query"' in body


def test_wrong_dim_rejected(settings: Settings) -> None:
    def handler(_request: httpx.Request) -> httpx.Response:
        return httpx.Response(200, json={"data": [{"index": 0, "embedding": [1.0, 0.0]}]})

    client = httpx.Client(transport=httpx.MockTransport(handler), base_url="http://llm.test/v1")
    embed = GemmaEmbed(settings, client=client)
    with pytest.raises(EmbedDimensionError):
        embed.embed(["x"], input_type="passage")