diff --git a/src/polygateway/sources.py b/src/polygateway/sources.py new file mode 100644 index 0000000..c33f439 --- /dev/null +++ b/src/polygateway/sources.py @@ -0,0 +1,61 @@ +"""选源策略与源冷却备忘(M1 设计 §2.3;蓝本 CHS app/providers/selector.py 与 governance.py:107)。 + +选源是端口(SourceSelector),两个首发实现逐字移植 CHS;冷却备忘是 +RetryMW 的进程本地状态——熔断开路的源在本地记冷却截止,选源时跳过, +避免每轮白烧 RPM 去探测已知开路的源。 +""" + +from __future__ import annotations + +import time +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Callable + + from polygateway.types import SourceConfig, SourceStats + + +class RoundRobinSelector: + """轮转起点后移(CHS selector.py:20 同款);单 client 内游标推进。""" + + def __init__(self) -> None: + self._n = 0 + + def order( + self, sources: list[SourceConfig], stats: dict[str, SourceStats] + ) -> list[SourceConfig]: + if not sources: + return [] + k = self._n % len(sources) + self._n += 1 + return sources[k:] + sources[:k] + + +class LeastInflightSelector: + """最少在途优先(CHS selector.py:36 同款);排序稳定,平局保持配置序。""" + + def order( + self, sources: list[SourceConfig], stats: dict[str, SourceStats] + ) -> list[SourceConfig]: + return sorted(sources, key=lambda s: stats[s.name].inflight if s.name in stats else 0) + + +class SourceCooldownMemo: + """进程本地的源冷却备忘(CHS governance.py:107 同款)。 + + 只记"冷却截止时刻";set_until 取更晚者,防止较早的提示回退已有备忘。 + """ + + def __init__(self, now: Callable[[], float] = time.monotonic) -> None: + self._now = now + self._until: dict[str, float] = {} + + def set_until(self, source_name: str, until: float) -> None: + self._until[source_name] = max(self._until.get(source_name, 0.0), until) + + def active(self, source_name: str) -> bool: + return self._until.get(source_name, 0.0) > self._now() + + def remaining(self, source_name: str) -> float: + return max(0.0, self._until.get(source_name, 0.0) - self._now()) diff --git a/tests/unit/test_sources.py b/tests/unit/test_sources.py new file mode 100644 index 0000000..30eef38 --- /dev/null +++ b/tests/unit/test_sources.py @@ -0,0 +1,58 @@ +"""sources.py 选源策略与冷却备忘测试(蓝本 CHS app/providers/selector.py:20,36)。""" + +from polygateway.sources import LeastInflightSelector, RoundRobinSelector, SourceCooldownMemo +from polygateway.types import SourceConfig, SourceStats + + +def _src(name): + return SourceConfig( + name=name, provider="openai", base_url="https://gw.example/v1", + api_key="sk", model="m", timeout_s=10.0, + ) + + +class TestRoundRobin: + def test_rotation_advances_each_call(self): + sources = [_src("a"), _src("b"), _src("c")] + sel = RoundRobinSelector() + orders = [tuple(s.name for s in sel.order(sources, {})) for _ in range(4)] + assert orders == [("a", "b", "c"), ("b", "c", "a"), ("c", "a", "b"), ("a", "b", "c")] + + def test_single_source_stable(self): + sources = [_src("only")] + sel = RoundRobinSelector() + assert [s.name for s in sel.order(sources, {})] == ["only"] + assert [s.name for s in sel.order(sources, {})] == ["only"] + + +class TestLeastInflight: + def test_orders_by_inflight_ascending(self): + sources = [_src("a"), _src("b"), _src("c")] + stats = {"a": SourceStats(5, 0, 0), "b": SourceStats(1, 0, 0), "c": SourceStats(3, 0, 0)} + sel = LeastInflightSelector() + assert [s.name for s in sel.order(sources, stats)] == ["b", "c", "a"] + + def test_stable_on_ties(self): + sources = [_src("a"), _src("b")] + stats = {"a": SourceStats(2, 0, 0), "b": SourceStats(2, 0, 0)} + assert [s.name for s in LeastInflightSelector().order(sources, stats)] == ["a", "b"] + + +class TestCooldownMemo: + def test_active_until_expiry(self): + t = {"now": 100.0} + memo = SourceCooldownMemo(now=lambda: t["now"]) + assert not memo.active("a") + memo.set_until("a", 130.0) + assert memo.active("a") + assert memo.remaining("a") == 30.0 + t["now"] = 131.0 + assert not memo.active("a") + assert memo.remaining("a") == 0.0 + + def test_extend_after_wait_hint(self): + t = {"now": 0.0} + memo = SourceCooldownMemo(now=lambda: t["now"]) + memo.set_until("a", 10.0) + memo.set_until("a", 5.0) # 更早的截止不回退已有备忘 + assert memo.remaining("a") == 10.0