refactor: rename flow enums and improve tool interaction logic
This commit is contained in:
@@ -1,20 +1,14 @@
|
||||
use crate::agent::run::Flow;
|
||||
use crate::agent::run::{TextFlow, ToolCallFlow};
|
||||
use crate::agent::tool::ToolCall;
|
||||
|
||||
pub trait Hook: Send + Sync {
|
||||
fn before_tool_call(&self, tool_call: &ToolCall) -> Flow;
|
||||
|
||||
/// Called after a tool call completes, before the JSON output is serialized to a string.
|
||||
/// The `output` value can be modified in place.
|
||||
/// Return `Flow::Stop` to abort the agent loop, `Flow::Skip` to discard this result,
|
||||
/// or `Flow::Continue` to proceed normally.
|
||||
fn after_tool_call(&self, _tool_name: &str, _output: &mut serde_json::Value) -> Flow {
|
||||
Flow::Continue
|
||||
}
|
||||
}
|
||||
|
||||
impl Hook for () {
|
||||
fn before_tool_call(&self, _tool_call: &ToolCall) -> Flow {
|
||||
Flow::Continue
|
||||
fn on_text_reply(&self, text: &str) -> TextFlow {
|
||||
TextFlow::Continue
|
||||
}
|
||||
fn before_tool_call(&self, tool_call: &ToolCall) -> ToolCallFlow {
|
||||
ToolCallFlow::Continue
|
||||
}
|
||||
fn after_tool_call(&self, tool_name: &str, output: &mut serde_json::Value) -> ToolCallFlow {
|
||||
ToolCallFlow::Continue
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,24 +2,29 @@ use crate::agent::agent::Agent;
|
||||
use crate::agent::agent::LlmApiAdaptor;
|
||||
use crate::agent::error::PromptError;
|
||||
use crate::agent::hook::Hook;
|
||||
use crate::agent::memory::{AgentMemory, VecMemory};
|
||||
use crate::agent::tool::{ToolCall, ToolDyn, ToolExecutionResult, ToolFailure, ToolOutcome};
|
||||
use serde::Deserialize;
|
||||
use serde_json::Value;
|
||||
use std::pin::Pin;
|
||||
use crate::agent::run::PromptOutput::{Text, Tool};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PromptOutput {
|
||||
pub text: String,
|
||||
pub tool_calls: Vec<ToolCall>,
|
||||
pub enum PromptOutput {
|
||||
Text(String),
|
||||
Tool(ToolCall),
|
||||
}
|
||||
|
||||
/// Flow control for tool execution within the agent loop.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum Flow {
|
||||
pub enum TextFlow {
|
||||
Continue,
|
||||
Stop,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ToolCallFlow {
|
||||
Continue,
|
||||
/// Skip this tool call but continue processing other responses.
|
||||
Skip,
|
||||
/// Stop the agent execution entirely.
|
||||
Stop,
|
||||
}
|
||||
|
||||
@@ -69,45 +74,35 @@ where
|
||||
M: LlmApiAdaptor<Message> + Send + Sync + 'a,
|
||||
Message: Send + Sync + Clone + 'a,
|
||||
{
|
||||
type Output = Result<PromptOutput, PromptError>;
|
||||
type Output = Result<Vec<PromptOutput>, PromptError>;
|
||||
type IntoFuture = Pin<Box<dyn Future<Output = Self::Output> + Send + 'a>>;
|
||||
|
||||
fn into_future(self) -> Self::IntoFuture {
|
||||
let agent = self.agent;
|
||||
let message = self.message;
|
||||
let per_prompt_system_prompt = self.system_prompt;
|
||||
let extra_hooks = self.hooks;
|
||||
let max_turns = self.max_turns;
|
||||
|
||||
Box::pin(async move {
|
||||
let mut memory: Vec<Message> = Vec::new();
|
||||
let mut memory = agent
|
||||
.memory
|
||||
.unwrap_or_else(|| VecMemory::new(Vec::new()).into());
|
||||
let mut prompt_result: Vec<PromptOutput> = Vec::new();
|
||||
|
||||
// Pre-populate from agent's memory backend if available
|
||||
if let Some(ref mem_backend) = agent.memory {
|
||||
memory.extend(mem_backend.get_memory().iter().cloned());
|
||||
}
|
||||
let history = match agent.system_prompt {
|
||||
Some(ref sp) => [sp].iter().chain(memory.iter().chain(message.into())),
|
||||
None => memory.iter().chain(message.into()),
|
||||
};
|
||||
for _ in 0..max_turns {
|
||||
// TODO perf issue for copy cost
|
||||
let text = agent.model.complete(&history).await?;
|
||||
|
||||
// Effective system prompt: per-prompt override > agent default
|
||||
match per_prompt_system_prompt {
|
||||
Some(sp) => memory.push(sp),
|
||||
None => {
|
||||
if let Some(ref sp) = agent.system_prompt {
|
||||
memory.push(sp.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
memory.push(message);
|
||||
|
||||
for turn in 0..max_turns {
|
||||
// Delegate completion to the model (rig's CompletionModel pattern)
|
||||
let text = agent.model.complete(&memory).await?;
|
||||
|
||||
// Try to parse the LLM reply as JSON (handle markdown fences)
|
||||
// Try to parse the LLM reply as JSON (handle Markdown fences)
|
||||
let cleaned = clean_json_response(&text);
|
||||
|
||||
match serde_json::from_str::<LlmJsonResponse>(&cleaned) {
|
||||
Ok(json_resp) => {
|
||||
// Check for <end/> marker — stop the turn immediately
|
||||
// Check for <end/> marker — stop the _turn immediately
|
||||
let reply_has_end = json_resp.reply_text.contains("<end/>");
|
||||
let clean_text = json_resp
|
||||
.reply_text
|
||||
@@ -115,7 +110,19 @@ where
|
||||
.trim()
|
||||
.to_string();
|
||||
|
||||
// Run on_text_reply hooks
|
||||
for hook in agent.hooks.iter().chain(extra_hooks.iter()) {
|
||||
match hook.on_text_reply(&clean_text) {
|
||||
TextFlow::Stop => {
|
||||
return Err(PromptError::ToolError(
|
||||
"text reply rejected by hook".to_string(),
|
||||
));
|
||||
}
|
||||
TextFlow::Continue => {}
|
||||
}
|
||||
}
|
||||
memory.push(agent.model.build_assistant_message(&clean_text));
|
||||
prompt_result.push(Text(clean_text));
|
||||
|
||||
if reply_has_end {
|
||||
return Ok(PromptOutput {
|
||||
@@ -129,7 +136,7 @@ where
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(idx, tc)| ToolCall {
|
||||
id: format!("call_{turn}_{idx}"),
|
||||
id: format!("{idx}"),
|
||||
name: tc.tool_name,
|
||||
args: tc.args,
|
||||
})
|
||||
@@ -143,20 +150,28 @@ where
|
||||
}
|
||||
|
||||
for (tc_id, tc) in tool_calls.iter().enumerate() {
|
||||
match run_hooks(&agent.hooks, &extra_hooks, tc) {
|
||||
Flow::Stop => {
|
||||
return Err(PromptError::ToolError(
|
||||
"tool call rejected by hook".to_string(),
|
||||
));
|
||||
// inline run_hooks: before_tool_call hook
|
||||
let mut should_skip = false;
|
||||
for hook in agent.hooks.iter().chain(extra_hooks.iter()) {
|
||||
match hook.before_tool_call(tc) {
|
||||
ToolCallFlow::Stop => {
|
||||
return Err(PromptError::ToolError(
|
||||
"tool call rejected by hook".to_string(),
|
||||
));
|
||||
}
|
||||
ToolCallFlow::Skip => {
|
||||
let msg = agent
|
||||
.model
|
||||
.tool_result_message(&tc.name, "(skipped by hook)");
|
||||
memory.push(msg);
|
||||
should_skip = true;
|
||||
break;
|
||||
}
|
||||
ToolCallFlow::Continue => {}
|
||||
}
|
||||
Flow::Skip => {
|
||||
let msg = agent
|
||||
.model
|
||||
.tool_result_message(&tc.name, "(skipped by hook)");
|
||||
memory.push(msg);
|
||||
continue;
|
||||
}
|
||||
Flow::Continue => {}
|
||||
}
|
||||
if should_skip {
|
||||
continue;
|
||||
}
|
||||
|
||||
let result = {
|
||||
@@ -181,21 +196,21 @@ where
|
||||
};
|
||||
|
||||
match result.outcome {
|
||||
ToolOutcome::Success => {
|
||||
ToolOutcome::InternalOk => {
|
||||
let mut json_output = result.output;
|
||||
// Run after_tool_call hooks to allow output modification
|
||||
for hook in agent.hooks.iter().chain(extra_hooks.iter()) {
|
||||
match hook.after_tool_call(&tc.name, &mut json_output) {
|
||||
Flow::Stop => {
|
||||
ToolCallFlow::Stop => {
|
||||
return Err(PromptError::ToolError(
|
||||
"tool call output caused this turn stopped by hook".to_string(),
|
||||
"tool call output caused this _turn stopped by hook".to_string(),
|
||||
));
|
||||
}
|
||||
Flow::Skip => {
|
||||
ToolCallFlow::Skip => {
|
||||
json_output = serde_json::json!({"message":"tool call is ignored by hook"});
|
||||
break;
|
||||
}
|
||||
Flow::Continue => {}
|
||||
ToolCallFlow::Continue => {}
|
||||
}
|
||||
}
|
||||
let arg_json = serde_json::to_string(&tc.args)
|
||||
@@ -209,11 +224,12 @@ where
|
||||
let msg =
|
||||
agent.model.tool_result_message(&tc.name, &output_str);
|
||||
memory.push(msg);
|
||||
prompt_result.push(Tool(tc))
|
||||
}
|
||||
ToolOutcome::Failure(failure) if failure.fatal => {
|
||||
return Err(PromptError::InternalError(failure.message));
|
||||
ToolOutcome::InternalError(failure) if failure.fatal => {
|
||||
return Err(PromptError::ToolError(failure.message));
|
||||
}
|
||||
ToolOutcome::Failure(failure) => {
|
||||
ToolOutcome::InternalError(failure) => {
|
||||
let msg = agent.model.tool_result_message(
|
||||
&tc.name,
|
||||
&format!("error: {}", failure.message),
|
||||
@@ -224,7 +240,7 @@ where
|
||||
}
|
||||
}
|
||||
Err(_) => {
|
||||
// Not valid JSON — push as assistant message and continue to next turn
|
||||
// Not valid JSON — push as assistant message and continue to next _turn
|
||||
memory.push(agent.model.build_assistant_message(&text));
|
||||
continue;
|
||||
}
|
||||
@@ -236,21 +252,6 @@ where
|
||||
}
|
||||
}
|
||||
|
||||
fn run_hooks(
|
||||
agent_hooks: &[Box<dyn Hook>],
|
||||
extra_hooks: &[Box<dyn Hook>],
|
||||
tool_call: &ToolCall,
|
||||
) -> Flow {
|
||||
for hook in agent_hooks.iter().chain(extra_hooks.iter()) {
|
||||
match hook.before_tool_call(tool_call) {
|
||||
Flow::Stop => return Flow::Stop,
|
||||
Flow::Skip => return Flow::Skip,
|
||||
Flow::Continue => {}
|
||||
}
|
||||
}
|
||||
Flow::Continue
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct LlmJsonResponse {
|
||||
reply_text: String,
|
||||
@@ -262,7 +263,7 @@ struct LlmJsonResponse {
|
||||
struct LlmToolCallRequest {
|
||||
tool_name: String,
|
||||
#[serde(default)]
|
||||
args: serde_json::Value,
|
||||
args: Value,
|
||||
}
|
||||
|
||||
pub(crate) fn clean_json_response(text: &str) -> String {
|
||||
|
||||
@@ -84,8 +84,8 @@ impl ToolFailure {
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub enum ToolOutcome {
|
||||
Success,
|
||||
Failure(ToolFailure),
|
||||
InternalOk,
|
||||
InternalError(ToolFailure),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
@@ -98,21 +98,21 @@ impl ToolExecutionResult {
|
||||
pub fn success(output: serde_json::Value) -> Self {
|
||||
Self {
|
||||
output,
|
||||
outcome: ToolOutcome::Success,
|
||||
outcome: ToolOutcome::InternalOk,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn failed(output: serde_json::Value, failure: ToolFailure) -> Self {
|
||||
Self {
|
||||
output,
|
||||
outcome: ToolOutcome::Failure(failure),
|
||||
outcome: ToolOutcome::InternalError(failure),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn failure(&self) -> Option<&ToolFailure> {
|
||||
match &self.outcome {
|
||||
ToolOutcome::Success => None,
|
||||
ToolOutcome::Failure(failure) => Some(failure),
|
||||
ToolOutcome::InternalOk => None,
|
||||
ToolOutcome::InternalError(failure) => Some(failure),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user