module type export Timestamp, # Abstract types messageContent, agentMessage, agent, # Model types modelCost, llmModel, llmUsage, # Message content types textContent, imageContent, # Message types userMessage, assistantMessage, toolResultMessage, # Tool types agentTool, validateRequiredArgs, # Context types agentContext, agentState, agentToolCall, prepareNextTurnContext, # Loop & execution types agentLoopConfig, abortSignal, agentToolResult, assistantMsgCtx, afterCtx, # Event types toolExecStartEvent, toolExecUpdateEvent, toolExecEndEvent, # Agent yiemAgent, # Tool call lifecycle types preparedToolCall, immediateOutcome, executedOutcome, finalizedOutcome, agentToolCallBatch, # Functions (defined elsewhere) run_agent, take_response, follow_up, stop_agent using Dates, UUIDs, DataStructures, JSON, NATS, Base.Threads using GeneralUtils const Timestamp = DateTime # ------------------------------------------------------------------------------------------------ # # LLM model info # # ------------------------------------------------------------------------------------------------ # struct modelCost # Model pricing per 1M tokens input::Float64 # Price per 1M input tokens output::Float64 # Price per 1M output tokens cache_read::Float64 # Price per 1M cached read tokens cache_write::Float64 # Price per 1M cache write tokens end struct llmModel # LLM model configuration id::String # Unique model identifier name::String # Human-readable model name provider::String # Provider name (e.g., "anthropic", "openai") baseUrl::String # API endpoint base URL reasoning::Bool # Whether the model supports chain-of-thought input::Vector{String} # Supported input modalities (e.g., "text", "image") cost::modelCost # Pricing information contextWindow::Int64 # Maximum context length in tokens maxTokens::Int64 # Maximum output tokens per completion end struct llmUsage inputTokens::Int64 outputTokens::Int64 end # ------------------------------------------------------------------------------------------------ # # Message content types # # ------------------------------------------------------------------------------------------------ # abstract type messageContent end # Base type for message content struct textContent <: messageContent # Plain text message content text::String # The text content end struct imageContent <: messageContent # Image message content data::String # Base64-encoded image data mimeType::String # MIME type (e.g., "image/png") end # ------------------------------------------------------------------------------------------------ # # Message types # # ------------------------------------------------------------------------------------------------ # abstract type agentMessage end # Base type for all agent messages struct userMessage <: agentMessage # Message from the user role::String # Always "user" content::Vector{messageContent} # Text and/or image content timestamp::Timestamp # When the message was sent end """ Create a new user message. # Arguments - `role::String`: Always "user" - `content::Vector{messageContent}`: Text and/or image content - `timestamp::Timestamp`: When the message was sent # Returns - A new `userMessage` instance # Examples ```julia julia> msg = userMessage(content=[textContent("Hello")]) userMessage("user", [textContent("Hello")], DateTime(...)) ``` """ function userMessage(; role="user", content=Vector{messageContent}(), timestamp=now()) return userMessage(role, content, timestamp) end struct assistantMessage <: agentMessage # Message from the AI assistant role::String # Always "assistant" content::Vector{messageContent} # Text and/or image content api::String # API name used (e.g., "openai") provider::String # Provider name (e.g., "anthropic") model::String # Model identifier usage::llmUsage # Token usage for this message stopReason::String # Why generation stopped (e.g., "end_turn") errorMessage::Union{String, Nothing} # Error if generation failed timestamp::Timestamp # When the message was received end """ Create a new assistant message. # Arguments - `role::String`: Always "assistant" - `content::Vector{messageContent}`: Text and/or image content - `api::String`: API name used (e.g., "openai") - `provider::String`: Provider name (e.g., "anthropic") - `model::String`: Model identifier - `usage::llmUsage`: Token usage for this message - `stopReason::String`: Why generation stopped (e.g., "end_turn") - `errorMessage::Union{String, Nothing}`: Error if generation failed - `timestamp::Timestamp`: When the message was received # Returns - A new `assistantMessage` instance # Examples ```julia julia> msg = assistantMessage(content=[textContent("Hello!")], model="gpt-4") assistantMessage("assistant", [textContent("Hello!")], "", "", "gpt-4", ..., "end_turn", nothing, DateTime(...)) ``` """ function assistantMessage(; role="assistant", content=Vector{messageContent}(), api="", provider="", model="", usage=llmUsage(0, 0), stopReason="end_turn", errorMessage=nothing, timestamp=now()) return assistantMessage(role, content, api, provider, model, usage, stopReason, errorMessage, timestamp) end struct toolResultMessage <: agentMessage # Result returned from a tool execution role::String # Always "tool" toolCallId::String # ID matching the tool call toolName::String # Name of the executed tool content::Vector{messageContent} # Tool output content details::Any # Additional tool-specific details usage::Union{llmUsage, Nothing} # Token usage if applicable addedToolNames::Union{Vector{String}, Nothing} # Tools added during execution isError::Bool # Whether the tool call resulted in an error timestamp::Timestamp # When the result was recorded end """ Create a new tool result message. # Arguments - `role::String`: Always "tool" - `toolCallId::String`: ID matching the tool call - `toolName::String`: Name of the executed tool - `content::Vector{messageContent}`: Tool output content - `details::Any`: Additional tool-specific details - `usage::Union{llmUsage, Nothing}`: Token usage if applicable - `addedToolNames::Union{Vector{String}, Nothing}`: Tools added during execution - `isError::Bool`: Whether the tool call resulted in an error - `timestamp::Timestamp`: When the result was recorded # Returns - A new `toolResultMessage` instance # Examples ```julia julia> msg = toolResultMessage(toolCallId="call_123", toolName="search", content=[textContent("results")]) toolResultMessage("tool", "call_123", "search", [textContent("results")], nothing, nothing, nothing, false, DateTime(...)) ``` """ function toolResultMessage(; role="tool", toolCallId="", toolName="", content=Vector{messageContent}(), details=nothing, usage=nothing, addedToolNames=nothing, isError=false, timestamp=now()) return toolResultMessage(role, toolCallId, toolName, content, details, usage, addedToolNames, isError, timestamp) end # ------------------------------------------------------------------------------------------------ # # Tool types # # ------------------------------------------------------------------------------------------------ # """ A tool available to the agent. Maps MCP server tool definitions to an executable Julia tool. # Arguments - `name::String`: Tool identifier (from MCP `name`) - `label::String`: Human-readable tool name (from MCP `title`) - `description::String`: What the tool does (from MCP `description`) - `inputSchema::Any`: Tool parameters schema (from MCP `inputSchema`, JSON Schema format) - `execute::Function`: Tool execution function, signature: `execute(toolCallId::String, args::Dict, signal::Union{Nothing,AbortSignal}, onPartialResult::Function)` - `prepareArguments::Union{Function, Nothing}`: Optional argument preparation callback - `validateRequiredArgs::Union{Function, Nothing}`: Optional validation hook, signature: `validateRequiredArgs(args::Dict) -> Union{Nothing, String}` where `String` is an error message - `parallelToolExecute::Bool`: Override: run tool calls sequentially or in parallel # Returns - A new `agentTool` instance # MCP Tool Example ``` { "name": "getWeather", "title": "Weather Lookup", "description": "Fetch current weather and forecast for a given city.", "inputSchema": { "type": "object", "properties": { "city": { "type": "string", "description": "City and state/country" }, "units": { "type": "string", "enum": ["celsius", "fahrenheit"], "default": "celsius" } }, "required": ["city"] } } ``` # Example ```julia tool = agentTool( name="getWeather", label="Weather Lookup", description="Fetch current weather and forecast for a given city.", inputSchema=Dict( "type" => "object", "properties" => Dict( "city" => Dict("type" => "string", "description" => "City name"), "units" => Dict("type" => "string", "enum" => ["celsius", "fahrenheit"]) ), "required" => ["city"] ), execute=(toolCallId, args, signal, onPartialResult) -> begin city = args["city"] return agentToolResult( [textContent("Sunny, 22C in Bangkok")], Dict{Any,Any}(), nothing, false ) end, prepareArguments = nothing, validateRequiredArgs = nothing, parallelToolExecute = false ) ``` """ struct agentTool # A tool available to the agent name::String # Tool identifier label::String # Human-readable tool name description::String # What the tool does inputSchema::Any # Tool parameters schema (JSON schema, MCP inputSchema format) execute::Function # Tool execution function prepareArguments::Union{Function, Nothing} # Optional argument preparation callback validateRequiredArgs::Union{Function, Nothing} # Optional validation hook for required args parallelToolExecute::Bool # Override: run tool calls sequentially or in parallel end # ------------------------------------------------------------------------------------------------ # # Agent context # # ------------------------------------------------------------------------------------------------ # """ Snapshot of the agent's conversation context. # Arguments - `systemPrompt::String`: System prompt for the agent - `messages::Vector{agentMessage}`: Conversation messages - `tools::Union{Vector{agentTool}, Nothing}`: Available tools # Returns - A new `agentContext` instance """ struct agentContext # Snapshot of the agent's conversation context systemPrompt::String # System prompt for the agent messages::Vector{agentMessage} # Conversation messages tools::Union{Vector{agentTool}, Nothing} # Available tools end # ------------------------------------------------------------------------------------------------ # # Agent state # # ------------------------------------------------------------------------------------------------ # mutable struct agentState # Mutable runtime state of an agent systemPrompt::String # System prompt text model::llmModel # LLM model to use tools::Vector{agentTool} # Available tools # messages history includes userMessage, assistantMessage, toolResultMessage. NO system prompt messages::Vector{agentMessage} pendingToolCalls::Vector{String} # Tool call IDs waiting for results activeRun::Bool # is agent processing user message? errorMessage::Union{String, Nothing} # Last error message end """ Create a new mutable agent state. Creates a deep copy of the provided tools and messages to isolate the new state from external references. # Arguments - `systemPrompt::String`: System prompt text - `model::llmModel`: LLM model to use (defaults to an unknown model) - `tools::Vector{agentTool}`: Available tools (deep copied) - `messages::Vector{agentMessage}`: Conversation messages (deep copied) # Returns - A new `agentState` instance with an empty pending tool calls list and no error # Examples ```julia julia> state = agentState(systemPrompt="You are a helpful assistant") agentState("You are a helpful assistant", ..., agentTool[], agentMessage[], String[], nothing) ``` """ function agentState( systemPrompt::String="", model::llmModel=llmModel{String}("", "", "unknown", "unknown", "", false, String[], modelCost(0.0, 0.0, 0.0, 0.0), 0, 0), tools::Vector{agentTool}=agentTool[], messages::Vector{agentMessage}=agentMessage[], ) agentState( systemPrompt, model, deepcopy(tools), deepcopy(messages), Vector{String}(), false, nothing, ) end struct agentToolCall # A tool invocation from the LLM type::String # Always "function" id::String # Unique tool call identifier name::String # Tool name arguments::Dict{String, Any} # Parsed tool arguments end """ Context for preparing the next conversation turn. # Arguments - `message::assistantMessage`: The assistant's message that just completed - `toolResults::Vector{toolResultMessage}`: Tool results from this turn - `context::agentContext`: Current conversation context - `newMessages::Vector{agentMessage}`: Messages to append to the context # Returns - A new `prepareNextTurnContext` instance """ struct prepareNextTurnContext # Context for preparing the next conversation turn message::assistantMessage # The assistant's message that just completed toolResults::Vector{toolResultMessage} # Tool results from this turn context::agentContext # Current conversation context newMessages::Vector{agentMessage} # Messages to append to the context end # ------------------------------------------------------------------------------------------------ # # Agent loop configuration & tool execution types # # ------------------------------------------------------------------------------------------------ # """ Configuration for the agent tool execution loop. # Arguments - `tools::Vector{agentTool}`: Available tools - `beforeToolCall::Union{Function, Nothing}`: Callback before tool execution - `afterToolCall::Union{Function, Nothing}`: Callback after tool execution - `toolExecution::String`: Execution mode — "sequential" or "parallel" """ struct agentLoopConfig tools::Vector{agentTool} beforeToolCall::Union{Function, Nothing} afterToolCall::Union{Function, Nothing} toolExecution::String end """ Signal for aborting ongoing operations. # Arguments - `aborted::Bool`: Whether the operation has been aborted """ struct abortSignal aborted::Bool end """ Result returned by tool execution before the `afterToolCall` hook. # Arguments - `content::Vector{messageContent}`: Tool output content - `details::Dict{Any,Any}`: Tool-specific details - `usage::Union{llmUsage, Nothing}`: Token usage if applicable - `terminate::Bool`: Whether tool requests termination of the agent loop """ struct agentToolResult content::Vector{messageContent} details::Dict{Any,Any} usage::Union{llmUsage, Nothing} terminate::Bool end """ Context passed to the `beforeToolCall` hook. # Arguments - `message::assistantMessage`: The assistant message containing the tool call - `toolCall::agentToolCall`: The tool call being prepared - `args::Dict{String,Any}`: Validated tool arguments - `context::agentContext`: Current conversation context """ struct assistantMsgCtx message::assistantMessage toolCall::agentToolCall args::Dict{String,Any} context::agentContext end """ Context passed to the `afterToolCall` hook. # Arguments - `message::assistantMessage`: The assistant message containing the tool call - `toolCall::agentToolCall`: The tool call that was executed - `args::Dict{String,Any}`: Tool arguments - `result::agentToolResult`: The raw tool result - `isError::Bool`: Whether execution resulted in an error - `context::agentContext`: Current conversation context """ struct afterCtx message::assistantMessage toolCall::agentToolCall args::Dict{String,Any} result::agentToolResult isError::Bool context::agentContext end """ Event emitted when a tool call execution starts. # Arguments - `toolCallId::String`: ID of the tool call - `toolName::String`: Name of the tool - `arguments::Dict{String,Any}`: Tool arguments """ struct toolExecStartEvent toolCallId::String toolName::String arguments::Dict{String,Any} end """ Event emitted with partial results during tool execution. # Arguments - `toolCallId::String`: ID of the tool call - `toolName::String`: Name of the tool - `arguments::Dict{String,Any}`: Tool arguments - `partialResult::Any`: The partial result data """ struct toolExecUpdateEvent toolCallId::String toolName::String arguments::Dict{String,Any} partialResult::Any end """ Event emitted when a tool call execution ends. # Arguments - `toolCallId::String`: ID of the tool call - `toolName::String`: Name of the tool - `result::agentToolResult`: The final tool result - `isError::Bool`: Whether execution resulted in an error """ struct toolExecEndEvent toolCallId::String toolName::String result::agentToolResult isError::Bool end # ------------------------------------------------------------------------------------------------ # # Agent struct # # ------------------------------------------------------------------------------------------------ # abstract type agent end """ docstring """ mutable struct yiemAgent <: agent # High-level agent wrapper _state::agentState # Current state (prompt, model, messages, tools, etc.) # user sends prompt message to agent. if agent is idle, it process user message right away. # if agent is running, it process user message after the current tool call finished. inputChannel::Channel # Buffers messages the user sends while the agent is busy. Processed after all inputChannel # messages are handled and the agent is idle (not using a tool call). followUpChannel::Channel # agent sends response message to user after processing all user messages in inputChannel # and all followUp messages. outputChannel::Channel _agent_loop::Union{Task, Nothing} # agent loop running in the background # Preprocess/transform messages and context (modify, filter, prune, inject context from memory, # reorder, ...) for a single LLM call in _process_message()'s loop. # returns new Vector{agentMessage} prepareContext ::Union{Function, Nothing} # Convert prepareContext()'s new Vector{agentMessage} to LLM message format formatMsgForLLM::Function # Actually invoke the LLM to get a completion response. The LLM response comes back as an # assistantMessage whose content is an array of content blocks. # Each block has a type — "text", "thinking", or "toolCall". # The code filters for type === "toolCall" blocks, then passes them to executeToolCalls(). llmCall::Function # Callback invoked before executing a tool call (ask for user permission/confirmation/abort, etc..) beforeToolCall::Union{Function, Nothing} executeToolCalls::Function # execute tool calls () # Callback invoked after executing a tool call to sanitize tools output so the output is ready # to be converted into toolResults message afterToolCall::Union{Function, Nothing} # prepareNextTurn::Union{Function, Nothing} # Callback to prepare the next conversation turn # prepareNextTurnWithContext::Union{Function, Nothing} # Same but receives context sessionId::Union{String, Nothing} # Optional session identifier maxRetryDelayMs::Union{Int64, Nothing} # Maximum delay between retries (ms) parallelToolExecute::Bool # Default: false agentEventSink::Function # agent emits its status via this function end """ Create a new yiemAgent instance with a background loop task. Spawns a background `@spawn` task that runs the agent loop, listening on `inputChannel` and `followUpChannel` channels concurrently. # Keyword Arguments - `systemPrompt::String`: System prompt for the agent - `model`: LLM model to use - `tools::Vector{agentTool}`: Available tools (default: empty) - `messages::Vector{agentMessage}`: Initial conversation messages (default: empty) - `formatMsgForLLM::Function`: Convert agent messages to LLM message format (default: `defaultformatMsgForLLM`) - `llmCall::Function`: Function to invoke the LLM (required) - `prepareContext::Union{Function, Nothing}`: Preprocess/transform messages before sending to LLM (default: `nothing`) - `beforeToolCall::Union{Function, Nothing}`: Callback invoked before executing a tool call (default: `nothing`) - `afterToolCall::Union{Function, Nothing}`: Callback invoked after executing a tool call (default: `nothing`) - `prepareNextTurn::Union{Function, Nothing}`: Callback to prepare the next conversation turn (default: `nothing`) - `prepareNextTurnWithContext::Union{Function, Nothing}`: Same but receives context (default: `nothing`) - `sessionId::Union{String, Nothing}`: Optional session identifier (default: `nothing`) - `maxRetryDelayMs::Union{Int64, Nothing}`: Maximum delay between retries in milliseconds (default: `nothing`) - `parallelToolExecute::Bool`: Run tool calls in parallel (default: `false`) - `agentEventSink::Function`: Callback to receive agent events # Returns - A new `yiemAgent` instance with an active background task # Examples ```julia julia> agent = yiemAgent(systemPrompt="You are a helpful assistant", model=my_model) yiemAgent(agentState(...), Channel(...), Channel(...), Channel(...), ..., ...) ``` """ function yiemAgent( ; systemPrompt::String="You are helpful assistant.", model=nothing, tools::Vector{agentTool}=agentTool[], messages::Vector{agentMessage}=agentMessage[], prepareContext::Union{Function, Nothing}=nothing, formatMsgForLLM::Function=defaultformatMsgForLLM, llmCall::Function, beforeToolCall::Union{Function, Nothing}=nothing, afterToolCall::Union{Function, Nothing}=nothing, # prepareNextTurn::Union{Function, Nothing}=nothing, # prepareNextTurnWithContext::Union{Function, Nothing}=nothing, sessionId::Union{String, Nothing}=nothing, maxRetryDelayMs::Union{Int64, Nothing}=nothing, parallelToolExecute::Bool=false, agentEventSink::Function, ) # Create channels: input (user -> agent), followUp (async queue), output (agent -> user) inputChannel = Channel(16) followUp = Channel(32) outputChannel = Channel(16) # Create struct with a placeholder task, then spawn and replace it agent = yiemAgent( agentState(systemPrompt, model, tools, messages), inputChannel, followUp, outputChannel, nothing, # placeholder — replaced below prepareContext, formatMsgForLLM, llmCall, beforeToolCall, afterToolCall, # prepareNextTurn, # prepareNextTurnWithContext, sessionId, maxRetryDelayMs, parallelToolExecute, agentEventSink, ) # Spawn the background loop and attach it agent._agent_loop = @spawn _agent_loop(agent) return agent end """ preparedToolCall(tool, toolCall, args) Intermediate state between tool call validation and execution. Created after `prepareToolCall` succeeds; serves as the bridge to the execution phase. Keeping the resolved tool, original call metadata, and validated args together avoids repeated lookups and allows the execution phase to access all necessary data without carrying the full context through the call chain. # Fields - `tool::agentTool`: The resolved tool definition from the context - `toolCall::agentToolCall`: The original tool call from the assistant - `args::any`: Validated (and coerced) argument values # Examples ```julia # After prepareToolCall succeeds, the agent holds a preparedToolCall prep = preparedToolCall( tool, # agentTool found in context.tools toolCall, # {id: "call_1", name: "search_wine", arguments: "{\"query\": \"red wine\"}"} validatedArgs # Dict("query" => "red wine") ) ``` """ struct preparedToolCall tool::agentTool # The resolved tool definition from the context toolCall::agentToolCall # The original tool call from the assistant args::Any # Validated (and coerced) argument values end """ immediateOutcome(result, isError) A tool call that was resolved without actual execution — either because the tool was not found, validation failed, or a `beforeToolCall` hook blocked the call. The result is produced immediately and emitted as a tool result message. Returning an outcome instead of throwing an exception is intentional: it lets the agent feed the error back to the LLM as a tool result so the model can recover — for example, by re-issuing a tool call with corrected arguments after a validation failure. # Fields - `result::agentToolResult`: The pre-computed tool result - `isError::Bool`: Whether this outcome represents an error # Examples ```julia # Tool not found — immediate error immediateOutcome( createErrorToolResult("Tool search_wine not found"), true ) # beforeToolCall hook blocked execution immediateOutcome( createErrorToolResult("Tool execution was blocked"), true ) ``` """ struct immediateOutcome result::agentToolResult # The pre-computed tool result isError::Bool # Whether this outcome represents an error end """ executedOutcome(result, isError) A tool call that has been executed by `tool.execute()` but has not yet been through the `afterToolCall` hook. This intermediate state is necessary because the hook may mutate the result (content, usage, termination, error status). Keeping execution and finalization separate allows the hook to inspect the raw result and decide whether to transform it or replace it entirely. # Fields - `result::agentToolResult`: The tool's execution result - `isError::Bool`: Whether execution raised an error # Examples ```julia # Successful execution executedOutcome( agentToolResult([textContent("text", "Found 3 wines")], dict{any,any}(), dict{any,any}()), false ) # Execution error executedOutcome( createErrorToolResult("Connection timeout"), true ) ``` """ struct executedOutcome result::agentToolResult # The tool's execution result isError::Bool # Whether execution raised an error end """ finalizedOutcome(toolCall, result, isError) The complete outcome of a tool call after both execution and the `afterToolCall` hook. This is the final form used to construct the `toolResultMessage` emitted to the agent loop. The three-phase design (prepare → execute → finalize) exists so that each phase has a single responsibility: preparation handles validation and gating, execution performs the actual work, and finalization applies post-processing hooks. This separation allows the agent loop to emit `toolExecutionEnd` events with the finalized data while keeping each phase independently testable and swappable. # Fields - `toolCall::agentToolCall`: The original tool call reference - `result::agentToolResult`: The final tool result (post-afterToolCall) - `isError::Bool`: Whether the call failed or was blocked # Examples ```julia # Normal successful finalization finalizedOutcome(tc, agentToolResult(content, details, usage, false), false) # afterToolCall mutated result and set terminate finalizedOutcome(tc, agentToolResult(content, details, usage, true), false) ``` """ struct finalizedOutcome toolCall::agentToolCall # The original tool call reference result::agentToolResult # The final tool result (post-afterToolCall) isError::Bool # Whether the call failed or was blocked end """ agentToolCallBatch(messages, terminate) A batch of tool result messages from executing one or more tool calls. The `terminate` flag indicates whether all tools in the batch requested termination, which causes the agent loop to stop processing further turns. This flag is set by the tool implementation (not the end user) to signal that the agent should not call the LLM again. Typical use cases: - Task completion: a tool like `deploy` or `submit` finishes its work and returns `terminate: true` so the agent stops instead of asking the LLM what to do next. - Unrecoverable error: a tool hits a fatal condition (e.g. database connection lost, auth token expired) and returns `terminate: true` so the agent stops with an error message rather than retrying. - Async handoff: a tool triggers a long-running external operation and wants the agent to stop now; the external system will later resume the agent via `continue()`. If `terminate` is `false` (default), the agent loop feeds the tool results back to the LLM for another turn. # Fields - `messages::Vector{toolResultMessage}`: Tool result messages for this batch - `terminate::Bool`: Whether the batch should terminate the loop # Examples ```julia # Batch of 3 tool results, no termination agentToolCallBatch(resultMessages, false) # All tools requested termination agentToolCallBatch(resultMessages, true) # Empty batch — terminate is false regardless agentToolCallBatch(toolResultMessage[], false) ``` """ struct agentToolCallBatch messages::Vector{toolResultMessage} # Tool result messages for this batch terminate::Bool # Whether the batch should terminate the loop end end # module type