From abb65c232435987328936ff3e64e8ffb6076ce9d Mon Sep 17 00:00:00 2001 From: iomgaa Date: Tue, 21 Jul 2026 00:50:33 -0400 Subject: [PATCH] feat: add postgres telemetry recorder with two-tier degradation --- src/polygateway/client.py | 54 +++++-- src/polygateway/telemetry/postgres.py | 136 ++++++++++++++++ tests/integration/test_postgres_telemetry.py | 159 +++++++++++++++++++ tests/unit/test_client.py | 16 ++ tests/unit/test_telemetry.py | 6 +- 5 files changed, 359 insertions(+), 12 deletions(-) create mode 100644 src/polygateway/telemetry/postgres.py create mode 100644 tests/integration/test_postgres_telemetry.py diff --git a/src/polygateway/client.py b/src/polygateway/client.py index 7ac5665..6e8a10f 100644 --- a/src/polygateway/client.py +++ b/src/polygateway/client.py @@ -162,9 +162,13 @@ class GatewayClient: transport_aclose = getattr(self._transport, "aclose", None) if transport_aclose is not None: await transport_aclose() - telemetry_close = getattr(self._telemetry, "close", None) - if telemetry_close is not None: - telemetry_close() + telemetry_aclose = getattr(self._telemetry, "aclose", None) + if telemetry_aclose is not None: + await telemetry_aclose() # Postgres 等异步后端 + else: + telemetry_close = getattr(self._telemetry, "close", None) + if telemetry_close is not None: + telemetry_close() cache_aclose = getattr(self._cache, "aclose", None) if cache_aclose is not None: await cache_aclose() @@ -196,14 +200,8 @@ class GatewayClient: scope=settings.scope, sources=sources, selector=_build_selector(settings.selector), - limiter=limiter - or InMemoryLimiter( - scope=settings.scope, - sources={s.name: s for s in sources}, - global_limits=settings.global_limits, - lease_ttl_s=settings.lease_ttl_s, - ), - breaker=breaker or InMemoryGate(config=settings.breaker), + limiter=limiter or _build_limiter(settings, sources), + breaker=breaker or _build_breaker(settings), transport=OpenAICompatTransport(registry=registry), retry=settings.retry, backpressure=settings.backpressure, @@ -240,6 +238,35 @@ class GatewayClient: ) +def _build_limiter(settings: GatewaySettings, sources: list[SourceConfig]) -> RateLimiter: + if settings.limiter_backend == "redis": + from polygateway.backends.redis.limiter import RedisLimiter + + assert settings.redis_url is not None # 内部不变量: config 已校验 + return RedisLimiter.from_url( + settings.redis_url, + scope=settings.scope, + sources={s.name: s for s in sources}, + global_limits=settings.global_limits, + lease_ttl_s=settings.lease_ttl_s, + ) + return InMemoryLimiter( + scope=settings.scope, + sources={s.name: s for s in sources}, + global_limits=settings.global_limits, + lease_ttl_s=settings.lease_ttl_s, + ) + + +def _build_breaker(settings: GatewaySettings) -> ProviderGate: + if settings.breaker_backend == "redis": + from polygateway.backends.redis.breaker import RedisGate + + assert settings.redis_url is not None # 内部不变量: config 已校验 + return RedisGate.from_url(settings.redis_url, config=settings.breaker, scope=settings.scope) + return InMemoryGate(config=settings.breaker) + + def _build_selector(name: str) -> SourceSelector: return RoundRobinSelector() if name == "round_robin" else LeastInflightSelector() @@ -258,6 +285,11 @@ def _build_cache(settings: GatewaySettings) -> CacheBackend | None: def _build_telemetry(settings: GatewaySettings) -> TelemetryRecorder | None: if settings.telemetry_backend == "none": return None + if settings.telemetry_backend == "postgres": + from polygateway.telemetry.postgres import PostgresRecorder + + assert settings.telemetry_pg_dsn is not None # 内部不变量: config 已校验 + return PostgresRecorder(settings.telemetry_pg_dsn) from polygateway.telemetry.sqlite import SQLiteRecorder assert settings.telemetry_sqlite_path is not None # 内部不变量: config 已校验 diff --git a/src/polygateway/telemetry/postgres.py b/src/polygateway/telemetry/postgres.py new file mode 100644 index 0000000..4e7fd81 --- /dev/null +++ b/src/polygateway/telemetry/postgres.py @@ -0,0 +1,136 @@ +"""Postgres 遥测后端(M2 设计 §5): asyncpg lazy 池 + 两级降级。 + +参考仓无先例(三项目遥测全 SQLite);asyncpg 工程写法取 GovDoc +`taskrun/postgres_store.py`($n 占位、`CREATE TABLE IF NOT EXISTS`、 +`ON CONFLICT DO NOTHING`),但其"失败冒泡"方向按遥测铁律**有意反转**: +① 结构性失败(建池/建表)→ warning 一次后永久降级(池置 None 短路); +② 运行时单条写失败 → 逐条 warning 丢弃,不降级不重试(连接抖动由 +asyncpg 池自恢复;避免浸泡开头一次抖动导致后续全程失遥测)。 +构造不连库(lazy),18 列 schema 与 SQLite 版同名同序。 +""" + +from __future__ import annotations + +import asyncio +from typing import TYPE_CHECKING + +from loguru import logger + +if TYPE_CHECKING: + import asyncpg + +_DDL = """ +CREATE TABLE IF NOT EXISTS llm_calls ( + call_id TEXT PRIMARY KEY, + parent_call_id TEXT, + session_id TEXT, + model TEXT NOT NULL, + provider TEXT NOT NULL, + source_name TEXT NOT NULL, + messages TEXT NOT NULL, + response TEXT NOT NULL, + thinking TEXT NOT NULL DEFAULT '', + prompt_tokens INTEGER NOT NULL, + completion_tokens INTEGER NOT NULL, + usage_source TEXT NOT NULL, + latency_ms INTEGER NOT NULL, + ttft_ms DOUBLE PRECISION, + max_inter_token_ms DOUBLE PRECISION, + cache_hit BOOLEAN NOT NULL DEFAULT FALSE, + error TEXT, + cost DOUBLE PRECISION, + created_at TIMESTAMPTZ NOT NULL DEFAULT now() +); +""" + +_COLUMNS = ( + "call_id", + "parent_call_id", + "session_id", + "model", + "provider", + "source_name", + "messages", + "response", + "thinking", + "prompt_tokens", + "completion_tokens", + "usage_source", + "latency_ms", + "ttft_ms", + "max_inter_token_ms", + "cache_hit", + "error", + "cost", +) + +_INSERT = ( + f"INSERT INTO llm_calls ({', '.join(_COLUMNS)}) " + f"VALUES ({', '.join(f'${i + 1}' for i in range(len(_COLUMNS)))}) " + "ON CONFLICT (call_id) DO NOTHING" +) + + +class PostgresRecorder: + """TelemetryRecorder 端口的 Postgres 实现;asyncpg 原生异步,无线程桥接。""" + + def __init__(self, dsn: str, *, pool: asyncpg.Pool | None = None) -> None: + try: + import asyncpg # noqa: F401 - 仅探测 extra 是否安装 + except ImportError as exc: + raise ImportError( + "Postgres 遥测未启用: 安装 pip install 'polygateway[postgres]' 后重试" + ) from exc + self._dsn = dsn + self._pool: asyncpg.Pool | None = pool + self._external_pool = pool is not None + self._schema_ready = False + self._failed = False # 结构性降级标志: 置位后所有写入短路 + self._init_lock = asyncio.Lock() + + async def _ensure_ready(self) -> asyncpg.Pool | None: + """lazy 建池+建表;结构性失败 warning 一次后永久降级(设计 §5 两级之一)。""" + if self._failed: + return None + if self._schema_ready: + return self._pool + async with self._init_lock: + if self._failed or self._schema_ready: + return None if self._failed else self._pool + try: + if self._pool is None: + import asyncpg + + self._pool = await asyncpg.create_pool(self._dsn, timeout=10) + async with self._pool.acquire() as conn: + await conn.execute(_DDL) + self._schema_ready = True + return self._pool + except asyncio.CancelledError: + raise + except Exception as exc: + self._failed = True + logger.warning("Postgres 遥测初始化失败,后续记录降级为 no-op: {}", exc) + return None + + async def record_llm_call(self, **fields: object) -> None: + """写一行遥测;单条失败逐条 warning 丢弃(两级降级之二),绝不冒泡。""" + pool = await self._ensure_ready() + if pool is None: + return + row = tuple(fields[col] for col in _COLUMNS) + try: + async with pool.acquire() as conn: + await conn.execute(_INSERT, *row) + except asyncio.CancelledError: + raise + except Exception as exc: + # 遥测铁律: 丢一条 < 拖垮调用;仅记 warning(非 pass),池自恢复 + logger.warning("Postgres 遥测写入失败(丢弃该行): {}", exc) + + async def aclose(self) -> None: + """幂等关闭自建池;注入的池归注入方管理。""" + pool, self._pool = self._pool, None + self._schema_ready = False + if pool is not None and not self._external_pool: + await pool.close() diff --git a/tests/integration/test_postgres_telemetry.py b/tests/integration/test_postgres_telemetry.py new file mode 100644 index 0000000..aa0bb82 --- /dev/null +++ b/tests/integration/test_postgres_telemetry.py @@ -0,0 +1,159 @@ +"""PostgresRecorder 集成测试(M2 设计 §5;真实实验室 Postgres,polygateway 专用库)。 + +DSN 走 .env `PGW_TELEMETRY_PG_DSN`,缺则 skip。该实例上有 app/chs_prod 等 +在用库——本测试只允许连 polygateway 专用库(fixture 里守卫)。 +""" + +from __future__ import annotations + +import asyncio +import os + +import pytest +from dotenv import dotenv_values + +from polygateway.telemetry.postgres import PostgresRecorder + +_EXPECTED_COLUMNS = [ + "call_id", + "parent_call_id", + "session_id", + "model", + "provider", + "source_name", + "messages", + "response", + "thinking", + "prompt_tokens", + "completion_tokens", + "usage_source", + "latency_ms", + "ttft_ms", + "max_inter_token_ms", + "cache_hit", + "error", + "cost", + "created_at", +] + + +def _dsn() -> str | None: + merged = {**dotenv_values(".env"), **os.environ} + raw = merged.get("PGW_TELEMETRY_PG_DSN") + if not raw: + return None + scheme, sep, rest = raw.partition("://") + return f"{scheme.partition('+')[0]}{sep}{rest}" + + +@pytest.fixture +async def dsn(): + value = _dsn() + if value is None: + pytest.skip("PGW_TELEMETRY_PG_DSN 未配置") + # 隔离守卫: 该实例有 app/chs_prod/mimiciv 等在用库,只许打 polygateway 专用库 + if not value.rstrip("/").endswith("/polygateway"): + pytest.fail(f"遥测测试只允许连 polygateway 专用库,当前 DSN 库名不符: {value!r}") + import asyncpg + + conn = await asyncpg.connect(value, timeout=10) + try: + await conn.execute("DROP TABLE IF EXISTS llm_calls") + finally: + await conn.close() + return value + + +async def _record_minimal(recorder: PostgresRecorder, call_id: str = "c1", **overrides) -> None: + fields = { + "call_id": call_id, + "parent_call_id": None, + "session_id": "sess-1", + "model": "m", + "provider": "p", + "source_name": "s1", + "messages": "[]", + "response": "ok", + "thinking": "", + "prompt_tokens": 1, + "completion_tokens": 2, + "usage_source": "measured", + "latency_ms": 10, + "ttft_ms": None, + "max_inter_token_ms": None, + "cache_hit": False, + "error": None, + "cost": None, + } + fields.update(overrides) + await recorder.record_llm_call(**fields) + + +async def _fetch(dsn: str, sql: str): + import asyncpg + + conn = await asyncpg.connect(dsn, timeout=10) + try: + return await conn.fetch(sql) + finally: + await conn.close() + + +class TestSchema: + async def test_schema_has_frozen_columns_in_order(self, dsn): + recorder = PostgresRecorder(dsn) + try: + await _record_minimal(recorder) + rows = await _fetch( + dsn, + "SELECT column_name FROM information_schema.columns " + "WHERE table_name='llm_calls' ORDER BY ordinal_position", + ) + assert [r["column_name"] for r in rows] == _EXPECTED_COLUMNS + finally: + await recorder.aclose() + + async def test_call_id_idempotent(self, dsn): + recorder = PostgresRecorder(dsn) + try: + await _record_minimal(recorder, call_id="dup") + await _record_minimal(recorder, call_id="dup", response="second") + rows = await _fetch(dsn, "SELECT response FROM llm_calls WHERE call_id='dup'") + assert [r["response"] for r in rows] == ["ok"] # ON CONFLICT DO NOTHING + finally: + await recorder.aclose() + + async def test_concurrent_writes_all_land(self, dsn): + recorder = PostgresRecorder(dsn) + try: + await asyncio.gather(*(_record_minimal(recorder, call_id=f"c{i}") for i in range(50))) + rows = await _fetch(dsn, "SELECT count(*) AS n FROM llm_calls") + assert rows[0]["n"] == 50 + finally: + await recorder.aclose() + + +class TestDegradation: + async def test_unreachable_server_degrades_silently(self): + """结构性失败(建池不通)→ warning 一次后永久降级,业务零感知。""" + recorder = PostgresRecorder("postgresql://u:p@127.0.0.1:1/x") + await _record_minimal(recorder) # 不抛 + await _record_minimal(recorder, call_id="c2") # 已降级短路,同样不抛 + await recorder.aclose() + + async def test_row_failure_does_not_poison_later_rows(self, dsn): + """运行时单条写失败(NUL 字节文本被 PG 拒)→ 丢该行,后续行照常落库。""" + recorder = PostgresRecorder(dsn) + try: + await _record_minimal(recorder, call_id="bad", response="nul\x00byte") + await _record_minimal(recorder, call_id="good") + rows = await _fetch(dsn, "SELECT call_id FROM llm_calls ORDER BY call_id") + assert [r["call_id"] for r in rows] == ["good"] + finally: + await recorder.aclose() + + async def test_aclose_idempotent(self, dsn): + recorder = PostgresRecorder(dsn) + await _record_minimal(recorder) + await recorder.aclose() + await recorder.aclose() diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 080112d..f7ceaec 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -130,6 +130,22 @@ class TestFactories: with pytest.raises(ValueError, match="glm"): GatewayClient.from_env("LLM", env=env) + def test_from_env_builds_redis_governance_backends(self): + """M2: 配置取 redis 时装配出 Redis 后端(构造不连库,unit 可测)。""" + from polygateway.backends.redis.breaker import RedisGate + from polygateway.backends.redis.limiter import RedisLimiter + from polygateway.client import _build_breaker, _build_limiter + + env = dict( + _ENV, + PGW_LIMITER_BACKEND="redis", + PGW_BREAKER_BACKEND="redis", + REDIS_URL="redis://:pw@10.0.0.1:6379/3", + ) + settings = GatewaySettings.from_env("LLM", env=env) + assert isinstance(_build_limiter(settings, list(settings.sources)), RedisLimiter) + assert isinstance(_build_breaker(settings), RedisGate) + def test_from_settings_respects_injection(self): settings = GatewaySettings.from_env("LLM", env=_ENV) shared = InMemoryLimiter( diff --git a/tests/unit/test_telemetry.py b/tests/unit/test_telemetry.py index be0245f..eaae1af 100644 --- a/tests/unit/test_telemetry.py +++ b/tests/unit/test_telemetry.py @@ -266,5 +266,9 @@ def test_single_emitter_discipline(): text=True, cwd=Path(__file__).resolve().parents[2], ).stdout.splitlines() - callers = [p for p in out if not p.endswith(("ports.py", "telemetry/sqlite.py"))] + callers = [ + p + for p in out + if not p.endswith(("ports.py", "telemetry/sqlite.py", "telemetry/postgres.py")) + ] assert callers == ["src/polygateway/middleware/telemetry.py"]