Coverage for agentos/core/middleware.py: 84%

1""" 

2AgentOS v1.1.4 Agent Runtime Middleware Pipeline — 可组合的执行生命周期中间件。 

3 

4在 Agent 执行的每个阶段（pre-LLM / post-LLM / pre-tool / post-tool） 

5插入策略检查、日志、脱敏、预算控制等拦截逻辑。 

6 

7灵感来自 Microsoft Agent Framework 1.0 的 Middleware Pipeline 和 CrewAI Runtime Hooks。 

8""" 

9 

10from __future__ import annotations 

11 

12from abc import ABC, abstractmethod 

13from dataclasses import dataclass, field 

14from enum import Enum 

15from typing import Any, Optional 

16 

17 

18class MiddlewarePhase(str, Enum): 

19 """中间件触发阶段。""" 

20 

21 PRE_LLM = "pre_llm" # LLM 调用前 

22 POST_LLM = "post_llm" # LLM 调用后、输出解析前 

23 PRE_TOOL = "pre_tool" # 工具调用前 

24 POST_TOOL = "post_tool" # 工具调用后 

25 ON_ERROR = "on_error" # 执行出错时 

26 ON_START = "on_start" # Agent 启动时 

27 ON_COMPLETE = "on_complete" # Agent 执行完成时 

28 

29 

30@dataclass 

31class MiddlewareContext: 

32 """中间件执行上下文。""" 

33 

34 phase: MiddlewarePhase 

35 agent_name: str = "" 

36 run_id: str = "" 

37 # LLM 阶段 

38 prompt: Optional[str] = None 

39 model_name: Optional[str] = None 

40 llm_output: Optional[str] = None 

41 # Tool 阶段 

42 tool_name: Optional[str] = None 

43 tool_args: Optional[dict] = None 

44 tool_result: Any = None 

45 # Error 

46 error: Optional[Exception] = None 

47 # 额外元数据 

48 metadata: dict[str, Any] = field(default_factory=dict) 

49 

50 

51@dataclass 

52class MiddlewareDecision: 

53 """中间件决策结果。""" 

54 

55 allow: bool = True 

56 """是否允许继续执行。""" 

57 

58 reason: str = "" 

59 """决策理由。""" 

60 

61 modified_context: Optional[MiddlewareContext] = None 

62 """修改后的上下文（如脱敏后的 prompt）。""" 

63 

64 action: str = "allow" # allow / warn / block / transform / escalate 

65 """决策动作。""" 

66 

67 blocked_by: str = "" 

68 """阻断方名称。""" 

69 

70 

71class AgentMiddleware(ABC): 

72 """Agent 运行时中间件基类。 

73 

74 每个中间件声明自己监听的阶段，通过 process() 返回决策。 

75 返回 MiddlewareDecision(allow=False) 阻断执行链。 

76 

77 __call__ 提供便捷调用：mw(ctx) → process(ctx)。 

78 """ 

79 

80 name: str = "base_middleware" 

81 

82 @property 

83 def phases(self) -> list[MiddlewarePhase]: 

84 """返回此中间件监听的阶段列表。""" 

85 return [MiddlewarePhase.PRE_LLM] 

86 

87 @abstractmethod 

88 async def process(self, ctx: MiddlewareContext) -> MiddlewareDecision: 

89 """处理中间件逻辑。返回决策。""" 

90 ... 

91 

92 async def __call__(self, ctx: MiddlewareContext) -> MiddlewareDecision: 

93 """便捷调用，等价于 process(ctx)。""" 

94 return await self.process(ctx) 

95 

96 

97# ── 内置中间件 ────────────────────────────────────────────────────────────── 

98 

99class PIIMaskingMiddleware(AgentMiddleware): 

100 """PII脱敏中间件：在 pre-LLM 阶段对 prompt 脱敏。""" 

101 

102 name = "pii_masking" 

103 

104 @property 

105 def phases(self) -> list[MiddlewarePhase]: 

