Coverage for agentos/core/loop.py: 33%

313 statements  

« prev     ^ index     » next       coverage.py v7.14.3, created at 2026-07-06 10:59 +0800

1""" 

2AgentOS v0.70 核心循环 — Gemini + Metrics + CostAnalytics 集成版。 

3v0.40: Swarm多Agent并行、Agent间通信、语义缓存、任务队列。 

4v0.70: MetricsCollector、CostAnalytics实时监控。 

5""" 

6 

7from __future__ import annotations 

8 

9import asyncio 

10import time 

11from dataclasses import dataclass, field 

12from enum import Enum 

13from typing import Callable 

14 

15from agentos.core.context import ContextManager 

16from agentos.tools.registry import ToolRegistry 

17from agentos.models.router import ModelRouter, AllModelsFailed 

18from agentos.security.sandbox import SandboxManager 

19from agentos.observability.tracer import Tracer 

20from agentos.observability.metrics import MetricsCollector 

21from agentos.observability.cost_analytics import CostAnalytics 

22from agentos.core.streaming import StreamChunk 

23from agentos.storage.base import CheckpointStore 

24from agentos.checkpoint.base import Checkpoint, CheckpointMetadata, CheckpointBackend 

25from agentos.cost.tracker import CostTracker 

26from agentos.swarm.coordinator import SwarmCoordinator, SwarmTopology, AgentRole, SwarmResult 

27from agentos.comm.layer import CommunicationLayer 

28from agentos.tools.audit_logger import AuditLogger, AuditEvent, Severity 

29from agentos.tools.rate_limiter import TokenBucket 

30 

31 

32class LoopState(str, Enum): 

33 

34 """主循环状态。""" 

35 

36 RUNNING = "running" 

37 PAUSED = "paused" 

38 WAITING_HUMAN = "waiting_human" 

39 COMPLETED = "completed" 

40 FAILED = "failed" 

41 CANCELLED = "cancelled" 

42 

43 

44@dataclass 

45class AgentResult: 

46 """Agent 主循环的最终运行结果。""" 

47 

48 output: str 

49 iterations: int 

50 tokens_used: dict[str, int] = field(default_factory=dict) 

51 cost_usd: float = 0.0 

52 duration_ms: float = 0.0 

53 tool_calls_total: int = 0 

54 reflections_count: int = 0 

55 human_interrupts: int = 0 

56 final_state: LoopState = LoopState.COMPLETED 

57 error: str | None = None 

58 # v0.40 

59 swarm_result: SwarmResult | None = None 

60 cache_hit: bool = False 

61 

62 

63@dataclass 

64class LoopConfig: 

65 """Agent 主循环的运行时配置。""" 

66 

67 max_iterations: int = 100 

68 max_retries_per_step: int = 2 

69 step_timeout_seconds: int = 120 

70 enable_streaming: bool = False 

71 enable_checkpoints: bool = True 

72 checkpoint_interval: int = 5 

73 # v0.30 

74 enable_reflection: bool = True 

75 reflection_frequency: int = 3 

76 max_reflection_loops: int = 3 

77 enable_self_critique: bool = True 

78 enable_human_in_the_loop: bool = False 

79 human_approval_trigger: str = "high_risk" 

80 enable_cost_tracking: bool = True 

81 auto_select_model: bool = True 

82 # v0.40 

83 enable_swarm: bool = False 

84 swarm_topology: str = "sequential" 

85 swarm_roles: list[AgentRole] = field(default_factory=list) 

86 max_parallel_agents: int = 4 

87 enable_comm_layer: bool = True 

88 enable_semantic_cache: bool = True 

89 # v1.11.0 — long-running task support 

90 checkpoint_backend: CheckpointBackend | None = None # Full checkpoint backend for crash recovery 

91 enable_auto_paging: bool = True # Auto-evict old memories when context fills 

92 auto_page_threshold: float = 0.85 # Page out at 85% context window usage 

93 

94 

95class MaxIterationsExceeded(Exception): 

96 

97 """超出最大迭代次数异常。""" 

98 

99 pass 

100 

101 

102class HumanInterruptNeeded(Exception): 

103 

104 """需要人工介入异常。""" 

105 

106 def __init__(self, message: str, context: dict | None = None): 

107 super().__init__(message) 

108 self.context = context or {} 

109 

110 

111@dataclass 

112class ReflectionResult: 

113 """反思结果。""" 

114 quality_score: float 

115 issues: list[str] 

116 suggestions: list[str] 

