feat: add governed embedding client with batching
This commit is contained in:
@@ -0,0 +1,83 @@
|
||||
"""真实网关 /embeddings 端点探测(M2 设计 §11.6;人类默认口径: 实现时探测)。
|
||||
|
||||
对 .env 的 LLM 源网关发一次真实 embeddings 请求: 支持则记录向量证据,
|
||||
不支持(404/翻译为领域错误)则 skip 并把响应记录进 tests/outputs/
|
||||
(降级证据)。无 EMBED scope 配置时复用 LLM 源的 base_url/api_key。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
import os
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from dotenv import dotenv_values
|
||||
|
||||
from polygateway.errors import PolyGatewayError
|
||||
from polygateway.transports.openai_compat import OpenAICompatTransport
|
||||
from polygateway.types import SourceConfig
|
||||
|
||||
_ENV = {k: v for k, v in {**dotenv_values(".env"), **os.environ}.items() if v is not None}
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
"LLM__MINIMAX__1__BASE_URL" not in _ENV, reason="缺真实网关配置(.env)"
|
||||
)
|
||||
|
||||
_OUT = Path("tests/outputs/embedding")
|
||||
|
||||
|
||||
def _record(name: str, lines: list[str]) -> Path:
|
||||
_OUT.mkdir(parents=True, exist_ok=True)
|
||||
path = _OUT / f"{name}_{datetime.now():%Y%m%d_%H%M%S}.md"
|
||||
path.write_text("\n".join(lines) + "\n", encoding="utf-8")
|
||||
return path
|
||||
|
||||
|
||||
async def test_probe_real_gateway_embeddings():
|
||||
source = SourceConfig(
|
||||
name="probe_1",
|
||||
provider="minimax",
|
||||
base_url=_ENV["LLM__MINIMAX__1__BASE_URL"],
|
||||
api_key=_ENV["LLM__MINIMAX__1__API_KEY"],
|
||||
model=_ENV.get("PGW_EMBED_PROBE_MODEL", "text-embedding-v1"),
|
||||
timeout_s=30.0,
|
||||
est_tokens=8,
|
||||
)
|
||||
transport = OpenAICompatTransport()
|
||||
try:
|
||||
result = await transport.embed(
|
||||
texts=["polygateway embedding probe"], source=source, call_id="probe"
|
||||
)
|
||||
except PolyGatewayError as exc:
|
||||
path = _record(
|
||||
"probe_unsupported",
|
||||
[
|
||||
"# Embedding 端点探测: 网关不支持",
|
||||
f"- base_url: {source.base_url}",
|
||||
f"- model: {source.model}",
|
||||
f"- 错误分类: {type(exc).__name__}",
|
||||
f"- status_code: {exc.status_code}",
|
||||
f"- 详情: {exc}",
|
||||
"",
|
||||
"结论: e2e 按设计 §11.6 降级,embedding 行为由 unit 全覆盖。",
|
||||
],
|
||||
)
|
||||
await transport.aclose()
|
||||
pytest.skip(f"网关不支持 embeddings({type(exc).__name__}),证据: {path}")
|
||||
else:
|
||||
await transport.aclose()
|
||||
assert result.dim > 0 and len(result.vectors) == 1
|
||||
_record(
|
||||
"probe_supported",
|
||||
[
|
||||
"# Embedding 端点探测: 网关支持",
|
||||
f"- base_url: {source.base_url}",
|
||||
f"- model: {source.model}",
|
||||
f"- dim: {result.dim}",
|
||||
f"- usage: {result.prompt_tokens}({result.usage_source})",
|
||||
f"- 向量前 5 维: {result.vectors[0][:5]}",
|
||||
f"- raw: {dataclasses.asdict(result)['raw']}",
|
||||
],
|
||||
)
|
||||
@@ -155,3 +155,214 @@ class TestEmbedTransport:
|
||||
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)
|
||||
|
||||
Reference in New Issue
Block a user