File size: 1,850 Bytes
2edb151 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 | 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")
|