117 should_continue: bool 

118 new_plan: str | None = None 

119 

120 

121class AgentLoop: 

122 """v0.30 核心循环 — Reflection + HITL + Self-Critique + 自动路由 + 成本追踪。""" 

123 

124 def __init__( 

125 self, 

126 model_router: ModelRouter, 

127 tool_registry: ToolRegistry, 

128 context_manager: ContextManager, 

129 sandbox_manager: SandboxManager | None = None, 

130 tracer: Tracer | None = None, 

131 checkpoint_store: CheckpointStore | None = None, 

132 checkpoint_backend: CheckpointBackend | None = None, 

133 cost_tracker: CostTracker | None = None, 

134 config: LoopConfig | None = None, 

135 on_iteration: Callable | None = None, 

136 on_stream: Callable[[StreamChunk], None] | None = None, 

137 on_human_interrupt: Callable[[str, dict], str | None] | None = None, 

138 on_reflection: Callable[[ReflectionResult], None] | None = None, 

139 metrics_collector: MetricsCollector | None = None, 

140 cost_analytics: CostAnalytics | None = None, 

141 audit_logger: AuditLogger | None = None, 

142 rate_limiter: TokenBucket | None = None, 

143 ): 

144 self.model_router = model_router 

145 self.tool_registry = tool_registry 

146 self.context_manager = context_manager 

147 self.sandbox_manager = sandbox_manager 

148 self.tracer = tracer or Tracer.noop() 

149 self.checkpoint_store = checkpoint_store 

150 self.cost_tracker = cost_tracker or CostTracker.noop() 

151 self.config = config or LoopConfig() 

152 self.checkpoint_backend = checkpoint_backend # v1.11.0 full checkpoint integration 

153 self._auto_page_callback: Callable | None = None # v1.11.0 auto-paging callback 

154 self.on_iteration = on_iteration 

155 self.on_stream = on_stream 

156 self.on_human_interrupt = on_human_interrupt 

157 self.on_reflection = on_reflection 

158 self.metrics = metrics_collector or MetricsCollector() 

159 self.cost_analytics = cost_analytics or CostAnalytics(self.cost_tracker) 

160 self.audit_logger = audit_logger 

161 self.rate_limiter = rate_limiter 

162 self._cancelled = False 

163 self._reflection_history: list[ReflectionResult] = [] 

164 self._human_interrupts = 0 

165 

166 # ── 运行入口 ────────────────────────────────── 

167 

168 async def run(self, task: str, session_id: str = "") -> AgentResult: 

169 start_time = time.time() 

170 await self.context_manager.init_session(session_id, task) 

171 

172 if self.audit_logger: 

173 self.audit_logger.log(event=AuditEvent( 

174 actor="agentos", 

175 action="loop.start", 

176 resource=session_id, 

177 outcome="initiated", 

178 details={"task": task[:200]}, 

179 )) 

180 

181 if self.config.auto_select_model: 

182 await self._auto_route_model(task) 

183 

184 iteration = await self._try_restore(session_id) 

185 tool_calls_total = 0 

186 reflection_loops = 0 

187 

188 while iteration < self.config.max_iterations and not self._cancelled: 

189 iteration += 1 

190 

191 if self.config.enable_reflection and iteration % self.config.reflection_frequency == 0: 

192 with self.tracer.step("reflection"): 

193 reflection = await self._reflect(session_id) 

194 self._reflection_history.append(reflection) 

195 if not reflection.should_continue and reflection_loops < self.config.max_reflection_loops: 

196 reflection_loops += 1 

197 if reflection.new_plan: 

198 self.context_manager.update_plan(reflection.new_plan) 

199 continue 

200 

201 try: 

202 with self.tracer.step(f"loop_{iteration}"): 

203 step_result = await self._execute_step_sync(iteration, session_id) 

204 

205 if step_result.is_terminal: 

206 duration_ms = (time.time() - start_time) * 1000 

207 if self.audit_logger: 

208 self.audit_logger.log(event=AuditEvent( 

209 actor="agentos", 

210 action="loop.complete", 

211 resource=session_id, 

212 outcome="success", 

213 details={"iterations": iteration, "duration_ms": duration_ms}, 

214 )) 

215 return AgentResult( 

216 output=step_result.content, 

217 iterations=iteration, 

218 tokens_used=self.tracer.token_summary(), 

219 cost_usd=self.cost_tracker.total_cost, 

220 duration_ms=duration_ms, 

221 tool_calls_total=tool_calls_total, 

222 reflections_count=len(self._reflection_history), 

223 human_interrupts=self._human_interrupts, 

224 ) 

