Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
2 changes: 1 addition & 1 deletion plugins/maxim/go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -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
Comment thread
Pratham-Mishra04 marked this conversation as resolved.
github.com/maximhq/maxim-go v0.1.3
)

Expand Down
4 changes: 2 additions & 2 deletions plugins/maxim/go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down
56 changes: 29 additions & 27 deletions plugins/maxim/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 != "" {
Expand All @@ -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 {
Expand All @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
}
Comment thread
Pratham-Mishra04 marked this conversation as resolved.
7 changes: 4 additions & 3 deletions plugins/maxim/plugin_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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?"),
},
Comment thread
Pratham-Mishra04 marked this conversation as resolved.
},
},
},
Expand All @@ -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()
}