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

316 statements  

« prev     ^ index     » next       coverage.py v7.14.3, created at 2026-07-06 08:01 +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 Any, AsyncIterator, 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, StreamEvent 

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, MessageBus 

27from agentos.comm.layer import CommunicationLayer 

28from agentos.cache.llm_cache import LLMCache 

29from agentos.multimodal.manager import MultimodalManager 

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

31from agentos.tools.rate_limiter import TokenBucket 

32 

33 

34class LoopState(str, Enum): 

35 

36 """主循环状态。""" 

37 

38 RUNNING = "running" 

39 PAUSED = "paused" 

40 WAITING_HUMAN = "waiting_human" 

41 COMPLETED = "completed" 

42 FAILED = "failed" 

43 CANCELLED = "cancelled" 

44 

45 

46@dataclass 

47class AgentResult: 

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

49 

50 output: str 

51 iterations: int 

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

53 cost_usd: float = 0.0 

54 duration_ms: float = 0.0 

55 tool_calls_total: int = 0 

56 reflections_count: int = 0 

57 human_interrupts: int = 0 

58 final_state: LoopState = LoopState.COMPLETED 

59 error: str | None = None 

60 # v0.40 

61 swarm_result: SwarmResult | None = None 

62 cache_hit: bool = False 

63 

64 

65@dataclass 

66class LoopConfig: 

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

68 

69 max_iterations: int = 100 

70 max_retries_per_step: int = 2 

71 step_timeout_seconds: int = 120 

72 enable_streaming: bool = False 

73 enable_checkpoints: bool = True 

74 checkpoint_interval: int = 5 

75 # v0.30 

76 enable_reflection: bool = True 

77 reflection_frequency: int = 3 

78 max_reflection_loops: int = 3 

79 enable_self_critique: bool = True 

80 enable_human_in_the_loop: bool = False 

81 human_approval_trigger: str = "high_risk" 

82 enable_cost_tracking: bool = True 

83 auto_select_model: bool = True 

84 # v0.40 

85 enable_swarm: bool = False 

86 swarm_topology: str = "sequential" 

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

88 max_parallel_agents: int = 4 

89 enable_comm_layer: bool = True 

90 enable_semantic_cache: bool = True 

91 # v1.11.0 — long-running task support 

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

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

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

95 

96 

97class MaxIterationsExceeded(Exception): 

98 

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

100 

101 pass 

102 

103 

104class HumanInterruptNeeded(Exception): 

105 

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

107 

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

109 super().__init__(message) 

110 self.context = context or {} 

111 

112 

113@dataclass 

114class ReflectionResult: 

115 """反思结果。""" 

116 quality_score: float 

117 issues: list[str] 

118 suggestions: list[str] 

119 should_continue: bool 

120 new_plan: str | None = None 

121 

122 

123class AgentLoop: 

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

125 

126 def __init__( 

127 self, 

128 model_router: ModelRouter, 

129 tool_registry: ToolRegistry, 

130 context_manager: ContextManager, 

131 sandbox_manager: SandboxManager | None = None, 

132 tracer: Tracer | None = None, 

133 checkpoint_store: CheckpointStore | None = None, 

134 checkpoint_backend: CheckpointBackend | None = None, 

135 cost_tracker: CostTracker | None = None, 

136 config: LoopConfig | None = None, 

137 on_iteration: Callable | None = None, 

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

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

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

141 metrics_collector: MetricsCollector | None = None, 

142 cost_analytics: CostAnalytics | None = None, 

143 audit_logger: AuditLogger | None = None, 

144 rate_limiter: TokenBucket | None = None, 

145 ): 

146 self.model_router = model_router 

147 self.tool_registry = tool_registry 

148 self.context_manager = context_manager 

149 self.sandbox_manager = sandbox_manager 

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

151 self.checkpoint_store = checkpoint_store 

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

153 self.config = config or LoopConfig() 

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

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

156 self.on_iteration = on_iteration 

157 self.on_stream = on_stream 

158 self.on_human_interrupt = on_human_interrupt 

159 self.on_reflection = on_reflection 

160 self.metrics = metrics_collector or MetricsCollector() 

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

162 self.audit_logger = audit_logger 

163 self.rate_limiter = rate_limiter 

164 self._cancelled = False 

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

