diff --git a/transports/bifrost-http/integrations/anthropic/types.go b/transports/bifrost-http/integrations/anthropic/types.go index 693232671e..c370c5d8ae 100644 --- a/transports/bifrost-http/integrations/anthropic/types.go +++ b/transports/bifrost-http/integrations/anthropic/types.go @@ -10,74 +10,47 @@ import ( var fnTypePtr = bifrost.Ptr(string(schemas.ToolChoiceTypeFunction)) -// AnthropicContent represents content in Anthropic message format -type AnthropicContent struct { +// AnthropicContentBlock represents content in Anthropic message format +type AnthropicContentBlock struct { Type string `json:"type"` // "text", "image", "tool_use", "tool_result" Text *string `json:"text,omitempty"` // For text content ToolUseID *string `json:"tool_use_id,omitempty"` // For tool_result content ID *string `json:"id,omitempty"` // For tool_use content Name *string `json:"name,omitempty"` // For tool_use content Input interface{} `json:"input,omitempty"` // For tool_use content - Content interface{} `json:"content,omitempty"` // For tool_result content + Content AnthropicContent `json:"content,omitempty"` // For tool_result content Source *AnthropicImageSource `json:"source,omitempty"` // For image content } // AnthropicImageSource represents image source in Anthropic format type AnthropicImageSource struct { - Type string `json:"type"` // "base64" or "url" - MediaType string `json:"media_type"` // "image/jpeg", "image/png", etc. - Data string `json:"data"` // Base64-encoded image data + Type string `json:"type"` // "base64" or "url" + MediaType *string `json:"media_type,omitempty"` // "image/jpeg", "image/png", etc. + Data *string `json:"data,omitempty"` // Base64-encoded image data + URL *string `json:"url,omitempty"` // URL of the image } // AnthropicMessage represents a message in Anthropic format type AnthropicMessage struct { - Role string `json:"role"` // "user", "assistant" - Content []AnthropicContent `json:"content"` // Array of content blocks + Role string `json:"role"` // "user", "assistant" + Content AnthropicContent `json:"content"` // Array of content blocks } -// UnmarshalJSON implements custom JSON unmarshaling for AnthropicMessage -// to handle both string and array content formats -func (m *AnthropicMessage) UnmarshalJSON(data []byte) error { - // First, try to unmarshal into a struct with Content as json.RawMessage - var temp struct { - Role string `json:"role"` - Content json.RawMessage `json:"content"` - } - - if err := json.Unmarshal(data, &temp); err != nil { - return err - } - - m.Role = temp.Role - - // Try to unmarshal content as string first - var contentStr string - if err := json.Unmarshal(temp.Content, &contentStr); err == nil { - // It's a string, convert to AnthropicContent array - m.Content = []AnthropicContent{ - { - Type: "text", - Text: &contentStr, - }, - } - return nil - } - - // Try to unmarshal as array of AnthropicContent - var contentArray []AnthropicContent - if err := json.Unmarshal(temp.Content, &contentArray); err == nil { - m.Content = contentArray - return nil - } - - return fmt.Errorf("content must be either a string or an array of content objects") +type AnthropicContent struct { + ContentStr *string + ContentBlocks *[]AnthropicContentBlock } // AnthropicTool represents a tool in Anthropic format type AnthropicTool struct { - Name string `json:"name"` - Description string `json:"description"` - InputSchema interface{} `json:"input_schema"` + Name string `json:"name"` + Type *string `json:"type,omitempty"` + Description string `json:"description"` + InputSchema *struct { + Type string `json:"type"` // "object" + Properties map[string]interface{} `json:"properties"` + Required []string `json:"required"` + } `json:"input_schema,omitempty"` } // AnthropicToolChoice represents tool choice in Anthropic format @@ -91,7 +64,7 @@ type AnthropicMessageRequest struct { Model string `json:"model"` MaxTokens int `json:"max_tokens"` Messages []AnthropicMessage `json:"messages"` - System *string `json:"system,omitempty"` + System *AnthropicContent `json:"system,omitempty"` Temperature *float64 `json:"temperature,omitempty"` TopP *float64 `json:"top_p,omitempty"` TopK *int `json:"top_k,omitempty"` @@ -101,23 +74,96 @@ type AnthropicMessageRequest struct { ToolChoice *AnthropicToolChoice `json:"tool_choice,omitempty"` } +// AnthropicMessageResponse represents an Anthropic messages API response +type AnthropicMessageResponse struct { + ID string `json:"id"` + Type string `json:"type"` + Role string `json:"role"` + Content []AnthropicContentBlock `json:"content"` + Model string `json:"model"` + StopReason *string `json:"stop_reason,omitempty"` + StopSequence *string `json:"stop_sequence,omitempty"` + Usage *AnthropicUsage `json:"usage,omitempty"` +} + +// AnthropicUsage represents usage information in Anthropic format +type AnthropicUsage struct { + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` +} + +// MarshalJSON implements custom JSON marshalling for MessageContent. +// It marshals either ContentStr or ContentBlocks directly without wrapping. +func (mc AnthropicContent) MarshalJSON() ([]byte, error) { + // Validation: ensure only one field is set at a time + if mc.ContentStr != nil && mc.ContentBlocks != nil { + return nil, fmt.Errorf("both ContentStr and ContentBlocks are set; only one should be non-nil") + } + + if mc.ContentStr != nil { + return json.Marshal(*mc.ContentStr) + } + if mc.ContentBlocks != nil { + return json.Marshal(*mc.ContentBlocks) + } + // If both are nil, return null + return json.Marshal(nil) +} + +// UnmarshalJSON implements custom JSON unmarshalling for MessageContent. +// It determines whether "content" is a string or array and assigns to the appropriate field. +// It also handles direct string/array content without a wrapper object. +func (mc *AnthropicContent) UnmarshalJSON(data []byte) error { + // First, try to unmarshal as a direct string + var stringContent string + if err := json.Unmarshal(data, &stringContent); err == nil { + mc.ContentStr = &stringContent + return nil + } + + // Try to unmarshal as a direct array of ContentBlock + var arrayContent []AnthropicContentBlock + if err := json.Unmarshal(data, &arrayContent); err == nil { + mc.ContentBlocks = &arrayContent + return nil + } + + return fmt.Errorf("content field is neither a string nor an array of ContentBlock") +} + // ConvertToBifrostRequest converts an Anthropic messages request to Bifrost format func (r *AnthropicMessageRequest) ConvertToBifrostRequest() *schemas.BifrostRequest { bifrostReq := &schemas.BifrostRequest{ Provider: schemas.Anthropic, Model: r.Model, - Input: schemas.RequestInput{ - ChatCompletionInput: &[]schemas.BifrostMessage{}, - }, } + messages := []schemas.BifrostMessage{} + // Add system message if present - if r.System != nil && *r.System != "" { - systemMsg := schemas.BifrostMessage{ - Role: schemas.ModelChatMessageRoleSystem, - Content: r.System, + if r.System != nil { + if r.System.ContentStr != nil && *r.System.ContentStr != "" { + messages = append(messages, schemas.BifrostMessage{ + Role: schemas.ModelChatMessageRoleSystem, + Content: schemas.MessageContent{ + ContentStr: r.System.ContentStr, + }, + }) + } else if r.System.ContentBlocks != nil { + contentBlocks := []schemas.ContentBlock{} + for _, block := range *r.System.ContentBlocks { + contentBlocks = append(contentBlocks, schemas.ContentBlock{ + Type: schemas.ContentBlockTypeText, + Text: block.Text, + }) + } + messages = append(messages, schemas.BifrostMessage{ + Role: schemas.ModelChatMessageRoleSystem, + Content: schemas.MessageContent{ + ContentBlocks: &contentBlocks, + }, + }) } - *bifrostReq.Input.ChatCompletionInput = append(*bifrostReq.Input.ChatCompletionInput, systemMsg) } // Convert messages @@ -125,76 +171,118 @@ func (r *AnthropicMessageRequest) ConvertToBifrostRequest() *schemas.BifrostRequ var bifrostMsg schemas.BifrostMessage bifrostMsg.Role = schemas.ModelChatMessageRole(msg.Role) - // Handle different content types - var toolCalls []schemas.ToolCall - var textContents []string - - for _, content := range msg.Content { - switch content.Type { - case "text": - if content.Text != nil { - textContents = append(textContents, *content.Text) - } - case "image": - if content.Source != nil { - bifrostMsg.UserMessage = &schemas.UserMessage{ - ImageContent: convertAnthropicImageSource(content.Source), - } - } - case "tool_use": - if content.ID != nil && content.Name != nil { - tc := schemas.ToolCall{ - Type: fnTypePtr, - ID: content.ID, - Function: schemas.FunctionCall{ - Name: content.Name, - Arguments: jsonifyInput(content.Input), - }, + if msg.Content.ContentStr != nil { + bifrostMsg.Content = schemas.MessageContent{ + ContentStr: msg.Content.ContentStr, + } + } else if msg.Content.ContentBlocks != nil { + // Handle different content types + var toolCalls []schemas.ToolCall + var contentBlocks []schemas.ContentBlock + + for _, content := range *msg.Content.ContentBlocks { + switch content.Type { + case "text": + if content.Text != nil { + contentBlocks = append(contentBlocks, schemas.ContentBlock{ + Type: schemas.ContentBlockTypeText, + Text: content.Text, + }) } - toolCalls = append(toolCalls, tc) - } - case "tool_result": - if content.ToolUseID != nil { - bifrostMsg.ToolMessage = &schemas.ToolMessage{ - ToolCallID: content.ToolUseID, + case "image": + if content.Source != nil { + contentBlocks = append(contentBlocks, schemas.ContentBlock{ + Type: schemas.ContentBlockTypeImage, + ImageURL: &schemas.ImageURLStruct{ + URL: func() string { + if content.Source.Data != nil { + mime := "image/png" + if content.Source.MediaType != nil && *content.Source.MediaType != "" { + mime = *content.Source.MediaType + } + return "data:" + mime + ";base64," + *content.Source.Data + } + if content.Source.URL != nil { + return *content.Source.URL + } + return "" + }(), + }, + }) } - if content.Content != nil { - if str, ok := content.Content.(string); ok { - bifrostMsg.Content = &str - } else { - jsonStr := jsonifyInput(content.Content) - bifrostMsg.Content = &jsonStr + case "tool_use": + if content.ID != nil && content.Name != nil { + tc := schemas.ToolCall{ + Type: fnTypePtr, + ID: content.ID, + Function: schemas.FunctionCall{ + Name: content.Name, + Arguments: jsonifyInput(content.Input), + }, } + toolCalls = append(toolCalls, tc) } - if content.Source != nil { - bifrostMsg.ToolMessage.ImageContent = convertAnthropicImageSource(content.Source) + case "tool_result": + if content.ToolUseID != nil { + bifrostMsg.ToolMessage = &schemas.ToolMessage{ + ToolCallID: content.ToolUseID, + } + if content.Content.ContentStr != nil { + contentBlocks = append(contentBlocks, schemas.ContentBlock{ + Type: schemas.ContentBlockTypeText, + Text: content.Content.ContentStr, + }) + } else if content.Content.ContentBlocks != nil { + for _, block := range *content.Content.ContentBlocks { + if block.Text != nil { + contentBlocks = append(contentBlocks, schemas.ContentBlock{ + Type: schemas.ContentBlockTypeText, + Text: block.Text, + }) + } else if block.Source != nil { + contentBlocks = append(contentBlocks, schemas.ContentBlock{ + Type: schemas.ContentBlockTypeImage, + ImageURL: &schemas.ImageURLStruct{ + URL: func() string { + if block.Source.Data != nil { + mime := "image/png" + if block.Source.MediaType != nil && *block.Source.MediaType != "" { + mime = *block.Source.MediaType + } + return "data:" + mime + ";base64," + *block.Source.Data + } + if block.Source.URL != nil { + return *block.Source.URL + } + return "" + }()}, + }) + } + } + } + bifrostMsg.Role = schemas.ModelChatMessageRoleTool } - bifrostMsg.Role = schemas.ModelChatMessageRoleTool } } - } - // Concatenate all text contents - if len(textContents) > 0 { - concatenatedText := "" - for i, text := range textContents { - if i > 0 { - concatenatedText += "\n" // Add newline separator between text blocks + // Concatenate all text contents + if len(contentBlocks) > 0 { + bifrostMsg.Content = schemas.MessageContent{ + ContentBlocks: &contentBlocks, } - concatenatedText += text } - bifrostMsg.Content = &concatenatedText - } - if len(toolCalls) > 0 { - bifrostMsg.AssistantMessage = &schemas.AssistantMessage{ - ToolCalls: &toolCalls, + if len(toolCalls) > 0 && msg.Role == string(schemas.ModelChatMessageRoleAssistant) { + bifrostMsg.AssistantMessage = &schemas.AssistantMessage{ + ToolCalls: &toolCalls, + } } } - - *bifrostReq.Input.ChatCompletionInput = append(*bifrostReq.Input.ChatCompletionInput, bifrostMsg) + messages = append(messages, bifrostMsg) } + bifrostReq.Input.ChatCompletionInput = &messages + // Convert parameters if r.MaxTokens > 0 || r.Temperature != nil || r.TopP != nil || r.TopK != nil || r.StopSequences != nil { params := &schemas.ModelParameters{} @@ -227,46 +315,19 @@ func (r *AnthropicMessageRequest) ConvertToBifrostRequest() *schemas.BifrostRequ Type: "object", } if tool.InputSchema != nil { - if schemaMap, ok := tool.InputSchema.(map[string]interface{}); ok { - if typeVal, ok := schemaMap["type"].(string); ok { - params.Type = typeVal - } - if desc, ok := schemaMap["description"].(string); ok { - params.Description = &desc - } - if required, ok := schemaMap["required"].([]interface{}); ok { - reqStrings := make([]string, len(required)) - for i, req := range required { - if reqStr, ok := req.(string); ok { - reqStrings[i] = reqStr - } - } - params.Required = reqStrings - } - if properties, ok := schemaMap["properties"].(map[string]interface{}); ok { - params.Properties = properties - } - if enum, ok := schemaMap["enum"].([]interface{}); ok { - enumStrings := make([]string, len(enum)) - for i, e := range enum { - if eStr, ok := e.(string); ok { - enumStrings[i] = eStr - } - } - params.Enum = &enumStrings - } - } + params.Type = tool.InputSchema.Type + params.Required = tool.InputSchema.Required + params.Properties = tool.InputSchema.Properties } - t := schemas.Tool{ + tools = append(tools, schemas.Tool{ Type: "function", Function: schemas.Function{ Name: tool.Name, Description: tool.Description, Parameters: params, }, - } - tools = append(tools, t) + }) } if bifrostReq.Params == nil { bifrostReq.Params = &schemas.ModelParameters{} @@ -310,24 +371,6 @@ func jsonifyInput(input interface{}) string { return string(jsonBytes) } -// AnthropicMessageResponse represents an Anthropic messages API response -type AnthropicMessageResponse struct { - ID string `json:"id"` - Type string `json:"type"` - Role string `json:"role"` - Content []AnthropicContent `json:"content"` - Model string `json:"model"` - StopReason *string `json:"stop_reason,omitempty"` - StopSequence *string `json:"stop_sequence,omitempty"` - Usage *AnthropicUsage `json:"usage,omitempty"` -} - -// AnthropicUsage represents usage information in Anthropic format -type AnthropicUsage struct { - InputTokens int `json:"input_tokens"` - OutputTokens int `json:"output_tokens"` -} - // DeriveAnthropicFromBifrostResponse converts a Bifrost response to Anthropic format func DeriveAnthropicFromBifrostResponse(bifrostResp *schemas.BifrostResponse) *AnthropicMessageResponse { if bifrostResp == nil { @@ -337,7 +380,7 @@ func DeriveAnthropicFromBifrostResponse(bifrostResp *schemas.BifrostResponse) *A anthropicResp := &AnthropicMessageResponse{ ID: bifrostResp.ID, Type: "message", - Role: "assistant", + Role: string(schemas.ModelChatMessageRoleAssistant), Model: bifrostResp.Model, } @@ -350,7 +393,7 @@ func DeriveAnthropicFromBifrostResponse(bifrostResp *schemas.BifrostResponse) *A } // Convert choices to content - var content []AnthropicContent + var content []AnthropicContentBlock if len(bifrostResp.Choices) > 0 { choice := bifrostResp.Choices[0] // Anthropic typically returns one choice @@ -363,18 +406,27 @@ func DeriveAnthropicFromBifrostResponse(bifrostResp *schemas.BifrostResponse) *A // Add thinking content if present if choice.Message.AssistantMessage != nil && choice.Message.AssistantMessage.Thought != nil && *choice.Message.AssistantMessage.Thought != "" { - content = append(content, AnthropicContent{ + content = append(content, AnthropicContentBlock{ Type: "thinking", Text: choice.Message.AssistantMessage.Thought, }) } // Add text content - if choice.Message.Content != nil && *choice.Message.Content != "" { - content = append(content, AnthropicContent{ + if choice.Message.Content.ContentStr != nil && *choice.Message.Content.ContentStr != "" { + content = append(content, AnthropicContentBlock{ Type: "text", - Text: choice.Message.Content, + Text: choice.Message.Content.ContentStr, }) + } else if choice.Message.Content.ContentBlocks != nil { + for _, block := range *choice.Message.Content.ContentBlocks { + if block.Text != nil { + content = append(content, AnthropicContentBlock{ + Type: "text", + Text: block.Text, + }) + } + } } // Add tool calls as tool_use content @@ -390,46 +442,20 @@ func DeriveAnthropicFromBifrostResponse(bifrostResp *schemas.BifrostResponse) *A input = map[string]interface{}{} } - tc := AnthropicContent{ + content = append(content, AnthropicContentBlock{ Type: "tool_use", ID: toolCall.ID, Name: toolCall.Function.Name, Input: input, - } - content = append(content, tc) + }) } } } if content == nil { - content = []AnthropicContent{} + content = []AnthropicContentBlock{} } anthropicResp.Content = content return anthropicResp } - -// convertAnthropicImageSource converts an Anthropic image source to Bifrost ImageContent format -func convertAnthropicImageSource(source *AnthropicImageSource) *schemas.ImageContent { - if source == nil { - return nil - } - - // Convert Anthropic source type to Bifrost ImageContentType - var contentType schemas.ImageContentType - switch source.Type { - case "base64": - contentType = schemas.ImageContentTypeBase64 - case "url": - contentType = schemas.ImageContentTypeURL - default: - // Default to base64 if unknown type, as this is more common in Anthropic - contentType = schemas.ImageContentTypeBase64 - } - - return &schemas.ImageContent{ - Type: contentType, - URL: source.Data, - MediaType: &source.MediaType, - } -} diff --git a/transports/bifrost-http/integrations/genai/types.go b/transports/bifrost-http/integrations/genai/types.go index de7a250103..2f5d1e6e88 100644 --- a/transports/bifrost-http/integrations/genai/types.go +++ b/transports/bifrost-http/integrations/genai/types.go @@ -47,18 +47,155 @@ func (r *GeminiChatRequest) ConvertToBifrostRequest() *schemas.BifrostRequest { }, } - // Convert system instruction if present + messages := []schemas.BifrostMessage{} + + allGenAiMessages := []genai_sdk.Content{} if r.SystemInstruction != nil { - systemMsgs := r.convertContentToBifrostMessages(*r.SystemInstruction, schemas.ModelChatMessageRoleSystem) - *bifrostReq.Input.ChatCompletionInput = append(*bifrostReq.Input.ChatCompletionInput, systemMsgs...) + allGenAiMessages = append(allGenAiMessages, *r.SystemInstruction) } + allGenAiMessages = append(allGenAiMessages, r.Contents...) + + for _, content := range allGenAiMessages { + if len(content.Parts) == 0 { + continue + } + + // Handle multiple parts - collect all content and tool calls + var toolCalls []schemas.ToolCall + var contentBlocks []schemas.ContentBlock + var thoughtStr string // Track thought content for assistant/model + + for _, part := range content.Parts { + switch { + case part.Text != "": + // Handle thought content specially for assistant messages + if part.Thought && + (content.Role == string(schemas.ModelChatMessageRoleAssistant) || content.Role == string(genai_sdk.RoleModel)) { + thoughtStr = thoughtStr + part.Text + "\n" + } else { + contentBlocks = append(contentBlocks, schemas.ContentBlock{ + Type: schemas.ContentBlockTypeText, + Text: &part.Text, + }) + } + + case part.FunctionCall != nil: + // Only add function calls for assistant messages + if content.Role == string(schemas.ModelChatMessageRoleAssistant) || content.Role == string(genai_sdk.RoleModel) { + jsonArgs, err := json.Marshal(part.FunctionCall.Args) + if err != nil { + jsonArgs = []byte(fmt.Sprintf("%v", part.FunctionCall.Args)) + } + id := part.FunctionCall.ID // create local copy + name := part.FunctionCall.Name // create local copy + toolCall := schemas.ToolCall{ + ID: bifrost.Ptr(id), + Type: fnTypePtr, + Function: schemas.FunctionCall{ + Name: &name, + Arguments: string(jsonArgs), + }, + } + toolCalls = append(toolCalls, toolCall) + } + + case part.FunctionResponse != nil: + // Create a separate tool response message + responseContent, err := json.Marshal(part.FunctionResponse.Response) + if err != nil { + responseContent = []byte(fmt.Sprintf("%v", part.FunctionResponse.Response)) + } + + toolResponseMsg := schemas.BifrostMessage{ + Role: schemas.ModelChatMessageRoleTool, + Content: schemas.MessageContent{ + ContentStr: bifrost.Ptr(string(responseContent)), + }, + ToolMessage: &schemas.ToolMessage{ + ToolCallID: &part.FunctionResponse.Name, + }, + } + + messages = append(messages, toolResponseMsg) + + case part.InlineData != nil: + // Handle inline images/media - only append if it's actually an image + if isImageMimeType(part.InlineData.MIMEType) { + contentBlocks = append(contentBlocks, schemas.ContentBlock{ + Type: schemas.ContentBlockTypeImage, + ImageURL: &schemas.ImageURLStruct{ + URL: fmt.Sprintf("data:%s;base64,%s", part.InlineData.MIMEType, base64.StdEncoding.EncodeToString(part.InlineData.Data)), + }, + }) + } + + case part.FileData != nil: + // Handle file data - only append if it's actually an image + if isImageMimeType(part.FileData.MIMEType) { + contentBlocks = append(contentBlocks, schemas.ContentBlock{ + Type: schemas.ContentBlockTypeImage, + ImageURL: &schemas.ImageURLStruct{ + URL: part.FileData.FileURI, + }, + }) + } + + case part.ExecutableCode != nil: + // Handle executable code as text content + codeText := fmt.Sprintf("```%s\n%s\n```", part.ExecutableCode.Language, part.ExecutableCode.Code) + contentBlocks = append(contentBlocks, schemas.ContentBlock{ + Type: schemas.ContentBlockTypeText, + Text: &codeText, + }) + + case part.CodeExecutionResult != nil: + // Handle code execution results as text content + resultText := fmt.Sprintf("Code execution result (%s):\n%s", part.CodeExecutionResult.Outcome, part.CodeExecutionResult.Output) + contentBlocks = append(contentBlocks, schemas.ContentBlock{ + Type: schemas.ContentBlockTypeText, + Text: &resultText, + }) + } + } + + // Only create message if there's actual content, tool calls, or thought content + if len(contentBlocks) > 0 || len(toolCalls) > 0 || thoughtStr != "" { + // Create main message with content blocks + bifrostMsg := schemas.BifrostMessage{ + Role: func(r string) schemas.ModelChatMessageRole { + if r == string(genai_sdk.RoleModel) { // GenAI's internal alias + return schemas.ModelChatMessageRoleAssistant + } + return schemas.ModelChatMessageRole(r) + }(content.Role), + } + + // Set content only if there are content blocks + if len(contentBlocks) > 0 { + bifrostMsg.Content = schemas.MessageContent{ + ContentBlocks: &contentBlocks, + } + } + + // Set assistant-specific fields for assistant/model messages + if content.Role == string(schemas.ModelChatMessageRoleAssistant) || content.Role == string(genai_sdk.RoleModel) { + if len(toolCalls) > 0 || thoughtStr != "" { + bifrostMsg.AssistantMessage = &schemas.AssistantMessage{} + if len(toolCalls) > 0 { + bifrostMsg.AssistantMessage.ToolCalls = &toolCalls + } + if thoughtStr != "" { + bifrostMsg.AssistantMessage.Thought = &thoughtStr + } + } + } - // Convert messages (contents) - for _, content := range r.Contents { - messages := r.convertContentToBifrostMessages(content, schemas.ModelChatMessageRole(content.Role)) - *bifrostReq.Input.ChatCompletionInput = append(*bifrostReq.Input.ChatCompletionInput, messages...) + messages = append(messages, bifrostMsg) + } } + bifrostReq.Input.ChatCompletionInput = &messages + // Convert generation config to parameters if params := r.convertGenerationConfigToParams(); params != nil { bifrostReq.Params = params @@ -136,125 +273,6 @@ func (r *GeminiChatRequest) ConvertToBifrostRequest() *schemas.BifrostRequest { return bifrostReq } -// convertContentToBifrostMessage converts a Gemini Content to BifrostMessage(s) -// Returns multiple messages when there are multiple images to ensure each image gets its own message -func (r *GeminiChatRequest) convertContentToBifrostMessages(content genai_sdk.Content, role schemas.ModelChatMessageRole) []schemas.BifrostMessage { - if len(content.Parts) == 0 { - return nil - } - - // Handle multiple parts - concatenate text parts and handle other types - var textParts []string - var toolCalls []schemas.ToolCall - var imageContents []schemas.ImageContent - - for _, part := range content.Parts { - switch { - case part.Text != "": - textParts = append(textParts, part.Text) - - case part.FunctionCall != nil: - jsonArgs, err := json.Marshal(part.FunctionCall.Args) - if err != nil { - jsonArgs = []byte(fmt.Sprintf("%v", part.FunctionCall.Args)) - } - toolCall := schemas.ToolCall{ - ID: bifrost.Ptr(part.FunctionCall.ID), - Type: fnTypePtr, - Function: schemas.FunctionCall{ - Name: &part.FunctionCall.Name, - Arguments: string(jsonArgs), - }, - } - - toolCalls = append(toolCalls, toolCall) - - case part.InlineData != nil: - // Handle inline images/media - imageContent := schemas.ImageContent{ - Type: schemas.ImageContentTypeBase64, - URL: fmt.Sprintf("data:%s;base64,%s", part.InlineData.MIMEType, base64.StdEncoding.EncodeToString(part.InlineData.Data)), - MediaType: &part.InlineData.MIMEType, - } - imageContents = append(imageContents, imageContent) - - case part.FileData != nil: - // Handle file references - imageContent := schemas.ImageContent{ - Type: schemas.ImageContentTypeURL, - URL: part.FileData.FileURI, - MediaType: &part.FileData.MIMEType, - } - imageContents = append(imageContents, imageContent) - - case part.FunctionResponse != nil: - responseContent, err := json.Marshal(part.FunctionResponse.Response) - if err != nil { - responseContent = []byte(fmt.Sprintf("%v", part.FunctionResponse.Response)) - } - - toolResponseMsg := schemas.BifrostMessage{ - Role: schemas.ModelChatMessageRoleTool, - Content: bifrost.Ptr(string(responseContent)), - ToolMessage: &schemas.ToolMessage{ - ToolCallID: &part.FunctionResponse.Name, - }, - } - - return []schemas.BifrostMessage{toolResponseMsg} - } - } - - var messages []schemas.BifrostMessage - - // Create main message with text content and tool calls - mainMsg := schemas.BifrostMessage{ - Role: role, - } - - // Set text content if we have any - if len(textParts) > 0 { - combinedText := strings.Join(textParts, "\n\n") - mainMsg.Content = &combinedText - } - - // Set tool calls if we have any - if len(toolCalls) > 0 && role == schemas.ModelChatMessageRoleAssistant { - mainMsg.AssistantMessage = &schemas.AssistantMessage{ - ToolCalls: &toolCalls, - } - } - - // Add main message if it has content or tool calls - if mainMsg.Content != nil || (mainMsg.AssistantMessage != nil && mainMsg.AssistantMessage.ToolCalls != nil) { - messages = append(messages, mainMsg) - } - - // Create separate messages for each image - for _, imageContent := range imageContents { - imageMsg := schemas.BifrostMessage{ - Role: role, - } - - // Set image content based on role - switch role { - case schemas.ModelChatMessageRoleUser: - imageMsg.UserMessage = &schemas.UserMessage{ - ImageContent: &imageContent, - } - messages = append(messages, imageMsg) - - case schemas.ModelChatMessageRoleTool: - imageMsg.ToolMessage = &schemas.ToolMessage{ - ImageContent: &imageContent, - } - messages = append(messages, imageMsg) - } - } - - return messages -} - // convertGenerationConfigToParams converts Gemini GenerationConfig to ModelParameters func (r *GeminiChatRequest) convertGenerationConfigToParams() *schemas.ModelParameters { params := &schemas.ModelParameters{ @@ -364,8 +382,14 @@ func DeriveGenAIFromBifrostResponse(bifrostResp *schemas.BifrostResponse) *genai } parts := []*genai_sdk.Part{} - if choice.Message.Content != nil && *choice.Message.Content != "" { - parts = append(parts, &genai_sdk.Part{Text: *choice.Message.Content}) + if choice.Message.Content.ContentStr != nil && *choice.Message.Content.ContentStr != "" { + parts = append(parts, &genai_sdk.Part{Text: *choice.Message.Content.ContentStr}) + } else if choice.Message.Content.ContentBlocks != nil { + for _, block := range *choice.Message.Content.ContentBlocks { + if block.Text != nil { + parts = append(parts, &genai_sdk.Part{Text: *block.Text}) + } + } } // Handle tool calls @@ -440,3 +464,46 @@ func DeriveGenAIFromBifrostResponse(bifrostResp *schemas.BifrostResponse) *genai return genaiResp } + +// isImageMimeType checks if a MIME type represents an image format +func isImageMimeType(mimeType string) bool { + if mimeType == "" { + return false + } + + // Convert to lowercase for case-insensitive comparison + mimeType = strings.ToLower(mimeType) + + // Remove any parameters (e.g., "image/jpeg; charset=utf-8" -> "image/jpeg") + if idx := strings.Index(mimeType, ";"); idx != -1 { + mimeType = strings.TrimSpace(mimeType[:idx]) + } + + // If it starts with "image/", it's an image + if strings.HasPrefix(mimeType, "image/") { + return true + } + + // Check for common image formats that might not have the "image/" prefix + commonImageTypes := []string{ + "jpeg", + "jpg", + "png", + "gif", + "webp", + "bmp", + "svg", + "tiff", + "ico", + "avif", + } + + // Check if the mimeType contains any of the common image type strings + for _, imageType := range commonImageTypes { + if strings.Contains(mimeType, imageType) { + return true + } + } + + return false +} diff --git a/transports/bifrost-http/integrations/openai/types.go b/transports/bifrost-http/integrations/openai/types.go index 3ac68e9819..cf37b01cb3 100644 --- a/transports/bifrost-http/integrations/openai/types.go +++ b/transports/bifrost-http/integrations/openai/types.go @@ -1,132 +1,41 @@ package openai import ( - "strings" - - bifrost "github.com/maximhq/bifrost/core" "github.com/maximhq/bifrost/core/schemas" ) -var fnTypePtr = bifrost.Ptr(string(schemas.ToolChoiceTypeFunction)) - -// OpenAIContentPart represents a part of the content (text or image) in OpenAI format -type OpenAIContentPart struct { - Type string `json:"type"` - Text *string `json:"text,omitempty"` - ImageURL *OpenAIImageURL `json:"image_url,omitempty"` -} - -// OpenAIImageURL represents an image URL with optional detail level in OpenAI format -type OpenAIImageURL struct { - URL string `json:"url"` - Detail *string `json:"detail,omitempty"` -} - -// OpenAIMessage represents a message in the OpenAI chat format -type OpenAIMessage struct { - Role string `json:"role"` - Content interface{} `json:"content,omitempty"` // Can be string or []OpenAIContentPart - Name *string `json:"name,omitempty"` - ToolCalls *[]schemas.ToolCall `json:"tool_calls,omitempty"` // Reuse schema type - ToolCallID *string `json:"tool_call_id,omitempty"` - FunctionCall *schemas.FunctionCall `json:"function_call,omitempty"` // Reuse schema type -} - // OpenAIChatRequest represents an OpenAI chat completion request type OpenAIChatRequest struct { - Model string `json:"model"` - Messages []OpenAIMessage `json:"messages"` - MaxTokens *int `json:"max_tokens,omitempty"` - Temperature *float64 `json:"temperature,omitempty"` - TopP *float64 `json:"top_p,omitempty"` - N *int `json:"n,omitempty"` - Stop interface{} `json:"stop,omitempty"` - PresencePenalty *float64 `json:"presence_penalty,omitempty"` - FrequencyPenalty *float64 `json:"frequency_penalty,omitempty"` - LogitBias map[string]float64 `json:"logit_bias,omitempty"` - User *string `json:"user,omitempty"` - Functions *[]schemas.Function `json:"functions,omitempty"` // Reuse schema type - FunctionCall interface{} `json:"function_call,omitempty"` - Tools *[]schemas.Tool `json:"tools,omitempty"` // Reuse schema type - ToolChoice interface{} `json:"tool_choice,omitempty"` - Stream *bool `json:"stream,omitempty"` - LogProbs *bool `json:"logprobs,omitempty"` - TopLogProbs *int `json:"top_logprobs,omitempty"` - ResponseFormat interface{} `json:"response_format,omitempty"` - Seed *int `json:"seed,omitempty"` + Model string `json:"model"` + Messages []schemas.BifrostMessage `json:"messages"` + MaxTokens *int `json:"max_tokens,omitempty"` + Temperature *float64 `json:"temperature,omitempty"` + TopP *float64 `json:"top_p,omitempty"` + N *int `json:"n,omitempty"` + Stop interface{} `json:"stop,omitempty"` + PresencePenalty *float64 `json:"presence_penalty,omitempty"` + FrequencyPenalty *float64 `json:"frequency_penalty,omitempty"` + LogitBias map[string]float64 `json:"logit_bias,omitempty"` + User *string `json:"user,omitempty"` + Tools *[]schemas.Tool `json:"tools,omitempty"` // Reuse schema type + ToolChoice *schemas.ToolChoice `json:"tool_choice,omitempty"` + Stream *bool `json:"stream,omitempty"` + LogProbs *bool `json:"logprobs,omitempty"` + TopLogProbs *int `json:"top_logprobs,omitempty"` + ResponseFormat interface{} `json:"response_format,omitempty"` + Seed *int `json:"seed,omitempty"` } // OpenAIChatResponse represents an OpenAI chat completion response type OpenAIChatResponse struct { - ID string `json:"id"` - Object string `json:"object"` - Created int `json:"created"` - Model string `json:"model"` - Choices []OpenAIChoice `json:"choices"` - Usage *schemas.LLMUsage `json:"usage,omitempty"` // Reuse schema type - SystemFingerprint *string `json:"system_fingerprint,omitempty"` -} - -// OpenAIChoice represents a choice in the OpenAI response -type OpenAIChoice struct { - Index int `json:"index"` - Message OpenAIMessage `json:"message"` - FinishReason *string `json:"finish_reason,omitempty"` - LogProbs interface{} `json:"logprobs,omitempty"` -} - -// convertOpenAIContent handles both string and structured content formats -// It returns the text content and any image content found -func convertOpenAIContent(content interface{}) (*string, *schemas.ImageContent) { - if content == nil { - return nil, nil - } - - switch c := content.(type) { - case string: - return &c, nil - case []interface{}: - // Handle array of content parts (for vision API) - var textParts []string - var imageContent *schemas.ImageContent - - for _, part := range c { - if partMap, ok := part.(map[string]interface{}); ok { - if partType, exists := partMap["type"].(string); exists { - switch partType { - case "text": - if text, textExists := partMap["text"].(string); textExists { - textParts = append(textParts, text) - } - case "image_url": - if imageURL, ok := partMap["image_url"].(map[string]interface{}); ok { - if url, urlExists := imageURL["url"].(string); urlExists { - // Initialize imageContent if we have a URL - imageContent = &schemas.ImageContent{ - URL: url, - } - - // Get detail level if specified - if detail, detailExists := imageURL["detail"].(string); detailExists { - imageContent.Detail = &detail - } - } - } - } - } - } - } - - var textContent *string - if len(textParts) > 0 { - combined := strings.Join(textParts, " ") - textContent = &combined - } - - return textContent, imageContent - } - - return nil, nil + ID string `json:"id"` + Object string `json:"object"` + Created int `json:"created"` + Model string `json:"model"` + Choices []schemas.BifrostResponseChoice `json:"choices"` + Usage *schemas.LLMUsage `json:"usage,omitempty"` // Reuse schema type + ServiceTier *string `json:"service_tier,omitempty"` + SystemFingerprint *string `json:"system_fingerprint,omitempty"` } // ConvertToBifrostRequest converts an OpenAI chat request to Bifrost format @@ -135,19 +44,13 @@ func (r *OpenAIChatRequest) ConvertToBifrostRequest() *schemas.BifrostRequest { Provider: schemas.OpenAI, Model: r.Model, Input: schemas.RequestInput{ - ChatCompletionInput: &[]schemas.BifrostMessage{}, + ChatCompletionInput: &r.Messages, }, } - // Convert messages - r.convertMessages(bifrostReq) - - // Convert parameters using dynamic reflection-based mapping + // Map extra parameters and tool settings bifrostReq.Params = r.convertParameters() - // Convert tools and tool choices - r.convertToolsAndChoices(bifrostReq) - return bifrostReq } @@ -158,7 +61,10 @@ func (r *OpenAIChatRequest) convertParameters() *schemas.ModelParameters { ExtraParams: make(map[string]interface{}), } - // Direct field mapping with type safety - much faster than reflection + params.Tools = r.Tools + params.ToolChoice = r.ToolChoice + + // Direct field mapping if r.MaxTokens != nil { params.ExtraParams["max_tokens"] = *r.MaxTokens } @@ -210,244 +116,6 @@ func (r *OpenAIChatRequest) convertParameters() *schemas.ModelParameters { return params } -// convertMessages handles message conversion from OpenAI to Bifrost format -func (r *OpenAIChatRequest) convertMessages(bifrostReq *schemas.BifrostRequest) { - for _, msg := range r.Messages { - bifrostMsg := r.convertSingleMessage(msg) - *bifrostReq.Input.ChatCompletionInput = append(*bifrostReq.Input.ChatCompletionInput, bifrostMsg) - } -} - -// convertSingleMessage converts a single OpenAI message to Bifrost format -func (r *OpenAIChatRequest) convertSingleMessage(msg OpenAIMessage) schemas.BifrostMessage { - // Handle different content formats (string vs array of content parts) - textContent, imageContent := convertOpenAIContent(msg.Content) - - // Create BifrostMessage with proper embedded struct setup - bifrostMsg := schemas.BifrostMessage{ - Role: schemas.ModelChatMessageRole(msg.Role), - Content: textContent, - } - - // Set embedded struct based on message role and image content - if imageContent != nil { - r.setImageContent(&bifrostMsg, msg.Role, imageContent) - } - - // Handle tool calls and function calls for assistant messages - r.setToolCalls(&bifrostMsg, msg) - - // Handle tool messages - r.setToolMessage(&bifrostMsg, msg) - - return bifrostMsg -} - -// setImageContent sets image content based on message role -func (r *OpenAIChatRequest) setImageContent(bifrostMsg *schemas.BifrostMessage, role string, imageContent *schemas.ImageContent) { - switch role { - case "user": - bifrostMsg.Role = schemas.ModelChatMessageRoleUser - bifrostMsg.UserMessage = &schemas.UserMessage{ - ImageContent: imageContent, - } - case "tool": - bifrostMsg.Role = schemas.ModelChatMessageRoleTool - bifrostMsg.ToolMessage = &schemas.ToolMessage{ - ImageContent: imageContent, - } - } -} - -// setToolCalls handles tool calls and function calls for assistant messages -func (r *OpenAIChatRequest) setToolCalls(bifrostMsg *schemas.BifrostMessage, msg OpenAIMessage) { - var toolCalls []schemas.ToolCall - - // Prioritize modern tool_calls format, only use legacy function_call if tool_calls is not present - if msg.ToolCalls != nil { - // Tool calls are already in the right format (schemas.ToolCall) - toolCalls = *msg.ToolCalls - } else if msg.FunctionCall != nil { - // Add legacy function call only if no modern tool calls exist - tc := schemas.ToolCall{ - Type: fnTypePtr, - Function: *msg.FunctionCall, // Already schemas.FunctionCall type - } - toolCalls = append(toolCalls, tc) - } - - // Assign AssistantMessage only if we have tool calls - if len(toolCalls) > 0 { - bifrostMsg.AssistantMessage = &schemas.AssistantMessage{ - ToolCalls: &toolCalls, - } - } -} - -// setToolMessage handles tool message-specific fields -func (r *OpenAIChatRequest) setToolMessage(bifrostMsg *schemas.BifrostMessage, msg OpenAIMessage) { - if msg.ToolCallID != nil { - if bifrostMsg.ToolMessage == nil { - bifrostMsg.ToolMessage = &schemas.ToolMessage{} - } - bifrostMsg.ToolMessage.ToolCallID = msg.ToolCallID - } -} - -// convertToolsAndChoices handles tools and tool choice conversion -func (r *OpenAIChatRequest) convertToolsAndChoices(bifrostReq *schemas.BifrostRequest) { - // Convert tools - allTools := r.convertTools() - if len(allTools) > 0 { - r.ensureParams(bifrostReq) - bifrostReq.Params.Tools = &allTools - } - - // Convert tool choice - toolChoice := r.convertToolChoice() - if toolChoice != nil { - r.ensureParams(bifrostReq) - bifrostReq.Params.ToolChoice = toolChoice - } -} - -// convertTools combines modern Tools and legacy Functions into a unified tool list -func (r *OpenAIChatRequest) convertTools() []schemas.Tool { - var allTools []schemas.Tool - - // Handle modern Tools field (already schemas.Tool type) - if r.Tools != nil { - allTools = append(allTools, *r.Tools...) - } - - // Handle legacy Functions field - if r.Functions != nil { - for _, function := range *r.Functions { - t := schemas.Tool{ - Type: string(schemas.ToolChoiceTypeFunction), - Function: function, // Already schemas.Function type - } - allTools = append(allTools, t) - } - } - - return allTools -} - -// convertToolChoice handles both modern tool_choice and legacy function_call -func (r *OpenAIChatRequest) convertToolChoice() *schemas.ToolChoice { - if r.ToolChoice == nil && r.FunctionCall == nil { - return nil - } - - // Handle ToolChoice (modern format) first - if r.ToolChoice != nil { - return r.parseToolChoice(r.ToolChoice) - } - - // Handle legacy FunctionCall - if r.FunctionCall != nil { - return r.parseFunctionCall(r.FunctionCall) - } - - return nil -} - -// parseToolChoice parses modern tool_choice format -func (r *OpenAIChatRequest) parseToolChoice(toolChoice interface{}) *schemas.ToolChoice { - tc := &schemas.ToolChoice{} - - switch v := toolChoice.(type) { - case string: - tc.Type = r.parseToolChoiceString(v) - case map[string]interface{}: - r.parseToolChoiceObject(tc, v) - } - - return tc -} - -// parseToolChoiceString handles string tool choice values -func (r *OpenAIChatRequest) parseToolChoiceString(value string) schemas.ToolChoiceType { - switch value { - case "none": - return schemas.ToolChoiceTypeNone - case "auto": - return schemas.ToolChoiceTypeAuto - case "required": - return schemas.ToolChoiceTypeRequired - default: - return schemas.ToolChoiceTypeAuto // fallback - } -} - -// parseToolChoiceObject handles object tool choice values -func (r *OpenAIChatRequest) parseToolChoiceObject(tc *schemas.ToolChoice, obj map[string]interface{}) { - typeVal, ok := obj["type"].(string) - if !ok { - tc.Type = schemas.ToolChoiceTypeAuto - return - } - - switch typeVal { - case "function": - tc.Type = schemas.ToolChoiceTypeFunction - if functionVal, ok := obj["function"].(map[string]interface{}); ok { - if name, ok := functionVal["name"].(string); ok { - tc.Function = schemas.ToolChoiceFunction{Name: name} - } - } - case "none": - tc.Type = schemas.ToolChoiceTypeNone - case "auto": - tc.Type = schemas.ToolChoiceTypeAuto - case "required": - tc.Type = schemas.ToolChoiceTypeRequired - default: - tc.Type = schemas.ToolChoiceTypeAuto // fallback - } -} - -// parseFunctionCall handles legacy function_call format -func (r *OpenAIChatRequest) parseFunctionCall(functionCall interface{}) *schemas.ToolChoice { - tc := &schemas.ToolChoice{} - - switch v := functionCall.(type) { - case string: - switch v { - case "none": - tc.Type = schemas.ToolChoiceTypeNone - case "auto": - tc.Type = schemas.ToolChoiceTypeAuto - default: - tc.Type = schemas.ToolChoiceTypeAuto // fallback - } - case map[string]interface{}: - if name, ok := v["name"].(string); ok { - tc.Type = schemas.ToolChoiceTypeFunction - tc.Function = schemas.ToolChoiceFunction{Name: name} - } - } - - return tc -} - -// ensureParams ensures bifrostReq.Params is initialized -func (r *OpenAIChatRequest) ensureParams(bifrostReq *schemas.BifrostRequest) { - if bifrostReq.Params == nil { - bifrostReq.Params = &schemas.ModelParameters{} - } -} - -// extractLegacyFunctionCall returns the FunctionCall for legacy compatibility -// when exactly one function tool-call is present, otherwise returns nil -func extractLegacyFunctionCall(toolCalls []schemas.ToolCall) *schemas.FunctionCall { - if len(toolCalls) == 1 && toolCalls[0].Type != nil && *toolCalls[0].Type == string(schemas.ToolChoiceTypeFunction) { - return &toolCalls[0].Function - } - return nil -} - // DeriveOpenAIFromBifrostResponse converts a Bifrost response to OpenAI format func DeriveOpenAIFromBifrostResponse(bifrostResp *schemas.BifrostResponse) *OpenAIChatResponse { if bifrostResp == nil { @@ -455,52 +123,14 @@ func DeriveOpenAIFromBifrostResponse(bifrostResp *schemas.BifrostResponse) *Open } openaiResp := &OpenAIChatResponse{ - ID: bifrostResp.ID, - Object: "chat.completion", - Created: bifrostResp.Created, - Model: bifrostResp.Model, - Choices: make([]OpenAIChoice, len(bifrostResp.Choices)), - } - - if bifrostResp.SystemFingerprint != nil { - openaiResp.SystemFingerprint = bifrostResp.SystemFingerprint - } - - // Convert usage information (using schemas.LLMUsage directly) - if bifrostResp.Usage != (schemas.LLMUsage{}) { - openaiResp.Usage = &bifrostResp.Usage - } - - // Convert choices - for i, choice := range bifrostResp.Choices { - openaiChoice := OpenAIChoice{ - Index: choice.Index, - FinishReason: choice.FinishReason, - } - - // Convert message - msg := OpenAIMessage{ - Role: string(choice.Message.Role), - } - - // Convert content back to proper format - if choice.Message.Content != nil { - msg.Content = *choice.Message.Content - } - - // Convert tool calls for assistant messages (already in schemas.ToolCall format) - if choice.Message.AssistantMessage != nil && choice.Message.AssistantMessage.ToolCalls != nil { - msg.ToolCalls = choice.Message.AssistantMessage.ToolCalls - msg.FunctionCall = extractLegacyFunctionCall(*choice.Message.AssistantMessage.ToolCalls) - } - - // Handle tool messages - propagate tool_call_id - if choice.Message.ToolMessage != nil && choice.Message.ToolMessage.ToolCallID != nil { - msg.ToolCallID = choice.Message.ToolMessage.ToolCallID - } - - openaiChoice.Message = msg - openaiResp.Choices[i] = openaiChoice + ID: bifrostResp.ID, + Object: bifrostResp.Object, + Created: bifrostResp.Created, + Model: bifrostResp.Model, + Choices: bifrostResp.Choices, + Usage: &bifrostResp.Usage, + ServiceTier: bifrostResp.ServiceTier, + SystemFingerprint: bifrostResp.SystemFingerprint, } return openaiResp diff --git a/transports/bifrost-http/tracking/plugin.go b/transports/bifrost-http/tracking/plugin.go index 9ab4047a45..417d4f5cf6 100644 --- a/transports/bifrost-http/tracking/plugin.go +++ b/transports/bifrost-http/tracking/plugin.go @@ -6,6 +6,7 @@ package tracking import ( "context" "fmt" + "log" "time" schemas "github.com/maximhq/bifrost/core/schemas" @@ -40,9 +41,14 @@ func NewPrometheusPlugin() *PrometheusPlugin { } } +// GetName returns the name of the plugin. +func (p *PrometheusPlugin) GetName() string { + return "bifrost-http-prometheus" +} + // PreHook records the start time of the request in the context. // This time is used later in PostHook to calculate request duration. -func (p *PrometheusPlugin) PreHook(ctx *context.Context, req *schemas.BifrostRequest) (*schemas.BifrostRequest, error) { +func (p *PrometheusPlugin) PreHook(ctx *context.Context, req *schemas.BifrostRequest) (*schemas.BifrostRequest, *schemas.BifrostResponse, error) { *ctx = context.WithValue(*ctx, startTimeKey, time.Now()) if req.Input.ChatCompletionInput != nil { @@ -51,24 +57,28 @@ func (p *PrometheusPlugin) PreHook(ctx *context.Context, req *schemas.BifrostReq *ctx = context.WithValue(*ctx, methodKey, "text") } - return req, nil + return req, nil, nil } // PostHook calculates duration and records upstream metrics for successful requests. // It records: // - Request latency // - Total request count -func (p *PrometheusPlugin) PostHook(ctx *context.Context, result *schemas.BifrostResponse) (*schemas.BifrostResponse, error) { +func (p *PrometheusPlugin) PostHook(ctx *context.Context, result *schemas.BifrostResponse, bifrostErr *schemas.BifrostError) (*schemas.BifrostResponse, *schemas.BifrostError, error) { + if result == nil { + return result, bifrostErr, nil + } + startTime, ok := (*ctx).Value(startTimeKey).(time.Time) if !ok { - fmt.Println("Warning: startTime not found in context for Prometheus PostHook") - return result, nil + log.Println("Warning: startTime not found in context for Prometheus PostHook") + return result, bifrostErr, nil } method, ok := (*ctx).Value(methodKey).(string) if !ok { - fmt.Println("Warning: method not found in context for Prometheus PostHook") - return result, nil + log.Println("Warning: method not found in context for Prometheus PostHook") + return result, bifrostErr, nil } // Collect prometheus labels from context @@ -93,5 +103,9 @@ func (p *PrometheusPlugin) PostHook(ctx *context.Context, result *schemas.Bifros p.UpstreamLatency.WithLabelValues(promLabelValues...).Observe(duration) p.UpstreamRequestsTotal.WithLabelValues(promLabelValues...).Inc() - return result, nil + return result, bifrostErr, nil +} + +func (p *PrometheusPlugin) Cleanup() error { + return nil } diff --git a/transports/go.mod b/transports/go.mod index da6c0f536e..94d73d966c 100644 --- a/transports/go.mod +++ b/transports/go.mod @@ -4,8 +4,8 @@ go 1.24.1 require ( github.com/fasthttp/router v1.5.4 - github.com/maximhq/bifrost/core v1.0.10 - github.com/maximhq/bifrost/plugins/maxim v1.0.3 + github.com/maximhq/bifrost/core v1.1.2 + github.com/maximhq/bifrost/plugins/maxim v1.0.5 github.com/prometheus/client_golang v1.22.0 github.com/valyala/fasthttp v1.62.0 google.golang.org/genai v1.4.0 diff --git a/transports/go.sum b/transports/go.sum index 8d7afb9170..4ad9363092 100644 --- a/transports/go.sum +++ b/transports/go.sum @@ -67,10 +67,10 @@ github.com/klauspost/compress v1.18.0 h1:c/Cqfb0r+Yi+JtIEq73FWXVkRonBlf0CRNYc8Zt github.com/klauspost/compress v1.18.0/go.mod h1:2Pp+KzxcywXVXMr50+X0Q/Lsb43OQHYWRCY2AiWywWQ= github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc= github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw= -github.com/maximhq/bifrost/core v1.0.10 h1:HLDOg11tFK7AwEHKqBhImUHbeWhFB775TiE/7BJMQhE= -github.com/maximhq/bifrost/core v1.0.10/go.mod h1:8ycaWQ9bjQezoUT/x6a82VmPjoqLzyGglQ0RnnlZjqo= -github.com/maximhq/bifrost/plugins/maxim v1.0.3 h1:3m3BGfC30pNVVYdon77etOBinEaD9B9RVgsTB8HtuDU= -github.com/maximhq/bifrost/plugins/maxim v1.0.3/go.mod h1:Zakfd201Id5uN368lFB09nrOJ3cCmGmzrKOFWq0KiAc= +github.com/maximhq/bifrost/core v1.1.2 h1:NR5zWD+2dMkj1ySmGqcE7VDJUhkvgrjoMoQikxsdXPU= +github.com/maximhq/bifrost/core v1.1.2/go.mod h1:8ycaWQ9bjQezoUT/x6a82VmPjoqLzyGglQ0RnnlZjqo= +github.com/maximhq/bifrost/plugins/maxim v1.0.5 h1:K67bqb49X0q07UbNCu1jlmdhmQG+gTdC8sxXN5ok5bY= +github.com/maximhq/bifrost/plugins/maxim v1.0.5/go.mod h1:Emik7JHo4BIa6kRWDEqOHgFp8M1BQv13bSOepoWw4aw= 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/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA=