131 lines
4.1 KiB
Rust
131 lines
4.1 KiB
Rust
//! Observability hooks for the ReAct agent loop.
|
|
//!
|
|
//! Hooks allow injecting custom behavior (logging, tracing, filtering, termination)
|
|
//! at each step of the reasoning loop without coupling to the core agent logic.
|
|
//!
|
|
//! Inspired by rig's `PromptHook` trait.
|
|
//!
|
|
//! # Example
|
|
//!
|
|
//! ```ignore
|
|
//! #[derive(Clone)]
|
|
//! struct MyHook;
|
|
//!
|
|
//! impl Hook for MyHook {
|
|
//! async fn on_thought(&self, step: usize, thought: &str) -> HookAction {
|
|
//! tracing::info!("[step {}] thinking: {}", step, thought);
|
|
//! HookAction::Continue
|
|
//! }
|
|
//! }
|
|
//!
|
|
//! let agent = ReactAgent::new(prompt, tools, config).with_hook(MyHook);
|
|
//! ```
|
|
|
|
use async_trait::async_trait;
|
|
|
|
/// Controls whether the agent loop continues after a hook callback.
|
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
|
pub enum HookAction {
|
|
/// Continue processing normally.
|
|
Continue,
|
|
/// Skip the current step and continue.
|
|
Skip,
|
|
/// Terminate the loop immediately with the given reason.
|
|
Terminate(&'static str),
|
|
}
|
|
|
|
/// Controls behavior after a tool call hook callback.
|
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
|
pub enum ToolCallAction {
|
|
/// Execute the tool normally.
|
|
Continue,
|
|
/// Skip tool execution and inject a custom result.
|
|
Skip(String),
|
|
/// Terminate the loop with the given reason.
|
|
Terminate(&'static str),
|
|
}
|
|
|
|
/// Default no-op hook that does nothing.
|
|
#[derive(Debug, Clone, Copy, Default)]
|
|
pub struct NoopHook;
|
|
|
|
impl Hook for NoopHook {}
|
|
|
|
impl Hook for () {}
|
|
|
|
/// A hook that logs everything to stderr using `eprintln`.
|
|
/// No external dependencies required.
|
|
#[derive(Debug, Clone, Copy, Default)]
|
|
pub struct TracingHook;
|
|
|
|
impl TracingHook {
|
|
pub fn new() -> Self {
|
|
Self
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl Hook for TracingHook {
|
|
async fn on_thought(&self, step: usize, thought: &str) -> HookAction {
|
|
eprintln!("[step {}] thought: {}", step, thought);
|
|
HookAction::Continue
|
|
}
|
|
|
|
async fn on_tool_call(&self, step: usize, name: &str, args_json: &str) -> ToolCallAction {
|
|
eprintln!("[step {}] tool_call: {}({})", step, name, args_json);
|
|
ToolCallAction::Continue
|
|
}
|
|
|
|
async fn on_observation(&self, step: usize, observation: &str) -> HookAction {
|
|
eprintln!("[step {}] observation: {}", step, observation);
|
|
HookAction::Continue
|
|
}
|
|
|
|
async fn on_answer(&self, step: usize, answer: &str) -> HookAction {
|
|
eprintln!("[step {}] answer: {}", step, answer);
|
|
HookAction::Continue
|
|
}
|
|
}
|
|
|
|
/// Hook trait for observing and controlling the ReAct agent loop.
|
|
///
|
|
/// Implement this trait to inject custom behavior at each step:
|
|
/// - Log thoughts, tool calls, observations, and final answers
|
|
/// - Filter or redact sensitive data
|
|
/// - Dynamically terminate the loop based on content
|
|
/// - Inject custom tool results (e.g., for testing or sandboxing)
|
|
///
|
|
/// All methods have default no-op implementations, so you only need to
|
|
/// override the ones you care about.
|
|
///
|
|
/// The hook is called synchronously during the agent loop. Keep hook
|
|
/// callbacks fast — avoid blocking I/O. For heavy work, spawn a task
|
|
/// and return immediately.
|
|
#[async_trait]
|
|
pub trait Hook: Send + Sync {
|
|
/// Called when the agent emits a thought/reasoning step.
|
|
///
|
|
/// Return `HookAction::Terminate` to stop the loop early.
|
|
async fn on_thought(&self, _step: usize, _thought: &str) -> HookAction {
|
|
HookAction::Continue
|
|
}
|
|
|
|
/// Called just before a tool is executed.
|
|
///
|
|
/// Return `ToolCallAction::Skip(result)` to skip execution and inject `result` instead.
|
|
/// Return `ToolCallAction::Terminate` to stop the loop without executing the tool.
|
|
async fn on_tool_call(&self, _step: usize, _name: &str, _args_json: &str) -> ToolCallAction {
|
|
ToolCallAction::Continue
|
|
}
|
|
|
|
/// Called after a tool returns an observation.
|
|
async fn on_observation(&self, _step: usize, _observation: &str) -> HookAction {
|
|
HookAction::Continue
|
|
}
|
|
|
|
/// Called when the agent produces a final answer.
|
|
async fn on_answer(&self, _step: usize, _answer: &str) -> HookAction {
|
|
HookAction::Continue
|
|
}
|
|
}
|