diff --git a/src/polygateway/ports.py b/src/polygateway/ports.py index c66180f..4ae42df 100644 --- a/src/polygateway/ports.py +++ b/src/polygateway/ports.py @@ -10,7 +10,14 @@ from dataclasses import dataclass from enum import StrEnum from typing import Any, Protocol, runtime_checkable -from .types import ChatRequest, LLMResponse, SourceConfig, SourceStats, TransportResult +from .types import ( + ChatRequest, + EmbeddingTransportResult, + LLMResponse, + SourceConfig, + SourceStats, + TransportResult, +) CallNext = Callable[[ChatRequest], Awaitable[LLMResponse]] @@ -37,6 +44,15 @@ class Transport(Protocol): ) -> TransportResult: ... +@runtime_checkable +class EmbeddingTransport(Protocol): + """一次原始 embedding 调用的协议细节(M2 §7);不含任何治理。""" + + async def embed( + self, *, texts: list[str], source: SourceConfig, call_id: str + ) -> EmbeddingTransportResult: ... + + @runtime_checkable class Permit(Protocol): """限流入场许可;settle/release 均幂等,finally 中必然执行。""" diff --git a/src/polygateway/transports/openai_compat.py b/src/polygateway/transports/openai_compat.py index 1dc56ac..e927bbe 100644 --- a/src/polygateway/transports/openai_compat.py +++ b/src/polygateway/transports/openai_compat.py @@ -15,10 +15,15 @@ from typing import TYPE_CHECKING, Any import httpx -from polygateway.errors import RequestRejectedError, SourceDeadError, TransientError +from polygateway.errors import ( + RequestRejectedError, + ResultInvalidError, + SourceDeadError, + TransientError, +) from polygateway.providers import ProviderProfile, get_provider from polygateway.streaming import StreamLivenessTimeout, stream_with_liveness_timeouts -from polygateway.types import SourceConfig, TransportResult +from polygateway.types import EmbeddingTransportResult, SourceConfig, TransportResult if TYPE_CHECKING: from collections.abc import AsyncIterator, Callable, Mapping @@ -141,6 +146,56 @@ def _resolve_usage(usage: dict[str, Any], source: SourceConfig) -> tuple[int, in return 0, source.est_tokens, "estimated" +def _extract_vectors( + data: dict[str, Any], source: SourceConfig, expected_count: int, ctx: dict[str, Any] +) -> list[list[float]]: + """按 data[].index 重排提取向量并校验条数/维度一致性。""" + try: + vectors = [ + [float(x) for x in item["embedding"]] + for item in sorted(data["data"], key=lambda it: int(it["index"])) + ] + except (KeyError, TypeError, ValueError) as exc: + raise ResultInvalidError(f"{source.name} embedding 响应形态异常: {exc}", **ctx) from exc + if len(vectors) != expected_count: + raise ResultInvalidError( + f"{source.name} 返回 {len(vectors)} 条向量,与输入 {expected_count} 条不符", **ctx + ) + if len({len(v) for v in vectors}) != 1 or not vectors[0]: + raise ResultInvalidError( + f"{source.name} 向量维度异常: {sorted({len(v) for v in vectors})}", **ctx + ) + return vectors + + +def _resolve_embedding_usage(data: dict[str, Any], source: SourceConfig) -> tuple[int, str]: + """usage 读取;缺失/非法按 est_tokens 保守兜底并标 estimated(与 chat 同口径)。""" + prompt = (data.get("usage") or {}).get("prompt_tokens") + if isinstance(prompt, int) and prompt > 0: + return prompt, "measured" + return source.est_tokens, "estimated" + + +def _parse_embedding_payload( + resp: httpx.Response, source: SourceConfig, expected_count: int +) -> EmbeddingTransportResult: + """解析 /embeddings 响应;一切形态异常归 ResultInvalidError(坏结果≠坏服务)。""" + ctx: dict[str, Any] = {"source_name": source.name, "operation": "embedding"} + try: + data = resp.json() + except json.JSONDecodeError as exc: + raise ResultInvalidError(f"{source.name} embedding 响应非 JSON: {exc}", **ctx) from exc + vectors = _extract_vectors(data, source, expected_count, ctx) + prompt_tokens, usage_source = _resolve_embedding_usage(data, source) + return EmbeddingTransportResult( + vectors=vectors, + dim=len(vectors[0]), + prompt_tokens=prompt_tokens, + usage_source=usage_source, + raw={"id": data.get("id")}, + ) + + def _default_client_factory(source: SourceConfig) -> httpx.AsyncClient: return httpx.AsyncClient( headers={"Authorization": f"Bearer {source.api_key}"}, @@ -217,6 +272,29 @@ class OpenAICompatTransport: # VT 宽集: 覆盖断连/协议错误/读写失败(设计 §9 行 8) raise TransientError(f"{source.name} 网络错误: {exc}", **ctx) from exc + async def embed( + self, *, texts: list[str], source: SourceConfig, call_id: str + ) -> EmbeddingTransportResult: + """一次原始 embedding 调用(M2 §7): POST /embeddings,错误翻译同 chat。 + + 响应按 data[].index 重排保序(GovDoc embedding.py:149 / VT :164 同款); + 空 data/长度不符/维度不一致 → ResultInvalidError(坏结果不熔断)。 + """ + if not texts: + raise ValueError("texts 不能为空(空输入由 EmbeddingClient 短路)") + url = source.base_url.rstrip("/") + "/embeddings" + client = self._client_for(source) + ctx: dict[str, Any] = {"source_name": source.name, "operation": "embedding"} + try: + resp = await client.post(url, json={"model": source.model, "input": texts}) + except httpx.TimeoutException as exc: + raise TransientError(f"{source.name} 超时: {exc}", **ctx) from exc + except httpx.TransportError as exc: + raise TransientError(f"{source.name} 网络错误: {exc}", **ctx) from exc + if resp.status_code != 200: + raise _status_to_error(source, resp.status_code, resp.text, resp.headers) + return _parse_embedding_payload(resp, source, len(texts)) + async def _complete_stream( self, client: httpx.AsyncClient, diff --git a/src/polygateway/types.py b/src/polygateway/types.py index c3f96ab..b70764f 100644 --- a/src/polygateway/types.py +++ b/src/polygateway/types.py @@ -186,3 +186,30 @@ class GlobalLimits: def __post_init__(self) -> None: if self.max_concurrency < 0 or self.rpm < 0 or self.tpm < 0: raise ValueError("全局限额不能为负(0 表示不启用)") + + +@dataclass(frozen=True) +class EmbeddingTransportResult: + """一次原始 embedding 调用的解析结果(M2 设计 §7.2;transport → client)。""" + + vectors: list[list[float]] + dim: int + prompt_tokens: int + usage_source: str # measured | estimated + raw: dict[str, Any] + + +@dataclass(frozen=True) +class EmbeddingResponse: + """一次治理 embedding 调用的统一响应(多批合并;与输入等长保序)。""" + + vectors: list[list[float]] + dim: int + model: str + provider: str + prompt_tokens: int + usage_source: str + latency_ms: int + call_id: str + source_name: str + cost: float | None = None diff --git a/tests/unit/test_embedding.py b/tests/unit/test_embedding.py new file mode 100644 index 0000000..c1f1bba --- /dev/null +++ b/tests/unit/test_embedding.py @@ -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")