225 

226 if step_result.tool_results: 

227 tool_calls_total += len(step_result.tool_results) 

228 

229 if self.on_iteration: 

230 self.on_iteration(iteration, step_result.tool_results or []) 

231 

232 except HumanInterruptNeeded as e: 

233 self._human_interrupts += 1 

234 if self.audit_logger: 

235 self.audit_logger.log(event=AuditEvent( 

236 actor="agentos", 

237 action="loop.human_interrupt", 

238 resource=session_id, 

239 outcome="paused", 

240 severity=Severity.WARNING, 

241 details={"interrupt_count": self._human_interrupts, "message": str(e)[:200]}, 

242 )) 

243 if self.on_human_interrupt: 

244 feedback = self.on_human_interrupt(str(e), e.context) 

245 if feedback: 

246 self.context_manager.append_user_message(feedback) 

247 continue 

248 

249 except StepTimeoutError: 

250 return AgentResult(output="", iterations=iteration, final_state=LoopState.FAILED, error="Step timeout") 

251 

252 if self.config.enable_checkpoints and iteration % self.config.checkpoint_interval == 0: 

253 await self._save_checkpoint(session_id, iteration) 

254 

255 raise MaxIterationsExceeded(f"超过 {self.config.max_iterations} 步") 

256 

257 # ── Reflection ──────────────────────────────── 

258 

259 async def _reflect(self, session_id: str) -> ReflectionResult: 

260 prompt = f"""你是一个反思者。审核以下Agent执行过程: 

261 

262任务: {self.context_manager.current_task} 

263已执行: {self.context_manager.step_count} 步 

264 

265评估并返回JSON: 

266{{"quality_score": 0.0-1.0, "issues": [...], "suggestions": [...], "should_continue": true/false, "new_plan": "如果调整,新计划"}}""" 

267 

268 resp = await self.model_router.call_simple(prompt) 

269 try: 

270 import json 

271 d = json.loads(resp) 

272 result = ReflectionResult( 

273 quality_score=d.get("quality_score", 0.5), 

274 issues=d.get("issues", []), 

275 suggestions=d.get("suggestions", []), 

276 should_continue=d.get("should_continue", True), 

277 new_plan=d.get("new_plan"), 

278 ) 

279 except Exception: 

280 result = ReflectionResult(0.5, [], [], True) 

281 if self.on_reflection: 

282 self.on_reflection(result) 

283 return result 

284 

285 # ── Self-Critique ───────────────────────────── 

286 

287 async def _self_critique(self, text: str) -> str: 

288 if not self.config.enable_self_critique: 

289 return text 

290 prompt = f"""审视以下回答,找出逻辑错误或不准确之处。如果已足够好就原样返回。 

291 

292{text[:3000]}""" 

293 improved = await self.model_router.call_simple(prompt) 

294 return improved or text 

295 

296 # ── Auto Route ──────────────────────────────── 

297 

298 async def _auto_route_model(self, task: str): 

299 score = self._estimate_complexity(task) 

300 if score > 0.7: 

301 self.model_router.set_preferred("deepseek-r1") 

302 elif score > 0.4: 

303 self.model_router.set_preferred("kimi-k2.6") 

304 else: 

305 self.model_router.set_preferred("deepseek-v3.1") 

306 

307 def _estimate_complexity(self, task: str) -> float: 

308 kw = ["分析", "对比", "设计", "架构", "review", "refactor", "实现", "优化", "诊断", "troubleshoot", "debug", "deploy", "migrate", "安全", "security"] 

309 score = sum(0.15 for k in kw if k in task.lower()) 

310 return min(score + min(len(task) / 2000, 0.3), 1.0) 

311 

312 # ── 步骤执行 ────────────────────────────────── 

313 

314 async def _execute_step_sync(self, iteration: int, session_id: str) -> "StepResult": 

315 last_error = None 

316 for attempt in range(self.config.max_retries_per_step + 1): 

317 try: 

318 return await asyncio.wait_for(self._do_step(iteration, session_id), timeout=self.config.step_timeout_seconds) 

319 except asyncio.TimeoutError: 

320 last_error = StepTimeoutError(f"Step {iteration} timeout") 

321 except AllModelsFailed as e: 

322 last_error = e 

323 await asyncio.sleep(2 ** attempt) 

324 raise last_error 

