pi_harness/packages/agent/julia_implementation/src/messages.jl

"""
    messages.jl - Custom message types and LLM conversion

This module provides custom message types and the convertToLlm function.
"""

module Messages

using ..Types: *

const COMPACTION_SUMMARY_PREFIX = """The conversation history before this point was compacted into the following summary:

<summary>
"""

const COMPACTION_SUMMARY_SUFFIX = """
</summary>"""

const BRANCH_SUMMARY_PREFIX = """The following is a summary of a branch that this conversation came back from:

<summary>
"""

const BRANCH_SUMMARY_SUFFIX = """</summary>"""

# ============================================================================
# Custom message types
# ============================================================================

mutable struct BashExecutionMessage
    role::String
    command::String
    output::String
    exit_code::Union{Int64, Nothing}
    cancelled::Bool
    truncated::Bool
    full_output_path::Union{String, Nothing}
    timestamp::Timestamp
    exclude_from_context::Bool
end

mutable struct CustomMessage{T}
    role::String
    custom_type::String
    content::Union{String, Vector{MessageContent}}
    display::Bool
    details::Union{T, Nothing}
    timestamp::Timestamp
end

mutable struct BranchSummaryMessage
    role::String
    summary::String
    from_id::String
    timestamp::Timestamp
end

mutable struct CompactionSummaryMessage
    role::String
    summary::String
    tokens_before::Int64
    timestamp::Timestamp
end

# ============================================================================
# Bash execution to text conversion
# ============================================================================

function bashExecutionToText(msg::BashExecutionMessage)::String
    text = "Ran `$(msg.command)`\n"
    if !isempty(msg.output)
        text *= "```\n$(msg.output)\n```"
    else
        text *= "(no output)"
    end
    if msg.cancelled
        text *= "\n\n(command cancelled)"
    elseif !isnothing(msg.exit_code) && msg.exit_code != 0
        text *= "\n\nCommand exited with code $(msg.exit_code)"
    end
    if msg.truncated && !isnothing(msg.full_output_path)
        text *= "\n\n[Output truncated. Full output: $(msg.full_output_path)]"
    end
    return text
end

# ============================================================================
# Message creation functions
# ============================================================================

function createBranchSummaryMessage(summary::String, from_id::String, timestamp::String)::BranchSummaryMessage
    return BranchSummaryMessage(
        "branchSummary",
        summary,
        from_id,
        Int64(Dates.now(Dates.UTC).datetime),
    )
end

function createCompactionSummaryMessage(summary::String, tokens_before::Int64, timestamp::String)::CompactionSummaryMessage
    return CompactionSummaryMessage(
        "compactionSummary",
        summary,
        tokens_before,
        Int64(Dates.now(Dates.UTC).datetime),
    )
end

function createCustomMessage(custom_type::String, content::Union{String, Vector{MessageContent}}, display::Bool, details::Union{Any, Nothing}, timestamp::String)::CustomMessage
    return CustomMessage(
        "custom",
        custom_type,
        content,
        display,
        details,
        Int64(Dates.now(Dates.UTC).datetime),
    )
end

# ============================================================================
# Convert to LLM messages
# ============================================================================

function convertToLlm(messages::Vector{AgentMessage})::Vector{Message}
    result::Vector{Message} = Message[]
    
    for m in messages
        converted = convertToLlmMessage(m)
        if !isnothing(converted)
            push!(result, converted)
        end
    end
    
    return result
end

function convertToLlmMessage(m::BashExecutionMessage)::Union{UserMessage, Nothing}
    if m.exclude_from_context
        return nothing
    end
    return UserMessage(
        "user",
        [TextContent(bashExecutionToText(m))],
        m.timestamp,
    )
end

function convertToLlmMessage(m::CustomMessage)::Union{UserMessage, Nothing}
    content = if m.content isa String
        [TextContent(m.content)]
    else
        m.content
    end
    return UserMessage("user", content, m.timestamp)
end

function convertToLlmMessage(m::BranchSummaryMessage)::UserMessage
    text = BRANCH_SUMMARY_PREFIX * m.summary * BRANCH_SUMMARY_SUFFIX
    return UserMessage("user", [TextContent(text)], m.timestamp)
end

function convertToLlmMessage(m::CompactionSummaryMessage)::UserMessage
    text = COMPACTION_SUMMARY_PREFIX * m.summary * COMPACTION_SUMMARY_SUFFIX
    return UserMessage("user", [TextContent(text)], m.timestamp)
end

function convertToLlmMessage(m::UserMessage)::UserMessage
    return m
end

function convertToLlmMessage(m::AssistantMessage)::AssistantMessage
    return m
end

function convertToLlmMessage(m::ToolResultMessage)::ToolResultMessage
    return m
end

function convertToLlmMessage(m::AgentMessage)::Union{Message, Nothing}
    return nothing
end

end