Add AnnotationMessage implementation

.github/workflows/CI.yml

@@ -37,6 +37,13 @@ jobs:
           arch: ${{ matrix.arch }}
       - uses: julia-actions/cache@v1
       - uses: julia-actions/julia-buildpkg@v1
+      - name: Install test dependencies
+        run: |
+          julia --project=test -e '
+            using Pkg;
+            Pkg.develop(PackageSpec(path=pwd()));
+            Pkg.instantiate();
+            Pkg.precompile()'
       - uses: julia-actions/julia-runtest@v1
       - uses: julia-actions/julia-processcoverage@v1
       - uses: codecov/codecov-action@v4
@@ -67,4 +74,4 @@ jobs:
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} # For authentication with GitHub Actions token
           GKSwstype: "100" # for Plots.jl plots (if you have them)
-          JULIA_DEBUG: "Documenter"
+          JULIA_DEBUG: "Documenter"

Project.toml

@@ -53,6 +53,7 @@ PrecompileTools = "1"
 Preferences = "1"
 REPL = "<0.0.1, 1"
 Random = "<0.0.1, 1"
+Snowball = "0.1"
 SparseArrays = "<0.0.1, 1"
 Statistics = "<0.0.1, 1"
 StreamCallbacks = "0.4, 0.5"
@@ -61,10 +62,13 @@ julia = "1.9, 1.10"
 
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+HTTP = "cd3eb016-35fb-5094-929b-558a96fad6f3"
+JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+OpenAI = "e9f21f70-7185-4079-aca2-91159181367c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Unicode = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"
 
 [targets]
-test = ["Aqua", "FlashRank", "SparseArrays", "Statistics", "LinearAlgebra", "Markdown", "Snowball", "Unicode"]
+test = ["Aqua", "FlashRank", "SparseArrays", "Statistics", "LinearAlgebra", "Markdown", "Snowball", "Unicode", "HTTP", "JSON3", "OpenAI"]

src/constants.jl

@@ -0,0 +1,23 @@
+# Constants used throughout the package
+
+# Model Registry
+const MODEL_REGISTRY = Dict{String,Any}()
+
+# Default Models
+const MODEL_CHAT = "gpt-3.5-turbo"
+const MODEL_COMPLETION = "gpt-3.5-turbo-instruct"
+
+# Reserved keywords that cannot be used as placeholders in templates
+const RESERVED_KWARGS = Symbol[
+    :model, :api_key, :verbose, :return_all, :dry_run, :conversation,
+    :streamcallback, :no_system_message, :name_user, :name_assistant,
+    :http_kwargs, :api_kwargs
+]
+
+# Default system message
+const DEFAULT_SYSTEM_MESSAGE = "You are a helpful AI assistant."
+
+# Default API Keys
+const OPENAI_API_KEY = get(ENV, "OPENAI_API_KEY", "")
+const ANTHROPIC_API_KEY = get(ENV, "ANTHROPIC_API_KEY", "")
+const GOOGLE_API_KEY = get(ENV, "GOOGLE_API_KEY", "")

src/llm_shared.jl

