feat: add embedding types, port and openai-compat transport

This commit is contained in:
2026-07-21 01:00:55 -04:00
parent d66299210b
commit 193d67da93
4 changed files with 281 additions and 3 deletions
+157
View File
@@ -0,0 +1,157 @@
"""Embedding 类型/端口/transport 测试(M2 设计 §7;T8)。
蓝本审计: GovDoc retrieval/embedding.py(分批/index 排序/维度校验)与
VT adapters/embedding.py(归一化);库裁决见设计 §7.3 表。
"""
import dataclasses
import json
import httpx
import pytest
from polygateway.errors import (
RequestRejectedError,
ResultInvalidError,
SourceDeadError,
TransientError,
)
from polygateway.ports import EmbeddingTransport
from polygateway.transports.openai_compat import OpenAICompatTransport
from polygateway.types import EmbeddingResponse, EmbeddingTransportResult, SourceConfig
def _src(**overrides):
base = {
"name": "e1",
"provider": "openai",
"base_url": "https://gw.example/v1",
"api_key": "sk",
"model": "embed-1",
"timeout_s": 10.0,
"est_tokens": 7,
}
base.update(overrides)
return SourceConfig(**base)
class TestTypes:
def test_embedding_response_frozen_with_defaults(self):
resp = EmbeddingResponse(
vectors=[[0.1, 0.2]],
dim=2,
model="m",
provider="p",
prompt_tokens=3,
usage_source="measured",
latency_ms=10,
call_id="c",
source_name="e1",
)
assert resp.cost is None
with pytest.raises(dataclasses.FrozenInstanceError):
resp.dim = 3
def test_transport_result_frozen(self):
r = EmbeddingTransportResult(
vectors=[[1.0]], dim=1, prompt_tokens=1, usage_source="measured", raw={}
)
with pytest.raises(dataclasses.FrozenInstanceError):
r.dim = 2
class _DummyEmbedTransport:
async def embed(self, *, texts, source, call_id):
raise NotImplementedError
def test_embedding_transport_protocol_runtime_checkable():
assert isinstance(_DummyEmbedTransport(), EmbeddingTransport)
assert isinstance(OpenAICompatTransport(), EmbeddingTransport)
def _transport_with(handler):
return OpenAICompatTransport(
client_factory=lambda source: httpx.AsyncClient(transport=httpx.MockTransport(handler))
)
def _ok_body(vectors, *, usage=None, shuffle=False):
data = [{"index": i, "embedding": v} for i, v in enumerate(vectors)]
if shuffle:
data = list(reversed(data))
body = {"data": data}
if usage is not None:
body["usage"] = usage
return body
class TestEmbedTransport:
async def test_sorts_by_index_and_measures_usage(self):
def handler(request):
assert request.url.path.endswith("/embeddings")
payload = json.loads(request.content)
assert payload == {"model": "embed-1", "input": ["a", "b"]}
return httpx.Response(
200,
json=_ok_body(
[[1.0, 0.0], [0.0, 1.0]], usage={"prompt_tokens": 5}, shuffle=True
),
)
result = await _transport_with(handler).embed(
texts=["a", "b"], source=_src(), call_id="c"
)
assert result.vectors == [[1.0, 0.0], [0.0, 1.0]] # 乱序响应按 index 重排
assert result.dim == 2
assert result.prompt_tokens == 5 and result.usage_source == "measured"
async def test_missing_usage_falls_back_estimated(self):
def handler(request):
return httpx.Response(200, json=_ok_body([[1.0]]))
result = await _transport_with(handler).embed(texts=["a"], source=_src(), call_id="c")
assert result.prompt_tokens == 7 and result.usage_source == "estimated" # est_tokens
@pytest.mark.parametrize(
("status", "exc_type"),
[(401, SourceDeadError), (400, RequestRejectedError), (500, TransientError)],
)
async def test_http_errors_translate(self, status, exc_type):
def handler(request):
return httpx.Response(status, text="boom")
with pytest.raises(exc_type):
await _transport_with(handler).embed(texts=["a"], source=_src(), call_id="c")
async def test_network_error_is_transient(self):
def handler(request):
raise httpx.ConnectError("refused")
with pytest.raises(TransientError):
await _transport_with(handler).embed(texts=["a"], source=_src(), call_id="c")
@pytest.mark.parametrize(
"body",
[
{"data": []}, # 空 data
{"data": [{"index": 0, "embedding": [1.0]}]}, # 数量与输入不符(输入 2 条)
{
"data": [
{"index": 0, "embedding": [1.0, 2.0]},
{"index": 1, "embedding": [1.0]}, # 维度不一致
]
},
{"nope": True}, # 缺 data
],
)
async def test_malformed_payload_is_result_invalid(self, body):
def handler(request):
return httpx.Response(200, json=body)
with pytest.raises(ResultInvalidError):
await _transport_with(handler).embed(texts=["a", "b"], source=_src(), call_id="c")
async def test_empty_texts_rejected(self):
with pytest.raises(ValueError):
await _transport_with(lambda r: None).embed(texts=[], source=_src(), call_id="c")