Skip to content

feat(providers/anthropic): add MessageSerializationCache to avoid O(N²) message serialization #25

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
johnstcn wants to merge 1 commit into
base: main
Choose a base branch
from
+91 −17
Draft

feat(providers/anthropic): add MessageSerializationCache to avoid O(N²) message serialization #25

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 50 additions & 4 deletions providers/anthropic/anthropic.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -286,7 +286,7 @@ func (a languageModel) prepareParams(call fantasy.Call) (
if providerOptions.SendReasoning != nil {
sendReasoning = *providerOptions.SendReasoning
}
systemBlocks, messages, warnings := toPrompt(call.Prompt, sendReasoning)
systemBlocks, messages, warnings := toPrompt(call.Prompt, sendReasoning, GetMessageCache(call.ProviderOptions))

if call.FrequencyPenalty != nil {
warnings = append(warnings, fantasy.CallWarning{
Expand Down Expand Up @@ -760,7 +760,7 @@ func (a languageModel) toTools(tools []fantasy.Tool, toolChoice *fantasy.ToolCho
return rawTools, anthropicToolChoice, warnings, betaFlags
}

func toPrompt(prompt fantasy.Prompt, sendReasoningData bool) ([]anthropic.TextBlockParam, []anthropic.MessageParam, []fantasy.CallWarning) {
func toPrompt(prompt fantasy.Prompt, sendReasoningData bool, cache MessageSerializationCache) ([]anthropic.TextBlockParam, []anthropic.MessageParam, []fantasy.CallWarning) {
var systemBlocks []anthropic.TextBlockParam
var messages []anthropic.MessageParam
var warnings []fantasy.CallWarning
Expand Down Expand Up @@ -799,6 +799,7 @@ func toPrompt(prompt fantasy.Prompt, sendReasoningData bool) ([]anthropic.TextBl

case fantasy.MessageRoleUser:
var anthropicContent []anthropic.ContentBlockParamUnion
blockHasCacheControl := false
for _, msg := range block.Messages {
if msg.Role == fantasy.MessageRoleUser {
for i, part := range msg.Content {
Expand All @@ -807,6 +808,9 @@ func toPrompt(prompt fantasy.Prompt, sendReasoningData bool) ([]anthropic.TextBl
if cacheControl == nil && isLastPart {
cacheControl = GetCacheControl(msg.ProviderOptions)
}
if cacheControl != nil {
blockHasCacheControl = true
}
switch part.GetType() {
case fantasy.ContentTypeText:
text, ok := fantasy.AsMessagePart[fantasy.TextPart](part)
Expand Down Expand Up @@ -855,6 +859,9 @@ func toPrompt(prompt fantasy.Prompt, sendReasoningData bool) ([]anthropic.TextBl
if cacheControl == nil && isLastPart {
cacheControl = GetCacheControl(msg.ProviderOptions)
}
if cacheControl != nil {
blockHasCacheControl = true
}
result, ok := fantasy.AsMessagePart[fantasy.ToolResultPart](part)
if !ok {
continue
Expand Down Expand Up @@ -923,16 +930,22 @@ func toPrompt(prompt fantasy.Prompt, sendReasoningData bool) ([]anthropic.TextBl
})
continue
}
messages = append(messages, anthropic.NewUserMessage(anthropicContent...))
outIdx := len(messages)
msgParam := anthropic.NewUserMessage(anthropicContent...)
messages = append(messages, applySerialisationCache(cache, outIdx, blockHasCacheControl, msgParam))
case fantasy.MessageRoleAssistant:
var anthropicContent []anthropic.ContentBlockParamUnion
blockHasCacheControl := false
for _, msg := range block.Messages {
for i, part := range msg.Content {
isLastPart := i == len(msg.Content)-1
cacheControl := GetCacheControl(part.Options())
if cacheControl == nil && isLastPart {
cacheControl = GetCacheControl(msg.ProviderOptions)
}
if cacheControl != nil {
blockHasCacheControl = true
}
switch part.GetType() {
case fantasy.ContentTypeText:
text, ok := fantasy.AsMessagePart[fantasy.TextPart](part)
Expand Down Expand Up @@ -1042,12 +1055,45 @@ func toPrompt(prompt fantasy.Prompt, sendReasoningData bool) ([]anthropic.TextBl
})
continue
}
messages = append(messages, anthropic.NewAssistantMessage(anthropicContent...))
outIdx := len(messages)
msgParam := anthropic.NewAssistantMessage(anthropicContent...)
messages = append(messages, applySerialisationCache(cache, outIdx, blockHasCacheControl, msgParam))
}
}
return systemBlocks, messages, warnings
}

// applySerialisationCache applies cache-based serialisation to a
// MessageParam. When the cache is nil or the block has cache
// control directives (which change between calls), the message is
// returned unmodified. On a cache hit the pre-serialised JSON is
// applied via param.SetJSON. On a miss the message is serialised,
// stored, and then param.SetJSON is applied so the SDK skips
// re-serialisation.
func applySerialisationCache(
cache MessageSerializationCache,
idx int,
hasCacheControl bool,
msg anthropic.MessageParam,
) anthropic.MessageParam {
if cache == nil || hasCacheControl {
return msg
}
if cached, ok := cache.Get(idx); ok {
param.SetJSON(cached, &msg)
return msg
}
data, err := msg.MarshalJSON()
if err != nil {
// Serialisation should never fail for well-formed params.
// Fall through without caching.
return msg
}
cache.Set(idx, data)
param.SetJSON(data, &msg)
return msg
}

func hasVisibleUserContent(content []anthropic.ContentBlockParamUnion) bool {
for _, block := range content {
if block.OfText != nil || block.OfImage != nil || block.OfToolResult != nil {
Expand Down
26 changes: 13 additions & 13 deletions providers/anthropic/anthropic_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ func TestToPrompt_DropsEmptyMessages(t *testing.T) {
},
}

systemBlocks, messages, warnings := toPrompt(prompt, true)
systemBlocks, messages, warnings := toPrompt(prompt, true, nil)

require.Empty(t, systemBlocks)
require.Len(t, messages, 1, "should only have user message, assistant message should be dropped")
Expand Down Expand Up @@ -86,7 +86,7 @@ func TestToPrompt_DropsEmptyMessages(t *testing.T) {
},
}

systemBlocks, messages, warnings := toPrompt(prompt, false)
systemBlocks, messages, warnings := toPrompt(prompt, false, nil)

require.Empty(t, systemBlocks)
require.Len(t, messages, 1, "should only have user message, assistant message should be dropped")
Expand All @@ -113,7 +113,7 @@ func TestToPrompt_DropsEmptyMessages(t *testing.T) {
},
}

systemBlocks, messages, warnings := toPrompt(prompt, true)
systemBlocks, messages, warnings := toPrompt(prompt, true, nil)

require.Empty(t, systemBlocks)
require.Len(t, messages, 1, "should only have user message")
Expand All @@ -140,7 +140,7 @@ func TestToPrompt_DropsEmptyMessages(t *testing.T) {
},
}

systemBlocks, messages, warnings := toPrompt(prompt, true)
systemBlocks, messages, warnings := toPrompt(prompt, true, nil)

require.Empty(t, systemBlocks)
require.Len(t, messages, 2, "should have both user and assistant messages")
Expand Down Expand Up @@ -169,7 +169,7 @@ func TestToPrompt_DropsEmptyMessages(t *testing.T) {
},
}

systemBlocks, messages, warnings := toPrompt(prompt, true)
systemBlocks, messages, warnings := toPrompt(prompt, true, nil)

require.Empty(t, systemBlocks)
require.Len(t, messages, 2, "should have both user and assistant messages")
Expand Down Expand Up @@ -198,7 +198,7 @@ func TestToPrompt_DropsEmptyMessages(t *testing.T) {
},
}

systemBlocks, messages, warnings := toPrompt(prompt, true)
systemBlocks, messages, warnings := toPrompt(prompt, true, nil)

require.Empty(t, systemBlocks)
require.Len(t, messages, 1, "should only have user message")
Expand Down Expand Up @@ -233,7 +233,7 @@ func TestToPrompt_DropsEmptyMessages(t *testing.T) {
},
}

systemBlocks, messages, warnings := toPrompt(prompt, true)
systemBlocks, messages, warnings := toPrompt(prompt, true, nil)

require.Empty(t, systemBlocks)
require.Len(t, messages, 2, "should have both user and assistant messages")
Expand All @@ -255,7 +255,7 @@ func TestToPrompt_DropsEmptyMessages(t *testing.T) {
},
}

systemBlocks, messages, warnings := toPrompt(prompt, true)
systemBlocks, messages, warnings := toPrompt(prompt, true, nil)

require.Empty(t, systemBlocks)
require.Len(t, messages, 1)
Expand All @@ -277,7 +277,7 @@ func TestToPrompt_DropsEmptyMessages(t *testing.T) {
},
}

systemBlocks, messages, warnings := toPrompt(prompt, true)
systemBlocks, messages, warnings := toPrompt(prompt, true, nil)

require.Empty(t, systemBlocks)
require.Empty(t, messages)
Expand All @@ -302,7 +302,7 @@ func TestToPrompt_DropsEmptyMessages(t *testing.T) {
},
}

systemBlocks, messages, warnings := toPrompt(prompt, true)
systemBlocks, messages, warnings := toPrompt(prompt, true, nil)

require.Empty(t, systemBlocks)
require.Len(t, messages, 1)
Expand All @@ -324,7 +324,7 @@ func TestToPrompt_DropsEmptyMessages(t *testing.T) {
},
}

systemBlocks, messages, warnings := toPrompt(prompt, true)
systemBlocks, messages, warnings := toPrompt(prompt, true, nil)

require.Empty(t, systemBlocks)
require.Len(t, messages, 1)
Expand All @@ -349,7 +349,7 @@ func TestToPrompt_DropsEmptyMessages(t *testing.T) {
},
}

systemBlocks, messages, warnings := toPrompt(prompt, true)
systemBlocks, messages, warnings := toPrompt(prompt, true, nil)

require.Empty(t, systemBlocks)
require.Len(t, messages, 1)
Expand Down Expand Up @@ -741,7 +741,7 @@ func TestToPrompt_WebSearchProviderExecutedToolResults(t *testing.T) {
},
}