106 return [MiddlewarePhase.PRE_LLM] 

107 

108 async def process(self, ctx: MiddlewareContext) -> MiddlewareDecision: 

109 if not ctx.prompt: 

110 return MiddlewareDecision(allow=True) 

111 from agentos.security.guard import PIIDetector 

112 detector = PIIDetector(auto_redact=True) 

113 sanitized, items = detector.redact(ctx.prompt) 

114 count = len(items) 

115 if count > 0: 

116 new_ctx = MiddlewareContext(**{**ctx.__dict__}) 

117 new_ctx.prompt = sanitized 

118 new_ctx.metadata["pii_count"] = count 

119 return MiddlewareDecision( 

120 allow=True, action="transform", 

121 reason=f"Masked {count} PII instances", 

122 modified_context=new_ctx, 

123 ) 

124 return MiddlewareDecision(allow=True) 

125 

126 

127class BudgetGuardMiddleware(AgentMiddleware): 

128 """预算守护中间件：pre-LLM 阶段检查预算。""" 

129 

130 name = "budget_guard" 

131 

132 def __init__(self, tracker=None, budget_limit: float = 0.0, warn_ratio: float = 0.8): 

133 self.tracker = tracker 

134 self.budget_limit = budget_limit 

135 self.warn_ratio = warn_ratio 

136 

137 @property 

138 def phases(self) -> list[MiddlewarePhase]: 

139 return [MiddlewarePhase.PRE_LLM, MiddlewarePhase.PRE_TOOL] 

140 

141 async def process(self, ctx: MiddlewareContext) -> MiddlewareDecision: 

142 if not self.tracker or self.budget_limit <= 0: 

143 return MiddlewareDecision(allow=True) 

144 spent = self.tracker.total_cost 

145 ratio = spent / self.budget_limit 

146 if ratio >= 1.0: 

147 return MiddlewareDecision( 

148 allow=False, action="block", 

149 reason=f"Budget exceeded: ${spent:.4f} / ${self.budget_limit:.2f}", 

150 blocked_by=self.name, 

151 ) 

152 if ratio >= self.warn_ratio: 

153 return MiddlewareDecision( 

154 allow=True, action="warn", 

155 reason=f"Budget warning: {ratio:.0%} used (${spent:.4f} / ${self.budget_limit:.2f})", 

156 ) 

157 return MiddlewareDecision(allow=True) 

158 

159 

160class ToolRiskGuardMiddleware(AgentMiddleware): 

161 """工具风险守护中间件：pre-tool 阶段根据风险等级决定是否阻断。""" 

162 

163 name = "tool_risk_guard" 

164 

165 def __init__(self, max_auto_level: str = "medium"): 

166 from agentos.tools.risk import ToolRiskLevel 

167 self.max_auto_level = ToolRiskLevel(max_auto_level) 

168 

169 @property 

170 def phases(self) -> list[MiddlewarePhase]: 

171 return [MiddlewarePhase.PRE_TOOL] 

172 

173 async def process(self, ctx: MiddlewareContext) -> MiddlewareDecision: 

174 if not ctx.tool_name: 

175 return MiddlewareDecision(allow=True) 

176 

177 from agentos.tools.risk import infer_risk_level 

178 risk = infer_risk_level(ctx.tool_name, tool_args=ctx.tool_args) 

179 

180 if risk.requires_user_confirm(): 

181 return MiddlewareDecision( 

182 allow=False, action="escalate", 

183 reason=f"Tool '{ctx.tool_name}' requires user approval: {risk.description}", 

184 blocked_by=self.name, 

185 ) 

186 

187 levels = ["low", "medium", "high", "critical"] 

188 if levels.index(risk.level.value) > levels.index(self.max_auto_level.value): 

189 return MiddlewareDecision( 

190 allow=False, action="block", 

191 reason=f"Tool '{ctx.tool_name}' risk {risk.level.value} exceeds auto limit {self.max_auto_level.value}", 

192 blocked_by=self.name, 

193 ) 

194 

195 return MiddlewareDecision(allow=True) 

196 

197 

198class AuditLogMiddleware(AgentMiddleware): 

199 """审计日志中间件：在所有阶段记录审计轨迹。""" 

200 

201 name = "audit_log" 

202 

203 @property 

204 def phases(self) -> list[MiddlewarePhase]: 

205 return [ 

206 MiddlewarePhase.ON_START, MiddlewarePhase.PRE_LLM, 

207 MiddlewarePhase.PRE_TOOL, MiddlewarePhase.POST_TOOL, 

208 MiddlewarePhase.ON_ERROR, MiddlewarePhase.ON_COMPLETE, 

209 ] 

210 

211 async def process(self, ctx: MiddlewareContext) -> MiddlewareDecision: 

212 import logging 

213 logger = logging.getLogger("agentos.audit") 

214 logger.info( 

215 f"[{ctx.phase.value}] agent={ctx.agent_name} run={ctx.run_id} " 

216 f"tool={ctx.tool_name or '-'}" 

217 ) 

218 return MiddlewareDecision(allow=True) 

219 

220 

221class TimingMiddleware(AgentMiddleware): 

222 """计时中间件：记录每个阶段的耗时。""" 

223 

224 name = "timing" 

225 

226 def __init__(self): 

227 self.timings: dict[str, float] = {} 

228 self._phase_start: dict[str, float] = {} 

229 

230 @property 

231 def phases(self) -> list[MiddlewarePhase]: 

232 return [ 

233 MiddlewarePhase.ON_START, MiddlewarePhase.PRE_LLM, 

234 MiddlewarePhase.POST_LLM, MiddlewarePhase.PRE_TOOL, 

235 MiddlewarePhase.POST_TOOL, MiddlewarePhase.ON_COMPLETE, 

236 MiddlewarePhase.ON_ERROR, 

237 ] 

238 

239 async def process(self, ctx: MiddlewareContext) -> MiddlewareDecision: 

240 import time 

241 phase_key = f"{ctx.phase.value}.{ctx.tool_name or ctx.agent_name or 'default'}" 

242 self._phase_start[phase_key] = time.monotonic() 

243 self.timings[phase_key] = self.timings.get(phase_key, 0.0) 

244 return MiddlewareDecision(allow=True) 

245 

246 def get_timings(self) -> dict[str, float]: 

247 return dict(self.timings) 

248 

249 def total_ms(self) -> float: 

250 return sum(self.timings.values()) * 1000 

251 

252 

253class RetryMiddleware(AgentMiddleware): 

254 """重试中间件：在 ON_ERROR 阶段自动重试。""" 

255 

256 name = "retry" 

257 

258 def __init__(self, max_retries: int = 2, backoff_base: float = 1.0): 

259 self.max_retries = max_retries 

260 self.backoff_base = backoff_base 

261 self._retry_counts: dict[str, int] = {} 

262 

263 @property 

264 def phases(self) -> list[MiddlewarePhase]: 

265 return [MiddlewarePhase.ON_ERROR] 

266 

267 async def process(self, ctx: MiddlewareContext) -> MiddlewareDecision: 

268 import asyncio 

269 run_key = ctx.run_id or "default" 

270 count = self._retry_counts.get(run_key, 0) 

271 

272 if count >= self.max_retries: 

273 return MiddlewareDecision( 

274 allow=True, action="warn", 

275 reason=f"Max retries ({self.max_retries}) exhausted", 

276 ) 

277 

278 self._retry_counts[run_key] = count + 1 

279 delay = self.backoff_base * (2 ** count) 

280 await asyncio.sleep(delay) 

281 

282 return MiddlewareDecision( 

283 allow=True, action="warn", 

284 reason=f"Retry {count + 1}/{self.max_retries} after {delay:.1f}s", 

285 ) 

286 

287 def reset(self, run_id: str = "") -> None: 

288 if run_id: 

289 self._retry_counts.pop(run_id, None) 

290 else: 

291 self._retry_counts.clear() 

