| 1 | # -*- coding: utf-8 -*- |
| 2 | """ |
| 3 | Doctor agent for generation failures. |
| 4 | |
| 5 | The doctor is intentionally not a workflow stage. It is called only at failure |
| 6 | boundaries to decide whether a prompt can be safely rewrite and retried. |
| 7 | """ |
| 8 | |
| 9 | import json |
| 10 | import logging |
| 11 | from typing import Any, Dict, Optional, Tuple |
| 12 | |
| 13 | from .base_agent import AgentInterface |
| 14 | from config import Config |
| 15 | from prompts.loader import load_prompt |
| 16 | |
| 17 | logger = logging.getLogger(__name__) |
| 18 | |
| 19 | |
| 20 | REASON_TYPE_PROMPT_SAFETY = "prompt_safety" |
| 21 | REASON_TYPE_PROMPT_TOO_LONG = "prompt_too_long" |
| 22 | REASON_TYPE_PROMPT_FORMAT = "prompt_format" |
| 23 | REASON_TYPE_MODEL_LIMITATION = "model_limitation" |
| 24 | REASON_TYPE_INPUT_ASSET = "input_asset" |
| 25 | REASON_TYPE_AUTH_CONFIG = "auth_config" |
| 26 | REASON_TYPE_NETWORK_TIMEOUT = "network_timeout" |
| 27 | REASON_TYPE_PROVIDER_INTERNAL = "provider_internal" |
| 28 | REASON_TYPE_UNKNOWN = "unknown" |
| 29 | |
| 30 | DOCTOR_REASON_TYPES = { |
| 31 | REASON_TYPE_PROMPT_SAFETY, |
| 32 | REASON_TYPE_PROMPT_TOO_LONG, |
| 33 | REASON_TYPE_PROMPT_FORMAT, |
| 34 | REASON_TYPE_MODEL_LIMITATION, |
| 35 | REASON_TYPE_INPUT_ASSET, |
| 36 | REASON_TYPE_AUTH_CONFIG, |
| 37 | REASON_TYPE_NETWORK_TIMEOUT, |
| 38 | REASON_TYPE_PROVIDER_INTERNAL, |
| 39 | REASON_TYPE_UNKNOWN, |
| 40 | } |
| 41 | |
| 42 | REWRITABLE_REASON_TYPES = { |
| 43 | REASON_TYPE_PROMPT_SAFETY, |
| 44 | REASON_TYPE_PROMPT_TOO_LONG, |
| 45 | REASON_TYPE_PROMPT_FORMAT, |
| 46 | } |
| 47 | |
| 48 | DOCTOR_TOOLS = {"rewrite_prompt", "none"} |
| 49 | MAX_DOCTOR_ATTEMPTS = 3 |
| 50 | |
| 51 | |
| 52 | class DoctorOutputError(ValueError): |
| 53 | """Raised when the doctor LLM returns invalid structured output.""" |
| 54 | |
| 55 | |
| 56 | class DoctorAgent(AgentInterface): |
| 57 | """Diagnose generation errors and optionally rewrite prompts.""" |
| 58 | |
| 59 | def __init__(self, llm_model: Optional[str] = None): |
| 60 | super().__init__(name="Doctor") |
| 61 | self.llm_model = llm_model or Config.LLM_MODEL |
| 62 | |
| 63 | async def process(self, input_data: Any, intervention: Optional[Dict] = None) -> Dict: |
| 64 | """Compatibility wrapper for the agent interface; doctor is normally called directly.""" |
| 65 | input_data = input_data if isinstance(input_data, dict) else {} |
| 66 | return { |
| 67 | "payload": self.diagnose_error( |
| 68 | stage=str(input_data.get("stage", "")), |
| 69 | model=str(input_data.get("model", "")), |
| 70 | prompt=str(input_data.get("prompt", "")), |
| 71 | error=str(input_data.get("error", "")), |
| 72 | context=input_data.get("context") if isinstance(input_data.get("context"), dict) else None, |
| 73 | ), |
| 74 | "requires_intervention": False, |
| 75 | "stage_completed": True, |
| 76 | } |
| 77 | |
| 78 | @staticmethod |
| 79 | def no_action(reason: str, reason_type: str = REASON_TYPE_UNKNOWN) -> Dict[str, Any]: |
| 80 | return { |
| 81 | "matched": False, |
| 82 | "reason_type": reason_type, |
| 83 | "tool": "none", |
| 84 | "should_retry": False, |
| 85 | "confidence": 0.0, |
| 86 | "reason": reason, |
| 87 | } |
| 88 | |
| 89 | @staticmethod |
| 90 | def _json_dumps(value: Any) -> str: |
| 91 | return json.dumps(value or {}, ensure_ascii=False, indent=2) |
| 92 | |
| 93 | @staticmethod |
| 94 | def _load_json_object(raw: str) -> Dict[str, Any]: |
| 95 | data = json.loads(raw) |
| 96 | if not isinstance(data, dict): |
| 97 | raise DoctorOutputError("doctor output must be a JSON object") |
| 98 | return data |
| 99 | |
| 100 | @staticmethod |
| 101 | def _validate_diagnosis(data: Dict[str, Any]) -> Dict[str, Any]: |
| 102 | reason_type = data.get("reason_type") |
| 103 | tool = data.get("tool") |
| 104 | should_retry = data.get("should_retry") |
| 105 | confidence = data.get("confidence") |
| 106 | reason = data.get("reason") |
| 107 | |
| 108 | if reason_type not in DOCTOR_REASON_TYPES: |
| 109 | raise DoctorOutputError(f"reason_type must be one of {sorted(DOCTOR_REASON_TYPES)}") |
| 110 | if tool not in DOCTOR_TOOLS: |
| 111 | raise DoctorOutputError(f"tool must be one of {sorted(DOCTOR_TOOLS)}") |
| 112 | if not isinstance(should_retry, bool): |
| 113 | raise DoctorOutputError("should_retry must be a boolean") |
| 114 | if not isinstance(confidence, (int, float)) or not 0 <= float(confidence) <= 1: |
| 115 | raise DoctorOutputError("confidence must be a number between 0 and 1") |
| 116 | if not isinstance(reason, str) or not reason.strip(): |
| 117 | raise DoctorOutputError("reason must be a non-empty string") |
| 118 | if tool == "rewrite_prompt" and reason_type not in REWRITABLE_REASON_TYPES: |
| 119 | raise DoctorOutputError("rewrite_prompt is only allowed for rewritable reason types") |
| 120 | |
| 121 | return { |
| 122 | "matched": bool(data.get("matched", tool != "none")), |
| 123 | "reason_type": reason_type, |
| 124 | "tool": tool, |
| 125 | "should_retry": should_retry, |
| 126 | "confidence": float(confidence), |
| 127 | "reason": reason.strip(), |
| 128 | } |
| 129 | |
| 130 | @staticmethod |
| 131 | def _validate_rewrite(data: Dict[str, Any], original_prompt: str) -> str: |
| 132 | rewrite = data.get("rewrite_prompt") |
| 133 | if not isinstance(rewrite, str) or not rewrite.strip(): |
| 134 | raise DoctorOutputError("rewrite_prompt must be a non-empty string") |
| 135 | rewrite = rewrite.strip() |
| 136 | if rewrite == (original_prompt or "").strip(): |
| 137 | raise DoctorOutputError("rewrite_prompt must differ from the original prompt") |
| 138 | return rewrite |
| 139 | |
| 140 | @staticmethod |
| 141 | def _rule_based_diagnosis(error: str) -> Optional[Dict[str, Any]]: |
| 142 | text = (error or "").lower() |
| 143 | safety_markers = ( |
| 144 | "datainspectionfailed", |
| 145 | "inappropriate content", |
| 146 | "content policy", |
| 147 | "safety", |
| 148 | "sensitive", |
| 149 | "risk control", |
| 150 | "审核", |
| 151 | "违规", |
| 152 | "敏感", |
| 153 | "不合规", |
| 154 | "安全", |
| 155 | ) |
| 156 | if any(marker in text for marker in safety_markers): |
| 157 | return { |
| 158 | "matched": True, |
| 159 | "reason_type": REASON_TYPE_PROMPT_SAFETY, |
| 160 | "tool": "rewrite_prompt", |
| 161 | "should_retry": True, |
| 162 | "confidence": 0.95, |
| 163 | "reason": "供应商返回内容安全或审核相关错误，疑似由生成提示词中的不文明、敏感或高风险表达触发。", |
| 164 | } |
| 165 | |
| 166 | length_markers = ("maximum context", "max tokens", "too long", "context length", "超长", "长度") |
| 167 | if any(marker in text for marker in length_markers): |
| 168 | return { |
| 169 | "matched": True, |
| 170 | "reason_type": REASON_TYPE_PROMPT_TOO_LONG, |
| 171 | "tool": "rewrite_prompt", |
| 172 | "should_retry": True, |
| 173 | "confidence": 0.85, |
| 174 | "reason": "供应商错误指向提示词或上下文过长，需要压缩提示词后重试。", |
| 175 | } |
| 176 | |
| 177 | asset_markers = ("file not found", "不存在", "missing", "not found", "路径") |
| 178 | if any(marker in text for marker in asset_markers): |
| 179 | return { |
| 180 | "matched": True, |
| 181 | "reason_type": REASON_TYPE_INPUT_ASSET, |
| 182 | "tool": "none", |
| 183 | "should_retry": False, |
| 184 | "confidence": 0.85, |
| 185 | "reason": "错误指向输入素材缺失或路径不可用，重写提示词无法修复。", |
| 186 | } |
| 187 | |
| 188 | auth_markers = ("api key", "apikey", "unauthorized", "forbidden", "permission", "权限", "密钥") |
| 189 | if any(marker in text for marker in auth_markers): |
| 190 | return { |
| 191 | "matched": True, |
| 192 | "reason_type": REASON_TYPE_AUTH_CONFIG, |
| 193 | "tool": "none", |
| 194 | "should_retry": False, |
| 195 | "confidence": 0.85, |
| 196 | "reason": "错误指向鉴权、权限或配置问题，重写提示词无法修复。", |
| 197 | } |
| 198 | |
| 199 | timeout_markers = ("timeout", "timed out", "超时") |
| 200 | if any(marker in text for marker in timeout_markers): |
| 201 | return { |
| 202 | "matched": True, |
| 203 | "reason_type": REASON_TYPE_NETWORK_TIMEOUT, |
| 204 | "tool": "none", |
| 205 | "should_retry": False, |
| 206 | "confidence": 0.8, |
| 207 | "reason": "错误指向网络或任务轮询超时，重写提示词不是合适修复方式。", |
| 208 | } |
| 209 | |
| 210 | return None |
| 211 | |
| 212 | def diagnose_error( |
| 213 | self, |
| 214 | *, |
| 215 | stage: str, |
| 216 | model: str, |
| 217 | prompt: str, |
| 218 | error: str, |
| 219 | context: Optional[Dict[str, Any]] = None, |
| 220 | ) -> Dict[str, Any]: |
| 221 | rule_result = self._rule_based_diagnosis(error) |
| 222 | if rule_result: |
| 223 | return self._validate_diagnosis(rule_result) |
| 224 | |
| 225 | from models.llm_client import LLM |
| 226 | |
| 227 | template = load_prompt("doctor", "diagnose", "zh") |
| 228 | validation_error = "" |
| 229 | for attempt in range(MAX_DOCTOR_ATTEMPTS): |
| 230 | doctor_prompt = template.format( |
| 231 | stage=stage, |
| 232 | model=model, |
| 233 | error=error, |
| 234 | original_prompt=prompt, |
| 235 | context=self._json_dumps(context), |
| 236 | reason_types=", ".join(sorted(DOCTOR_REASON_TYPES)), |
| 237 | rewritable_reason_types=", ".join(sorted(REWRITABLE_REASON_TYPES)), |
| 238 | tools=", ".join(sorted(DOCTOR_TOOLS)), |
| 239 | validation_error=validation_error or "无", |
| 240 | ) |
| 241 | try: |
| 242 | raw = LLM().query(doctor_prompt, model=self.llm_model, safe_content=False) |
| 243 | data = self._load_json_object(raw) |
| 244 | return self._validate_diagnosis(data) |
| 245 | except Exception as exc: |
| 246 | validation_error = str(exc) |
| 247 | logger.warning( |
| 248 | "Doctor diagnosis output invalid; retrying attempt=%s/%s error=%s", |
| 249 | attempt + 1, |
| 250 | MAX_DOCTOR_ATTEMPTS, |
| 251 | validation_error, |
| 252 | ) |
| 253 | |
| 254 | return self.no_action(f"doctor 输出多次校验失败，跳过自动修复: {validation_error}") |
| 255 | |
| 256 | def rewrite_prompt( |
| 257 | self, |
| 258 | *, |
| 259 | prompt: str, |
| 260 | reason: str, |
| 261 | reason_type: str, |
| 262 | stage: str, |
| 263 | context: Optional[Dict[str, Any]] = None, |
| 264 | ) -> Optional[str]: |
| 265 | if reason_type not in REWRITABLE_REASON_TYPES: |
| 266 | return None |
| 267 | |
| 268 | from models.llm_client import LLM |
| 269 | |
| 270 | template = load_prompt("doctor", "rewrite_prompt", "zh") |
| 271 | validation_error = "" |
| 272 | for attempt in range(MAX_DOCTOR_ATTEMPTS): |
| 273 | rewrite_prompt = template.format( |
| 274 | stage=stage, |
| 275 | reason_type=reason_type, |
| 276 | reason=reason, |
| 277 | original_prompt=prompt, |
| 278 | context=self._json_dumps(context), |
| 279 | validation_error=validation_error or "无", |
| 280 | ) |
| 281 | try: |
| 282 | raw = LLM().query(rewrite_prompt, model=self.llm_model, safe_content=False) |
| 283 | data = self._load_json_object(raw) |
| 284 | return self._validate_rewrite(data, prompt) |
| 285 | except Exception as exc: |
| 286 | validation_error = str(exc) |
| 287 | logger.warning( |
| 288 | "Doctor rewrite output invalid; retrying attempt=%s/%s error=%s", |
| 289 | attempt + 1, |
| 290 | MAX_DOCTOR_ATTEMPTS, |
| 291 | validation_error, |
| 292 | ) |
| 293 | |
| 294 | return None |
| 295 | |
| 296 | @staticmethod |
| 297 | def build_rewrite_result( |
| 298 | *, |
| 299 | stage: str, |
| 300 | model: str, |
| 301 | original_prompt: str, |
| 302 | optimized_prompt: str, |
| 303 | error: str, |
| 304 | diagnosis: Dict[str, Any], |
| 305 | ) -> Dict[str, Any]: |
| 306 | """Build the persisted rewrite record in one place so all stages stay consistent.""" |
| 307 | return { |
| 308 | "stage": stage, |
| 309 | "model": model, |
| 310 | "error": error, |
| 311 | "original_prompt": original_prompt, |
| 312 | "optimized_prompt": optimized_prompt, |
| 313 | "doctor_reason_type": diagnosis.get("reason_type", REASON_TYPE_UNKNOWN), |
| 314 | "doctor_reason": diagnosis.get("reason", ""), |
| 315 | "confidence": diagnosis.get("confidence", 0.0), |
| 316 | } |
| 317 | |
| 318 | def maybe_rewrite_prompt( |
| 319 | self, |
| 320 | *, |
| 321 | stage: str, |
| 322 | model: str, |
| 323 | prompt: str, |
| 324 | error: str, |
| 325 | context: Optional[Dict[str, Any]] = None, |
| 326 | ) -> Tuple[Optional[str], Dict[str, Any], Optional[Dict[str, Any]]]: |
| 327 | """Return a rewrite prompt when doctor finds a safe rewrite path.""" |
| 328 | try: |
| 329 | diagnosis = self.diagnose_error( |
| 330 | stage=stage, |
| 331 | model=model, |
| 332 | prompt=prompt, |
| 333 | error=error, |
| 334 | context=context, |
| 335 | ) |
| 336 | if diagnosis.get("tool") != "rewrite_prompt" or not diagnosis.get("should_retry"): |
| 337 | return None, diagnosis, None |
| 338 | rewrite = self.rewrite_prompt( |
| 339 | prompt=prompt, |
| 340 | reason=diagnosis["reason"], |
| 341 | reason_type=diagnosis["reason_type"], |
| 342 | stage=stage, |
| 343 | context=context, |
| 344 | ) |
| 345 | if not rewrite: |
| 346 | return None, diagnosis, None |
| 347 | return rewrite, diagnosis, self.build_rewrite_result( |
| 348 | stage=stage, |
| 349 | model=model, |
| 350 | original_prompt=prompt, |
| 351 | optimized_prompt=rewrite, |
| 352 | error=error, |
| 353 | diagnosis=diagnosis, |
| 354 | ) |
| 355 | except Exception as exc: |
| 356 | logger.warning("Doctor failed and will be skipped: %s", exc, exc_info=True) |
| 357 | return None, self.no_action(f"doctor 执行失败，跳过自动修复: {exc}"), None |
| 358 |