Coverage for agentos/validation/schema_enforcer.py: 100%
156 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-08 13:14 +0800
1"""AgentOS v1.3.9 - Schema Enforcer 模块。
3对 Agent 输出执行 Pydantic schema 校验，校验失败时自动修复/重试。
4支持 JSON 修复、字段回退、LLM 辅助修正三种修复策略。
5"""
7from __future__ import annotations
9import asyncio
10import json
11import logging
12from dataclasses import dataclass, field
13from enum import Enum, auto
14from typing import Any
16logger = logging.getLogger(__name__)
19class FixStrategy(Enum):
20 """修复策略枚举。"""
22 JSON_REPAIR = auto()
23 FIELD_FALLBACK = auto()
24 LLM_ASSISTED = auto()
25 RAISE = auto()
28@dataclass
29class EnforcerResult:
30 """校验执行结果。"""
32 is_valid: bool
33 original_output: Any
34 repaired_output: Any | None = None
35 errors: list[str] = field(default_factory=list)
36 fix_strategy_used: FixStrategy | None = None
37 fix_attempts: int = 0
40@dataclass
41class EnforcerConfig:
42 """Schema Enforcer 配置。"""
44 max_retries: int = 3
45 strategy_order: list[FixStrategy] = field(
46 default_factory=lambda: [
47 FixStrategy.JSON_REPAIR,
48 FixStrategy.FIELD_FALLBACK,
49 FixStrategy.LLM_ASSISTED,
50 ]
51 )
52 llm_fix_prompt_template: str = ""
53 default_value_fallback: bool = True
54 log_rejections: bool = True
57@dataclass
58class EnforcerStats:
59 """校验统计。"""
61 total_checks: int = 0
62 total_rejections: int = 0
63 total_repairs: int = 0
64 repairs_by_strategy: dict[str, int] = field(default_factory=dict)
67class SchemaEnforcer:
68 """对 Agent 输出执行 Pydantic schema 校验与自动修复。
70 核心流程：
71 1. 尝试直接 model_validate
72 2. 失败时按 strategy_order 依次尝试修复
73 3. 所有策略耗尽仍失败则降级为 FIELD_FALLBACK（最佳努力）
74 """
76 def __init__(self, config: EnforcerConfig | None = None):
77 self.config = config or EnforcerConfig()
78 self.stats = EnforcerStats()
80 async def enforce(
81 self,
82 output: dict | str | Any,
83 schema_model: type,
84 context: dict | None = None,
85 ) -> EnforcerResult:
86 """对单次输出执行 schema 校验。"""
87 self.stats.total_checks += 1
88 errors: list[str] = []
90 try:
91 validated = schema_model.model_validate(output)
92 return EnforcerResult(is_valid=True, original_output=output, repaired_output=validated)
93 except Exception as e:
94 errors.append(str(e))
95 self.stats.total_rejections += 1
97 result = EnforcerResult(is_valid=False, original_output=output, errors=errors)
99 for attempt in range(self.config.max_retries):
100 for strategy in self.config.strategy_order:
101 try:
102 repaired = await self._apply_fix(
103 strategy, output, schema_model, errors, context
104 )
105 if repaired is not None:
106 validated = schema_model.model_validate(repaired)
107 self.stats.total_repairs += 1
108 strat_key = strategy.name
109 self.stats.repairs_by_strategy[strat_key] = (
110 self.stats.repairs_by_strategy.get(strat_key, 0) + 1
111 )
112 result.is_valid = True
113 result.repaired_output = validated
114 result.fix_strategy_used = strategy
115 result.fix_attempts = attempt + 1
116 if self.config.log_rejections:
117 logger.info(
118 "Schema fixed via %s (attempt %d/%d)",
119 strategy.name,
120 attempt + 1,
121 self.config.max_retries,
122 )
123 return result
124 except Exception as fix_error:
125 errors.append(f"[{strategy.name}] {fix_error}")
127 if self.config.default_value_fallback:
128 try:
129 fallback = self._build_fallback(schema_model)
130 self.stats.total_repairs += 1
131 self.stats.repairs_by_strategy["FALLBACK"] = (
132 self.stats.repairs_by_strategy.get("FALLBACK", 0) + 1
133 )
134 result.is_valid = True
135 result.repaired_output = fallback
136 result.fix_strategy_used = FixStrategy.FIELD_FALLBACK
137 result.fix_attempts = self.config.max_retries
138 return result
139 except Exception:
140 pass
142 return result
144 async def _apply_fix(
145 self,
146 strategy: FixStrategy,
147 output: Any,
148 model: type,
149 errors: list[str],
150 context: dict | None,
151 ) -> dict | None:
152 if strategy == FixStrategy.JSON_REPAIR:
153 return self._json_repair(output)
154 elif strategy == FixStrategy.FIELD_FALLBACK:
155 return self._field_fallback(output, model, errors)
156 elif strategy == FixStrategy.LLM_ASSISTED:
157 return await self._llm_fix(output, model, errors, context)
158 return None
160 def _json_repair(self, output: Any) -> dict | None:
161 """尝试修复 JSON 格式问题（尾部逗号、单引号、截断等）。"""
162 if isinstance(output, dict):
163 return output
164 if isinstance(output, str):
165 s = output.strip()
166 # 去除 markdown 代码块包裹
167 if s.startswith("```"):
168 lines = s.split("\n")
169 if lines[0].startswith("```"):
170 lines = lines[1:]
171 if lines and lines[-1].strip() == "```":
172 lines = lines[:-1]
173 s = "\n".join(lines)
174 # 修复常见 JSON 问题
175 s = s.replace("'", '"')
176 # 修复尾部多余逗号
177 import re
179 s = re.sub(r",(\s*[}\]])", r"\1", s)
180 try:
181 return json.loads(s)
182 except json.JSONDecodeError:
183 pass
184 return None
186 def _field_fallback(self, output: Any, model: type, errors: list[str]) -> dict | None:
187 """从原始输出中尽力提取有效字段，缺失字段填默认值。"""
188 from pydantic_core import PydanticUndefined
190 try:
191 if not isinstance(output, dict):
192 return None
193 fields_info = model.model_fields
194 clean: dict = {}
195 for key, finfo in fields_info.items():
196 if key in output:
197 clean[key] = output[key]
198 elif finfo.default is not PydanticUndefined:
199 clean[key] = finfo.default
200 elif finfo.default_factory is not None:
201 clean[key] = finfo.default_factory()
202 return clean if clean else None
203 except Exception:
204 return None
206 async def _llm_fix(
207 self, output: Any, model: type, errors: list[str], context: dict | None
208 ) -> dict | None:
209 """通过 LLM 辅助修复（调用方需注入 llm_call 回调）。"""
210 if self.config.llm_fix_prompt_template:
211 logger.warning("LLM-assisted fix requires llm_call callback (not implemented inline).")
212 return None
214 def _build_fallback(self, model: type) -> Any:
215 """使用全默认值构建回退对象。"""
216 from pydantic_core import PydanticUndefined
218 fields_info = model.model_fields
219 kwargs: dict = {}
220 for key, finfo in fields_info.items():
221 if finfo.default is not PydanticUndefined:
222 kwargs[key] = finfo.default
223 elif finfo.default_factory is not None:
224 kwargs[key] = finfo.default_factory()
225 else:
226 annotation = finfo.annotation
227 origin = getattr(annotation, "__origin__", None)
228 if annotation is str:
229 kwargs[key] = ""
230 elif annotation is int:
231 kwargs[key] = 0
232 elif annotation is float:
233 kwargs[key] = 0.0
234 elif annotation is bool:
235 kwargs[key] = False
236 elif annotation is list or origin is list:
237 kwargs[key] = []
238 elif annotation is dict or origin is dict:
239 kwargs[key] = {}
240 return model(**kwargs)
242 async def enforce_batch(
243 self,
244 outputs: list[dict | str],
245 schema_model: type,
246 context: dict | None = None,
247 ) -> list[EnforcerResult]:
248 """批量校验，利用异步并发。"""
249 tasks = [self.enforce(out, schema_model, context) for out in outputs]
250 return await asyncio.gather(*tasks)
253__all__ = [
254 "SchemaEnforcer",
255 "EnforcerConfig",
256 "EnforcerResult",
257 "EnforcerStats",
258 "FixStrategy",
259]