pi harness reimplement
This commit is contained in:
@@ -0,0 +1,861 @@
|
||||
"""
|
||||
agent_loop.jl - Low-level agent loop implementation
|
||||
|
||||
This module implements the core agentLoop functionality that works with AgentMessage
|
||||
throughout, transforming to Message[] only at the LLM call boundary.
|
||||
"""
|
||||
|
||||
module AgentLoop
|
||||
|
||||
using ..Types: *
|
||||
using ..StreamFn: *
|
||||
|
||||
# ============================================================================
|
||||
# Event sink type
|
||||
# ============================================================================
|
||||
|
||||
const AgentEventSink = Function
|
||||
|
||||
# ============================================================================
|
||||
# Main agent loop function
|
||||
# ============================================================================
|
||||
|
||||
function agentLoop(
|
||||
prompts::Vector{AgentMessage},
|
||||
context::AgentContext,
|
||||
config::AgentLoopConfig,
|
||||
signal::Union{Nothing, AbortSignal},
|
||||
stream_fn::StreamFn,
|
||||
)::EventStream
|
||||
stream = createAgentStream()
|
||||
|
||||
Threads.@spawn begin
|
||||
messages = runAgentLoop(
|
||||
prompts,
|
||||
context,
|
||||
config,
|
||||
(event) -> push!(stream, event),
|
||||
signal,
|
||||
stream_fn,
|
||||
)
|
||||
end(stream, messages)
|
||||
end
|
||||
|
||||
return stream
|
||||
end
|
||||
|
||||
# ============================================================================
|
||||
# Continue agent loop function
|
||||
# ============================================================================
|
||||
|
||||
function agentLoopContinue(
|
||||
context::AgentContext,
|
||||
config::AgentLoopConfig,
|
||||
signal::Union{Nothing, AbortSignal},
|
||||
stream_fn::StreamFn,
|
||||
)::EventStream
|
||||
if isempty(context.messages)
|
||||
throw(ErrorException("Cannot continue: no messages in context"))
|
||||
end
|
||||
|
||||
if context.messages[end].role == "assistant"
|
||||
throw(ErrorException("Cannot continue from message role: assistant"))
|
||||
end
|
||||
|
||||
stream = createAgentStream()
|
||||
|
||||
Threads.@spawn begin
|
||||
messages = runAgentLoopContinue(
|
||||
context,
|
||||
config,
|
||||
(event) -> push!(stream, event),
|
||||
signal,
|
||||
stream_fn,
|
||||
)
|
||||
end(stream, messages)
|
||||
end
|
||||
|
||||
return stream
|
||||
end
|
||||
|
||||
# ============================================================================
|
||||
# Run agent loop function
|
||||
# ============================================================================
|
||||
|
||||
function runAgentLoop(
|
||||
prompts::Vector{AgentMessage},
|
||||
context::AgentContext,
|
||||
config::AgentLoopConfig,
|
||||
emit::AgentEventSink,
|
||||
signal::Union{Nothing, AbortSignal},
|
||||
stream_fn::StreamFn,
|
||||
)::Vector{AgentMessage}
|
||||
new_messages::Vector{AgentMessage} = copy(prompts)
|
||||
current_context::AgentContext = AgentContext(
|
||||
context.system_prompt,
|
||||
vcat(context.messages, copy(prompts)),
|
||||
context.tools,
|
||||
)
|
||||
|
||||
emit(AgentStartEvent())
|
||||
emit(TurnStartEvent())
|
||||
for prompt in prompts
|
||||
emit(MessageStartEvent(prompt))
|
||||
emit(MessageEndEvent(prompt))
|
||||
end
|
||||
|
||||
runLoop(
|
||||
current_context,
|
||||
new_messages,
|
||||
config,
|
||||
signal,
|
||||
emit,
|
||||
stream_fn,
|
||||
)
|
||||
return new_messages
|
||||
end
|
||||
|
||||
# ============================================================================
|
||||
# Run agent loop continue function
|
||||
# ============================================================================
|
||||
|
||||
function runAgentLoopContinue(
|
||||
context::AgentContext,
|
||||
config::AgentLoopConfig,
|
||||
emit::AgentEventSink,
|
||||
signal::Union{Nothing, AbortSignal},
|
||||
stream_fn::StreamFn,
|
||||
)::Vector{AgentMessage}
|
||||
if isempty(context.messages)
|
||||
throw(ErrorException("Cannot continue: no messages in context"))
|
||||
end
|
||||
|
||||
if context.messages[end].role == "assistant"
|
||||
throw(ErrorException("Cannot continue from message role: assistant"))
|
||||
end
|
||||
|
||||
new_messages::Vector{AgentMessage} = []
|
||||
current_context::AgentContext = context
|
||||
|
||||
emit(AgentStartEvent())
|
||||
emit(TurnStartEvent())
|
||||
|
||||
runLoop(
|
||||
current_context,
|
||||
new_messages,
|
||||
config,
|
||||
signal,
|
||||
emit,
|
||||
stream_fn,
|
||||
)
|
||||
return new_messages
|
||||
end
|
||||
|
||||
# ============================================================================
|
||||
# Create agent stream function
|
||||
# ============================================================================
|
||||
|
||||
function createAgentStream()::EventStream
|
||||
return EventStream(
|
||||
(event::AgentEvent) -> event isa AgentEndEvent,
|
||||
(event::AgentEvent) -> event isa AgentEndEvent ? event.messages : AgentMessage[],
|
||||
)
|
||||
end
|
||||
|
||||
# ============================================================================
|
||||
# Main loop logic shared by agentLoop and agentLoopContinue
|
||||
# ============================================================================
|
||||
|
||||
function runLoop(
|
||||
initial_context::AgentContext,
|
||||
new_messages::Vector{AgentMessage},
|
||||
initial_config::AgentLoopConfig,
|
||||
signal::Union{Nothing, AbortSignal},
|
||||
emit::AgentEventSink,
|
||||
stream_function::StreamFn,
|
||||
)::Nothing
|
||||
current_context::AgentContext = initial_context
|
||||
config::AgentLoopConfig = initial_config
|
||||
first_turn::Bool = true
|
||||
pending_messages::Vector{AgentMessage} = getSteeringMessages(config) do
|
||||
get_steering_messages(config)
|
||||
end
|
||||
|
||||
while true
|
||||
has_more_tool_calls::Bool = true
|
||||
|
||||
while has_more_tool_calls || !isempty(pending_messages)
|
||||
if !first_turn
|
||||
emit(TurnStartEvent())
|
||||
else
|
||||
first_turn = false
|
||||
end
|
||||
|
||||
if !isempty(pending_messages)
|
||||
for message in pending_messages
|
||||
emit(MessageStartEvent(message))
|
||||
emit(MessageEndEvent(message))
|
||||
push!(current_context.messages, message)
|
||||
push!(new_messages, message)
|
||||
end
|
||||
pending_messages = AgentMessage[]
|
||||
end
|
||||
|
||||
message = streamAssistantResponse(
|
||||
current_context,
|
||||
config,
|
||||
signal,
|
||||
emit,
|
||||
stream_function,
|
||||
)
|
||||
push!(new_messages, message)
|
||||
|
||||
if message.stop_reason in ("error", "aborted")
|
||||
emit(TurnEndEvent(message, ToolResultMessage[]))
|
||||
emit(AgentEndEvent(new_messages))
|
||||
return
|
||||
end
|
||||
|
||||
tool_calls = filter(
|
||||
(c) -> c isa ToolCall,
|
||||
message.content,
|
||||
)
|
||||
|
||||
tool_results::Vector{ToolResultMessage} = []
|
||||
has_more_tool_calls = false
|
||||
if !isempty(tool_calls)
|
||||
executed_tool_batch =
|
||||
message.stop_reason == "length"
|
||||
? failToolCallsFromTruncatedMessage(tool_calls, emit)
|
||||
: executeToolCalls(
|
||||
current_context,
|
||||
message,
|
||||
config,
|
||||
signal,
|
||||
emit,
|
||||
)
|
||||
append!(tool_results, executed_tool_batch.messages)
|
||||
has_more_tool_calls = !executed_tool_batch.terminate
|
||||
|
||||
for result in tool_results
|
||||
push!(current_context.messages, result)
|
||||
push!(new_messages, result)
|
||||
end
|
||||
end
|
||||
|
||||
emit(TurnEndEvent(message, tool_results))
|
||||
|
||||
next_turn_context = PrepareNextTurnContext(
|
||||
message,
|
||||
tool_results,
|
||||
current_context,
|
||||
new_messages,
|
||||
)
|
||||
next_turn_snapshot = prepare_next_turn(config, next_turn_context)
|
||||
|
||||
if !isnothing(next_turn_snapshot)
|
||||
current_context = next_turn_snapshot.context
|
||||
config = AgentLoopConfig(
|
||||
model = next_turn_snapshot.model,
|
||||
reasoning = next_turn_snapshot.thinking_level,
|
||||
convert_to_llm = config.convert_to_llm,
|
||||
transform_context = config.transform_context,
|
||||
get_api_key = config.get_api_key,
|
||||
should_stop_after_turn = config.should_stop_after_turn,
|
||||
prepare_next_turn = config.prepare_next_turn,
|
||||
get_steering_messages = config.get_steering_messages,
|
||||
get_follow_up_messages = config.get_follow_up_messages,
|
||||
tool_execution = config.tool_execution,
|
||||
before_tool_call = config.before_tool_call,
|
||||
after_tool_call = config.after_tool_call,
|
||||
max_tokens = config.max_tokens,
|
||||
temperature = config.temperature,
|
||||
reasoning = config.reasoning,
|
||||
cache_retention = config.cache_retention,
|
||||
session_id = config.session_id,
|
||||
headers = config.headers,
|
||||
metadata = config.metadata,
|
||||
transport = config.transport,
|
||||
signal = signal,
|
||||
api_key = config.api_key,
|
||||
on_payload = config.on_payload,
|
||||
on_response = config.on_response,
|
||||
max_retry_delay_ms = config.max_retry_delay_ms,
|
||||
)
|
||||
end
|
||||
|
||||
if should_stop_after_turn(config, next_turn_context)
|
||||
emit(AgentEndEvent(new_messages))
|
||||
return
|
||||
end
|
||||
|
||||
pending_messages = getSteeringMessages(config) do
|
||||
get_steering_messages(config)
|
||||
end
|
||||
end
|
||||
|
||||
follow_up_messages = getFollowUpMessages(config) do
|
||||
get_follow_up_messages(config)
|
||||
end
|
||||
|
||||
if !isempty(follow_up_messages)
|
||||
pending_messages = follow_up_messages
|
||||
continue
|
||||
end
|
||||
|
||||
break
|
||||
end
|
||||
|
||||
emit(AgentEndEvent(new_messages))
|
||||
end
|
||||
|
||||
# ============================================================================
|
||||
# Helper types
|
||||
# ============================================================================
|
||||
|
||||
struct PrepareNextTurnContext
|
||||
message::AssistantMessage
|
||||
tool_results::Vector{ToolResultMessage}
|
||||
context::AgentContext
|
||||
new_messages::Vector{AgentMessage}
|
||||
end
|
||||
|
||||
struct AgentLoopTurnUpdate
|
||||
context::Union{AgentContext, Nothing}
|
||||
model::Union{Model, Nothing}
|
||||
thinking_level::Union{ThinkingLevel, Nothing}
|
||||
end
|
||||
|
||||
# ============================================================================
|
||||
# Helper functions for getting messages from queues
|
||||
# ============================================================================
|
||||
|
||||
macro getSteeringMessages(config)
|
||||
:(get_steering_messages($(esc(config))))
|
||||
end
|
||||
|
||||
macro getFollowUpMessages(config)
|
||||
:(get_follow_up_messages($(esc(config))))
|
||||
end
|
||||
|
||||
function get_steering_messages(config::AgentLoopConfig)::Vector{AgentMessage}
|
||||
return isnothing(config.get_steering_messages) ? AgentMessage[] : config.get_steering_messages()
|
||||
end
|
||||
|
||||
function get_follow_up_messages(config::AgentLoopConfig)::Vector{AgentMessage}
|
||||
return isnothing(config.get_follow_up_messages) ? AgentMessage[] : config.get_follow_up_messages()
|
||||
end
|
||||
|
||||
function prepare_next_turn(config::AgentLoopConfig, context::PrepareNextTurnContext)::Union{AgentLoopTurnUpdate, Nothing}
|
||||
return isnothing(config.prepare_next_turn) ? nothing : config.prepare_next_turn(context)
|
||||
end
|
||||
|
||||
function should_stop_after_turn(config::AgentLoopConfig, context::PrepareNextTurnContext)::Bool
|
||||
return isnothing(config.should_stop_after_turn) ? false : config.should_stop_after_turn(context)
|
||||
end
|
||||
|
||||
# ============================================================================
|
||||
# Stream assistant response function
|
||||
# ============================================================================
|
||||
|
||||
function streamAssistantResponse(
|
||||
context::AgentContext,
|
||||
config::AgentLoopConfig,
|
||||
signal::Union{Nothing, AbortSignal},
|
||||
emit::AgentEventSink,
|
||||
stream_function::StreamFn,
|
||||
)::AssistantMessage
|
||||
messages::Vector{AgentMessage} = context.messages
|
||||
|
||||
if !isnothing(config.transform_context)
|
||||
messages = config.transform_context(messages, signal)
|
||||
end
|
||||
|
||||
llm_messages::Vector{Message} = config.convert_to_llm(messages)
|
||||
|
||||
llm_context::Context = Context(
|
||||
context.system_prompt,
|
||||
llm_messages,
|
||||
context.tools,
|
||||
)
|
||||
|
||||
resolved_api_key::Union{String, Nothing} =
|
||||
!isnothing(config.get_api_key)
|
||||
? config.get_api_key(config.model.provider)
|
||||
: nothing
|
||||
|
||||
response = stream_function(
|
||||
config.model,
|
||||
llm_context,
|
||||
merge(
|
||||
config,
|
||||
Dict(:apiKey => resolved_api_key, :signal => signal),
|
||||
),
|
||||
)
|
||||
|
||||
partial_message::Union{AssistantMessage, Nothing} = nothing
|
||||
added_partial::Bool = false
|
||||
|
||||
for event in response
|
||||
if event.type == "start"
|
||||
partial_message = event.partial
|
||||
push!(context.messages, partial_message)
|
||||
added_partial = true
|
||||
emit(MessageStartEvent(copy(partial_message)))
|
||||
elseif event.type in ("text_start", "text_delta", "text_end", "thinking_start", "thinking_delta", "thinking_end", "toolcall_start", "toolcall_delta", "toolcall_end")
|
||||
if !isnothing(partial_message)
|
||||
partial_message = event.partial
|
||||
context.messages[end] = partial_message
|
||||
emit(MessageUpdateEvent(copy(partial_message), event))
|
||||
end
|
||||
elseif event.type in ("done", "error")
|
||||
final_message = response.result()
|
||||
if added_partial
|
||||
context.messages[end] = final_message
|
||||
else
|
||||
push!(context.messages, final_message)
|
||||
end
|
||||
if !added_partial
|
||||
emit(MessageStartEvent(copy(final_message)))
|
||||
end
|
||||
emit(MessageEndEvent(final_message))
|
||||
return final_message
|
||||
end
|
||||
end
|
||||
|
||||
final_message = response.result()
|
||||
if added_partial
|
||||
context.messages[end] = final_message
|
||||
else
|
||||
push!(context.messages, final_message)
|
||||
emit(MessageStartEvent(copy(final_message)))
|
||||
end
|
||||
emit(MessageEndEvent(final_message))
|
||||
return final_message
|
||||
end
|
||||
|
||||
# ============================================================================
|
||||
# Fail tool calls from truncated message
|
||||
# ============================================================================
|
||||
|
||||
struct ExecutedToolCallBatch
|
||||
messages::Vector{ToolResultMessage}
|
||||
terminate::Bool
|
||||
end
|
||||
|
||||
function failToolCallsFromTruncatedMessage(
|
||||
tool_calls::Vector{ToolCall},
|
||||
emit::AgentEventSink,
|
||||
)::ExecutedToolCallBatch
|
||||
messages::Vector{ToolResultMessage} = []
|
||||
|
||||
for tool_call in tool_calls
|
||||
emit(ToolExecutionStartEvent(tool_call.id, tool_call.name, tool_call.arguments))
|
||||
|
||||
finalized = FinalizedToolCallOutcome(
|
||||
tool_call,
|
||||
createErrorToolResult(
|
||||
"Tool call \"$(tool_call.name)\" was not executed: the response hit the output token limit, so its arguments may be truncated. Re-issue the tool call with complete arguments.",
|
||||
),
|
||||
true,
|
||||
)
|
||||
|
||||
emitToolExecutionEnd(finalized, emit)
|
||||
tool_result_message = createToolResultMessage(finalized)
|
||||
emitToolResultMessage(tool_result_message, emit)
|
||||
push!(messages, tool_result_message)
|
||||
end
|
||||
|
||||
return ExecutedToolCallBatch(messages, false)
|
||||
end
|
||||
|
||||
# ============================================================================
|
||||
# Execute tool calls
|
||||
# ============================================================================
|
||||
|
||||
function executeToolCalls(
|
||||
current_context::AgentContext,
|
||||
assistant_message::AssistantMessage,
|
||||
config::AgentLoopConfig,
|
||||
signal::Union{Nothing, AbortSignal},
|
||||
emit::AgentEventSink,
|
||||
)::ExecutedToolCallBatch
|
||||
tool_calls = filter(
|
||||
(c) -> c isa ToolCall,
|
||||
assistant_message.content,
|
||||
)
|
||||
|
||||
has_sequential_tool_call = any(
|
||||
(tc) -> begin
|
||||
tool = findfirst((t) -> t.name == tc.name, current_context.tools)
|
||||
!isnothing(tool) && tool.execution_mode == EXECUTION_SEQUENTIAL
|
||||
end,
|
||||
tool_calls,
|
||||
)
|
||||
|
||||
if config.tool_execution == EXECUTION_SEQUENTIAL || has_sequential_tool_call
|
||||
return executeToolCallsSequential(
|
||||
current_context,
|
||||
assistant_message,
|
||||
tool_calls,
|
||||
config,
|
||||
signal,
|
||||
emit,
|
||||
)
|
||||
end
|
||||
return executeToolCallsParallel(
|
||||
current_context,
|
||||
assistant_message,
|
||||
tool_calls,
|
||||
config,
|
||||
signal,
|
||||
emit,
|
||||
)
|
||||
end
|
||||
|
||||
# ============================================================================
|
||||
# Execute tool calls sequentially
|
||||
# ============================================================================
|
||||
|
||||
function executeToolCallsSequential(
|
||||
current_context::AgentContext,
|
||||
assistant_message::AssistantMessage,
|
||||
tool_calls::Vector{ToolCall},
|
||||
config::AgentLoopConfig,
|
||||
signal::Union{Nothing, AbortSignal},
|
||||
emit::AgentEventSink,
|
||||
)::ExecutedToolCallBatch
|
||||
finalized_calls::Vector{FinalizedToolCallOutcome} = []
|
||||
messages::Vector{ToolResultMessage} = []
|
||||
|
||||
for tool_call in tool_calls
|
||||
emit(ToolExecutionStartEvent(tool_call.id, tool_call.name, tool_call.arguments))
|
||||
|
||||
preparation = prepareToolCall(current_context, assistant_message, tool_call, config, signal)
|
||||
|
||||
finalized = if preparation.kind == "immediate"
|
||||
FinalizedToolCallOutcome(tool_call, preparation.result, preparation.is_error)
|
||||
else
|
||||
executed = executePreparedToolCall(preparation, signal, emit)
|
||||
finalizeExecutedToolCall(
|
||||
current_context,
|
||||
assistant_message,
|
||||
preparation,
|
||||
executed,
|
||||
config,
|
||||
signal,
|
||||
)
|
||||
end
|
||||
|
||||
emitToolExecutionEnd(finalized, emit)
|
||||
tool_result_message = createToolResultMessage(finalized)
|
||||
emitToolResultMessage(tool_result_message, emit)
|
||||
push!(finalized_calls, finalized)
|
||||
push!(messages, tool_result_message)
|
||||
|
||||
if !isnothing(signal) && signal.aborted
|
||||
break
|
||||
end
|
||||
end
|
||||
|
||||
return ExecutedToolCallBatch(messages, shouldTerminateToolBatch(finalized_calls))
|
||||
end
|
||||
|
||||
# ============================================================================
|
||||
# Execute tool calls in parallel
|
||||
# ============================================================================
|
||||
|
||||
function executeToolCallsParallel(
|
||||
current_context::AgentContext,
|
||||
assistant_message::AssistantMessage,
|
||||
tool_calls::Vector{ToolCall},
|
||||
config::AgentLoopConfig,
|
||||
signal::Union{Nothing, AbortSignal},
|
||||
emit::AgentEventSink,
|
||||
)::ExecutedToolCallBatch
|
||||
finalized_calls::Vector{Union{FinalizedToolCallOutcome, Function}} = []
|
||||
|
||||
for tool_call in tool_calls
|
||||
emit(ToolExecutionStartEvent(tool_call.id, tool_call.name, tool_call.arguments))
|
||||
|
||||
preparation = prepareToolCall(current_context, assistant_message, tool_call, config, signal)
|
||||
|
||||
if preparation.kind == "immediate"
|
||||
finalized = FinalizedToolCallOutcome(
|
||||
tool_call,
|
||||
preparation.result,
|
||||
preparation.is_error,
|
||||
)
|
||||
emitToolExecutionEnd(finalized, emit)
|
||||
push!(finalized_calls, finalized)
|
||||
if !isnothing(signal) && signal.aborted
|
||||
break
|
||||
end
|
||||
continue
|
||||
end
|
||||
|
||||
push!(finalized_calls, () -> begin
|
||||
executed = executePreparedToolCall(preparation, signal, emit)
|
||||
finalized = finalizeExecutedToolCall(
|
||||
current_context,
|
||||
assistant_message,
|
||||
preparation,
|
||||
executed,
|
||||
config,
|
||||
signal,
|
||||
)
|
||||
emitToolExecutionEnd(finalized, emit)
|
||||
return finalized
|
||||
end)
|
||||
|
||||
if !isnothing(signal) && signal.aborted
|
||||
break
|
||||
end
|
||||
end
|
||||
|
||||
ordered_finalized_calls = map(
|
||||
(entry) -> if entry isa Function
|
||||
entry()
|
||||
else
|
||||
entry
|
||||
end,
|
||||
finalized_calls,
|
||||
)
|
||||
|
||||
messages::Vector{ToolResultMessage} = []
|
||||
for finalized in ordered_finalized_calls
|
||||
tool_result_message = createToolResultMessage(finalized)
|
||||
emitToolResultMessage(tool_result_message, emit)
|
||||
push!(messages, tool_result_message)
|
||||
end
|
||||
|
||||
return ExecutedToolCallBatch(messages, shouldTerminateToolBatch(ordered_finalized_calls))
|
||||
end
|
||||
|
||||
# ============================================================================
|
||||
# Prepared tool call types
|
||||
# ============================================================================
|
||||
|
||||
struct PreparedToolCall
|
||||
kind::String
|
||||
tool_call::ToolCall
|
||||
tool::AgentTool
|
||||
args::Any
|
||||
end
|
||||
|
||||
struct ImmediateToolCallOutcome
|
||||
kind::String
|
||||
result::AgentToolResultMutable
|
||||
is_error::Bool
|
||||
end
|
||||
|
||||
struct ExecutedToolCallOutcome
|
||||
result::AgentToolResultMutable
|
||||
is_error::Bool
|
||||
end
|
||||
|
||||
struct FinalizedToolCallOutcome
|
||||
tool_call::ToolCall
|
||||
result::AgentToolResultMutable
|
||||
is_error::Bool
|
||||
end
|
||||
|
||||
# ============================================================================
|
||||
# Helper functions
|
||||
# ============================================================================
|
||||
|
||||
function shouldTerminateToolBatch(finalized_calls::Vector{FinalizedToolCallOutcome})::Bool
|
||||
return !isempty(finalized_calls) && all(
|
||||
(finalized) -> finalized.result.terminate === true,
|
||||
finalized_calls,
|
||||
)
|
||||
end
|
||||
|
||||
function prepareToolCallArguments(tool::AgentTool, tool_call::ToolCall)::ToolCall
|
||||
if isnothing(tool.prepare_arguments)
|
||||
return tool_call
|
||||
end
|
||||
prepared_arguments = tool.prepare_arguments(tool_call.arguments)
|
||||
if prepared_arguments === tool_call.arguments
|
||||
return tool_call
|
||||
end
|
||||
return ToolCall(
|
||||
tool_call.type,
|
||||
tool_call.id,
|
||||
tool_call.name,
|
||||
prepared_arguments,
|
||||
tool_call.partial_json,
|
||||
)
|
||||
end
|
||||
|
||||
function prepareToolCall(
|
||||
current_context::AgentContext,
|
||||
assistant_message::AssistantMessage,
|
||||
tool_call::ToolCall,
|
||||
config::AgentLoopConfig,
|
||||
signal::Union{Nothing, AbortSignal},
|
||||
)::Union{PreparedToolCall, ImmediateToolCallOutcome}
|
||||
tool = findfirst((t) -> t.name == tool_call.name, current_context.tools)
|
||||
if isnothing(tool)
|
||||
return ImmediateToolCallOutcome("immediate", createErrorToolResult("Tool $(tool_call.name) not found"), true)
|
||||
end
|
||||
|
||||
try
|
||||
prepared_tool_call = prepareToolCallArguments(tool, tool_call)
|
||||
validated_args = validateToolArguments(tool, prepared_tool_call)
|
||||
|
||||
if !isnothing(config.before_tool_call)
|
||||
before_result = config.before_tool_call(
|
||||
BeforeToolCallContext(assistant_message, tool_call, validated_args, current_context),
|
||||
signal,
|
||||
)
|
||||
if !isnothing(signal) && signal.aborted
|
||||
return ImmediateToolCallOutcome("immediate", createErrorToolResult("Operation aborted"), true)
|
||||
end
|
||||
if !isnothing(before_result) && before_result.block
|
||||
reason = isnothing(before_result.reason) ? "Tool execution was blocked" : before_result.reason
|
||||
return ImmediateToolCallOutcome("immediate", createErrorToolResult(reason), true)
|
||||
end
|
||||
end
|
||||
|
||||
if !isnothing(signal) && signal.aborted
|
||||
return ImmediateToolCallOutcome("immediate", createErrorToolResult("Operation aborted"), true)
|
||||
end
|
||||
|
||||
return PreparedToolCall("prepared", tool_call, tool, validated_args)
|
||||
catch error
|
||||
return ImmediateToolCallOutcome("immediate", createErrorToolResult(string(error)), true)
|
||||
end
|
||||
end
|
||||
|
||||
function executePreparedToolCall(
|
||||
prepared::PreparedToolCall,
|
||||
signal::Union{Nothing, AbortSignal},
|
||||
emit::AgentEventSink,
|
||||
)::ExecutedToolCallOutcome
|
||||
update_events::Vector{Future} = []
|
||||
accepting_updates::Bool = true
|
||||
|
||||
try
|
||||
result = prepared.tool.execute(
|
||||
prepared.tool_call.id,
|
||||
prepared.args,
|
||||
signal,
|
||||
(partial_result) -> begin
|
||||
if !accepting_updates
|
||||
return
|
||||
end
|
||||
push!(
|
||||
update_events,
|
||||
Threads.@spawn begin
|
||||
emit(
|
||||
ToolExecutionUpdateEvent(
|
||||
prepared.tool_call.id,
|
||||
prepared.tool_call.name,
|
||||
prepared.tool_call.arguments,
|
||||
partial_result,
|
||||
),
|
||||
)
|
||||
end,
|
||||
)
|
||||
end,
|
||||
)
|
||||
accepting_updates = false
|
||||
wait.(update_events)
|
||||
return ExecutedToolCallOutcome(result, false)
|
||||
catch error
|
||||
accepting_updates = false
|
||||
wait.(update_events)
|
||||
return ExecutedToolCallOutcome(createErrorToolResult(string(error)), true)
|
||||
finally
|
||||
accepting_updates = false
|
||||
end
|
||||
end
|
||||
|
||||
function finalizeExecutedToolCall(
|
||||
current_context::AgentContext,
|
||||
assistant_message::AssistantMessage,
|
||||
prepared::PreparedToolCall,
|
||||
executed::ExecutedToolCallOutcome,
|
||||
config::AgentLoopConfig,
|
||||
signal::Union{Nothing, AbortSignal},
|
||||
)::FinalizedToolCallOutcome
|
||||
result = executed.result
|
||||
is_error = executed.is_error
|
||||
|
||||
if !isnothing(config.after_tool_call)
|
||||
try
|
||||
after_result = config.after_tool_call(
|
||||
AfterToolCallContext(
|
||||
assistant_message,
|
||||
prepared.tool_call,
|
||||
prepared.args,
|
||||
result,
|
||||
is_error,
|
||||
current_context,
|
||||
),
|
||||
signal,
|
||||
)
|
||||
if !isnothing(after_result)
|
||||
result = AgentToolResultMutable(
|
||||
isnothing(after_result.content) ? result.content : after_result.content,
|
||||
isnothing(after_result.details) ? result.details : after_result.details,
|
||||
isnothing(after_result.usage) ? result.usage : after_result.usage,
|
||||
result.added_tool_names,
|
||||
isnothing(after_result.terminate) ? result.terminate : after_result.terminate,
|
||||
)
|
||||
is_error = isnothing(after_result.is_error) ? is_error : after_result.is_error
|
||||
end
|
||||
catch error
|
||||
result = createErrorToolResult(string(error))
|
||||
is_error = true
|
||||
end
|
||||
end
|
||||
|
||||
return FinalizedToolCallOutcome(prepared.tool_call, result, is_error)
|
||||
end
|
||||
|
||||
function createErrorToolResult(message::String)::AgentToolResultMutable
|
||||
return AgentToolResultMutable([TextContent(message)], Dict{String, Any}(), nothing, nothing, nothing)
|
||||
end
|
||||
|
||||
function emitToolExecutionEnd(finalized::FinalizedToolCallOutcome, emit::AgentEventSink)::Nothing
|
||||
emit(ToolExecutionEndEvent(
|
||||
finalized.tool_call.id,
|
||||
finalized.tool_call.name,
|
||||
finalized.result,
|
||||
finalized.is_error,
|
||||
))
|
||||
return nothing
|
||||
end
|
||||
|
||||
function createToolResultMessage(finalized::FinalizedToolCallOutcome)::ToolResultMessage
|
||||
return ToolResultMessage(
|
||||
"toolResult",
|
||||
finalized.tool_call.id,
|
||||
finalized.tool_call.name,
|
||||
isnothing(finalized.result.content) ? MessageContent[] : finalized.result.content,
|
||||
finalized.result.details,
|
||||
finalized.result.usage,
|
||||
finalized.result.added_tool_names,
|
||||
finalized.is_error,
|
||||
Dates.now(Dates.UTC).datetime,
|
||||
)
|
||||
end
|
||||
|
||||
function emitToolResultMessage(tool_result_message::ToolResultMessage, emit::AgentEventSink)::Nothing
|
||||
emit(MessageStartEvent(tool_result_message))
|
||||
emit(MessageEndEvent(tool_result_message))
|
||||
return nothing
|
||||
end
|
||||
|
||||
# ============================================================================
|
||||
# Validation helper
|
||||
# ============================================================================
|
||||
|
||||
function validateToolArguments(tool::AgentTool, tool_call::ToolCall)::Any
|
||||
# Simplified validation - in a full implementation, this would use TypeBox-like validation
|
||||
return tool_call.arguments
|
||||
end
|
||||
|
||||
end
|
||||
Reference in New Issue
Block a user