166 self._human_interrupts = 0 

167 

168 # ── 运行入口 ────────────────────────────────── 

169 

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

171 start_time = time.time() 

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

173 

174 if self.audit_logger: 

175 self.audit_logger.log(event=AuditEvent( 

176 actor="agentos", 

177 action="loop.start", 

178 resource=session_id, 

179 outcome="initiated", 

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

181 )) 

182 

183 if self.config.auto_select_model: 

184 await self._auto_route_model(task) 

185 

186 iteration = await self._try_restore(session_id) 

187 tool_calls_total = 0 

188 reflection_loops = 0 

189 

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

191 iteration += 1 

192 

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

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

195 reflection = await self._reflect(session_id) 

196 self._reflection_history.append(reflection) 

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

198 reflection_loops += 1 

199 if reflection.new_plan: 

200 self.context_manager.update_plan(reflection.new_plan) 

201 continue 

202 

203 try: 

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

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

206 

207 if step_result.is_terminal: 

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

209 if self.audit_logger: 

210 self.audit_logger.log(event=AuditEvent( 

211 actor="agentos", 

212 action="loop.complete", 

213 resource=session_id, 

214 outcome="success", 

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

216 )) 

217 return AgentResult( 

218 output=step_result.content, 

219 iterations=iteration, 

220 tokens_used=self.tracer.token_summary(), 

221 cost_usd=self.cost_tracker.total_cost, 

222 duration_ms=duration_ms, 

223 tool_calls_total=tool_calls_total, 

224 reflections_count=len(self._reflection_history), 

225 human_interrupts=self._human_interrupts, 

226 ) 

227 

228 if step_result.tool_results: 

229 tool_calls_total += len(step_result.tool_results) 

230 

231 if self.on_iteration: 

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

233 

234 except HumanInterruptNeeded as e: 

235 self._human_interrupts += 1 

236 if self.audit_logger: 

237 self.audit_logger.log(event=AuditEvent( 

238 actor="agentos", 

239 action="loop.human_interrupt", 

240 resource=session_id, 

241 outcome="paused", 

242 severity=Severity.WARNING, 

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

244 )) 

245 if self.on_human_interrupt: 

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

247 if feedback: 

248 self.context_manager.append_user_message(feedback) 

249 continue 

250 

251 except StepTimeoutError: 

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

253 

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

255 await self._save_checkpoint(session_id, iteration) 

256 

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

258 

259 # ── Reflection ──────────────────────────────── 

260 

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

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

263 

264任务: {self.context_manager.current_task} 

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

266 

267评估并返回JSON: 

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

269 

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

271 try: 

272 import json 

273 d = json.loads(resp) 

274 result = ReflectionResult( 

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

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

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

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

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

280 ) 

281 except Exception: 

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

283 if self.on_reflection: 

284 self.on_reflection(result) 

285 return result 

286 

287 # ── Self-Critique ───────────────────────────── 

288 

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

290 if not self.config.enable_self_critique: 

291 return text 

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

293 

