304 lines
11 KiB
Python
304 lines
11 KiB
Python
"""Pi agent-core 风格的有状态 Agent 门面。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections.abc import Awaitable, Callable
|
|
from contextlib import suppress
|
|
from dataclasses import dataclass, field, replace
|
|
from enum import StrEnum
|
|
from typing import Any
|
|
|
|
from app.ai_provider.types.chat import ChatChunk, TokenUsage
|
|
from app.ai_provider.types.messages import BaseMessage, HumanMessage
|
|
from app.application.agent.confirmation import ConfirmationRegistry
|
|
from app.application.agent.extensions import (
|
|
AgentEvent,
|
|
AgentRunContext,
|
|
CancellationToken,
|
|
ExtensionRegistry,
|
|
)
|
|
from app.application.agent.runtime import AgentTurnResult, run_agent_turn
|
|
from app.modules.tool_runtime.definition import ToolRuntimeLike
|
|
|
|
|
|
class AgentStatus(StrEnum):
|
|
idle = "idle"
|
|
running = "running"
|
|
completed = "completed"
|
|
failed = "failed"
|
|
cancelled = "cancelled"
|
|
|
|
|
|
class QueueMode(StrEnum):
|
|
all = "all"
|
|
one_at_a_time = "one-at-a-time"
|
|
|
|
|
|
@dataclass
|
|
class AgentState:
|
|
"""有状态 Agent 的进程内快照。"""
|
|
|
|
messages: list[BaseMessage] = field(default_factory=list)
|
|
status: AgentStatus = AgentStatus.idle
|
|
streaming_message: str = ""
|
|
pending_tool_calls: set[str] = field(default_factory=set)
|
|
error_message: str | None = None
|
|
usage: TokenUsage = field(default_factory=TokenUsage)
|
|
|
|
|
|
class PendingMessageQueue:
|
|
"""steering/follow-up 共用的小型 FIFO。"""
|
|
|
|
def __init__(self, mode: QueueMode = QueueMode.one_at_a_time) -> None:
|
|
self.mode = mode
|
|
self._messages: list[BaseMessage] = []
|
|
|
|
def put(self, message: BaseMessage) -> None:
|
|
self._messages.append(message)
|
|
|
|
def drain(self) -> list[BaseMessage]:
|
|
if not self._messages:
|
|
return []
|
|
if self.mode is QueueMode.all:
|
|
messages, self._messages = self._messages, []
|
|
return messages
|
|
return [self._messages.pop(0)]
|
|
|
|
def clear(self) -> None:
|
|
self._messages.clear()
|
|
|
|
def __bool__(self) -> bool:
|
|
return bool(self._messages)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AgentConfig:
|
|
"""Stateful Agent 的稳定运行配置。"""
|
|
|
|
tools: tuple[dict[str, Any], ...]
|
|
tool_runtime: ToolRuntimeLike | None
|
|
model_hint: str | None = None
|
|
thinking_enabled: bool | None = None
|
|
reasoning_effort: str | None = None
|
|
max_iterations: int = 8
|
|
confirmation_registry: ConfirmationRegistry | None = None
|
|
confirmation_timeout: float = 300.0
|
|
on_turn: Callable[[], Awaitable[None]] | None = None
|
|
on_chunk: Callable[[ChatChunk], Awaitable[None]] | None = None
|
|
on_tool_event: Callable[[str, dict[str, Any]], Awaitable[None]] | None = None
|
|
on_confirm_timeout: Callable[[str], Awaitable[None]] | None = None
|
|
|
|
|
|
class Agent:
|
|
"""管理状态、消息队列、取消与扩展 Hook；低层循环仍由 run_agent_turn 负责。"""
|
|
|
|
def __init__(
|
|
self,
|
|
config: AgentConfig,
|
|
*,
|
|
messages: list[BaseMessage] | None = None,
|
|
run_context: AgentRunContext | None = None,
|
|
extensions: ExtensionRegistry | None = None,
|
|
runner: Callable[..., Awaitable[AgentTurnResult]] = run_agent_turn,
|
|
) -> None:
|
|
self.config = config
|
|
self.state = AgentState(messages=list(messages or []))
|
|
self.run_context = run_context or AgentRunContext()
|
|
self.extensions = extensions or ExtensionRegistry()
|
|
self._runner = runner
|
|
self.steering = PendingMessageQueue()
|
|
self.follow_ups = PendingMessageQueue()
|
|
self._active_task: asyncio.Task[AgentTurnResult] | None = None
|
|
self._accepting_messages = False
|
|
self._idle = asyncio.Event()
|
|
self._idle.set()
|
|
|
|
def steer(self, message: str | BaseMessage) -> bool:
|
|
if not self._accepting_messages:
|
|
return False
|
|
self.steering.put(self._as_message(message))
|
|
return True
|
|
|
|
def follow_up(self, message: str | BaseMessage) -> bool:
|
|
if not self._accepting_messages:
|
|
return False
|
|
self.follow_ups.put(self._as_message(message))
|
|
return True
|
|
|
|
async def prompt(self, message: str | BaseMessage) -> AgentTurnResult:
|
|
if self.state.status is AgentStatus.running:
|
|
raise RuntimeError("Agent 正在运行，请使用 steer 或 follow_up")
|
|
self.state.messages.append(self._as_message(message))
|
|
return await self._run()
|
|
|
|
async def continue_run(self) -> AgentTurnResult:
|
|
if self.state.status is AgentStatus.running:
|
|
raise RuntimeError("Agent 正在运行")
|
|
return await self._run()
|
|
|
|
async def abort(self) -> bool:
|
|
task = self._active_task
|
|
if task is None or task.done():
|
|
return False
|
|
self.run_context.cancellation.cancel()
|
|
task.cancel()
|
|
return True
|
|
|
|
async def wait_for_idle(self) -> None:
|
|
await self._idle.wait()
|
|
|
|
async def reset(self) -> None:
|
|
await self.abort()
|
|
await self.wait_for_idle()
|
|
self.state = AgentState()
|
|
self.steering.clear()
|
|
self.follow_ups.clear()
|
|
|
|
async def _run(self) -> AgentTurnResult:
|
|
if self.run_context.cancellation.cancelled:
|
|
self.run_context = replace(self.run_context, cancellation=CancellationToken())
|
|
self.state.status = AgentStatus.running
|
|
self.state.streaming_message = ""
|
|
self.state.pending_tool_calls.clear()
|
|
self.state.error_message = None
|
|
self._accepting_messages = True
|
|
self._idle.clear()
|
|
try:
|
|
await self.extensions.publish(AgentEvent("agent.started"), self.run_context)
|
|
task = asyncio.create_task(self._run_loop())
|
|
self._active_task = task
|
|
result = await task
|
|
self._accepting_messages = False
|
|
self.state.status = AgentStatus.completed
|
|
self.state.streaming_message = result.content
|
|
self.state.usage = result.usage
|
|
await self.extensions.publish(
|
|
AgentEvent("agent.ended", {"iterations": result.iterations}), self.run_context
|
|
)
|
|
return result
|
|
except asyncio.CancelledError:
|
|
self._accepting_messages = False
|
|
self.state.status = AgentStatus.cancelled
|
|
await self.extensions.publish(AgentEvent("agent.cancelled"), self.run_context)
|
|
raise
|
|
except Exception as exc:
|
|
self._accepting_messages = False
|
|
self.state.status = AgentStatus.failed
|
|
self.state.error_message = str(exc)
|
|
with suppress(Exception):
|
|
await self.extensions.publish(
|
|
AgentEvent("agent.failed", {"error": str(exc)}), self.run_context
|
|
)
|
|
raise
|
|
finally:
|
|
self._accepting_messages = False
|
|
with suppress(Exception):
|
|
await self.extensions.publish(
|
|
AgentEvent("agent.settled", {"status": self.state.status.value}),
|
|
self.run_context,
|
|
)
|
|
self._active_task = None
|
|
self._idle.set()
|
|
|
|
async def _run_loop(self) -> AgentTurnResult:
|
|
remaining = self.config.max_iterations
|
|
total_iterations = 0
|
|
total_usage = TokenUsage()
|
|
total_tool_history: list[dict[str, Any]] = []
|
|
final = AgentTurnResult(content="", reasoning="")
|
|
while remaining > 0:
|
|
working_messages = list(self.state.messages)
|
|
final = await self._runner(
|
|
messages=working_messages,
|
|
tools=list(self.config.tools),
|
|
tool_runtime=self.config.tool_runtime,
|
|
model_hint=self.config.model_hint,
|
|
thinking_enabled=self.config.thinking_enabled,
|
|
reasoning_effort=self.config.reasoning_effort,
|
|
max_iterations=remaining,
|
|
run_id=self.run_context.run_id,
|
|
character_id=self.run_context.character_id,
|
|
trace_id=self.run_context.trace_id,
|
|
on_turn=self.config.on_turn,
|
|
on_chunk=self._on_chunk,
|
|
on_tool_event=self._on_tool_event,
|
|
confirmation_registry=self.config.confirmation_registry,
|
|
confirmation_timeout=self.config.confirmation_timeout,
|
|
on_confirm_timeout=self.config.on_confirm_timeout,
|
|
extensions=self.extensions,
|
|
run_context=self.run_context,
|
|
get_steering_messages=self.steering.drain,
|
|
)
|
|
self.state.messages = working_messages
|
|
total_iterations += final.iterations
|
|
total_usage += final.usage
|
|
total_tool_history.extend(final.tool_history)
|
|
remaining -= final.iterations
|
|
if final.truncated:
|
|
return AgentTurnResult(
|
|
content=final.content,
|
|
reasoning=final.reasoning,
|
|
tool_history=total_tool_history,
|
|
iterations=total_iterations,
|
|
truncated=True,
|
|
usage=total_usage,
|
|
)
|
|
if remaining <= 0:
|
|
return AgentTurnResult(
|
|
content=final.content,
|
|
reasoning=final.reasoning,
|
|
tool_history=total_tool_history,
|
|
iterations=total_iterations,
|
|
truncated=bool(self.follow_ups),
|
|
usage=total_usage,
|
|
)
|
|
follow_ups = self.follow_ups.drain()
|
|
if not follow_ups:
|
|
return AgentTurnResult(
|
|
content=final.content,
|
|
reasoning=final.reasoning,
|
|
tool_history=total_tool_history,
|
|
iterations=total_iterations,
|
|
truncated=final.truncated,
|
|
usage=total_usage,
|
|
)
|
|
self.state.messages.extend(follow_ups)
|
|
return final
|
|
|
|
async def _on_chunk(self, chunk: ChatChunk) -> None:
|
|
if chunk.delta_content:
|
|
self.state.streaming_message += chunk.delta_content
|
|
if self.config.on_chunk is not None:
|
|
await self.config.on_chunk(chunk)
|
|
|
|
async def _on_tool_event(self, event_type: str, payload: dict[str, Any]) -> None:
|
|
call_id = str(payload.get("tool_call_id", ""))
|
|
if event_type == "tool.started" and call_id:
|
|
self.state.pending_tool_calls.add(call_id)
|
|
elif event_type in ("tool.completed", "tool.failed") and call_id:
|
|
self.state.pending_tool_calls.discard(call_id)
|
|
# 先转发 SSE（用户可见流优先），再把同一事件发布进扩展事件流，
|
|
# 让 AgentEventListener 也能观测工具级生命周期。
|
|
# 扩展 listener 异常被吞掉，避免观测副菜阻断工具事件主流程。
|
|
if self.config.on_tool_event is not None:
|
|
await self.config.on_tool_event(event_type, payload)
|
|
with suppress(Exception):
|
|
await self.extensions.publish(
|
|
AgentEvent(event_type, payload), self.run_context
|
|
)
|
|
|
|
@staticmethod
|
|
def _as_message(message: str | BaseMessage) -> BaseMessage:
|
|
return HumanMessage(content=message) if isinstance(message, str) else message
|
|
|
|
|
|
__all__ = [
|
|
"Agent",
|
|
"AgentConfig",
|
|
"AgentState",
|
|
"AgentStatus",
|
|
"PendingMessageQueue",
|
|
"QueueMode",
|
|
]
|