325 

326 async def _do_step(self, iteration: int, session_id: str) -> "StepResult": 

327 ctx = self.context_manager.build_context( 

328 model_type=self.model_router.model_type, 

329 tools=self.tool_registry.get_schemas_for_model(self.model_router.model_type), 

330 ) 

331 

332 # v1.11.0 — auto-page old memories if context nearing limit 

333 if self.config.enable_auto_paging and self._auto_page_callback: 

334 usage_ratio = self.context_manager.estimate_context_usage() 

335 if usage_ratio > self.config.auto_page_threshold: 

336 page_count = await self._auto_page_callback(usage_ratio) 

337 

338 # v1.16.6 — rate limiting before model calls 

339 if self.rate_limiter and not self.rate_limiter.try_acquire("model_call"): 

340 raise StepTimeoutError(f"Rate limit exceeded for model call at step {iteration}") 

341 

342 resp = await self.model_router.call(ctx) 

343 

344 # 成本记录 

345 if self.config.enable_cost_tracking and hasattr(resp, "usage"): 

346 self.cost_tracker.record(self.model_router.current_model, resp.usage) 

347 

348 if not resp.tool_calls: 

349 if self.config.enable_self_critique: 

350 improved = await self._self_critique(resp.content) 

351 return StepResult(content=improved, is_terminal=True) 

352 return StepResult(content=resp.content, is_terminal=True) 

353 

354 # HITL 检查 

355 if self.config.enable_human_in_the_loop: 

356 for tc in resp.tool_calls: 

357 if self._is_high_risk(tc): 

358 raise HumanInterruptNeeded(f"高风险操作需确认: {tc.name}", {"tool": tc.name, "args": tc.arguments}) 

359 

360 groups = self._group_independent_calls(resp.tool_calls) 

361 all_results = [] 

362 for group in groups: 

363 sandbox = self.sandbox_manager.get_sandbox(session_id) if self.sandbox_manager else None 

364 batch_results = await self.tool_registry.execute_batch(group, sandbox=sandbox) 

365 all_results.extend(batch_results) 

366 

367 self.context_manager.append_tool_results(all_results) 

368 return StepResult(content="", is_terminal=False, tool_results=all_results) 

369 

370 def _is_high_risk(self, tc) -> bool: 

371 risky = ["delete", "rm", "uninstall", "format", "sudo", "kill", "drop"] 

372 name = tc.name.lower() if hasattr(tc, "name") else tc.get("name", "").lower() 

373 return any(r in name for r in risky) 

374 

375 def _group_independent_calls(self, tool_calls: list) -> list[list]: 

376 if len(tool_calls) <= 1: 

377 return [tool_calls] if tool_calls else [] 

378 groups: list[list] = [] 

379 for call in tool_calls: 

380 for group in groups: 

381 if not self._has_conflict(call, group): 

382 group.append(call) 

383 break 

384 else: 

385 groups.append([call]) 

386 return groups 

387 

388 def _has_conflict(self, call, group: list) -> bool: 

389 write_paths = set() 

390 for tc in group: 

391 tool = self.tool_registry.get(tc.name) 

392 if tool and tool.is_write_operation(tc.arguments): 

393 if p := tool.extract_target_path(tc.arguments): 

394 write_paths.add(p) 

395 cur = self.tool_registry.get(call.name) 

396 if cur and cur.is_read_operation(call.arguments): 

397 return cur.extract_target_path(call.arguments) in write_paths 

398 return False 

399 

400 # ── v1.11.0 全量 Checkpoint (完整状态快照) ──── 

401 

402 async def _save_checkpoint(self, session_id: str, iteration: int): 

403 """Save full runtime state snapshot via CheckpointBackend.""" 

404 backend = self.checkpoint_backend 

405 if not backend: 

406 # Fallback to thin CheckpointStore 

407 if not self.checkpoint_store: 

408 return 

409 snap = { 

410 "session_id": session_id, "iteration": iteration, 

411 "messages": [{"role": m.role, "content": m.content} for m in self.context_manager._messages], 

412 "timestamp": time.time(), 

413 } 

414 await self.checkpoint_store.save(session_id, snap) 

415 return 

416 

417 # Full checkpoint via CheckpointBackend 

418 try: 

419 from datetime import datetime, timezone 

420 

421 checkpoint_id = f"ckpt-{session_id}-{iteration:06d}" 

422 parent_id = getattr(self, '_last_checkpoint_id', None) 

423 

