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

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 中必然执行。"""

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,

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

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")