diff --git a/src/polygateway/embedding.py b/src/polygateway/embedding.py index d7e9b70..18ad334 100644 --- a/src/polygateway/embedding.py +++ b/src/polygateway/embedding.py @@ -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: diff --git a/tests/unit/test_embedding.py b/tests/unit/test_embedding.py index 8519583..1880f4d 100644 --- a/tests/unit/test_embedding.py +++ b/tests/unit/test_embedding.py @@ -412,6 +412,45 @@ class TestEmbedTelemetry: assert len(rec.rows[1]["messages"]) < 1000 # 长文本截断后入库 +class TestEmbedCallerDimensions: + """issue #11: 调用方自定义维度必须沿 embed 链四层透传到每一行遥测。""" + + async def test_single_batch_row_carries_dimensions(self): + rec = _MemoryRecorder() + client, _ = _embed_client([_src()], ["ok"], batch_size=2, telemetry=rec) + await client.embed(["a"], tenant_id="t1", meta={"batch": "b-42"}) + assert rec.rows[0]["tenant_id"] == "t1" + assert rec.rows[0]["meta"] == '{"batch": "b-42"}' + + async def test_every_batch_row_carries_the_same_dimensions(self): + """维度属于本次 `embed()` 调用,不随批次变化。 + + 只断言首行会漏掉"只有第一批带维度"的实现——那正是逐层透传最容易漏的形态。 + """ + rec = _MemoryRecorder() + client, _ = _embed_client([_src()], ["ok", "ok", "ok"], batch_size=1, telemetry=rec) + await client.embed(["a", "b", "c"], tenant_id="t1", meta={"batch": "b-42"}) + assert len(rec.rows) == 3 # 切成三批,每批一行 + assert [r["tenant_id"] for r in rec.rows] == ["t1", "t1", "t1"] + assert [r["meta"] for r in rec.rows] == ['{"batch": "b-42"}'] * 3 + + async def test_invalid_meta_rejected_before_any_telemetry(self): + """校验在切批之前: 遥测层的失败都被降级成 warning,放下游等于没有校验。""" + rec = _MemoryRecorder() + client, _ = _embed_client([_src()], ["ok"], batch_size=2, telemetry=rec) + with pytest.raises(ValueError, match="meta"): + await client.embed(["a"], meta={"Bad Key": 1}) + assert rec.rows == [] + assert client._transport.calls == [] # 连调用都没发出 + + async def test_defaults_land_as_sentinels(self): + rec = _MemoryRecorder() + client, _ = _embed_client([_src()], ["ok"], batch_size=2, telemetry=rec) + await client.embed(["a"]) + assert rec.rows[0]["tenant_id"] == "" # 空串哨兵,不是 None + assert rec.rows[0]["meta"] == "{}" + + @contextlib.contextmanager def _captured_warnings(): """捕获库发出的 WARNING;loguru 不经标准 logging,pytest 的 caplog 抓不到。"""