424 cp = Checkpoint( 

425 metadata=CheckpointMetadata( 

426 thread_id=session_id, 

427 checkpoint_id=checkpoint_id, 

428 step=iteration, 

429 parent_checkpoint_id=parent_id, 

430 created_at=datetime.now(timezone.utc).isoformat(), 

431 tags=["auto", f"iter_{iteration}"], 

432 ), 

433 messages=[{"role": m.role, "content": m.content} for m in self.context_manager._messages], 

434 state={ 

435 "iteration": iteration, 

436 "task": self.context_manager.current_task, 

437 "session_id": session_id, 

438 "cost_usd": self.cost_tracker.total_cost, 

439 "reflections": len(self._reflection_history), 

440 "human_interrupts": self._human_interrupts, 

441 "loop_state": self.context_manager.current_state if hasattr(self.context_manager, 'current_state') else "running", 

442 }, 

443 tools_result={}, 

444 next_node="loop", 

445 ) 

446 await backend.put(cp) 

447 self._last_checkpoint_id = checkpoint_id 

448 

449 except Exception: 

450 pass # Checkpoint failure must not crash the loop 

451 

452 async def _try_restore(self, session_id: str) -> int: 

453 """Restore full state from last checkpoint. Returns iteration to resume from.""" 

454 backend = self.checkpoint_backend 

455 if not backend: 

456 # Fallback to thin CheckpointStore 

457 if not self.checkpoint_store or not self.config.enable_checkpoints: 

458 return 0 

459 snap = await self.checkpoint_store.load(session_id) 

460 if not snap: 

461 return 0 

462 iter_count = snap.get("iteration", 0) 

463 if iter_count > 0: 

464 msgs = snap.get("messages", []) 

465 for msg in msgs: 

466 self.context_manager.append_message(msg["role"], msg["content"]) 

467 return iter_count 

468 

469 if not self.config.enable_checkpoints: 

470 return 0 

471 

472 try: 

473 latest = await backend.get_latest(session_id) 

474 if not latest: 

475 return 0 

476 self._last_checkpoint_id = latest.metadata.checkpoint_id 

477 iter_count = latest.metadata.step 

478 

479 # Restore messages 

480 for msg in latest.messages: 

481 self.context_manager.append_message(msg.get("role", "user"), msg.get("content", "")) 

482 

483 # Restore state 

484 state = latest.state 

485 self._human_interrupts = state.get("human_interrupts", 0) 

486 

487 return iter_count 

488 except Exception: 

489 return 0 

490 

491 def set_auto_paging(self, callback: Callable): 

492 """Register callback for automatic memory paging (v1.11.0).""" 

493 self._auto_page_callback = callback 

494 

495 def cancel(self): 

496 self._cancelled = True 

497 

498 # ── v0.40 Swarm执行 ────────────────────────── 

499 

500 async def run_swarm(self, task: str, roles: list[AgentRole] | None = None) -> AgentResult: 

501 """以Swarm模式执行任务 — 多Agent协作。""" 

502 start_time = time.time() 

503 roles = roles or self.config.swarm_roles 

504 if not roles: 

505 return AgentResult(output="[Swarm] No roles defined", iterations=0, final_state=LoopState.FAILED, error="No roles") 

506 

507 topology = SwarmTopology(self.config.swarm_topology) 

508 comm_layer = CommunicationLayer() if self.config.enable_comm_layer else None 

509 

510 swarm = SwarmCoordinator( 

511 router=self.model_router, 

512 tool_registry=self.tool_registry, 

513 topology=topology, 

514 max_parallel=self.config.max_parallel_agents, 

515 ) 

516 swarm.register_roles(roles) 

517 

518 swarm_result = await swarm.execute(task, roles) 

519 duration_ms = (time.time() - start_time) * 1000 

520 

521 return AgentResult( 

522 output=swarm_result.combined_output, 

523 iterations=1, 

524 cost_usd=self.cost_tracker.total_cost, 

525 duration_ms=duration_ms, 

526 tool_calls_total=0, 

527 reflections_count=0, 

528 human_interrupts=0, 

529 final_state=LoopState.COMPLETED, 

530 swarm_result=swarm_result, 

531 ) 

532 

533 

534class StepTimeoutError(Exception): 

535 

536 """步骤超时异常。""" 

537 

538 pass 

539 

540 

541@dataclass 

542class StepResult: 

543 """步骤执行结果。""" 

544 content: str 

545 is_terminal: bool = False 

546 tool_results: list | None = None