Coverage for agentos/core/loop.py: 32%
312 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 17:01 +0800
1""" # noqa: E501
2AgentOS v0.70 核心循环 — Gemini + Metrics + CostAnalytics 集成版。
3v0.40: Swarm多Agent并行、Agent间通信、语义缓存、任务队列。
4v0.70: MetricsCollector、CostAnalytics实时监控。
5"""
7from __future__ import annotations
9import asyncio
10import time
11from collections.abc import Callable
12from dataclasses import dataclass, field
13from datetime import UTC
14from enum import StrEnum
16from agentos.checkpoint.base import Checkpoint, CheckpointBackend, CheckpointMetadata
17from agentos.comm.layer import CommunicationLayer
18from agentos.core.context import ContextManager
19from agentos.core.streaming import StreamChunk
20from agentos.cost.tracker import CostTracker
21from agentos.models.router import AllModelsFailed, ModelRouter
22from agentos.observability.cost_analytics import CostAnalytics
23from agentos.observability.metrics import MetricsCollector
24from agentos.observability.tracer import Tracer
25from agentos.security.sandbox import SandboxManager
26from agentos.storage.base import CheckpointStore
27from agentos.swarm.coordinator import AgentRole, SwarmCoordinator, SwarmResult, SwarmTopology
28from agentos.tools.audit_logger import AuditEvent, AuditLogger, Severity
29from agentos.tools.rate_limiter import TokenBucket
30from agentos.tools.registry import ToolRegistry
33class LoopState(StrEnum):
34 """主循环状态。"""
36 RUNNING = "running"
37 PAUSED = "paused"
38 WAITING_HUMAN = "waiting_human"
39 COMPLETED = "completed"
40 FAILED = "failed"
41 CANCELLED = "cancelled"
44@dataclass
45class AgentResult:
46 """Agent 主循环的最终运行结果。"""
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
63@dataclass
64class LoopConfig:
65 """Agent 主循环的运行时配置。"""
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 = (
91 None # Full checkpoint backend for crash recovery
92 )
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
97class MaxIterationsExceeded(Exception): # noqa: N818
98 """超出最大迭代次数异常。"""
102class HumanInterruptNeeded(Exception): # noqa: N818
103 """需要人工介入异常。"""
105 def __init__(self, message: str, context: dict | None = None):
106 super().__init__(message)
107 self.context = context or {}
110@dataclass
111class ReflectionResult:
112 """反思结果。"""
114 quality_score: float
115 issues: list[str]
116 suggestions: list[str]
117 should_continue: bool
118 new_plan: str | None = None
121class AgentLoop:
122 """v0.30 核心循环 — Reflection + HITL + Self-Critique + 自动路由 + 成本追踪。"""
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
166 # ── 运行入口 ──────────────────────────────────
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)
172 if self.audit_logger:
173 self.audit_logger.log(
174 event=AuditEvent(
175 actor="agentos",
176 action="loop.start",
177 resource=session_id,
178 outcome="initiated",
179 details={"task": task[:200]},
180 )
181 )
183 if self.config.auto_select_model:
184 await self._auto_route_model(task)
186 iteration = await self._try_restore(session_id)
187 tool_calls_total = 0
188 reflection_loops = 0
190 while iteration < self.config.max_iterations and not self._cancelled:
191 iteration += 1
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 (
198 not reflection.should_continue
199 and reflection_loops < self.config.max_reflection_loops
200 ):
201 reflection_loops += 1
202 if reflection.new_plan:
203 self.context_manager.update_plan(reflection.new_plan)
204 continue
206 try:
207 with self.tracer.step(f"loop_{iteration}"):
208 step_result = await self._execute_step_sync(iteration, session_id)
210 if step_result.is_terminal:
211 duration_ms = (time.time() - start_time) * 1000
212 if self.audit_logger:
213 self.audit_logger.log(
214 event=AuditEvent(
215 actor="agentos",
216 action="loop.complete",
217 resource=session_id,
218 outcome="success",
219 details={"iterations": iteration, "duration_ms": duration_ms},
220 )
221 )
222 return AgentResult(
223 output=step_result.content,
224 iterations=iteration,
225 tokens_used=self.tracer.token_summary(),
226 cost_usd=self.cost_tracker.total_cost,
227 duration_ms=duration_ms,
228 tool_calls_total=tool_calls_total,
229 reflections_count=len(self._reflection_history),
230 human_interrupts=self._human_interrupts,
231 )
233 if step_result.tool_results:
234 tool_calls_total += len(step_result.tool_results)
236 if self.on_iteration:
237 self.on_iteration(iteration, step_result.tool_results or [])
239 except HumanInterruptNeeded as e:
240 self._human_interrupts += 1
241 if self.audit_logger:
242 self.audit_logger.log(
243 event=AuditEvent(
244 actor="agentos",
245 action="loop.human_interrupt",
246 resource=session_id,
247 outcome="paused",
248 severity=Severity.WARNING,
249 details={
250 "interrupt_count": self._human_interrupts,
251 "message": str(e)[:200],
252 },
253 )
254 )
255 if self.on_human_interrupt:
256 feedback = self.on_human_interrupt(str(e), e.context)
257 if feedback:
258 self.context_manager.append_user_message(feedback)
259 continue
261 except StepTimeoutError:
262 return AgentResult(
263 output="",
264 iterations=iteration,
265 final_state=LoopState.FAILED,
266 error="Step timeout",
267 )
269 if self.config.enable_checkpoints and iteration % self.config.checkpoint_interval == 0:
270 await self._save_checkpoint(session_id, iteration)
272 raise MaxIterationsExceeded(f"超过 {self.config.max_iterations} 步")
274 # ── Reflection ────────────────────────────────
276 async def _reflect(self, session_id: str) -> ReflectionResult:
277 prompt = f"""你是一个反思者。审核以下Agent执行过程：
279任务: {self.context_manager.current_task}
280已执行: {self.context_manager.step_count} 步
282评估并返回JSON:
283{{"quality_score": 0.0-1.0, "issues": [...], "suggestions": [...], "should_continue": true/false, "new_plan": "如果调整，新计划"}}""" # noqa: E501
285 resp = await self.model_router.call_simple(prompt)
286 try:
287 import json
289 d = json.loads(resp)
290 result = ReflectionResult(
291 quality_score=d.get("quality_score", 0.5),
292 issues=d.get("issues", []),
293 suggestions=d.get("suggestions", []),
294 should_continue=d.get("should_continue", True),
295 new_plan=d.get("new_plan"),
296 )
297 except Exception:
298 result = ReflectionResult(0.5, [], [], True)
299 if self.on_reflection:
300 self.on_reflection(result)
301 return result
303 # ── Self-Critique ─────────────────────────────
305 async def _self_critique(self, text: str) -> str:
306 if not self.config.enable_self_critique:
307 return text
308 prompt = f"""审视以下回答，找出逻辑错误或不准确之处。如果已足够好就原样返回。
310{text[:3000]}"""
311 improved = await self.model_router.call_simple(prompt)
312 return improved or text
314 # ── Auto Route ────────────────────────────────
316 async def _auto_route_model(self, task: str):
317 score = self._estimate_complexity(task)
318 if score > 0.7:
319 self.model_router.set_preferred("deepseek-r1")
320 elif score > 0.4:
321 self.model_router.set_preferred("kimi-k2.6")
322 else:
323 self.model_router.set_preferred("deepseek-v3.1")
325 def _estimate_complexity(self, task: str) -> float:
326 kw = [
327 "分析",
328 "对比",
329 "设计",
330 "架构",
331 "review",
332 "refactor",
333 "实现",
334 "优化",
335 "诊断",
336 "troubleshoot",
337 "debug",
338 "deploy",
339 "migrate",
340 "安全",
341 "security",
342 ]
343 score = sum(0.15 for k in kw if k in task.lower())
344 return min(score + min(len(task) / 2000, 0.3), 1.0)
346 # ── 步骤执行 ──────────────────────────────────
348 async def _execute_step_sync(self, iteration: int, session_id: str) -> StepResult:
349 last_error = None
350 for attempt in range(self.config.max_retries_per_step + 1):
351 try:
352 return await asyncio.wait_for(
353 self._do_step(iteration, session_id), timeout=self.config.step_timeout_seconds
354 )
355 except TimeoutError:
356 last_error = StepTimeoutError(f"Step {iteration} timeout")
357 except AllModelsFailed as e:
358 last_error = e
359 await asyncio.sleep(2**attempt)
360 raise last_error
362 async def _do_step(self, iteration: int, session_id: str) -> StepResult:
363 ctx = self.context_manager.build_context(
364 model_type=self.model_router.model_type,
365 tools=self.tool_registry.get_schemas_for_model(self.model_router.model_type),
366 )
368 # v1.11.0 — auto-page old memories if context nearing limit
369 if self.config.enable_auto_paging and self._auto_page_callback:
370 usage_ratio = self.context_manager.estimate_context_usage()
371 if usage_ratio > self.config.auto_page_threshold:
372 await self._auto_page_callback(usage_ratio)
374 # v1.16.6 — rate limiting before model calls
375 if self.rate_limiter and not self.rate_limiter.try_acquire("model_call"):
376 raise StepTimeoutError(f"Rate limit exceeded for model call at step {iteration}")
378 resp = await self.model_router.call(ctx)
380 # 成本记录
381 if self.config.enable_cost_tracking and hasattr(resp, "usage"):
382 self.cost_tracker.record(self.model_router.current_model, resp.usage)
384 if not resp.tool_calls:
385 if self.config.enable_self_critique:
386 improved = await self._self_critique(resp.content)
387 return StepResult(content=improved, is_terminal=True)
388 return StepResult(content=resp.content, is_terminal=True)
390 # HITL 检查
391 if self.config.enable_human_in_the_loop:
392 for tc in resp.tool_calls:
393 if self._is_high_risk(tc):
394 raise HumanInterruptNeeded(
395 f"高风险操作需确认: {tc.name}", {"tool": tc.name, "args": tc.arguments}
396 )
398 groups = self._group_independent_calls(resp.tool_calls)
399 all_results = []
400 for group in groups:
401 sandbox = self.sandbox_manager.get_sandbox(session_id) if self.sandbox_manager else None
402 batch_results = await self.tool_registry.execute_batch(group, sandbox=sandbox)
403 all_results.extend(batch_results)
405 self.context_manager.append_tool_results(all_results)
406 return StepResult(content="", is_terminal=False, tool_results=all_results)
408 def _is_high_risk(self, tc) -> bool:
409 risky = ["delete", "rm", "uninstall", "format", "sudo", "kill", "drop"]
410 name = tc.name.lower() if hasattr(tc, "name") else tc.get("name", "").lower()
411 return any(r in name for r in risky)
413 def _group_independent_calls(self, tool_calls: list) -> list[list]:
414 if len(tool_calls) <= 1:
415 return [tool_calls] if tool_calls else []
416 groups: list[list] = []
417 for call in tool_calls:
418 for group in groups:
419 if not self._has_conflict(call, group):
420 group.append(call)
421 break
422 else:
423 groups.append([call])
424 return groups
426 def _has_conflict(self, call, group: list) -> bool:
427 write_paths = set()
428 for tc in group:
429 tool = self.tool_registry.get(tc.name)
430 if tool and tool.is_write_operation(tc.arguments):
431 if p := tool.extract_target_path(tc.arguments):
432 write_paths.add(p)
433 cur = self.tool_registry.get(call.name)
434 if cur and cur.is_read_operation(call.arguments):
435 return cur.extract_target_path(call.arguments) in write_paths
436 return False
438 # ── v1.11.0 全量 Checkpoint (完整状态快照) ────
440 async def _save_checkpoint(self, session_id: str, iteration: int):
441 """Save full runtime state snapshot via CheckpointBackend."""
442 backend = self.checkpoint_backend
443 if not backend:
444 # Fallback to thin CheckpointStore
445 if not self.checkpoint_store:
446 return
447 snap = {
448 "session_id": session_id,
449 "iteration": iteration,
450 "messages": [
451 {"role": m.role, "content": m.content} for m in self.context_manager._messages
452 ],
453 "timestamp": time.time(),
454 }
455 await self.checkpoint_store.save(session_id, snap)
456 return
458 # Full checkpoint via CheckpointBackend
459 try:
460 from datetime import datetime
462 checkpoint_id = f"ckpt-{session_id}-{iteration:06d}"
463 parent_id = getattr(self, "_last_checkpoint_id", None)
465 cp = Checkpoint(
466 metadata=CheckpointMetadata(
467 thread_id=session_id,
468 checkpoint_id=checkpoint_id,
469 step=iteration,
470 parent_checkpoint_id=parent_id,
471 created_at=datetime.now(UTC).isoformat(),
472 tags=["auto", f"iter_{iteration}"],
473 ),
474 messages=[
475 {"role": m.role, "content": m.content} for m in self.context_manager._messages
476 ],
477 state={
478 "iteration": iteration,
479 "task": self.context_manager.current_task,
480 "session_id": session_id,
481 "cost_usd": self.cost_tracker.total_cost,
482 "reflections": len(self._reflection_history),
483 "human_interrupts": self._human_interrupts,
484 "loop_state": (
485 self.context_manager.current_state
486 if hasattr(self.context_manager, "current_state")
487 else "running"
488 ),
489 },
490 tools_result={},
491 next_node="loop",
492 )
493 await backend.put(cp)
494 self._last_checkpoint_id = checkpoint_id
496 except Exception:
497 pass # Checkpoint failure must not crash the loop
499 async def _try_restore(self, session_id: str) -> int:
500 """Restore full state from last checkpoint. Returns iteration to resume from."""
501 backend = self.checkpoint_backend
502 if not backend:
503 # Fallback to thin CheckpointStore
504 if not self.checkpoint_store or not self.config.enable_checkpoints:
505 return 0
506 snap = await self.checkpoint_store.load(session_id)
507 if not snap:
508 return 0
509 iter_count = snap.get("iteration", 0)
510 if iter_count > 0:
511 msgs = snap.get("messages", [])
512 for msg in msgs:
513 self.context_manager.append_message(msg["role"], msg["content"])
514 return iter_count
516 if not self.config.enable_checkpoints:
517 return 0
519 try:
520 latest = await backend.get_latest(session_id)
521 if not latest:
522 return 0
523 self._last_checkpoint_id = latest.metadata.checkpoint_id
524 iter_count = latest.metadata.step
526 # Restore messages
527 for msg in latest.messages:
528 self.context_manager.append_message(msg.get("role", "user"), msg.get("content", ""))
530 # Restore state
531 state = latest.state
532 self._human_interrupts = state.get("human_interrupts", 0)
534 return iter_count
535 except Exception:
536 return 0
538 def set_auto_paging(self, callback: Callable):
539 """Register callback for automatic memory paging (v1.11.0)."""
540 self._auto_page_callback = callback
542 def cancel(self):
543 self._cancelled = True
545 # ── v0.40 Swarm执行 ──────────────────────────
547 async def run_swarm(self, task: str, roles: list[AgentRole] | None = None) -> AgentResult:
548 """以Swarm模式执行任务 — 多Agent协作。"""
549 start_time = time.time()
550 roles = roles or self.config.swarm_roles
551 if not roles:
552 return AgentResult(
553 output="[Swarm] No roles defined",
554 iterations=0,
555 final_state=LoopState.FAILED,
556 error="No roles",
557 )
559 topology = SwarmTopology(self.config.swarm_topology)
560 CommunicationLayer() if self.config.enable_comm_layer else None
562 swarm = SwarmCoordinator(
563 router=self.model_router,
564 tool_registry=self.tool_registry,
565 topology=topology,
566 max_parallel=self.config.max_parallel_agents,
567 )
568 swarm.register_roles(roles)
570 swarm_result = await swarm.execute(task, roles)
571 duration_ms = (time.time() - start_time) * 1000
573 return AgentResult(
574 output=swarm_result.combined_output,
575 iterations=1,
576 cost_usd=self.cost_tracker.total_cost,
577 duration_ms=duration_ms,
578 tool_calls_total=0,
579 reflections_count=0,
580 human_interrupts=0,
581 final_state=LoopState.COMPLETED,
582 swarm_result=swarm_result,
583 )
586class StepTimeoutError(Exception):
587 """步骤超时异常。"""
591@dataclass
592class StepResult:
593 """步骤执行结果。"""
595 content: str
596 is_terminal: bool = False
597 tool_results: list | None = None