Improve tracer schema / automated logging (#151)

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

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).
+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)
+```
+

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

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

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}[]

src/llm_interface.jl

@@ -7,6 +7,7 @@
 # Ideally, each new interface would be defined in a separate `llm_<interface>.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",

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

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

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."

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