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

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 StrEnum 

15from typing import Any 

16 

17 

18class MiddlewarePhase(StrEnum): 

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: str | None = None 

39 model_name: str | None = None 

40 llm_output: str | None = None 

41 # Tool 阶段 

42 tool_name: str | None = None 

43 tool_args: dict | None = None 

44 tool_result: Any = None 

45 # Error 

46 error: Exception | None = 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: MiddlewareContext | None = 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 

99 

100class PIIMaskingMiddleware(AgentMiddleware): 

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

102 

103 name = "pii_masking" 

104 

105 @property 

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

107 return [MiddlewarePhase.PRE_LLM] 

108 

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

110 if not ctx.prompt: 

111 return MiddlewareDecision(allow=True) 

112 from agentos.security.guard import PIIDetector 

113 

114 detector = PIIDetector(auto_redact=True) 

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

116 count = len(items) 

117 if count > 0: 

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

119 new_ctx.prompt = sanitized 

120 new_ctx.metadata["pii_count"] = count 

121 return MiddlewareDecision( 

122 allow=True, 

123 action="transform", 

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

125 modified_context=new_ctx, 

126 ) 

127 return MiddlewareDecision(allow=True) 

128 

129 

130class BudgetGuardMiddleware(AgentMiddleware): 

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

132 

133 name = "budget_guard" 

134 

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

136 self.tracker = tracker 

137 self.budget_limit = budget_limit 

138 self.warn_ratio = warn_ratio 

139 

140 @property 

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

142 return [MiddlewarePhase.PRE_LLM, MiddlewarePhase.PRE_TOOL] 

143 

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

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

146 return MiddlewareDecision(allow=True) 

147 spent = self.tracker.total_cost 

148 ratio = spent / self.budget_limit 

149 if ratio >= 1.0: 

150 return MiddlewareDecision( 

151 allow=False, 

152 action="block", 

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

154 blocked_by=self.name, 

155 ) 

156 if ratio >= self.warn_ratio: 

157 return MiddlewareDecision( 

158 allow=True, 

159 action="warn", 

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

161 ) 

162 return MiddlewareDecision(allow=True) 

163 

164 

165class ToolRiskGuardMiddleware(AgentMiddleware): 

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

167 

168 name = "tool_risk_guard" 

169 

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

171 from agentos.tools.risk import ToolRiskLevel 

172 

173 self.max_auto_level = ToolRiskLevel(max_auto_level) 

174 

175 @property 

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

177 return [MiddlewarePhase.PRE_TOOL] 

178 

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

180 if not ctx.tool_name: 

181 return MiddlewareDecision(allow=True) 

182 

183 from agentos.tools.risk import infer_risk_level 

184 

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

186 

187 if risk.requires_user_confirm(): 

188 return MiddlewareDecision( 

189 allow=False, 

190 action="escalate", 

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

192 blocked_by=self.name, 

193 ) 

194 

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

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

197 return MiddlewareDecision( 

198 allow=False, 

199 action="block", 

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

201 blocked_by=self.name, 

202 ) 

203 

204 return MiddlewareDecision(allow=True) 

205 

206 

207class AuditLogMiddleware(AgentMiddleware): 

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

209 

210 name = "audit_log" 

211 

212 @property 

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

214 return [ 

215 MiddlewarePhase.ON_START, 

216 MiddlewarePhase.PRE_LLM, 

217 MiddlewarePhase.PRE_TOOL, 

218 MiddlewarePhase.POST_TOOL, 

219 MiddlewarePhase.ON_ERROR, 

220 MiddlewarePhase.ON_COMPLETE, 

221 ] 

222 

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

224 import logging 

225 

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

227 logger.info( 

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

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

230 ) 

231 return MiddlewareDecision(allow=True) 

232 

233 

234class TimingMiddleware(AgentMiddleware): 

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

236 

237 name = "timing" 

238 

239 def __init__(self): 

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

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

242 

243 @property 

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

245 return [ 

246 MiddlewarePhase.ON_START, 

247 MiddlewarePhase.PRE_LLM, 

248 MiddlewarePhase.POST_LLM, 

249 MiddlewarePhase.PRE_TOOL, 

250 MiddlewarePhase.POST_TOOL, 

251 MiddlewarePhase.ON_COMPLETE, 

252 MiddlewarePhase.ON_ERROR, 

253 ] 

254 

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

256 import time 

257 

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

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

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

261 return MiddlewareDecision(allow=True) 

262 

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

264 return dict(self.timings) 

265 

266 def total_ms(self) -> float: 

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

268 

269 

270class RetryMiddleware(AgentMiddleware): 

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

272 

273 name = "retry" 

274 

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

276 self.max_retries = max_retries 

277 self.backoff_base = backoff_base 

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

279 

280 @property 

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

282 return [MiddlewarePhase.ON_ERROR] 

283 

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

285 import asyncio 

286 

287 run_key = ctx.run_id or "default" 

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

289 

290 if count >= self.max_retries: 

291 return MiddlewareDecision( 

292 allow=True, 

293 action="warn", 

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

295 ) 

296 

297 self._retry_counts[run_key] = count + 1 

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

299 await asyncio.sleep(delay) 

300 

301 return MiddlewareDecision( 

302 allow=True, 

303 action="warn", 

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

305 ) 

306 

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

308 if run_id: 

309 self._retry_counts.pop(run_id, None) 

310 else: 

311 self._retry_counts.clear() 

312 

313 

314# ── 中间件管道 ────────────────────────────────────────────────────────────── 

315 

316 

317class MiddlewarePipeline: 

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

319 

320 每个阶段： 

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

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

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

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

325 """ 

326 

327 def __init__(self, middlewares: list[AgentMiddleware] | None = None): 

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

329 

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

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

332 self._middlewares.append(middleware) 

333 return self 

334 

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

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

337 

338 @property 

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

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

341 

342 async def execute_phase( 

343 self, 

344 phase: MiddlewarePhase, 

345 ctx: MiddlewareContext, 

346 ) -> MiddlewareDecision: 

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

348 current_ctx = ctx 

349 for mw in self._middlewares: 

350 if phase not in mw.phases: 

351 continue 

352 decision = await mw.process(current_ctx) 

353 if not decision.allow: 

354 decision.blocked_by = mw.name 

355 return decision 

356 if decision.modified_context: 

357 current_ctx = decision.modified_context 

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

359 

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

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

362 

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

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

365 

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

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

368 

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

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

371 

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

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

374 

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

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

377 

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

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

380 

381 async def run( 

382 self, 

383 ctx: MiddlewareContext, 

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

385 ) -> MiddlewareDecision: 

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

387 

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

389 """ 

390 if phases is None: 

391 phases = [ 

392 MiddlewarePhase.ON_START, 

393 MiddlewarePhase.PRE_LLM, 

394 MiddlewarePhase.POST_LLM, 

395 MiddlewarePhase.ON_COMPLETE, 

396 ] 

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

398 for phase in phases: 

399 current_ctx = decision.modified_context or ctx 

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

401 if not decision.allow: 

402 return decision 

403 return decision