From 27e4301a7d2b461ccd85394719b6084c43212ffa Mon Sep 17 00:00:00 2001 From: J S <49557684+svilupp@users.noreply.github.com> Date: Sun, 19 May 2024 10:50:01 +0100 Subject: [PATCH] Improve tracer schema / automated logging (#151) --- CHANGELOG.md | 10 +- docs/src/frequently_asked_questions.md | 55 ++++++++++- src/PromptingTools.jl | 3 +- src/llm_anthropic.jl | 9 +- src/llm_google.jl | 19 ++-- src/llm_interface.jl | 48 ++++++++- src/llm_ollama.jl | 11 +-- src/llm_openai.jl | 9 +- src/llm_shared.jl | 13 +++ src/llm_sharegpt.jl | 15 +-- src/llm_tracer.jl | 132 ++++++++++++++++++++++--- src/messages.jl | 25 +++-- src/templates.jl | 20 +++- src/user_preferences.jl | 10 +- test/llm_openai.jl | 6 ++ test/llm_shared.jl | 10 +- test/llm_sharegpt.jl | 7 +- test/llm_tracer.jl | 82 ++++++++++++++- test/messages.jl | 9 ++ 19 files changed, 415 insertions(+), 78 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e48cb7ba..271eff83 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,12 +12,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [0.26.0] +### BREAKING CHANGES +- Added new field `meta` to `TracerMessage` and `TracerMessageLike` to hold metadata in a simply dictionary. Change is backward-compatible. +- Changed behaviour of `aitemplates(name::Symbol)` to look for the exact match on the template name, not just a partial match. This is a breaking change for the `aitemplates` function only. Motivation is that having multiple matches could have introduced subtle bugs when looking up valid placeholders for a template. + + ### Added - Improved support for `aiclassify` with OpenAI models (you can now encode upto 40 choices). - Added a template for routing questions `:QuestionRouter` (to be used with `aiclassify`) - -### Changed -- [BREAKING] Changed behaviour of `aitemplates(name::Symbol)` to look for the exact match on the template name, not just a partial match. This is a breaking change for the `aitemplates` function only. Motivation is that having multiple matches could have introduced subtle bugs when looking up valid placeholders for a template. +- Improved tracing by `TracerSchema` to automatically capture crucial metadata such as any LLM API kwargs (`api_kwargs`), use of prompt templates and its versions. Information is captured in `meta(tracer)` dictionary. See `?TracerSchema` for more information. +- New tracing schema `SaverSchema` allows to automatically serialize all conversations. It can be composed with other tracing schemas, eg, `TracerSchema` to automatically capture necessary metadata and serialize. See `?SaverSchema` for more information. ### Fixed - Fixed a bug where `aiclassify` would not work when returning the full conversation for choices with extra descriptions diff --git a/docs/src/frequently_asked_questions.md b/docs/src/frequently_asked_questions.md index c4a724df..d9a89d43 100644 --- a/docs/src/frequently_asked_questions.md +++ b/docs/src/frequently_asked_questions.md @@ -469,4 +469,57 @@ PT.render(PT.NoSchema(), tpl; system = "I exist", user = "say hi") PromptingTools.UserMessage("say hi") ``` -For more information about the rendering pipeline and examples refer to [Walkthrough Example for aigenerate](@ref). \ No newline at end of file +For more information about the rendering pipeline and examples refer to [Walkthrough Example for aigenerate](@ref). + + +## Automatic Logging / Tracing + +If you would like to automatically capture metadata about your conversations, you can use the `TracerSchema`. It automatically captures the necessary metadata such as model, task (`parent_id`), current thread (`thread_id`), API kwargs used and any prompt templates (and its versions). +```julia +using PromptingTools: TracerSchema, OpenAISchema + +wrap_schema = TracerSchema(OpenAISchema()) +msg = aigenerate(wrap_schema, "Say hi!"; model="gpt-4") +# output type should be TracerMessage +msg isa TracerMessage +``` + +You can work with the message like any other message (properties of the inner `object` are overloaded). +You can extract the original message with `unwrap`: +```julia +unwrap(msg) isa String +``` +You can extract the metadata with `meta`: +```julia +meta(msg) isa Dict +``` + + +If you would like to automatically save the conversations, you can use the `SaverSchema`. It automatically serializes the conversation to a file in the directory specified by the environment variable `LOG_DIR`. + +```julia +using PromptingTools: SaverSchema + +wrap_schema = SaverSchema(OpenAISchema()) +msg = aigenerate(wrap_schema, "Say hi!"; model="gpt-4") +``` +See `LOG_DIR` location to find the serialized conversation. + + +You can also compose multiple tracing schemas. For example, you can capture metadata with `TracerSchema` and then save everything automatically with `SaverSchema`: +```julia +using PromptingTools: TracerSchema, SaverSchema, OpenAISchema + +wrap_schema = OpenAISchema() |> TracerSchema |> SaverSchema +conv = aigenerate(wrap_schema,:BlankSystemUser; system="You're a French-speaking assistant!", + user="Say hi!"; model="gpt-4", api_kwargs=(;temperature=0.1), return_all=true) +``` + +`conv` is a vector of tracing messages that will be saved to a JSON together with metadata about the template and `api_kwargs`. + +If you would like to enable this behavior automatically, you can register your favorite model (or re-register existing models) with the "wrapped" schema: + +```julia +PT.register_model!(; name= "gpt-3.5-turbo", schema=OpenAISchema() |> TracerSchema |> SaverSchema) +``` + diff --git a/src/PromptingTools.jl b/src/PromptingTools.jl index 672f2d5b..33d2a0e4 100644 --- a/src/PromptingTools.jl +++ b/src/PromptingTools.jl @@ -2,7 +2,8 @@ module PromptingTools import AbstractTrees using Base64: base64encode -using Dates: now, DateTime +import Dates +using Dates: now, DateTime, @dateformat_str using Logging using OpenAI using JSON3 diff --git a/src/llm_anthropic.jl b/src/llm_anthropic.jl index 049bce70..e2368ddb 100644 --- a/src/llm_anthropic.jl +++ b/src/llm_anthropic.jl @@ -34,17 +34,12 @@ function render(schema::AbstractAnthropicSchema, conversation = Dict{String, Any}[] for msg in messages_replaced - role = if msg isa UserMessage || msg isa UserMessageWithImages - "user" - elseif msg isa AIMessage - "assistant" - end - if msg isa SystemMessage system = msg.content elseif msg isa UserMessage || msg isa AIMessage content = msg.content - push!(conversation, Dict("role" => role, "content" => content)) + push!(conversation, + Dict("role" => role4render(schema, msg), "content" => content)) elseif msg isa UserMessageWithImages error("AbstractAnthropicSchema does not yet support UserMessageWithImages. Please use OpenAISchema instead.") end diff --git a/src/llm_google.jl b/src/llm_google.jl index 165e6a0e..0a9f0fa2 100644 --- a/src/llm_google.jl +++ b/src/llm_google.jl @@ -1,4 +1,11 @@ ## Rendering of converation history for the OpenAI API +## No system message, we need to merge with UserMessage, see below +function role4render(schema::AbstractGoogleSchema, msg::SystemMessage) + "user" +end +function role4render(schema::AbstractGoogleSchema, msg::AIMessage) + "model" +end """ render(schema::AbstractGoogleSchema, messages::Vector{<:AbstractMessage}; @@ -24,15 +31,9 @@ function render(schema::AbstractGoogleSchema, # replace any handlebar variables in the messages for msg in messages_replaced - role = if msg isa SystemMessage - ## No system message, we need to merge with UserMessage, see below - "user" - elseif msg isa UserMessage - "user" - elseif msg isa AIMessage - "model" - end - push!(conversation, Dict(:role => role, :parts => [Dict("text" => msg.content)])) + push!(conversation, + Dict( + :role => role4render(schema, msg), :parts => [Dict("text" => msg.content)])) end ## Merge any subsequent UserMessages merged_conversation = Dict{Symbol, Any}[] diff --git a/src/llm_interface.jl b/src/llm_interface.jl index 650fb633..eaaab9ed 100644 --- a/src/llm_interface.jl +++ b/src/llm_interface.jl @@ -7,6 +7,7 @@ # Ideally, each new interface would be defined in a separate `llm_.jl` file (eg, `llm_chatml.jl`). ## Main Functions +function role4render end function render end function aigenerate end function aiembed end @@ -323,7 +324,15 @@ A schema designed to wrap another schema, enabling pre- and post-execution callb The `TracerSchema` acts as a middleware, allowing developers to insert custom logic before and after the execution of the primary schema's functionality. This can include logging, performance measurement, or any other form of tracing required to understand or improve the execution flow. -# Usage +`TracerSchema` automatically wraps messages in `TracerMessage` type, which has several important fields, eg, +- `object`: the original message - unwrap with utility `unwrap` +- `meta`: a dictionary with metadata about the tracing process (eg, prompt templates, LLM API kwargs) - extract with utility `meta` +- `parent_id`: an identifier for the overall job / high-level conversation with the user where the current conversation `thread` originated. It should be the same for objects in the same thread. +- `thread_id`: an identifier for the current thread or execution context (sub-task, sub-process, CURRENT CONVERSATION or vector of messages) within the broader parent task. It should be the same for objects in the same thread. + +See also: `meta`, `unwrap`, `SaverSchema`, `initialize_tracer`, `finalize_tracer` + +# Example ```julia wrap_schema = TracerSchema(OpenAISchema()) msg = aigenerate(wrap_schema, "Say hi!"; model="gpt-4") @@ -336,6 +345,43 @@ struct TracerSchema <: AbstractTracerSchema schema::AbstractPromptSchema end +""" + SaverSchema <: AbstractTracerSchema + +SaverSchema is a schema that automatically saves the conversation to the disk. +It's useful for debugging and for persistent logging. + +It can be composed with any other schema, eg, `TracerSchema` to save additional metadata. + +Set environment variable `LOG_DIR` to the directory where you want to save the conversation (see `?PREFERENCES`). +Conversations are named by the hash of the first message in the conversation to naturally group subsequent conversations together. + +To use it automatically, re-register the models you use with the schema wrapped in `SaverSchema` + +See also: `meta`, `unwrap`, `TracerSchema`, `initialize_tracer`, `finalize_tracer` + +# Example +```julia +using PromptingTools: TracerSchema, OpenAISchema, SaverSchema +# This schema will first trace the metadata (change to TraceMessage) and then save the conversation to the disk + +wrap_schema = OpenAISchema() |> TracerSchema |> SaverSchema +conv = aigenerate(wrap_schema,:BlankSystemUser; system="You're a French-speaking assistant!", + user="Say hi!"; model="gpt-4", api_kwargs=(;temperature=0.1), return_all=true) + +# conv is a vector of messages that will be saved to a JSON together with metadata about the template and api_kwargs +``` + +If you wanted to enable this automatically for models you use, you can do it like this: +```julia +PT.register_model!(; name= "gpt-3.5-turbo", schema=OpenAISchema() |> TracerSchema |> SaverSchema) +``` +Any subsequent calls `model="gpt-3.5-turbo"` will automatically capture metadata and save the conversation to the disk. +""" +struct SaverSchema <: AbstractTracerSchema + schema::AbstractPromptSchema +end + ## Dispatch into a default schema (can be set by Preferences.jl) # Since we load it as strings, we need to convert it to a symbol and instantiate it global PROMPT_SCHEMA::AbstractPromptSchema = @load_preference("PROMPT_SCHEMA", diff --git a/src/llm_ollama.jl b/src/llm_ollama.jl index ae9b95a8..5d6bdd24 100644 --- a/src/llm_ollama.jl +++ b/src/llm_ollama.jl @@ -7,6 +7,7 @@ ## Schema dedicated to [Ollama's models](https://ollama.ai/), which also managed the prompt templates # ## Rendering of converation history for the Ollama API (similar to OpenAI but not for the images) + """ render(schema::AbstractOllamaSchema, messages::Vector{<:AbstractMessage}; @@ -32,14 +33,8 @@ function render(schema::AbstractOllamaSchema, # replace any handlebar variables in the messages for msg in messages_replaced - role = if msg isa SystemMessage - "system" - elseif msg isa UserMessage || msg isa UserMessageWithImages - "user" - elseif msg isa AIMessage - "assistant" - end - new_message = Dict{String, Any}("role" => role, "content" => msg.content) + new_message = Dict{String, Any}( + "role" => role4render(schema, msg), "content" => msg.content) ## Special case for images if msg isa UserMessageWithImages new_message["images"] = msg.image_url diff --git a/src/llm_openai.jl b/src/llm_openai.jl index 980da7e8..4bee1ced 100644 --- a/src/llm_openai.jl +++ b/src/llm_openai.jl @@ -28,13 +28,6 @@ function render(schema::AbstractOpenAISchema, # replace any handlebar variables in the messages for msg in messages_replaced - role = if msg isa SystemMessage - "system" - elseif msg isa UserMessage || msg isa UserMessageWithImages - "user" - elseif msg isa AIMessage - "assistant" - end ## Special case for images if msg isa UserMessageWithImages # Build message content @@ -50,7 +43,7 @@ function render(schema::AbstractOpenAISchema, else content = msg.content end - push!(conversation, Dict("role" => role, "content" => content)) + push!(conversation, Dict("role" => role4render(schema, msg), "content" => content)) end return conversation diff --git a/src/llm_shared.jl b/src/llm_shared.jl index 06354d80..38682fde 100644 --- a/src/llm_shared.jl +++ b/src/llm_shared.jl @@ -1,4 +1,12 @@ # Reusable functionality across different schemas +function role4render(schema::AbstractPromptSchema, msg::AbstractMessage) + throw(ArgumentError("Function `role4render` is not implemented for the provided schema ($(typeof(schema))) and $(typeof(msg)).")) +end +role4render(schema::AbstractPromptSchema, msg::SystemMessage) = "system" +role4render(schema::AbstractPromptSchema, msg::UserMessage) = "user" +role4render(schema::AbstractPromptSchema, msg::UserMessageWithImages) = "user" +role4render(schema::AbstractPromptSchema, msg::AIMessage) = "assistant" + """ render(schema::NoSchema, messages::Vector{<:AbstractMessage}; @@ -49,6 +57,11 @@ function render(schema::NoSchema, elseif msg isa AIMessage # no replacements push!(conversation, msg) + elseif istracermessage(msg) && issystemmessage(msg.object) + # Look for tracers + count_system_msg += 1 + # move to the front + pushfirst!(conversation, msg) else # Note: Ignores any DataMessage or other types for the prompt/conversation history @warn "Unexpected message type: $(typeof(msg)). Skipping." diff --git a/src/llm_sharegpt.jl b/src/llm_sharegpt.jl index fc935497..52701cfc 100644 --- a/src/llm_sharegpt.jl +++ b/src/llm_sharegpt.jl @@ -1,13 +1,14 @@ ### RENDERING -function sharegpt_role(::AbstractMessage) - throw(ArgumentError("Unsupported message type $(typeof(msg))")) +role4render(::AbstractShareGPTSchema, ::AIMessage) = "gpt" +role4render(::AbstractShareGPTSchema, ::UserMessage) = "human" +role4render(::AbstractShareGPTSchema, ::SystemMessage) = "system" +function role4render(::AbstractShareGPTSchema, ::UserMessageWithImages) + throw(ArgumentError("UserMessageWithImages is not supported in ShareGPT schema")) end -sharegpt_role(::AIMessage) = "gpt" -sharegpt_role(::UserMessage) = "human" -sharegpt_role(::SystemMessage) = "system" -function render(::AbstractShareGPTSchema, conv::AbstractVector{<:AbstractMessage}) - Dict("conversations" => [Dict("from" => sharegpt_role(msg), "value" => msg.content) +function render(schema::AbstractShareGPTSchema, conv::AbstractVector{<:AbstractMessage}) + Dict("conversations" => [Dict("from" => role4render(schema, msg), + "value" => msg.content) for msg in conv]) end diff --git a/src/llm_tracer.jl b/src/llm_tracer.jl index 1d63349a..6f01390f 100644 --- a/src/llm_tracer.jl +++ b/src/llm_tracer.jl @@ -4,6 +4,18 @@ # - Call your ai* function with the tracer schema as usual # Simple passthrough, do nothing +function role4render(schema::AbstractTracerSchema, msg::SystemMessage) + role4render(schema.schema, msg) +end +function role4render(schema::AbstractTracerSchema, msg::UserMessage) + role4render(schema.schema, msg) +end +function role4render(schema::AbstractTracerSchema, msg::UserMessageWithImages) + role4render(schema.schema, msg) +end +function role4render(schema::AbstractTracerSchema, msg::AIMessage) + role4render(schema.schema, msg) +end """ render(tracer_schema::AbstractTracerSchema, conv::AbstractVector{<:AbstractMessage}; kwargs...) @@ -17,28 +29,63 @@ end """ initialize_tracer( - tracer_schema::AbstractTracerSchema; model = "", tracer_kwargs = NamedTuple(), kwargs...) + tracer_schema::AbstractTracerSchema; model = "", tracer_kwargs = NamedTuple(), + prompt::ALLOWED_PROMPT_TYPE = "", kwargs...) Initializes `tracer`/callback (if necessary). Can provide any keyword arguments in `tracer_kwargs` (eg, `parent_id`, `thread_id`, `run_id`). Is executed prior to the `ai*` calls. +By default it captures: +- `time_sent`: the time the request was sent +- `model`: the model to use +- `meta`: a dictionary of additional metadata that is not part of the tracer itself + - `template_name`: the template to use if any + - `template_version`: the template version to use if any + - expanded `api_kwargs`, ie, the keyword arguments to pass to the API call + In the default implementation, we just collect the necessary data to build the tracer object in `finalize_tracer`. + +See also: `meta`, `unwrap`, `TracerSchema`, `SaverSchema`, `finalize_tracer` """ function initialize_tracer( - tracer_schema::AbstractTracerSchema; model = "", tracer_kwargs = NamedTuple(), kwargs...) - return (; time_sent = now(), model, tracer_kwargs...) + tracer_schema::AbstractTracerSchema; model = "", tracer_kwargs = NamedTuple(), + prompt::ALLOWED_PROMPT_TYPE = "", api_kwargs::NamedTuple = NamedTuple(), + kwargs...) + meta = Dict{Symbol, Any}(k => v for (k, v) in pairs(api_kwargs)) + if haskey(kwargs, :_tracer_template) + tpl = get(kwargs, :_tracer_template, nothing) + meta[:template_name] = tpl.name + metadata = aitemplates(tpl.name) + if !isempty(metadata) + meta[:template_version] = metadata[1].version + end + end + return (; time_sent = now(), model, meta, + tracer_kwargs...) end +function finalize_tracer( + tracer_schema::AbstractTracerSchema, tracer, msg_or_conv; + tracer_kwargs = NamedTuple(), model = "", kwargs...) + # default is a passthrough + return msg_or_conv +end """ finalize_tracer( - tracer_schema::AbstractTracerSchema, tracer, msg_or_conv; tracer_kwargs = NamedTuple(), model = "", kwargs...) + tracer_schema::AbstractTracerSchema, tracer, msg_or_conv::Union{ + AbstractMessage, AbstractVector{<:AbstractMessage}}; + tracer_kwargs = NamedTuple(), model = "", kwargs...) Finalizes the calltracer of whatever is nedeed after the `ai*` calls. Use `tracer_kwargs` to provide any information necessary (eg, `parent_id`, `thread_id`, `run_id`). In the default implementation, we convert all non-tracer messages into `TracerMessage`. + +See also: `meta`, `unwrap`, `SaverSchema`, `initialize_tracer` """ function finalize_tracer( - tracer_schema::AbstractTracerSchema, tracer, msg_or_conv; tracer_kwargs = NamedTuple(), model = "", kwargs...) + tracer_schema::AbstractTracerSchema, tracer, msg_or_conv::Union{ + AbstractMessage, AbstractVector{<:AbstractMessage}}; + tracer_kwargs = NamedTuple(), model = "", kwargs...) # We already captured all kwargs, they are already in `tracer`, we can ignore them in this implementation time_received = now() # work with arrays for unified processing @@ -46,22 +93,74 @@ function finalize_tracer( conv = msg_or_conv isa AbstractVector{<:AbstractMessage} ? convert(Vector{AbstractMessage}, msg_or_conv) : AbstractMessage[msg_or_conv] + # extract the relevant properties from the tracer + tracer_subset = [f => get(tracer, f, nothing) + for f in fieldnames(TracerMessage) if haskey(tracer, f)] # all msg non-traced, set times for i in eachindex(conv) msg = conv[i] # change into TracerMessage if not already, use the current kwargs if !istracermessage(msg) # we saved our data for `tracer` - conv[i] = TracerMessage(; object = msg, tracer..., time_received) + conv[i] = TracerMessage(; object = msg, tracer_subset..., time_received) end end return is_vector ? conv : first(conv) end +## Specialized finalizer to save the response to the disk """ - aigenerate(tracer_schema::AbstractTracerSchema, prompt::ALLOWED_PROMPT_TYPE; + finalize_tracer( + tracer_schema::SaverSchema, tracer, msg_or_conv::Union{ + AbstractMessage, AbstractVector{<:AbstractMessage}}; tracer_kwargs = NamedTuple(), model = "", kwargs...) +Finalizes the calltracer by saving the provided conversation `msg_or_conv` to the disk. + +Path is `LOG_DIR/conversation____.json`, + where `LOG_DIR` is set by user preferences or ENV variable (defaults to `log/` in current working directory). + +It can be composed with `TracerSchema` to also attach necessary metadata (see below). + +# Example +```julia +wrap_schema = PT.SaverSchema(PT.TracerSchema(PT.OpenAISchema())) +conv = aigenerate(wrap_schema,:BlankSystemUser; system="You're a French-speaking assistant!", + user="Say hi!"; model="gpt-4", api_kwargs=(;temperature=0.1), return_all=true) + +# conv is a vector of messages that will be saved to a JSON together with metadata about the template and api_kwargs +``` + +See also: `meta`, `unwrap`, `TracerSchema`, `initialize_tracer` +""" +function finalize_tracer( + tracer_schema::SaverSchema, tracer, msg_or_conv::Union{ + AbstractMessage, AbstractVector{<:AbstractMessage}}; + tracer_kwargs = NamedTuple(), model = "", kwargs...) + # We already captured all kwargs, they are already in `tracer`, we can ignore them in this implementation + time_received = now() + # work with arrays for unified processing + is_vector = msg_or_conv isa AbstractVector + conv = msg_or_conv isa AbstractVector{<:AbstractMessage} ? + convert(Vector{AbstractMessage}, msg_or_conv) : + AbstractMessage[msg_or_conv] + + # Log the conversation to disk, save by hash of the first convo message + timestamp + first_msg_hash = hash(first(conv).content) + time_received_str = Dates.format( + time_received, dateformat"YYYYmmdd_HHMMSS") + path = joinpath( + LOG_DIR, + "conversation__$(first_msg_hash)__$(time_received_str).json") + mkpath(dirname(path)) + save_conversation(path, conv) + return is_vector ? conv : first(conv) +end + +""" + aigenerate(tracer_schema::AbstractTracerSchema, prompt::ALLOWED_PROMPT_TYPE; + tracer_kwargs = NamedTuple(), model = "", return_all::Bool = false, kwargs...) + Wraps the normal `aigenerate` call in a tracing/callback system. Use `tracer_kwargs` to provide any information necessary to the tracer/callback system only (eg, `parent_id`, `thread_id`, `run_id`). Logic: @@ -86,12 +185,15 @@ all(PT.istracermessage, conv) #true ``` """ function aigenerate(tracer_schema::AbstractTracerSchema, prompt::ALLOWED_PROMPT_TYPE; - tracer_kwargs = NamedTuple(), model = "", kwargs...) - tracer = initialize_tracer(tracer_schema; model, tracer_kwargs, kwargs...) + tracer_kwargs = NamedTuple(), model = "", return_all::Bool = false, kwargs...) + tracer = initialize_tracer(tracer_schema; model, tracer_kwargs, prompt, kwargs...) + # Force to return all convo and then subset as necessary merged_kwargs = isempty(model) ? kwargs : (; model, kwargs...) # to not override default model for each schema if not provided - msg_or_conv = aigenerate(tracer_schema.schema, prompt; merged_kwargs...) - return finalize_tracer( + msg_or_conv = aigenerate( + tracer_schema.schema, prompt; tracer_kwargs, return_all = true, merged_kwargs...) + output = finalize_tracer( tracer_schema, tracer, msg_or_conv; model, tracer_kwargs, kwargs...) + return return_all ? output : last(output) end """ @@ -130,7 +232,7 @@ Logic: """ function aiclassify(tracer_schema::AbstractTracerSchema, prompt::ALLOWED_PROMPT_TYPE; tracer_kwargs = NamedTuple(), model = "", kwargs...) - tracer = initialize_tracer(tracer_schema; model, tracer_kwargs..., kwargs...) + tracer = initialize_tracer(tracer_schema; model, prompt, tracer_kwargs..., kwargs...) merged_kwargs = isempty(model) ? kwargs : (; model, kwargs...) # to not override default model for each schema if not provided classify_or_conv = aiclassify(tracer_schema.schema, prompt; merged_kwargs...) return finalize_tracer( @@ -150,7 +252,7 @@ Logic: """ function aiextract(tracer_schema::AbstractTracerSchema, prompt::ALLOWED_PROMPT_TYPE; tracer_kwargs = NamedTuple(), model = "", kwargs...) - tracer = initialize_tracer(tracer_schema; model, tracer_kwargs..., kwargs...) + tracer = initialize_tracer(tracer_schema; model, prompt, tracer_kwargs..., kwargs...) merged_kwargs = isempty(model) ? kwargs : (; model, kwargs...) # to not override default model for each schema if not provided extract_or_conv = aiextract(tracer_schema.schema, prompt; merged_kwargs...) return finalize_tracer( @@ -170,7 +272,7 @@ Logic: """ function aiscan(tracer_schema::AbstractTracerSchema, prompt::ALLOWED_PROMPT_TYPE; tracer_kwargs = NamedTuple(), model = "", kwargs...) - tracer = initialize_tracer(tracer_schema; model, tracer_kwargs..., kwargs...) + tracer = initialize_tracer(tracer_schema; model, prompt, tracer_kwargs..., kwargs...) merged_kwargs = isempty(model) ? kwargs : (; model, kwargs...) # to not override default model for each schema if not provided scan_or_conv = aiscan(tracer_schema.schema, prompt; merged_kwargs...) return finalize_tracer( @@ -190,7 +292,7 @@ Logic: """ function aiimage(tracer_schema::AbstractTracerSchema, prompt::ALLOWED_PROMPT_TYPE; tracer_kwargs = NamedTuple(), model = "", kwargs...) - tracer = initialize_tracer(tracer_schema; model, tracer_kwargs..., kwargs...) + tracer = initialize_tracer(tracer_schema; model, prompt, tracer_kwargs..., kwargs...) merged_kwargs = isempty(model) ? kwargs : (; model, kwargs...) # to not override default model for each schema if not provided image_or_conv = aiimage(tracer_schema.schema, prompt; merged_kwargs...) return finalize_tracer( diff --git a/src/messages.jl b/src/messages.jl index c6fecbb9..fdc93120 100644 --- a/src/messages.jl +++ b/src/messages.jl @@ -142,6 +142,10 @@ issystemmessage(m::AbstractMessage) = m isa SystemMessage isdatamessage(m::AbstractMessage) = m isa DataMessage isaimessage(m::AbstractMessage) = m isa AIMessage istracermessage(m::AbstractMessage) = m isa AbstractTracerMessage +isusermessage(m::AbstractTracerMessage) = isusermessage(m.object) +issystemmessage(m::AbstractTracerMessage) = issystemmessage(m.object) +isdatamessage(m::AbstractTracerMessage) = isdatamessage(m.object) +isaimessage(m::AbstractTracerMessage) = isaimessage(m.object) # equality check for testing, only equal if all fields are equal and type is the same Base.var"=="(m1::AbstractMessage, m2::AbstractMessage) = false @@ -225,6 +229,7 @@ A mutable wrapper message designed for tracing the flow of messages through the - `model::String`: The name of the model that generated the message. Defaults to empty. - `parent_id::Symbol`: An identifier for the job or process that the message is associated with. Higher-level tracing ID. - `thread_id::Symbol`: An identifier for the thread (series of messages for one model/agent) or execution context within the job where the message originated. It should be the same for messages in the same thread. +- `meta::Union{Nothing, Dict{Symbol, Any}}`: A dictionary for additional metadata that is not part of the message itself. Try to limit to a small number of items and singletons to be serializable. - `_type::Symbol`: A fixed symbol identifying the type of the message as `:eventmessage`, used for type discrimination. This structure is particularly useful for debugging, monitoring, and auditing the flow of messages in systems that involve complex interactions or asynchronous processing. @@ -254,6 +259,7 @@ Base.@kwdef mutable struct TracerMessage{T <: parent_id::Symbol = gensym("parent") thread_id::Symbol = gensym("thread") run_id::Union{Nothing, Int} = Int(rand(Int32)) + meta::Union{Nothing, Dict{Symbol, Any}} = Dict{Symbol, Any}() _type::Symbol = :tracermessage end function TracerMessage(msg::Union{AbstractChatMessage, AbstractDataMessage}; kwargs...) @@ -275,8 +281,9 @@ It provides a flexible way to track and annotate objects as they move through di - `time_sent::Union{Nothing, DateTime}`: The timestamp when the object was originally sent, if available. - `model::String`: The name of the model or process that generated or is associated with the object. Defaults to empty. - `parent_id::Symbol`: An identifier for the job or process that the object is associated with. Higher-level tracing ID. -- `thread_id::Symbol`: An identifier for the thread or execution context within the job where the object originated. It should be the same for objects in the same thread. -- `run_id::Union{Nothing, Int}`: A unique identifier for the run or instance of the process that generated the object. Defaults to a random integer. +- `thread_id::Symbol`: An identifier for the thread or execution context (sub-task, sub-process) within the job where the object originated. It should be the same for objects in the same thread. +- `run_id::Union{Nothing, Int}`: A unique identifier for the run or instance of the process (ie, a single call to the LLM) that generated the object. Defaults to a random integer. +- `meta::Union{Nothing, Dict{Symbol, Any}}`: A dictionary for additional metadata that is not part of the object itself. Try to limit to a small number of items and singletons to be serializable. - `_type::Symbol`: A fixed symbol identifying the type of the tracer as `:tracermessage`, used for type discrimination. This structure is particularly useful for systems that involve complex interactions or asynchronous processing, where tracking the flow and transformation of objects is crucial. @@ -294,12 +301,10 @@ All fields are optional besides the `object`. parent_id::Symbol = gensym("parent") thread_id::Symbol = gensym("thread") run_id::Union{Nothing, Int} = Int(rand(Int32)) + meta::Union{Nothing, Dict{Symbol, Any}} = Dict{Symbol, Any}() _type::Symbol = :tracermessagelike ## TracerMessageLike() = new() end -## function TracerMessageLike() -## TracerMessageLike(; object = undef) -## end function TracerMessageLike( object; kwargs...) TracerMessageLike(; object, kwargs...) @@ -325,11 +330,16 @@ function Base.copy(t::T) where {T <: Union{AbstractTracerMessage, AbstractTracer T([deepcopy(getfield(t, f)) for f in fieldnames(T)]...) end -"Unwraps the tracer message, returning the original `object`." +"Unwraps the tracer message or tracer-like object, returning the original `object`." function unwrap(t::Union{AbstractTracerMessage, AbstractTracer}) getfield(t, :object) end +"Extracts the metadata dictionary from the tracer message or tracer-like object." +function meta(t::Union{AbstractTracerMessage, AbstractTracer}) + getfield(t, :meta) +end + "Aligns the tracer message, updating the `parent_id`, `thread_id`. Often used to align multiple tracers in the vector to have the same IDs." function align_tracer!( t::Union{AbstractTracerMessage, AbstractTracer}; parent_id::Symbol = t.parent_id, @@ -406,6 +416,9 @@ end # kwargs...) # render(schema, messages; kwargs...) # end +function role4render(schema::AbstractPromptSchema, msg::AbstractTracerMessage) + role4render(schema, msg.object) +end function render(schema::AbstractPromptSchema, msg::AbstractMessage; kwargs...) render(schema, [msg]; kwargs...) end diff --git a/src/templates.jl b/src/templates.jl index d189d255..97259868 100644 --- a/src/templates.jl +++ b/src/templates.jl @@ -377,19 +377,29 @@ function render(schema::AbstractTracerSchema, template::AITemplate; kwargs...) render(schema.schema, template; kwargs...) end function aigenerate(schema::AbstractTracerSchema, template::Symbol; kwargs...) - aigenerate(schema, render(schema, AITemplate(template)); kwargs...) + tpl = AITemplate(template) + aigenerate(schema, render(schema, tpl); + _tracer_template = tpl, kwargs...) end function aiclassify(schema::AbstractTracerSchema, template::Symbol; kwargs...) - aiclassify(schema, render(schema, AITemplate(template)); kwargs...) + tpl = AITemplate(template) + aiclassify(schema, render(schema, tpl); + _tracer_template = tpl, kwargs...) end function aiextract(schema::AbstractTracerSchema, template::Symbol; kwargs...) - aiextract(schema, render(schema, AITemplate(template)); kwargs...) + tpl = AITemplate(template) + aiextract(schema, render(schema, tpl); + _tracer_template = tpl, kwargs...) end function aiscan(schema::AbstractTracerSchema, template::Symbol; kwargs...) - aiscan(schema, render(schema, AITemplate(template)); kwargs...) + tpl = AITemplate(template) + aiscan(schema, render(schema, tpl); + _tracer_template = tpl, kwargs...) end function aiimage(schema::AbstractTracerSchema, template::Symbol; kwargs...) - aiimage(schema, render(schema, AITemplate(template)); kwargs...) + tpl = AITemplate(template) + aiimage(schema, render(schema, tpl); + _tracer_template = tpl, kwargs...) end ## Utility for creating templates diff --git a/src/user_preferences.jl b/src/user_preferences.jl index 89ad4e0b..cd0d0ae6 100644 --- a/src/user_preferences.jl +++ b/src/user_preferences.jl @@ -31,6 +31,7 @@ Check your preferences by calling `get_preferences(key::String)`. See `CONV_HISTORY` for more information. - `LOCAL_SERVER`: The URL of the local server to use for `ai*` calls. Defaults to `http://localhost:10897/v1`. This server is called when you call `model="local"` See `?LocalServerOpenAISchema` for more information and examples. +- `LOG_DIR`: The directory to save the logs to, eg, when using `SaverSchema <: AbstractTracerSchema`. Defaults to `joinpath(pwd(), "log")`. Refer to `?SaverSchema` for more information on how it works and examples. At the moment it is not possible to persist changes to `MODEL_REGISTRY` across sessions. Define your `register_model!()` calls in your `startup.jl` file to make them available across sessions or put them at the top of your script. @@ -48,6 +49,7 @@ Define your `register_model!()` calls in your `startup.jl` file to make them ava - `VOYAGE_API_KEY`: The API key for the Voyage API. Free tier is upto 50M tokens! Get yours from [here](https://dash.voyageai.com/api-keys). - `GROQ_API_KEY`: The API key for the Groq API. Free in beta! Get yours from [here](https://console.groq.com/keys). - `DEEPSEEK_API_KEY`: The API key for the DeepSeek API. Get \$5 credit when you join. Get yours from [here](https://platform.deepseek.com/api_keys). +- `LOG_DIR`: The directory to save the logs to, eg, when using `SaverSchema <: AbstractTracerSchema`. Defaults to `joinpath(pwd(), "log")`. Refer to `?SaverSchema` for more information on how it works and examples. Preferences.jl takes priority over ENV variables, so if you set a preference, it will take precedence over the ENV variable. @@ -72,7 +74,8 @@ const ALLOWED_PREFERENCES = ["MISTRALAI_API_KEY", "MODEL_ALIASES", "PROMPT_SCHEMA", "MAX_HISTORY_LENGTH", - "LOCAL_SERVER"] + "LOCAL_SERVER", + "LOG_DIR"] """ set_preferences!(pairs::Pair{String, <:Any}...) @@ -191,6 +194,11 @@ _temp = get(ENV, "LOCAL_SERVER", "http://localhost:10897/v1") const LOCAL_SERVER::String = @load_preference("LOCAL_SERVER", default=_temp); +_temp = get(ENV, "LOG_DIR", joinpath(pwd(), "log")) +## Address of the local server +const LOG_DIR::String = @load_preference("LOG_DIR", + default=_temp); + ## CONVERSATION HISTORY """ CONV_HISTORY diff --git a/test/llm_openai.jl b/test/llm_openai.jl index 1020f7c3..4dc55454 100644 --- a/test/llm_openai.jl +++ b/test/llm_openai.jl @@ -8,6 +8,12 @@ using PromptingTools: encode_choices, decode_choices, response_to_message, call_ @testset "render-OpenAI" begin schema = OpenAISchema() + + role4render(schema, SystemMessage("System message 1")) == "system" + role4render(schema, UserMessage("User message 1")) == "user" + role4render(schema, UserMessageWithImages("User message 1"; image_url = "")) == "user" + role4render(schema, AIMessage("AI message 1")) == "assistant" + # Given a schema and a vector of messages with handlebar variables, it should replace the variables with the correct values in the conversation dictionary. messages = [ SystemMessage("Act as a helpful AI assistant"), diff --git a/test/llm_shared.jl b/test/llm_shared.jl index ec36599c..82df905a 100644 --- a/test/llm_shared.jl +++ b/test/llm_shared.jl @@ -1,10 +1,18 @@ using PromptingTools: render, NoSchema using PromptingTools: AIMessage, SystemMessage, AbstractMessage, AbstractChatMessage using PromptingTools: UserMessage, UserMessageWithImages -using PromptingTools: finalize_outputs +using PromptingTools: finalize_outputs, role4render @testset "render-NoSchema" begin schema = NoSchema() + + @test role4render(schema, SystemMessage("System message 1")) == "system" + @test role4render(schema, UserMessage("User message 1")) == "user" + @test role4render(schema, UserMessageWithImages("User message 1"; image_url = "")) == + "user" + @test role4render(schema, AIMessage("AI message 1")) == "assistant" + @test_throws ArgumentError role4render(schema, DataMessage(; content = ones(3, 3))) + # Given a schema and a vector of messages with handlebar variables, it should replace the variables with the correct values in the conversation dictionary. messages = [ SystemMessage("Act as a helpful AI assistant"), diff --git a/test/llm_sharegpt.jl b/test/llm_sharegpt.jl index 812269d3..d6a449e1 100644 --- a/test/llm_sharegpt.jl +++ b/test/llm_sharegpt.jl @@ -1,9 +1,14 @@ -using PromptingTools: render, ShareGPTSchema +using PromptingTools: role4render, render, ShareGPTSchema using PromptingTools: AIMessage, SystemMessage, AbstractMessage using PromptingTools: UserMessage, UserMessageWithImages, DataMessage @testset "render-ShareGPT" begin schema = ShareGPTSchema() + + role4render(schema, SystemMessage("System message 1")) == "system" + role4render(schema, UserMessage("User message 1")) == "human" + role4render(schema, AIMessage("AI message 1")) == "gpt" + # Ignores any handlebar replacement, takes conversations as is messages = [ SystemMessage("Act as a helpful AI assistant"), diff --git a/test/llm_tracer.jl b/test/llm_tracer.jl index aa69dcd2..5b9d2162 100644 --- a/test/llm_tracer.jl +++ b/test/llm_tracer.jl @@ -1,12 +1,30 @@ -using PromptingTools: TestEchoOpenAISchema, render, OpenAISchema, TracerSchema +using PromptingTools: TestEchoOpenAISchema, render, OpenAISchema, TracerSchema, SaverSchema using PromptingTools: AIMessage, SystemMessage, AbstractMessage using PromptingTools: UserMessage, UserMessageWithImages, DataMessage, TracerMessage using PromptingTools: CustomProvider, CustomOpenAISchema, MistralOpenAISchema, MODEL_EMBEDDING, MODEL_IMAGE_GENERATION using PromptingTools: initialize_tracer, finalize_tracer, isaimessage, istracermessage, - unwrap, AITemplate + unwrap, meta, AITemplate, render, role4render +@testset "role4render-Tracer" begin + schema = TracerSchema(OpenAISchema()) + + # unwrapping schema + @test role4render(schema, SystemMessage("System message 1")) == "system" + @test role4render(schema, UserMessage("User message 1")) == "user" + @test role4render(schema, UserMessageWithImages("User message 1"; image_url = "")) == + "user" + @test role4render(schema, AIMessage("AI message 1")) == "assistant" + + # unwrapping TracerMessage + @test role4render(OpenAISchema(), TracerMessage(SystemMessage("Abc123"))) == "system" + @test role4render(OpenAISchema(), TracerMessage(UserMessage("Abc123"))) == "user" + @test role4render( + OpenAISchema(), TracerMessage(UserMessageWithImages("Abc123"; image_url = ""))) == + "user" + @test role4render(OpenAISchema(), TracerMessage(AIMessage("Abc123"))) == "assistant" +end @testset "render-Tracer" begin schema = TracerSchema(OpenAISchema()) # Given a schema and a vector of messages with handlebar variables, it should replace the variables with the correct values in the conversation dictionary. @@ -19,6 +37,11 @@ using PromptingTools: initialize_tracer, finalize_tracer, isaimessage, istracerm conv = render(schema, AITemplate(:InputClassifier)) @test conv isa Vector + + ## other schema + schema = SaverSchema(OpenAISchema()) + conv = render(schema, messages) + @test conv == messages end @testset "initialize_tracer" begin @@ -30,22 +53,28 @@ end @test tracer.time_sent >= time_before @test tracer.model == "" @test tracer.a == 1 + @test isempty(tracer.meta) ## custom model and tracer_kwargs custom_model = "custom_model" custom_tracer_kwargs = (parent_id = :parent, thread_id = :thread, run_id = 1) tracer = initialize_tracer( - schema; model = custom_model, tracer_kwargs = custom_tracer_kwargs) + schema; model = custom_model, api_kwargs = (; temperature = 1.0), + tracer_kwargs = custom_tracer_kwargs, _tracer_template = AITemplate(:BlankSystemUser)) @test tracer.time_sent >= time_before @test tracer.model == custom_model @test tracer.parent_id == :parent @test tracer.thread_id == :thread @test tracer.run_id == 1 + @test tracer.meta[:temperature] == 1.0 + @test tracer.meta[:template_name] == :BlankSystemUser + @test tracer.meta[:template_version] == aitemplates(:BlankSystemUser)[1].version end @testset "finalize_tracer" begin schema = TracerSchema(OpenAISchema()) tracer = initialize_tracer(schema; model = "test_model", + api_kwargs = (; temperature = 1.0), tracer_kwargs = (parent_id = :parent, thread_id = :thread, run_id = 1)) time_before = now() @@ -59,6 +88,8 @@ end @test finalized_msg.thread_id == :thread @test finalized_msg.run_id == 1 @test finalized_msg.time_received >= time_before + @test finalized_msg.meta[:temperature] == 1.0 + @test meta(finalized_msg)[:temperature] == 1.0 # vector of non-tracer messages msgs = [SystemMessage("Test message 1"), SystemMessage("Test message 2")] @@ -83,6 +114,26 @@ end @test length(finalized_msgs) == 2 @test finalized_msgs[1] isa TracerMessage @test finalized_msgs[2] === tracer_msg # should be the same object, not a new one + @test meta(finalized_msgs[2])[:temperature] == 1.0 + + ## other schema -- SaverSchema + schema = SaverSchema(OpenAISchema()) + tracer = initialize_tracer(schema) + msgs = [SystemMessage("Test message 1"), SystemMessage("Test message 2")] + conv = finalize_tracer(schema, tracer, msgs) + fn = filter( + x -> occursin("conversation__$(hash(msgs[1].content))", x), readdir( + PT.LOG_DIR; join = true)) |> + first + @test isfile(fn) + @test PT.load_conversation(fn) == conv + ## clean up + isfile(fn) && rm(fn) + + # Passthrough for non-messages (dry-runs) + schema = TracerSchema(OpenAISchema()) + conv = finalize_tracer(schema, tracer, [1, 2, 3, 4, 5]) + @test conv == [1, 2, 3, 4, 5] end @testset "aigenerate-Tracer" begin @@ -104,8 +155,31 @@ end @test msg.model == "xyz" @test msg.thread_id == :ABC1 - msg = aigenerate(schema1, :BlankSystemUser) + msg = aigenerate(schema1, :BlankSystemUser; system = "abc", user = "xyz") + @test istracermessage(msg) + @test msg.meta[:template_name] == :BlankSystemUser + @test msg.meta[:template_version] == aitemplates(:BlankSystemUser)[1].version + + ## other schema -- SaverSchema + schema2 = schema1 |> SaverSchema + msgs = [TracerMessage(SystemMessage("Test message 1")), UserMessage("Hello World")] + msg = aigenerate( + schema2, msgs; model = "xyz", tracer_kwargs = (; thread_id = :ABC1)) @test istracermessage(msg) + fn = filter( + x -> occursin("conversation__$(hash(msgs[1].content))", x), readdir( + PT.LOG_DIR; join = true)) |> + last + @test isfile(fn) + load_conv = PT.load_conversation(fn) + @test length(load_conv) == 3 + loaded_msg = load_conv[end] + @test unwrap(loaded_msg) |> isaimessage + @test loaded_msg.content == "Hello!" + @test loaded_msg.model == "xyz" + @test loaded_msg.thread_id == :ABC1 + ## clean up + isfile(fn) && rm(fn) end @testset "aiembed-Tracer" begin diff --git a/test/messages.jl b/test/messages.jl index b0c2c5b7..93e2b2de 100644 --- a/test/messages.jl +++ b/test/messages.jl @@ -6,6 +6,7 @@ using PromptingTools: isusermessage, issystemmessage, isdatamessage, isaimessage istracermessage using PromptingTools: TracerMessageLike, TracerMessage, align_tracer!, unwrap, AbstractTracerMessage, AbstractTracer, pprint +using PromptingTools: TracerSchema, SaverSchema @testset "Message constructors" begin # Creates an instance of MSG with the given content string. @@ -140,6 +141,14 @@ end @test copy(tr2) !== tr2 # Specific methods + # type trait passthrough to the underlying message + content = "say hi" + @test TracerMessage(UserMessage(content)) |> isusermessage + @test TracerMessage(SystemMessage(content)) |> issystemmessage + @test TracerMessage(DataMessage(; content)) |> isdatamessage + @test TracerMessage(AIMessage(; content)) |> isaimessage + @test TracerMessage(UserMessage(content)) |> AIMessage |> isaimessage + # unwrap the tracer @test unwrap(tr1) == msg1