iomgaa/PolyGateway
1
0
Fork 0
Code Issues 3 Pull Requests Actions Packages Projects Releases 9 Wiki Activity

feat: carry caller dimensions down the embedding chain

Browse Source 
EmbeddingClient does not go through the chat onion: it builds its own
ChatRequest inside _emit purely to reuse the shared TelemetryEmitter, so
wiring chat() alone left every embed row without a tenant. Validate the
dimensions at the embed() entry (before batching, since anything failing
further down is degraded to a warning) and thread them through
_embed_batch -> _attempt -> _emit so every batch row carries the same
pair.
This commit is contained in:
iomgaa
2026-08-17 09:53:37 -04:00
parent 4be2b4f287
commit 702040d1a3
2 changed files with 118 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/polygateway/embedding.py
+79 -9
View File
@@ -21,7 +21,7 @@ import random
import time
import uuid
from dataclasses import dataclass
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Any
from loguru import logger
@@ -47,6 +47,7 @@ from polygateway.types import (
EmbeddingResponse,
LLMResponse,
strip_unsupported_extra_body,
validate_caller_dimensions,
)
if TYPE_CHECKING:
@@ -145,10 +146,22 @@ class EmbeddingClient:
*,
session_id: str | None = None,
parent_call_id: str | None = None,
tenant_id: str | None = None,
meta: Mapping[str, Any] | None = None,
) -> EmbeddingResponse:
"""一次治理 embedding 调用: 按 batch_size 切批,批间串行,全批合并返回。"""
"""一次治理 embedding 调用: 按 batch_size 切批,批间串行,全批合并返回。
`tenant_id` 与 `meta` 是调用方自定义维度,只进遥测(issue #11);它们属于
本次调用而非某一批,故每批的遥测行都带同一份维度。
"""
if not isinstance(texts, list) or any(not isinstance(t, str) for t in texts):
raise TypeError("texts 必须是 list[str](显式优于隐式,不收单条 str)")
# 必须在切批之前校验: 洋葱/链路内的一切失败都被遥测层降级成 warning
# (库铁律「遥测写失败降级不冒泡」),校验放下游等于没有校验——非法维度
# 会变成静默丢失的遥测行,而调用照常发出(issue #11 §4.2)
dimension_tenant_id, dimensions = validate_caller_dimensions(
tenant_id, meta, origin="embed(tenant_id=..., meta=...)"
)
if not texts:
return EmbeddingResponse(
vectors=[],
@@ -167,7 +180,11 @@ class EmbeddingClient:
for start in range(0, len(texts), self._batch_size):
outcomes.append(
await self._embed_batch(
texts[start : start + self._batch_size], session_id, parent_call_id
texts[start : start + self._batch_size],
session_id,
parent_call_id,
dimension_tenant_id,
dimensions,
)
)
return self._merge(outcomes)
@@ -175,7 +192,12 @@ class EmbeddingClient:
# —— 治理循环(与 RetryMW 同构;设计 §7.1 已声明的有限重复)——
async def _embed_batch(
self, batch: list[str], session_id: str | None, parent_call_id: str | None
self,
batch: list[str],
session_id: str | None,
parent_call_id: str | None,
tenant_id: str | None,
meta: dict[str, Any],
) -> _BatchOutcome:
fails = 0
reasons: dict[str, str] = {}
@@ -187,7 +209,9 @@ class EmbeddingClient:
await self._on_no_runnable(gate_rejections, reasons, clock)
continue
async with clock.attempting():
outcome = await self._attempt(batch, *picked, reasons, session_id, parent_call_id)
outcome = await self._attempt(
batch, *picked, reasons, session_id, parent_call_id, tenant_id, meta
)
if isinstance(outcome, _BatchOutcome):
return outcome
fails += 1
@@ -266,6 +290,8 @@ class EmbeddingClient:
reasons: dict[str, str],
session_id: str | None,
parent_call_id: str | None,
tenant_id: str | None,
meta: dict[str, Any],
) -> _BatchOutcome | _FailedBatch:
call_id = str(uuid.uuid4())
started = self._now()
@@ -286,17 +312,45 @@ class EmbeddingClient:
await self._record_quietly(self._breaker.record_success(entry))
await self._record_quietly(self._quota.mark_progress())
latency_ms = int((self._now() - started) * 1000)
await self._emit(batch, source, call_id, started, session_id, parent_call_id, result)
await self._emit(
batch,
source,
call_id,
started,
session_id,
parent_call_id,
tenant_id,
meta,
result,
)
return _BatchOutcome(result, source, call_id, latency_ms)
except (RequestRejectedError, ResultInvalidError) as exc:
await self._gate_on_terminal(exc, entry)
await self._emit(batch, source, call_id, started, session_id, parent_call_id, error=exc)
await self._emit(
batch,
source,
call_id,
started,
session_id,
parent_call_id,
tenant_id,
meta,
error=exc,
)
raise
except asyncio.CancelledError:
if entry.is_probe:
await self._record_quietly(self._breaker.release_probe(entry))
await self._emit(
batch, source, call_id, started, session_id, parent_call_id, error="cancelled"
batch,
source,
call_id,
started,
session_id,
parent_call_id,
tenant_id,
meta,
error="cancelled",
)
raise
except (SourceDeadError, TransientError) as exc:
@@ -307,7 +361,17 @@ class EmbeddingClient:
if not dead:
# 保守: 失败请求可能已被网关计费(CHS 同款);与入场预扣同源取值
actual = source.effective_est_tokens()
await self._emit(batch, source, call_id, started, session_id, parent_call_id, error=exc)
await self._emit(
batch,
source,
call_id,
started,
session_id,
parent_call_id,
tenant_id,
meta,
error=exc,
)
return _FailedBatch(exc, immediate=dead)
finally:
await self._settle_and_release(permit, actual)
@@ -351,16 +415,22 @@ class EmbeddingClient:
started: float,
session_id: str | None,
parent_call_id: str | None,
tenant_id: str | None,
meta: dict[str, Any],
result: EmbeddingTransportResult | None = None,
error: object | None = None,
) -> None:
"""逐批遥测(经同一 Emitter): messages=截断 texts、向量绝不入库。"""
if self._emitter is None:
return
# 这个 ChatRequest 只为复用同一个 Emitter 而现场构造(embedding 不走 chat
# 洋葱),故调用方维度必须在这里显式填回,否则 embed 行的维度恒为空
request = ChatRequest(
messages=[{"role": "user", "content": t[:_TELEMETRY_TEXT_CAP]} for t in batch],
session_id=session_id,
parent_call_id=parent_call_id,
tenant_id=tenant_id,
meta=meta,
)
response = None
if result is not None: