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")