230 lines
7.1 KiB
Python
230 lines
7.1 KiB
Python
"""Agent 扩展契约与注册表。
|
|
|
|
扩展只接触本次运行的快照与 ToolResult，不持有数据库会话、密钥或安全治理对象。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections.abc import Awaitable, Callable
|
|
from dataclasses import dataclass, field, replace
|
|
from typing import Any
|
|
|
|
from app.ai_provider.types.chat import TokenUsage
|
|
from app.ai_provider.types.messages import BaseMessage
|
|
from app.modules.tool_runtime.definition import ToolResult
|
|
|
|
|
|
class CancellationToken:
|
|
"""协作式取消信号；阻塞 await 仍由活动 asyncio.Task.cancel() 负责中断。"""
|
|
|
|
def __init__(self) -> None:
|
|
self._cancelled = False
|
|
|
|
@property
|
|
def cancelled(self) -> bool:
|
|
return self._cancelled
|
|
|
|
def cancel(self) -> None:
|
|
self._cancelled = True
|
|
|
|
def raise_if_cancelled(self) -> None:
|
|
if self._cancelled:
|
|
raise asyncio.CancelledError
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AgentRunContext:
|
|
"""扩展可见的最小运行上下文。"""
|
|
|
|
run_id: str | None = None
|
|
session_id: str | None = None
|
|
character_id: str | None = None
|
|
trace_id: str | None = None
|
|
metadata: dict[str, Any] = field(default_factory=dict)
|
|
cancellation: CancellationToken = field(default_factory=CancellationToken)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AgentContext:
|
|
"""每轮发送给模型前可转换的上下文快照。"""
|
|
|
|
messages: tuple[BaseMessage, ...]
|
|
tools: tuple[dict[str, Any], ...]
|
|
model_hint: str | None = None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AgentEvent:
|
|
"""Agent Core 内部生命周期事件。"""
|
|
|
|
type: str
|
|
payload: dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ToolCallContext:
|
|
"""工具调用 Hook 的稳定输入。"""
|
|
|
|
tool_call_id: str
|
|
name: str
|
|
input: dict[str, Any]
|
|
messages: tuple[BaseMessage, ...]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class BeforeToolCallDecision:
|
|
"""执行前决策；扩展只能收紧或改参，不能绕过 ToolRuntime。"""
|
|
|
|
allowed: bool = True
|
|
reason: str | None = None
|
|
input: dict[str, Any] | None = None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AfterToolCallDecision:
|
|
"""执行后结果覆盖与提前终止提示。"""
|
|
|
|
result: ToolResult | None = None
|
|
terminate: bool = False
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TurnSnapshot:
|
|
"""一轮模型调用及其工具执行结果摘要。"""
|
|
|
|
iteration: int
|
|
content: str
|
|
reasoning: str
|
|
tool_calls: tuple[dict[str, Any], ...]
|
|
tool_results: tuple[dict[str, Any], ...]
|
|
usage: TokenUsage = field(default_factory=TokenUsage)
|
|
|
|
|
|
ContextTransformer = Callable[[AgentContext, AgentRunContext], Awaitable[AgentContext]]
|
|
BeforeToolCallHook = Callable[
|
|
[ToolCallContext, AgentRunContext], Awaitable[BeforeToolCallDecision | None]
|
|
]
|
|
AfterToolCallHook = Callable[
|
|
[ToolCallContext, ToolResult, AgentRunContext], Awaitable[AfterToolCallDecision | None]
|
|
]
|
|
TurnPolicy = Callable[[TurnSnapshot, AgentRunContext], Awaitable[bool]]
|
|
AgentEventListener = Callable[[AgentEvent, AgentRunContext], Awaitable[None]]
|
|
|
|
|
|
class ExtensionRegistry:
|
|
"""受约束的 Agent 扩展注册表；按注册顺序执行 Hook。"""
|
|
|
|
def __init__(self) -> None:
|
|
self._context_transformers: list[ContextTransformer] = []
|
|
self._before_tool_hooks: list[BeforeToolCallHook] = []
|
|
self._after_tool_hooks: list[AfterToolCallHook] = []
|
|
self._turn_policies: list[TurnPolicy] = []
|
|
self._event_listeners: list[AgentEventListener] = []
|
|
|
|
def add_context_transformer(self, transformer: ContextTransformer) -> None:
|
|
self._context_transformers.append(transformer)
|
|
|
|
def add_before_tool_hook(self, hook: BeforeToolCallHook) -> None:
|
|
self._before_tool_hooks.append(hook)
|
|
|
|
def add_after_tool_hook(self, hook: AfterToolCallHook) -> None:
|
|
self._after_tool_hooks.append(hook)
|
|
|
|
def add_turn_policy(self, policy: TurnPolicy) -> None:
|
|
self._turn_policies.append(policy)
|
|
|
|
def add_event_listener(self, listener: AgentEventListener) -> Callable[[], None]:
|
|
self._event_listeners.append(listener)
|
|
|
|
def unsubscribe() -> None:
|
|
if listener in self._event_listeners:
|
|
self._event_listeners.remove(listener)
|
|
|
|
return unsubscribe
|
|
|
|
def extend(self, registry: ExtensionRegistry) -> None:
|
|
"""按原注册顺序合并另一组扩展。"""
|
|
self._context_transformers.extend(registry._context_transformers)
|
|
self._before_tool_hooks.extend(registry._before_tool_hooks)
|
|
self._after_tool_hooks.extend(registry._after_tool_hooks)
|
|
self._turn_policies.extend(registry._turn_policies)
|
|
self._event_listeners.extend(registry._event_listeners)
|
|
|
|
async def transform_context(
|
|
self, context: AgentContext, run: AgentRunContext
|
|
) -> AgentContext:
|
|
for transformer in self._context_transformers:
|
|
run.cancellation.raise_if_cancelled()
|
|
context = await transformer(context, run)
|
|
return context
|
|
|
|
async def before_tool_call(
|
|
self, call: ToolCallContext, run: AgentRunContext
|
|
) -> BeforeToolCallDecision:
|
|
current = call
|
|
for hook in self._before_tool_hooks:
|
|
run.cancellation.raise_if_cancelled()
|
|
decision = await hook(current, run)
|
|
if decision is None:
|
|
continue
|
|
if not decision.allowed:
|
|
return decision
|
|
if decision.input is not None:
|
|
current = replace(current, input=decision.input)
|
|
return BeforeToolCallDecision(input=current.input)
|
|
|
|
async def after_tool_call(
|
|
self,
|
|
call: ToolCallContext,
|
|
result: ToolResult,
|
|
run: AgentRunContext,
|
|
) -> AfterToolCallDecision:
|
|
current = result
|
|
terminate = False
|
|
for hook in self._after_tool_hooks:
|
|
run.cancellation.raise_if_cancelled()
|
|
decision = await hook(call, current, run)
|
|
if decision is None:
|
|
continue
|
|
if decision.result is not None:
|
|
current = decision.result
|
|
terminate = terminate or decision.terminate
|
|
return AfterToolCallDecision(result=current, terminate=terminate)
|
|
|
|
async def should_stop(self, turn: TurnSnapshot, run: AgentRunContext) -> bool:
|
|
for policy in self._turn_policies:
|
|
run.cancellation.raise_if_cancelled()
|
|
if await policy(turn, run):
|
|
return True
|
|
return False
|
|
|
|
async def publish(self, event: AgentEvent, run: AgentRunContext) -> None:
|
|
for listener in tuple(self._event_listeners):
|
|
await listener(event, run)
|
|
|
|
@property
|
|
def enabled(self) -> bool:
|
|
return any(
|
|
(
|
|
self._context_transformers,
|
|
self._before_tool_hooks,
|
|
self._after_tool_hooks,
|
|
self._turn_policies,
|
|
self._event_listeners,
|
|
)
|
|
)
|
|
|
|
|
|
__all__ = [
|
|
"AfterToolCallDecision",
|
|
"AgentContext",
|
|
"AgentEvent",
|
|
"AgentRunContext",
|
|
"BeforeToolCallDecision",
|
|
"CancellationToken",
|
|
"ExtensionRegistry",
|
|
"ToolCallContext",
|
|
"TurnSnapshot",
|
|
]
|