Skip to main content

lellm_agent/
hook.rs

1//! Agent-level Hook — AgentLoop 执行扩展点。
2//!
3//! 与 graph 层 `AgentHook`(观测型)不同,agent 层 hook 可以:
4//! - 在 agent loop 执行前后注入 StateDelta
5//! - 观测 agent 执行过程
6//!
7//! # 设计原则
8//!
9//! - Hook 方法**可以修改 State**(通过返回 `Vec<StateDelta>`)
10//! - Hook 修改经过 Reducer,进入 ExecutionTrace — 无审计盲区
11//! - Hook 失败不影响 Agent 执行(降级为 ObservedError)
12
13use lellm_graph::StateDelta;
14
15use crate::runtime::{AgentEvent, ToolUseResult};
16
17/// Agent 执行上下文 — 传递给 Hook 的执行信息。
18pub struct AgentHookContext {
19    /// Agent 节点名称
20    pub node_name: String,
21    /// 输入消息数量
22    pub input_message_count: usize,
23}
24
25/// Agent 执行快照 — after_agent 收到的执行结果。
26pub struct AgentHookSnapshot {
27    /// 执行结果
28    pub result: ToolUseResult,
29    /// 收到的事件流(用于审计)
30    pub events: Vec<AgentEvent>,
31}
32
33/// Agent-level Hook trait。
34///
35/// 在 AgentFlowNode 执行 Agent Loop 前后调用。
36/// Hook 返回的 `Vec<StateDelta>` 会经过 Reducer 合并到 State。
37///
38/// # 示例
39///
40/// ```rust,ignore
41/// use lellm_agent::hook::{AgentHook, AgentHookContext, AgentHookSnapshot};
42/// use lellm_graph::StateDelta;
43///
44/// struct AuditHook;
45///
46/// impl AgentHook for AuditHook {
47///     fn before_agent(&self, ctx: &AgentHookContext) -> Vec<StateDelta> {
48///         vec![StateDelta::put("audit_log", serde_json::json!([
49///             format!("agent '{}' started", ctx.node_name)
50///         ]))]
51///     }
52///
53///     fn after_agent(&self, snapshot: &AgentHookSnapshot) -> Vec<StateDelta> {
54///         vec![StateDelta::put(
55///             "audit_log",
56///             serde_json::json!([format!("agent completed: {} iterations", snapshot.result.iterations)]),
57///         )]
58///     }
59/// }
60/// ```
61pub trait AgentHook: Send + Sync {
62    /// Agent loop 执行前调用。
63    /// 返回的 Deltas 会在 agent loop 启动前 apply 到 State。
64    fn before_agent(&self, _ctx: &AgentHookContext) -> Vec<StateDelta> {
65        Vec::new()
66    }
67
68    /// Agent loop 执行后调用。
69    /// 返回的 Deltas 会在 agent loop 完成后 apply 到 State。
70    fn after_agent(&self, _snapshot: &AgentHookSnapshot) -> Vec<StateDelta> {
71        Vec::new()
72    }
73}
74
75/// 无操作 Hook — 默认行为。
76#[derive(Debug, Clone, Default)]
77pub struct NoOpAgentHook;
78
79impl AgentHook for NoOpAgentHook {}
80
81/// 日志 Hook — 将 agent 执行事件输出为 tracing 日志。
82#[derive(Debug, Clone)]
83pub struct TracingAgentHook;
84
85impl AgentHook for TracingAgentHook {
86    fn before_agent(&self, ctx: &AgentHookContext) -> Vec<StateDelta> {
87        tracing::debug!(
88            node = %ctx.node_name,
89            input_messages = ctx.input_message_count,
90            "agent loop starting"
91        );
92        Vec::new()
93    }
94
95    fn after_agent(&self, snapshot: &AgentHookSnapshot) -> Vec<StateDelta> {
96        tracing::debug!(
97            iterations = snapshot.result.iterations,
98            tool_calls = snapshot.result.tool_calls_executed,
99            stop_reason = ?snapshot.result.stop_reason,
100            "agent loop completed"
101        );
102        Vec::new()
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109
110    #[test]
111    fn test_noop_agent_hook() {
112        let hook = NoOpAgentHook;
113        let ctx = AgentHookContext {
114            node_name: "test".to_string(),
115            input_message_count: 0,
116        };
117        let deltas = hook.before_agent(&ctx);
118        assert!(deltas.is_empty());
119    }
120
121    #[test]
122    fn test_tracing_agent_hook() {
123        let hook = TracingAgentHook;
124        let ctx = AgentHookContext {
125            node_name: "test".to_string(),
126            input_message_count: 5,
127        };
128        let deltas = hook.before_agent(&ctx);
129        assert!(deltas.is_empty());
130    }
131}