iomgaa
afd6101c08
Mutation testing showed the negative structured-retries and expected-dim checks in the env parsing path could be deleted with every test still passing. Their value is the env key name in the message, so they need tests that assert it. Changelog now states the real scope of this release and warns that normalising scope moves the Redis keys, the one change here that silently relocates runtime state. Records the breaker threshold derivation as deliberately env-only so it does not resurface as another round.
370 lines
14 KiB
Python
370 lines
14 KiB
Python
"""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")
|
|
|
|
|
|
# ═══════════ T9: EmbeddingClient 治理循环 ═══════════
|
|
|
|
import asyncio # noqa: E402
|
|
|
|
from polygateway.backends.memory.breaker import InMemoryGate # noqa: E402
|
|
from polygateway.backends.memory.limiter import InMemoryLimiter # noqa: E402
|
|
from polygateway.config import EmbeddingSettings # noqa: E402
|
|
from polygateway.embedding import EmbeddingClient # noqa: E402
|
|
from polygateway.sources import RoundRobinSelector # noqa: E402
|
|
from polygateway.types import ( # noqa: E402
|
|
BackpressurePolicy,
|
|
BreakerConfig,
|
|
GlobalLimits,
|
|
RetryPolicy,
|
|
)
|
|
|
|
_BREAKER = BreakerConfig(fail_threshold=3, cooldown_s=60.0, probe_ttl_s=120.0)
|
|
_NO_GLOBAL = GlobalLimits(max_concurrency=0, rpm=0, tpm=0)
|
|
|
|
|
|
def _vec_for(texts):
|
|
"""确定性向量: 每条 text 一个 [len(text)] 一维向量,便于断言保序。"""
|
|
return EmbeddingTransportResult(
|
|
vectors=[[float(len(t))] for t in texts],
|
|
dim=1,
|
|
prompt_tokens=len(texts),
|
|
usage_source="measured",
|
|
raw={},
|
|
)
|
|
|
|
|
|
class ScriptedEmbedTransport:
|
|
"""按脚本响应: 条目为 Exception / "ok"(按输入生成) / EmbeddingTransportResult / "hang"。"""
|
|
|
|
def __init__(self, script):
|
|
self.script = list(script)
|
|
self.calls = []
|
|
|
|
async def embed(self, *, texts, source, call_id):
|
|
self.calls.append((source.name, list(texts), call_id))
|
|
action = self.script.pop(0)
|
|
if isinstance(action, Exception):
|
|
raise action
|
|
if action == "hang":
|
|
await asyncio.Event().wait()
|
|
if action == "ok":
|
|
return _vec_for(texts)
|
|
return action
|
|
|
|
|
|
class _MemoryRecorder:
|
|
def __init__(self):
|
|
self.rows = []
|
|
|
|
async def record_llm_call(self, **fields):
|
|
self.rows.append(fields)
|
|
|
|
|
|
def _embed_client(sources, script, *, batch_size=2, telemetry=None, **overrides):
|
|
limiter = InMemoryLimiter(
|
|
scope="embed",
|
|
sources={s.name: s for s in sources},
|
|
global_limits=_NO_GLOBAL,
|
|
lease_ttl_s=100.0,
|
|
)
|
|
kwargs = {
|
|
"scope": "embed",
|
|
"sources": sources,
|
|
"selector": RoundRobinSelector(),
|
|
"limiter": limiter,
|
|
"breaker": InMemoryGate(config=_BREAKER),
|
|
"transport": ScriptedEmbedTransport(script),
|
|
"retry": RetryPolicy(max_attempts=3, backoff_base_s=0.001, backoff_max_s=0.01),
|
|
"backpressure": BackpressurePolicy(stall_window_s=300.0, poll_interval_s=0.001),
|
|
"batch_size": batch_size,
|
|
"telemetry": telemetry,
|
|
}
|
|
kwargs.update(overrides)
|
|
client = EmbeddingClient(**kwargs)
|
|
return client, limiter
|
|
|
|
|
|
class TestEmbedBatching:
|
|
async def test_batches_sequential_and_order_preserved(self):
|
|
texts = ["a", "bb", "ccc", "dddd", "eeeee"]
|
|
client, _ = _embed_client([_src()], ["ok", "ok", "ok"], batch_size=2)
|
|
resp = await client.embed(texts)
|
|
transport = client._transport
|
|
assert [len(batch) for _, batch, _ in transport.calls] == [2, 2, 1]
|
|
assert resp.vectors == [[1.0], [2.0], [3.0], [4.0], [5.0]] # 全批拼接保序
|
|
assert resp.prompt_tokens == 5 and resp.dim == 1
|
|
|
|
async def test_empty_input_short_circuits(self):
|
|
client, _ = _embed_client([_src()], [])
|
|
resp = await client.embed([])
|
|
assert resp.vectors == [] and resp.prompt_tokens == 0
|
|
assert client._transport.calls == []
|
|
|
|
async def test_usage_source_aggregates_conservatively(self):
|
|
estimated = EmbeddingTransportResult(
|
|
vectors=[[1.0], [1.0]], dim=1, prompt_tokens=9, usage_source="estimated", raw={}
|
|
)
|
|
client, _ = _embed_client([_src()], ["ok", estimated], batch_size=2)
|
|
resp = await client.embed(["a", "b", "c", "d"])
|
|
assert resp.usage_source == "estimated" # 任一批 estimated 则整体 estimated
|
|
assert resp.prompt_tokens == 2 + 9
|
|
|
|
|
|
class TestEmbedPostProcess:
|
|
async def test_normalize_l2(self):
|
|
raw = EmbeddingTransportResult(
|
|
vectors=[[3.0, 4.0]], dim=2, prompt_tokens=1, usage_source="measured", raw={}
|
|
)
|
|
client, _ = _embed_client([_src()], [raw], normalize=True)
|
|
resp = await client.embed(["x"])
|
|
assert resp.vectors[0] == pytest.approx([0.6, 0.8])
|
|
|
|
async def test_zero_vector_normalize_no_nan(self):
|
|
raw = EmbeddingTransportResult(
|
|
vectors=[[0.0, 0.0]], dim=2, prompt_tokens=1, usage_source="measured", raw={}
|
|
)
|
|
client, _ = _embed_client([_src()], [raw], normalize=True)
|
|
resp = await client.embed(["x"])
|
|
assert resp.vectors[0] == [0.0, 0.0] # max(norm, 1e-12) 防除零(VT 语义)
|
|
|
|
async def test_expected_dim_violation_is_result_invalid(self):
|
|
client, _ = _embed_client([_src()], ["ok"], expected_dim=768)
|
|
with pytest.raises(ResultInvalidError):
|
|
await client.embed(["x"])
|
|
|
|
|
|
class TestEmbedGovernance:
|
|
async def test_transient_retries_then_succeeds(self):
|
|
client, _ = _embed_client(
|
|
[_src()], [TransientError("boom", status_code=500), "ok"], batch_size=8
|
|
)
|
|
resp = await client.embed(["a", "b"])
|
|
assert resp.vectors == [[1.0], [1.0]]
|
|
assert len(client._transport.calls) == 2
|
|
|
|
async def test_source_dead_switches_source(self):
|
|
s1, s2 = _src(name="e1"), _src(name="e2")
|
|
client, _ = _embed_client(
|
|
[s1, s2], [SourceDeadError("401", status_code=401), "ok"], batch_size=8
|
|
)
|
|
await client.embed(["a"])
|
|
assert [name for name, _, _ in client._transport.calls] == ["e1", "e2"]
|
|
|
|
async def test_cancel_releases_permit(self):
|
|
client, limiter = _embed_client([_src(max_concurrency=1)], ["hang"])
|
|
task = asyncio.create_task(client.embed(["a"]))
|
|
while not (await limiter.source_stats("e1")).inflight:
|
|
await asyncio.sleep(0.01)
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
assert (await limiter.source_stats("e1")).inflight == 0
|
|
|
|
|
|
class TestEmbedTelemetry:
|
|
async def test_per_batch_rows_with_digest(self):
|
|
rec = _MemoryRecorder()
|
|
client, _ = _embed_client([_src()], ["ok", "ok"], batch_size=1, telemetry=rec)
|
|
await client.embed(["hello", "x" * 5000], session_id="sess", parent_call_id="pc")
|
|
assert len(rec.rows) == 2 # 每批一行
|
|
row = rec.rows[0]
|
|
assert row["session_id"] == "sess" and row["parent_call_id"] == "pc"
|
|
assert row["completion_tokens"] == 0
|
|
assert row["response"] == "<vectors n=1 dim=1>" # 向量绝不入库
|
|
assert len(rec.rows[1]["messages"]) < 1000 # 长文本截断后入库
|
|
|
|
|
|
class TestEmbeddingSettings:
|
|
_ENV = {
|
|
"EMBED__QWEN__1__BASE_URL": "https://gw.example/v1",
|
|
"EMBED__QWEN__1__API_KEY": "sk-a",
|
|
"EMBED__QWEN__1__MODEL": "text-embedding-v3",
|
|
"EMBED__QWEN__1__TIMEOUT_S": "60",
|
|
"EMBED__RETRY__MAX_ATTEMPTS": "3",
|
|
"EMBED__RETRY__BACKOFF_BASE_S": "1.0",
|
|
"EMBED__RETRY__BACKOFF_MAX_S": "10.0",
|
|
"EMBED__BREAKER__FAIL_THRESHOLD": "5",
|
|
"EMBED__BREAKER__COOLDOWN_S": "60",
|
|
"PGW_CACHE_BACKEND": "none",
|
|
"PGW_TELEMETRY_BACKEND": "none",
|
|
"EMBED__BATCH_SIZE": "64",
|
|
}
|
|
|
|
def test_loads_scope_and_batch(self):
|
|
s = EmbeddingSettings.from_env("EMBED", env=self._ENV)
|
|
assert s.gateway.sources[0].model == "text-embedding-v3"
|
|
assert s.batch_size == 64 and s.normalize is False and s.expected_dim is None
|
|
|
|
def test_batch_size_required_and_positive(self):
|
|
env = {k: v for k, v in self._ENV.items() if k != "EMBED__BATCH_SIZE"}
|
|
with pytest.raises(ValueError, match="BATCH_SIZE"):
|
|
EmbeddingSettings.from_env("EMBED", env=env)
|
|
with pytest.raises(ValueError, match="BATCH_SIZE"):
|
|
EmbeddingSettings.from_env("EMBED", env={**self._ENV, "EMBED__BATCH_SIZE": "0"})
|
|
|
|
def test_optional_normalize_and_dim(self):
|
|
env = {**self._ENV, "EMBED__NORMALIZE": "true", "EMBED__EXPECTED_DIM": "768"}
|
|
s = EmbeddingSettings.from_env("EMBED", env=env)
|
|
assert s.normalize is True and s.expected_dim == 768
|
|
|
|
def test_expected_dim_must_be_positive(self):
|
|
"""env 层的检查保留是为了报错能点出键名(构造期那道点的是字段名)。"""
|
|
with pytest.raises(ValueError, match="EXPECTED_DIM"):
|
|
EmbeddingSettings.from_env("EMBED", env={**self._ENV, "EMBED__EXPECTED_DIM": "0"})
|
|
|
|
def test_from_settings_assembles_client(self):
|
|
s = EmbeddingSettings.from_env("EMBED", env=self._ENV)
|
|
client = EmbeddingClient.from_settings(s)
|
|
assert isinstance(client, EmbeddingClient)
|