@@ -73,6 +73,9 @@ function render(schema::NoSchema,
             count_system_msg += 1
             # move to the front
             pushfirst!(conversation, msg)
+        elseif isabstractannotationmessage(msg)
+            # Silently skip annotation messages - they are not meant for LLM consumption
+            continue
         else
             # Note: Ignores any DataMessage or other types for the prompt/conversation history
             @warn "Unexpected message type: $(typeof(msg)). Skipping."

src/messages.jl

@@ -5,6 +5,7 @@ abstract type AbstractMessage end
 abstract type AbstractChatMessage <: AbstractMessage end # with text-based content
 abstract type AbstractDataMessage <: AbstractMessage end # with data-based content, eg, embeddings
 abstract type AbstractTracerMessage{T <: AbstractMessage} <: AbstractMessage end # message with annotation that exposes the underlying message
+abstract type AbstractAnnotationMessage <: AbstractMessage end # message with metadata that is never sent to LLMs
 # Complementary type for tracing, follows the same API as TracerMessage
 abstract type AbstractTracer{T <: Any} end
 
@@ -41,9 +42,9 @@ end
 """
     UserMessage
 
-A message type for user-generated text-based responses. 
+A message type for user-generated text-based responses.
 Consumed by `ai*` functions to generate responses.
-    
+
 # Fields
 - `content::T`: The content of the message.
 - `variables::Vector{Symbol}`: The variables in the message.
@@ -68,9 +69,9 @@ end
 """
     UserMessageWithImages
 
-A message type for user-generated text-based responses with images. 
+A message type for user-generated text-based responses with images.
 Consumed by `ai*` functions to generate responses.
-    
+
 # Fields
 - `content::T`: The content of the message.
 - `image_url::Vector{String}`: The URLs of the images.
@@ -97,9 +98,9 @@ end
 """
     AIMessage
 
-A message type for AI-generated text-based responses. 
+A message type for AI-generated text-based responses.
 Returned by `aigenerate`, `aiclassify`, and `aiscan` functions.
-    
+
 # Fields
 - `content::Union{AbstractString, Nothing}`: The content of the message.
 - `status::Union{Int, Nothing}`: The status of the message from the API.
@@ -131,9 +132,9 @@ end
 """
     DataMessage
 
-A message type for AI-generated data-based responses, ie, different `content` than text. 
+A message type for AI-generated data-based responses, ie, different `content` than text.
 Returned by `aiextract`, and `aiextract` functions.
-    
+
 # Fields
 - `content::Union{AbstractString, Nothing}`: The content of the message.
 - `status::Union{Int, Nothing}`: The status of the message from the API.
@@ -163,8 +164,8 @@ end
 """
     ToolMessage
 
-A message type for tool calls. 
-    
+A message type for tool calls.
+
 It represents both the request (fields `args`, `name`) and the response (field `content`).
 
 # Fields
@@ -188,9 +189,9 @@ end
 """
     AIToolRequest
 
-A message type for AI-generated tool requests. 
+A message type for AI-generated tool requests.
 Returned by `aitools` functions.
-    
+
 # Fields
 - `content::Union{AbstractString, Nothing}`: The content of the message.
 - `tool_calls::Vector{ToolMessage}`: The vector of tool call requests.
@@ -224,11 +225,100 @@ Base.@kwdef struct AIToolRequest{T <: Union{AbstractString, Nothing}} <: Abstrac
     sample_id::Union{Nothing, Int} = nothing
     _type::Symbol = :aitoolrequest
 end
-"Get the vector of tool call requests from an AIToolRequest/message."
-tool_calls(msg::AIToolRequest) = msg.tool_calls
-tool_calls(msg::AbstractMessage) = ToolMessage[]
-tool_calls(msg::ToolMessage) = [msg]
-tool_calls(msg::AbstractTracerMessage) = tool_calls(msg.object)
+
+"""
+    AnnotationMessage{T} <: AbstractAnnotationMessage
+
+A message type for adding metadata and documentation to conversation histories without affecting LLM context.
+
+Fields:
+- `content::T`: The main content of the annotation
+- `extras::Dict{Symbol,Any}`: Additional metadata as key-value pairs
+- `tags::Vector{Symbol}`: Tags for categorizing the annotation
+- `comment::String`: Human-readable comment (not used for automatic operations)
+- `run_id::Union{Nothing,Int}`: Optional run ID for tracking message provenance
+- `_type::Symbol`: Internal type identifier
+
+The AnnotationMessage is designed to bundle key information and documentation together with conversation data
+without sending any of it to large language models.
+"""
+Base.@kwdef mutable struct AnnotationMessage{T} <: AbstractAnnotationMessage
+    content::T
+    extras::Dict{Symbol,Any} = Dict{Symbol,Any}()
+    tags::Vector{Symbol} = Symbol[]
+    comment::String = ""
+    run_id::Union{Nothing,Int} = get_run_id()
+    _type::Symbol = :annotationmessage
+end
+
+# Constructor that accepts content as first argument
+AnnotationMessage(content; kwargs...) = AnnotationMessage(; content=content, kwargs...)
+
+"""
+    annotate!(messages::Vector{<:AbstractMessage}, content; kwargs...)
+    annotate!(message::AbstractMessage, content; kwargs...)
+
+Add an annotation message to a vector of messages or wrap a single message in a vector with an annotation.
+
+The annotation is inserted after any existing annotation messages to maintain logical grouping.
+All kwargs are passed to the AnnotationMessage constructor.
+
+Returns the modified vector of messages.
+"""
+function annotate!(messages::Vector{<:AbstractMessage}, content; kwargs...)
+    msg = AnnotationMessage(content; kwargs...)
+    # Find the last annotation message
+    last_annotation_idx = findlast(isabstractannotationmessage, messages)
+    if isnothing(last_annotation_idx)
+        # No existing annotations, insert at start
+        insert!(messages, 1, msg)
+    else
+        # Insert after last annotation
+        insert!(messages, last_annotation_idx + 1, msg)
+    end
+    return messages
+end
+
+"""
+    annotate!(msgs::Vector{<:AbstractMessage}, content::String; kwargs...)
+
+Add an AnnotationMessage to a vector of messages, placing it after any existing annotations.
+Returns the modified vector.
+
+# Arguments
+- `msgs`: Vector of messages to annotate
+- `content`: Content of the annotation
+- `kwargs...`: Additional fields for AnnotationMessage (extras, tags, comment)
+
+# Example
+```julia
+msgs = [UserMessage("Hello"), AIMessage("Hi")]
+annotate!(msgs, "Conversation start"; tags=[:greeting])
+```
+"""
+function annotate!(msgs::Vector{<:AbstractMessage}, content::String; kwargs...)
+    annotation = AnnotationMessage(content; kwargs...)
+    # Find the last annotation message index
+    last_annotation_idx = findlast(isabstractannotationmessage, msgs)
+
+    if isnothing(last_annotation_idx)
+        # No existing annotations, insert at start
+        insert!(msgs, 1, annotation)
+    else
+        # Insert after the last annotation
+        insert!(msgs, last_annotation_idx + 1, annotation)
+    end
+    return msgs
+end
+
+"""
+    annotate!(message::AbstractMessage, content; kwargs...)
+
+Convenience method to annotate a single message by wrapping it in a vector.
+"""
+function annotate!(message::AbstractMessage, content; kwargs...)
+    annotate!([message], content; kwargs...)
+end
 
 ### Other Message methods
 # content-only constructor
@@ -251,6 +341,7 @@ isaimessage(m::Any) = m isa AIMessage
 istoolmessage(m::Any) = m isa ToolMessage
 isaitoolrequest(m::Any) = m isa AIToolRequest
 istracermessage(m::Any) = m isa AbstractTracerMessage
+isabstractannotationmessage(m::Any) = m isa AbstractAnnotationMessage
 isusermessage(m::AbstractTracerMessage) = isusermessage(m.object)
 isusermessagewithimages(m::AbstractTracerMessage) = isusermessagewithimages(m.object)
 issystemmessage(m::AbstractTracerMessage) = issystemmessage(m.object)
@@ -502,6 +593,34 @@ function Base.show(io::IO, ::MIME"text/plain", m::AbstractChatMessage)
     end
     print(io, "(\"", m.content, "\")")
 end
+
+function Base.show(io::IO, ::MIME"text/plain", m::AnnotationMessage)
+    printstyled(io, "AnnotationMessage"; color = :cyan)
+    print(io, "(\"", m.content[1:min(30, length(m.content))])
+    if length(m.content) > 30
+        print(io, "...")
+    end
+    if !isempty(m.tags)
+        print(io, "\" [", join(m.tags, ", "), "])")
+    else
+        print(io, "\")")
+    end
+end
+
+function pprint(io::IO, msg::AnnotationMessage)
+    println(io, "Annotation Message:")
+    println(io, "Content: ", msg.content)
+    if !isempty(msg.tags)
+        println(io, "Tags: ", msg.tags)
+    end
+    if !isempty(msg.comment)
+        println(io, "Comment: ", msg.comment)
+    end
+    if !isempty(msg.extras)
+        println(io, "Extras: ", msg.extras)
+    end
+end
+
 function Base.show(io::IO, ::MIME"text/plain", m::AbstractDataMessage)
     type_ = string(typeof(m)) |> x -> split(x, "{")[begin]
     printstyled(io, type_; color = :light_yellow)
@@ -557,7 +676,8 @@ function StructTypes.subtypes(::Type{AbstractMessage})
         systemmessage = SystemMessage,
         metadatamessage = MetadataMessage,
         datamessage = DataMessage,
-        tracermessage = TracerMessage)
+        tracermessage = TracerMessage,
+        annotationmessage = AnnotationMessage)
 end
 
 StructTypes.StructType(::Type{AbstractChatMessage}) = StructTypes.AbstractType()
