diff --git a/src/polygateway/middleware/structured.py b/src/polygateway/middleware/structured.py new file mode 100644 index 0000000..2a70b63 --- /dev/null +++ b/src/polygateway/middleware/structured.py @@ -0,0 +1,112 @@ +"""StructuredMW: D14 五级阶梯的编排层(遥测→缓存→**结构化**→重试)。 + +重问 = 再次调用 call_next(内层重试循环)——天然照过限流/熔断门并逐次 +遥测;`ResultInvalidError` 不会进入重试循环的失败计数(RetryMW 对其记 +成功后上抛,由本层决定是否带反馈重问)。缓存在本层之外,只固化阶梯 +通过的最终响应。 +""" + +from __future__ import annotations + +import dataclasses +from typing import TYPE_CHECKING + +from polygateway.errors import ResultInvalidError + +if TYPE_CHECKING: + from polygateway.ports import CallNext, StructuredOutputStrategy + from polygateway.types import ChatRequest, LLMResponse + +# 反馈模板: 库内常量,英文、零业务词(设计 §5 细则 1) +_FEEDBACK_TEMPLATE = ( + "Your previous reply was not valid JSON matching the required schema. " + "Errors: {errors}. Reply with ONLY the corrected JSON object." +) +_MAX_FEEDBACK_ERRORS = 3 +_MAX_ERROR_CHARS = 200 + + +def _format_errors(errors: list[str]) -> str: + clipped = [e[:_MAX_ERROR_CHARS] for e in errors[:_MAX_FEEDBACK_ERRORS]] + return "; ".join(clipped) if clipped else "output could not be parsed" + + +class StructuredMW: + """三档分派: 不传直通 / "json" 仅修复 / pydantic 模型走完整阶梯。""" + + def __init__( + self, + *, + strategy: StructuredOutputStrategy, + max_retries: int = 1, + escalation: StructuredOutputStrategy | None = None, + ) -> None: + if max_retries < 0: + raise ValueError("max_structured_retries 不能为负(0 = 不重问转人工)") + self._strategy = strategy + self._max_retries = max_retries + self._escalation = escalation + + async def __call__(self, request: ChatRequest, call_next: CallNext) -> LLMResponse: + if request.structured is None: + return await call_next(request) + if request.structured == "json": + response = await call_next(self._shape(request, self._strategy, schema=None)) + return dataclasses.replace( + response, structured_data=self._strategy.parse(response.content) + ) + return await self._run_ladder(request, call_next) + + async def _run_ladder(self, request: ChatRequest, call_next: CallNext) -> LLMResponse: + model_cls = request.structured + schema = model_cls.model_json_schema() + current = self._shape(request, self._strategy, schema=schema) + response = await call_next(current) + reasks = 0 + while True: + errors: list[str] = [] + repair_error: str | None = None + try: + parsed = self._strategy.parse(response.content) + except ResultInvalidError as exc: + repair_error = exc.repair_error + errors.append(exc.repair_error or "not parseable as JSON") + else: + try: + validated = model_cls.model_validate(parsed) + except ValueError as exc: # pydantic ValidationError 继承 ValueError + errors.append(str(exc)) + else: + return dataclasses.replace(response, structured_data=validated) + if reasks >= self._max_retries: + raise ResultInvalidError( + "结构化输出阶梯耗尽", + raw_text=response.content, + repair_error=repair_error, + validation_errors=tuple(errors), + ) + reasks += 1 + current = self._with_feedback(current, response.content, errors, schema) + response = await call_next(current) + + def _shape( + self, request: ChatRequest, strategy: StructuredOutputStrategy, *, schema: dict | None + ) -> ChatRequest: + overlay = strategy.request_overlay(schema) + if not overlay: + return request + return dataclasses.replace(request, overlay={**request.overlay, **overlay}) + + def _with_feedback( + self, current: ChatRequest, bad_content: str, errors: list[str], schema: dict | None + ) -> ChatRequest: + """构造带反馈的重问请求;可升级到原生 schema 策略(设计 §5 细则 2)。""" + messages = [ + *current.messages, + {"role": "assistant", "content": bad_content}, + {"role": "user", "content": _FEEDBACK_TEMPLATE.format(errors=_format_errors(errors))}, + ] + reask = dataclasses.replace(current, messages=messages) + if self._escalation is not None: + reask = self._shape(reask, self._escalation, schema=schema) + return reask diff --git a/src/polygateway/structured/json_repair.py b/src/polygateway/structured/json_repair.py new file mode 100644 index 0000000..3b4b7e1 --- /dev/null +++ b/src/polygateway/structured/json_repair.py @@ -0,0 +1,57 @@ +"""JsonRepairStrategy: 阶梯②修复(D7/D14;蓝本 VT core/agent/loop.py:341-377)。 + +处理链: 围栏剥离 → json_repair → json.loads → 可选 normalize 钩子。 +业务 schema 相关的变体归一化(如 VT `_normalize_action` 的 DeepSeek 平铺 +收拢)**不入库**——零业务假设铁律;业务侧经 `normalize` 注入自带函数。 +""" + +from __future__ import annotations + +import json +import re +from typing import TYPE_CHECKING, Any + +from polygateway.errors import ResultInvalidError + +if TYPE_CHECKING: + from collections.abc import Callable + +try: + from json_repair import repair_json +except ImportError as _exc: # pragma: no cover - 依赖缺失路径 + repair_json = None + _IMPORT_ERROR = _exc +else: + _IMPORT_ERROR = None + +# VT loop.py:37 同款围栏正则 +_CODE_FENCE_RE = re.compile(r"^\s*```(?:json)?\s*\n?|\n?\s*```\s*$") + + +class JsonRepairStrategy: + """prompt 约定 + 事后修复策略;不改请求体(request_overlay 恒空)。""" + + def __init__(self, normalize: Callable[[Any], Any] | None = None) -> None: + if repair_json is None: + raise ImportError( + "结构化输出需要 json_repair 包: pip install 'polygateway[structured]'" + ) from _IMPORT_ERROR + self._normalize = normalize + + def request_overlay(self, schema: dict[str, Any] | None) -> dict[str, Any]: + return {} + + def parse(self, text: str) -> Any: + """修复并解析;失败抛 ResultInvalidError(坏结果 ≠ 坏服务)。""" + stripped = _CODE_FENCE_RE.sub("", text).strip() + if not stripped: + raise ResultInvalidError("结构化输出为空", raw_text=text, repair_error="empty content") + try: + data = json.loads(repair_json(stripped)) + except (json.JSONDecodeError, ValueError) as exc: + raise ResultInvalidError( + "JSON 修复失败", raw_text=text, repair_error=str(exc) + ) from exc + if self._normalize is not None: + data = self._normalize(data) + return data diff --git a/src/polygateway/structured/native_schema.py b/src/polygateway/structured/native_schema.py new file mode 100644 index 0000000..202c3fd --- /dev/null +++ b/src/polygateway/structured/native_schema.py @@ -0,0 +1,32 @@ +"""NativeSchemaStrategy: 阶梯①预防(D7/D14)——网关支持时用协议级约束。 + +请求侧注入 response_format(有 schema 用 json_schema 严格模式,无 schema +退为 json_object);响应侧仍复用修复链兜底——原生约束下模型偶发的围栏/ +噪声不至于直接判死。 +""" + +from __future__ import annotations + +from typing import Any + +from polygateway.structured.json_repair import JsonRepairStrategy + + +class NativeSchemaStrategy: + """response_format 注入策略;由装配层按 provider 注册表能力选择。""" + + def __init__(self) -> None: + self._repair = JsonRepairStrategy() + + def request_overlay(self, schema: dict[str, Any] | None) -> dict[str, Any]: + if schema is None: + return {"response_format": {"type": "json_object"}} + return { + "response_format": { + "type": "json_schema", + "json_schema": {"name": "structured_output", "strict": True, "schema": schema}, + } + } + + def parse(self, text: str) -> Any: + return self._repair.parse(text) diff --git a/tests/unit/test_structured.py b/tests/unit/test_structured.py new file mode 100644 index 0000000..dbcc0bf --- /dev/null +++ b/tests/unit/test_structured.py @@ -0,0 +1,169 @@ +"""结构化输出阶梯测试(D14/设计 §5): 三档分派、修复链、有界带反馈重问。""" + +import dataclasses + +import pytest +from pydantic import BaseModel + +from polygateway.errors import ResultInvalidError +from polygateway.middleware.structured import StructuredMW +from polygateway.structured.json_repair import JsonRepairStrategy +from polygateway.structured.native_schema import NativeSchemaStrategy +from polygateway.types import ChatRequest, LLMResponse + +_MSGS = [{"role": "user", "content": "give json"}] + + +class Verdict(BaseModel): + answer: int + reason: str + + +def _resp(content): + return LLMResponse( + content=content, thinking="", model="m", provider="p", prompt_tokens=1, + completion_tokens=2, latency_ms=10, ttft_ms=None, max_inter_token_ms=None, + cache_hit=False, call_id="cid", source_name="s1", usage_source="measured", + ) + + +class ScriptedTerminal: + """按脚本逐次返回 content;记录收到的 ChatRequest 序列。""" + + def __init__(self, contents): + self.contents = list(contents) + self.requests = [] + + async def __call__(self, request): + self.requests.append(request) + return _resp(self.contents.pop(0)) + + +class TestJsonRepairStrategy: + @pytest.mark.parametrize( + "dirty", + [ + '```json\n{"answer": 1, "reason": "ok"}\n```', # 围栏 + '{"answer": 1, "reason": "ok",}', # 尾逗号 + "{'answer': 1, 'reason': 'ok'}", # 单引号 + '{"answer": 1, "reason": "ok"', # 缺右括号 + ], + ) + def test_repairs_real_world_dirt(self, dirty): + assert JsonRepairStrategy().parse(dirty) == {"answer": 1, "reason": "ok"} + + def test_unrepairable_raises_result_invalid(self): + with pytest.raises(ResultInvalidError) as ei: + JsonRepairStrategy().parse("I refuse to answer in JSON.") + assert ei.value.raw_text + + def test_normalize_hook_applied(self): + strategy = JsonRepairStrategy(normalize=lambda d: {**d, "tagged": True}) + assert strategy.parse('{"a": 1}') == {"a": 1, "tagged": True} + + def test_request_overlay_empty(self): + assert JsonRepairStrategy().request_overlay({"type": "object"}) == {} + + +class TestNativeSchemaStrategy: + def test_overlay_with_schema(self): + overlay = NativeSchemaStrategy().request_overlay(Verdict.model_json_schema()) + rf = overlay["response_format"] + assert rf["type"] == "json_schema" + assert rf["json_schema"]["schema"]["required"] == ["answer", "reason"] + + def test_overlay_without_schema_is_json_object(self): + assert NativeSchemaStrategy().request_overlay(None) == { + "response_format": {"type": "json_object"} + } + + +def _mw(**kwargs): + defaults = dict(strategy=JsonRepairStrategy(), max_retries=1, escalation=None) + defaults.update(kwargs) + return StructuredMW(**defaults) + + +class TestThreeTiers: + async def test_tier_none_passthrough(self): + terminal = ScriptedTerminal(["free text"]) + resp = await _mw()(ChatRequest(messages=_MSGS), terminal) + assert resp.structured_data is None + assert terminal.requests[0].overlay == {} + + async def test_tier_json_repair_only_no_retry(self): + terminal = ScriptedTerminal(["not json at all"]) + with pytest.raises(ResultInvalidError): + await _mw()(ChatRequest(messages=_MSGS, structured="json"), terminal) + assert len(terminal.requests) == 1 # "json" 档失败不重问(CHS 语义由 0 档覆盖) + + async def test_tier_json_success(self): + terminal = ScriptedTerminal(['```json\n{"x": 1}\n```']) + resp = await _mw()(ChatRequest(messages=_MSGS, structured="json"), terminal) + assert resp.structured_data == {"x": 1} + + async def test_tier_model_full_ladder_success(self): + terminal = ScriptedTerminal(['{"answer": 7, "reason": "sure"}']) + resp = await _mw()(ChatRequest(messages=_MSGS, structured=Verdict), terminal) + assert isinstance(resp.structured_data, Verdict) + assert resp.structured_data.answer == 7 + + +class TestFeedbackRetry: + async def test_validation_failure_triggers_feedback_reask(self): + terminal = ScriptedTerminal( + ['{"answer": "not-an-int"}', '{"answer": 7, "reason": "fixed"}'] + ) + resp = await _mw()(ChatRequest(messages=_MSGS, structured=Verdict), terminal) + assert resp.structured_data.answer == 7 + assert len(terminal.requests) == 2 + reask = terminal.requests[1].messages + # 反馈模板: 原 messages + assistant 坏输出 + user 纠错指令(设计 §5 细则 1) + assert reask[0] == _MSGS[0] + assert reask[1]["role"] == "assistant" and "not-an-int" in reask[1]["content"] + assert reask[2]["role"] == "user" and "valid JSON" in reask[2]["content"] + assert "answer" in reask[2]["content"] # 校验错误进入反馈 + + async def test_exhaustion_raises_with_diagnosis(self): + terminal = ScriptedTerminal(['{"answer": "a"}', '{"answer": "b"}']) + with pytest.raises(ResultInvalidError) as ei: + await _mw()(ChatRequest(messages=_MSGS, structured=Verdict), terminal) + assert len(terminal.requests) == 2 # 首次 + 1 次重问 + assert ei.value.raw_text == '{"answer": "b"}' + assert ei.value.validation_errors + + async def test_zero_retries_is_chs_policy(self): + terminal = ScriptedTerminal(['{"answer": "bad"}']) + with pytest.raises(ResultInvalidError): + await _mw(max_retries=0)(ChatRequest(messages=_MSGS, structured=Verdict), terminal) + assert len(terminal.requests) == 1 + + async def test_reask_escalates_to_native_schema(self): + terminal = ScriptedTerminal(['{"answer": "bad"}', '{"answer": 1, "reason": "r"}']) + mw = _mw(escalation=NativeSchemaStrategy()) + await mw(ChatRequest(messages=_MSGS, structured=Verdict), terminal) + assert "response_format" not in terminal.requests[0].overlay # 首发 JsonRepair 无 overlay + assert terminal.requests[1].overlay["response_format"]["type"] == "json_schema" + + async def test_error_feedback_truncated(self): + huge_reason = "x" * 5000 + terminal = ScriptedTerminal( + [f'{{"answer": "{huge_reason}"}}', '{"answer": 1, "reason": "r"}'] + ) + await _mw()(ChatRequest(messages=_MSGS, structured=Verdict), terminal) + feedback = terminal.requests[1].messages[-1]["content"] + assert len(feedback) < 2000 # 每条错误截断 200 字符,防 prompt 膨胀 + + +class TestNativeOverlayFirstAttempt: + async def test_native_strategy_shapes_first_request(self): + terminal = ScriptedTerminal(['{"answer": 1, "reason": "r"}']) + mw = _mw(strategy=NativeSchemaStrategy()) + await mw(ChatRequest(messages=_MSGS, structured=Verdict), terminal) + assert terminal.requests[0].overlay["response_format"]["type"] == "json_schema" + + async def test_response_immutability_preserved(self): + terminal = ScriptedTerminal(['{"answer": 1, "reason": "r"}']) + resp = await _mw()(ChatRequest(messages=_MSGS, structured=Verdict), terminal) + with pytest.raises(dataclasses.FrozenInstanceError): + resp.structured_data = None