YiemAgent/src/type.jl

module type
   export agent, sommelier, companion, virtualcustomer, agentContext


using Dates, UUIDs, DataStructures, JSON, NATS
using GeneralUtils

# ---------------------------------------------- 100 --------------------------------------------- #


# ============================================================================
# Message types
# ============================================================================
abstract type agentMessage end

struct userMessage <: agentMessage
    role::String
    content::Vector{messageContent}
    timestamp::Timestamp
end

struct assistantMessage <: agentMessage
    role::String
    content::Vector{messageContent}
    api::String
    provider::String
    model::String
    usage::Usage
    stop_reason::String
    error_message::Union{String, Nothing}
    timestamp::Timestamp
end

struct toolResultMessage <: agentMessage

    role::String
    tool_call_id::String
    tool_name::String
    content::Vector{messageContent}
    details::Any
    usage::Union{Usage, Nothing}
    added_tool_names::Union{Vector{String}, Nothing}
    is_error::Bool
    timestamp::Timestamp
end


# ============================================================================
# Message content types
# ============================================================================

abstract type messageContent end

struct textContent <: messageContent
    text::String
end

struct imageContent <: messageContent
    data::String
    mime_type::String
end


# ============================================================================
# Tool types
# ============================================================================

struct agentTool{TParameters, TDetails}
    name::String
    label::String
    description::String
    parameters::TParameters
    execute::Function
    prepare_arguments::Union{Function, Nothing}
    execution_mode::Union{ToolExecutionMode, Nothing}
end


# ============================================================================
# Agent context
# ============================================================================

struct agentContext
    system_prompt::String    
    messages::Vector{agentMessage}
    tools::Union{Vector{agentTool}, Nothing}
end



# ============================================================================
# Assistant message event types
# ============================================================================

abstract type assistantMessageEvent end

struct startEvent <: assistantMessageEvent
    partial::assistantMessage
end
struct textStartEvent <: assistantMessageEvent
    content_index::Int64
    partial::assistantMessage
end
struct textDeltaEvent <: assistantMessageEvent
    content_index::Int64
    delta::String
    partial::assistantMessage
end
struct textEndEvent <: assistantMessageEvent
    content_index::Int64
    content::String
    partial::assistantMessage
end
struct doneEvent <: assistantMessageEvent
    reason::String
    usage::Usage
    message::assistantMessage
end
struct errorEvent <: assistantMessageEvent
    reason::String
    error_message::Union{String, Nothing}
    usage::Usage
    error::assistantMessage
end



# ============================================================================
# Agent state
# ============================================================================

mutable struct agentState
    system_prompt::String
    model::Model
    thinking_level::ThinkingLevel
    tools::Vector{agentTool}
    messages::Vector{agentMessage}
    is_streaming::Bool
    streaming_message::Union{agentMessage, Nothing}
    pending_tool_calls::Set{String}
    error_message::Union{String, Nothing}

    function agentState(
        system_prompt::String="",
        model::Model=Model("", "", "unknown", "unknown", "", false, String[], ModelCost(0.0, 0.0, 0.0, 0.0), 0, 0),
        thinking_level::ThinkingLevel=THINKING_OFF,
        tools::Vector{agentTool}=agentTool[],
        messages::Vector{agentMessage}=agentMessage[],
    )
        new(
            system_prompt,
            model,
            thinking_level,
            copy(tools),
            copy(messages),
            false,
            nothing,
            Set{String}(),
            nothing,
        )
    end
end

# ============================================================================
# Tool call types
# ============================================================================

struct toolCall
    type::String
    id::String
    name::String
    arguments::Dict{String, Any}
    partial_json::Union{String, Nothing}
end


# ============================================================================
# Next turn context
# ============================================================================

struct nextTurnContext
    message::assistantMessage
    tool_results::Vector{toolResultMessage}
    context::agentContext
    new_messages::Vector{agentMessage}
end


end # module type