292 

293 

294# ── 中间件管道 ────────────────────────────────────────────────────────────── 

295 

296class MiddlewarePipeline: 

297 """编排多个中间件按阶段执行。 

298 

299 每个阶段： 

300 1. 筛选监听该阶段的中间件 

301 2. 按注册顺序依次执行 

302 3. 任一返回 allow=False 即阻断 

303 4. 若返回 modified_context 则传递给后续中间件 

304 """ 

305 

306 def __init__(self, middlewares: Optional[list[AgentMiddleware]] = None): 

307 self._middlewares: list[AgentMiddleware] = list(middlewares or []) 

308 

309 def add(self, middleware: AgentMiddleware) -> MiddlewarePipeline: 

310 """添加中间件，返回自身以支持链式调用。""" 

311 self._middlewares.append(middleware) 

312 return self 

313 

314 def remove(self, name: str) -> None: 

315 self._middlewares = [m for m in self._middlewares if m.name != name] 

316 

317 @property 

318 def middleware_names(self) -> list[str]: 

319 return [m.name for m in self._middlewares] 

320 

321 async def execute_phase( 

322 self, 

323 phase: MiddlewarePhase, 

324 ctx: MiddlewareContext, 

325 ) -> MiddlewareDecision: 

326 """执行指定阶段的所有中间件。""" 

327 current_ctx = ctx 

328 for mw in self._middlewares: 

329 if phase not in mw.phases: 

330 continue 

331 decision = await mw.process(current_ctx) 

332 if not decision.allow: 

333 decision.blocked_by = mw.name 

334 return decision 

335 if decision.modified_context: 

336 current_ctx = decision.modified_context 

337 return MiddlewareDecision(allow=True, modified_context=current_ctx) 

338 

339 async def on_start(self, ctx: MiddlewareContext) -> MiddlewareDecision: 

340 return await self.execute_phase(MiddlewarePhase.ON_START, ctx) 

341 

342 async def pre_llm(self, ctx: MiddlewareContext) -> MiddlewareDecision: 

343 return await self.execute_phase(MiddlewarePhase.PRE_LLM, ctx) 

344 

345 async def post_llm(self, ctx: MiddlewareContext) -> MiddlewareDecision: 

346 return await self.execute_phase(MiddlewarePhase.POST_LLM, ctx) 

347 

348 async def pre_tool(self, ctx: MiddlewareContext) -> MiddlewareDecision: 

349 return await self.execute_phase(MiddlewarePhase.PRE_TOOL, ctx) 

350 

351 async def post_tool(self, ctx: MiddlewareContext) -> MiddlewareDecision: 

352 return await self.execute_phase(MiddlewarePhase.POST_TOOL, ctx) 

353 

354 async def on_error(self, ctx: MiddlewareContext) -> MiddlewareDecision: 

355 return await self.execute_phase(MiddlewarePhase.ON_ERROR, ctx) 

356 

357 async def on_complete(self, ctx: MiddlewareContext) -> MiddlewareDecision: 

358 return await self.execute_phase(MiddlewarePhase.ON_COMPLETE, ctx) 

359 

360 async def run( 

361 self, 

362 ctx: MiddlewareContext, 

363 phases: list[MiddlewarePhase] | None = None, 

364 ) -> MiddlewareDecision: 

365 """便捷方法：按阶段列表依次执行管道。 

366 

367 phases 默认为完整的生命周期序列。 

368 """ 

369 if phases is None: 

370 phases = [ 

371 MiddlewarePhase.ON_START, MiddlewarePhase.PRE_LLM, 

372 MiddlewarePhase.POST_LLM, MiddlewarePhase.ON_COMPLETE, 

373 ] 

374 decision = MiddlewareDecision(allow=True, modified_context=ctx) 

375 for phase in phases: 

376 current_ctx = decision.modified_context or ctx 

377 decision = await self.execute_phase(phase, current_ctx) 

378 if not decision.allow: 

379 return decision 

380 return decision