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

This commit is contained in:
2026-07-21 01:00:55 -04:00
parent d66299210b
commit 193d67da93
4 changed files with 281 additions and 3 deletions
+17 -1
View File
@@ -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 中必然执行。"""
+80 -2
View File
@@ -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,
+27
View File
@@ -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