369 lines
13 KiB
Python
369 lines
13 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_from_settings_assembles_client(self):
|
|
s = EmbeddingSettings.from_env("EMBED", env=self._ENV)
|
|
client = EmbeddingClient.from_settings(s)
|
|
assert isinstance(client, EmbeddingClient)
|