294{text[:3000]}""" 

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

296 return improved or text 

297 

298 # ── Auto Route ──────────────────────────────── 

299 

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

301 score = self._estimate_complexity(task) 

302 if score > 0.7: 

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

304 elif score > 0.4: 

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

306 else: 

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

308 

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

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

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

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

313 

314 # ── 步骤执行 ────────────────────────────────── 

315 

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

317 last_error = None 

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

319 try: 

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

321 except asyncio.TimeoutError: 

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

323 except AllModelsFailed as e: 

324 last_error = e 

325 await asyncio.sleep(2 ** attempt) 

326 raise last_error 

327 

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

329 ctx = self.context_manager.build_context( 

330 model_type=self.model_router.model_type, 

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

332 ) 

333 

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

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

336 usage_ratio = self.context_manager.estimate_context_usage() 

337 if usage_ratio > self.config.auto_page_threshold: 

338 page_count = await self._auto_page_callback(usage_ratio) 

339 

340 # v1.16.6 — rate limiting before model calls 

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

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

343 

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

345 

346 # 成本记录 

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

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

349 

350 if not resp.tool_calls: 

351 if self.config.enable_self_critique: 

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

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

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

355 

356 # HITL 检查 

357 if self.config.enable_human_in_the_loop: 

358 for tc in resp.tool_calls: 

359 if self._is_high_risk(tc): 

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

361 

362 groups = self._group_independent_calls(resp.tool_calls) 

363 all_results = [] 

364 for group in groups: 

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

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

367 all_results.extend(batch_results) 

368 

369 self.context_manager.append_tool_results(all_results) 

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

371 

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

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

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

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

376 

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

378 if len(tool_calls) <= 1: 

379 return [tool_calls] if tool_calls else [] 

380 groups: list[list] = [] 

381 for call in tool_calls: 

382 for group in groups: 

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

384 group.append(call) 

385 break 

386 else: 

387 groups.append([call]) 

388 return groups 

389 

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

391 write_paths = set() 

392 for tc in group: 

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

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

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

396 write_paths.add(p) 

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

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

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

400 return False 

401 

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

403 

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

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

406 backend = self.checkpoint_backend 

407 if not backend: 

408 # Fallback to thin CheckpointStore 

409 if not self.checkpoint_store: 

410 return 

411 snap = { 

412 "session_id": session_id, "iteration": iteration, 

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

414 "timestamp": time.time(), 

415 } 

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

417 return 

418 

419 # Full checkpoint via CheckpointBackend 

420 try: 

421 from datetime import datetime, timezone 

422 import uuid 

423 

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

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

426 

427 cp = Checkpoint( 

428 metadata=CheckpointMetadata( 

429 thread_id=session_id, 

430 checkpoint_id=checkpoint_id, 

431 step=iteration, 

432 parent_checkpoint_id=parent_id, 

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

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

435 ), 

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

437 state={ 

438 "iteration": iteration, 

439 "task": self.context_manager.current_task, 

440 "session_id": session_id, 

441 "cost_usd": self.cost_tracker.total_cost, 

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

443 "human_interrupts": self._human_interrupts, 

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

445 }, 

446 tools_result={}, 

447 next_node="loop", 

448 ) 

449 await backend.put(cp) 

450 self._last_checkpoint_id = checkpoint_id 

451 

452 except Exception as e: 

453 pass # Checkpoint failure must not crash the loop 

454 

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

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

457 backend = self.checkpoint_backend 

458 if not backend: 

459 # Fallback to thin CheckpointStore 

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

461 return 0 

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

463 if not snap: 

464 return 0 

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

466 if iter_count > 0: 

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

468 for msg in msgs: 

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

470 return iter_count 

471 

472 if not self.config.enable_checkpoints: 

473 return 0 

474 

475 try: 

476 latest = await backend.get_latest(session_id) 

477 if not latest: 

478 return 0 

479 self._last_checkpoint_id = latest.metadata.checkpoint_id 

480 iter_count = latest.metadata.step 

481 

482 # Restore messages 

483 for msg in latest.messages: 

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

485 

486 # Restore state 

487 state = latest.state 

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

489 

490 return iter_count 

491 except Exception: 

492 return 0 

493 

494 def set_auto_paging(self, callback: Callable): 

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

496 self._auto_page_callback = callback 

497 

498 def cancel(self): 

499 self._cancelled = True 

500 

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

502 

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

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

505 start_time = time.time() 

506 roles = roles or self.config.swarm_roles 

507 if not roles: 

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

509 

510 topology = SwarmTopology(self.config.swarm_topology) 

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

512 

513 swarm = SwarmCoordinator( 

514 router=self.model_router, 

515 tool_registry=self.tool_registry, 

516 topology=topology, 

517 max_parallel=self.config.max_parallel_agents, 

518 ) 

519 swarm.register_roles(roles) 

520 

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

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

523 

524 return AgentResult( 

525 output=swarm_result.combined_output, 

526 iterations=1, 

527 cost_usd=self.cost_tracker.total_cost, 

528 duration_ms=duration_ms, 

529 tool_calls_total=0, 

530 reflections_count=0, 

531 human_interrupts=0, 

532 final_state=LoopState.COMPLETED, 

533 swarm_result=swarm_result, 

534 ) 

535 

536 

537class StepTimeoutError(Exception): 

538 

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

540 

541 pass 

542 

543 

544@dataclass 

545class StepResult: 

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

547 content: str 

548 is_terminal: bool = False 

549 tool_results: list | None = None