diff --git a/plugins/maxim/go.mod b/plugins/maxim/go.mod index 0832c55425..8976df5562 100644 --- a/plugins/maxim/go.mod +++ b/plugins/maxim/go.mod @@ -3,7 +3,7 @@ module github.com/maximhq/bifrost/plugins/maxim go 1.24.1 require ( - github.com/maximhq/bifrost/core v1.0.9 + github.com/maximhq/bifrost/core v1.1.0 github.com/maximhq/maxim-go v0.1.3 ) diff --git a/plugins/maxim/go.sum b/plugins/maxim/go.sum index c9e51dcc3e..5d234ea0ae 100644 --- a/plugins/maxim/go.sum +++ b/plugins/maxim/go.sum @@ -36,8 +36,8 @@ github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/klauspost/compress v1.18.0 h1:c/Cqfb0r+Yi+JtIEq73FWXVkRonBlf0CRNYc8Zttxdo= github.com/klauspost/compress v1.18.0/go.mod h1:2Pp+KzxcywXVXMr50+X0Q/Lsb43OQHYWRCY2AiWywWQ= -github.com/maximhq/bifrost/core v1.0.9 h1:OWUHCWCQsBH43YPIy2AsqNMZhoFXYe/qhJSCLbw5su8= -github.com/maximhq/bifrost/core v1.0.9/go.mod h1:8ycaWQ9bjQezoUT/x6a82VmPjoqLzyGglQ0RnnlZjqo= +github.com/maximhq/bifrost/core v1.1.0 h1:qZ3rm6SjZqx4R/5Wu9Fen7yLkL3SFK8dXyiQ2k08Oro= +github.com/maximhq/bifrost/core v1.1.0/go.mod h1:8ycaWQ9bjQezoUT/x6a82VmPjoqLzyGglQ0RnnlZjqo= github.com/maximhq/maxim-go v0.1.3 h1:nVzdz3hEjZVxmWHARWIM+Yrn1Jp50qrsK4BA/sz2jj8= github.com/maximhq/maxim-go v0.1.3/go.mod h1:0+UTWM7UZwNNE5VnljLtr/vpRGtYP8r/2q9WDwlLWFw= github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw= diff --git a/plugins/maxim/main.go b/plugins/maxim/main.go index c50ad930d5..a7ca677727 100644 --- a/plugins/maxim/main.go +++ b/plugins/maxim/main.go @@ -103,7 +103,7 @@ type Plugin struct { // Returns: // - *schemas.BifrostRequest: The original request, unmodified // - error: Any error that occurred during trace/generation creation -func (plugin *Plugin) PreHook(ctx *context.Context, req *schemas.BifrostRequest) (*schemas.BifrostRequest, error) { +func (plugin *Plugin) PreHook(ctx *context.Context, req *schemas.BifrostRequest) (*schemas.BifrostRequest, *schemas.BifrostResponse, error) { var traceID string var sessionID string @@ -111,7 +111,7 @@ func (plugin *Plugin) PreHook(ctx *context.Context, req *schemas.BifrostRequest) if ctx != nil { if existingGenerationID, ok := (*ctx).Value(GenerationIDKey).(string); ok && existingGenerationID != "" { // If generationID exists, return early - return req, nil + return req, nil, nil } if existingTraceID, ok := (*ctx).Value(TraceIDKey).(string); ok && existingTraceID != "" { @@ -137,32 +137,28 @@ func (plugin *Plugin) PreHook(ctx *context.Context, req *schemas.BifrostRequest) "model": req.Model, } for _, message := range *req.Input.ChatCompletionInput { - if message.Content != nil { - messages = append(messages, logging.CompletionRequest{ - Role: string(message.Role), - Content: message.Content, - }) - } else if message.UserMessage != nil && message.UserMessage.ImageContent != nil { - messages = append(messages, logging.CompletionRequest{ - Role: string(message.Role), - Content: message.UserMessage.ImageContent, - }) - } else if message.ToolCalls != nil { - messages = append(messages, logging.CompletionRequest{ - Role: string(message.Role), - Content: message.ToolCalls, - }) - } else if message.ToolMessage != nil && message.ToolMessage.ToolCallID != nil { - messages = append(messages, logging.CompletionRequest{ - Role: string(message.Role), - Content: message.ToolMessage, - }) - } + messages = append(messages, logging.CompletionRequest{ + Role: string(message.Role), + Content: message.Content, + }) } if len(*req.Input.ChatCompletionInput) > 0 { lastMsg := (*req.Input.ChatCompletionInput)[len(*req.Input.ChatCompletionInput)-1] - if lastMsg.Content != nil { - latestMessage = *lastMsg.Content + if lastMsg.Content.ContentStr != nil { + latestMessage = *lastMsg.Content.ContentStr + } else if lastMsg.Content.ContentBlocks != nil { + // Find the last text content block + for i := len(*lastMsg.Content.ContentBlocks) - 1; i >= 0; i-- { + block := (*lastMsg.Content.ContentBlocks)[i] + if block.Type == "text" && block.Text != nil { + latestMessage = *block.Text + break + } + } + // If no text block found, use placeholder + if latestMessage == "" { + latestMessage = "-" + } } } } else if req.Input.TextCompletionInput != nil { @@ -172,7 +168,7 @@ func (plugin *Plugin) PreHook(ctx *context.Context, req *schemas.BifrostRequest) "model": req.Model, } messages = append(messages, logging.CompletionRequest{ - Role: "user", + Role: string(schemas.ModelChatMessageRoleUser), Content: req.Input.TextCompletionInput, }) latestMessage = *req.Input.TextCompletionInput @@ -225,7 +221,7 @@ func (plugin *Plugin) PreHook(ctx *context.Context, req *schemas.BifrostRequest) *ctx = context.WithValue(*ctx, GenerationIDKey, generationID) } - return req, nil + return req, nil, nil } // PostHook is called after a request has been processed by Bifrost. @@ -265,3 +261,9 @@ func (plugin *Plugin) PostHook(ctxRef *context.Context, res *schemas.BifrostResp return res, nil } + +func (plugin *Plugin) Cleanup() error { + plugin.logger.Flush() + + return nil +} diff --git a/plugins/maxim/plugin_test.go b/plugins/maxim/plugin_test.go index 4cfc2068e8..459d3def06 100644 --- a/plugins/maxim/plugin_test.go +++ b/plugins/maxim/plugin_test.go @@ -109,8 +109,10 @@ func TestMaximLoggerPlugin(t *testing.T) { Input: schemas.RequestInput{ ChatCompletionInput: &[]schemas.BifrostMessage{ { - Role: "user", - Content: bifrost.Ptr("Hello, how are you?"), + Role: "user", + Content: schemas.MessageContent{ + ContentStr: bifrost.Ptr("Hello, how are you?"), + }, }, }, }, @@ -122,6 +124,5 @@ func TestMaximLoggerPlugin(t *testing.T) { log.Println("Bifrost request completed, check your Maxim Dashboard for the trace") - // Clean up resources client.Cleanup() }