chore: maxim plugin refactor to support core/v1.1.0

Author: Pratham-Mishra04

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
 )
 

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=

plugins/maxim/main.go

@@ -103,15 +103,15 @@ 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
 
 	// Check if context already has traceID and generationID
 	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
+}

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()
 }