_, messages, warnings := toPrompt(prompt, true)
_, messages, warnings := toPrompt(prompt, true, nil)

// No warnings expected; the provider-executed result is in the
// assistant message so there is no empty tool message to drop.
Expand Down
28 changes: 28 additions & 0 deletions providers/anthropic/provider_options.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -188,6 +188,34 @@ type CacheControl struct {
Type string `json:"type"`
}

// MessageCacheKey is the provider-options key used to pass a
// MessageSerializationCache through fantasy.Call.
const MessageCacheKey = Name + ".message_cache"

// MessageSerializationCache allows pre-serialized message JSON to
// be reused across agentic loop steps. The cache is keyed by output
// MessageParam index. Implementations need not be safe for
// concurrent use.
type MessageSerializationCache interface {
Get(index int) (json.RawMessage, bool)
Set(index int, data json.RawMessage)
Clear()
}

// GetMessageCache extracts a MessageSerializationCache from
// provider options, if present.
func GetMessageCache(providerOptions fantasy.ProviderOptions) MessageSerializationCache {
v, ok := providerOptions[MessageCacheKey]
if !ok {
return nil
}
cache, ok := v.(MessageSerializationCache)
if !ok {
return nil
}
return cache
}

// NewProviderOptions creates new provider options for the Anthropic provider.
func NewProviderOptions(opts *ProviderOptions) fantasy.ProviderOptions {
return fantasy.ProviderOptions{
Expand Down