@@ -592,6 +712,7 @@ StructTypes.StructType(::Type{AIMessage}) = StructTypes.Struct()
 StructTypes.StructType(::Type{DataMessage}) = StructTypes.Struct()
 StructTypes.StructType(::Type{TracerMessage}) = StructTypes.Struct() # Ignore mutability once we serialize
 StructTypes.StructType(::Type{TracerMessageLike}) = StructTypes.Struct() # Ignore mutability once we serialize
+StructTypes.StructType(::Type{AnnotationMessage}) = StructTypes.Struct()
 
 ### Utilities for Pretty Printing
 """
@@ -615,6 +736,8 @@ function pprint(io::IO, msg::AbstractMessage; text_width::Int = displaysize(io)[
         "AI Tool Request"
     elseif msg isa ToolMessage
         "Tool Message"
+    elseif msg isa AnnotationMessage
+        "Annotation Message"
     else
         "Unknown Message"
     end
@@ -633,6 +756,11 @@ function pprint(io::IO, msg::AbstractMessage; text_width::Int = displaysize(io)[
     elseif istoolmessage(msg)
         isnothing(msg.content) ? string("Name: ", msg.name, ", Args: ", msg.raw) :
         string(msg.content)
+    elseif msg isa AnnotationMessage
+        content_str = wrap_string(msg.content, text_width)
+        tags_str = isempty(msg.tags) ? "" : "\nTags: " * join(msg.tags, ", ")
+        comment_str = isempty(msg.comment) ? "" : "\nComment: " * msg.comment
+        string(content_str, tags_str, comment_str)
     else
         wrap_string(msg.content, text_width)
     end