450 lines
14 KiB
Go
450 lines
14 KiB
Go
package bedrock
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/maximhq/bifrost/core/schemas"
|
|
)
|
|
|
|
// ToBedrockChatCompletionRequest converts a Bifrost request to Bedrock Converse API format
|
|
func ToBedrockChatCompletionRequest(ctx *schemas.BifrostContext, bifrostReq *schemas.BifrostChatRequest) (*BedrockConverseRequest, error) {
|
|
if bifrostReq == nil {
|
|
return nil, fmt.Errorf("bifrost request is nil")
|
|
}
|
|
|
|
if bifrostReq.Input == nil {
|
|
return nil, fmt.Errorf("only chat completion requests are supported for Bedrock Converse API")
|
|
}
|
|
|
|
bedrockReq := &BedrockConverseRequest{
|
|
ModelID: bifrostReq.Model,
|
|
}
|
|
|
|
// Convert messages and system messages
|
|
messages, systemMessages, err := convertMessages(bifrostReq.Input)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to convert messages: %w", err)
|
|
}
|
|
bedrockReq.Messages = messages
|
|
if len(systemMessages) > 0 {
|
|
bedrockReq.System = systemMessages
|
|
}
|
|
|
|
// Convert parameters and configurations
|
|
if err := convertChatParameters(ctx, bifrostReq, bedrockReq); err != nil {
|
|
return nil, fmt.Errorf("failed to convert chat parameters: %w", err)
|
|
}
|
|
|
|
// Ensure tool config is present when needed
|
|
ensureChatToolConfigForConversation(bifrostReq, bedrockReq)
|
|
|
|
return bedrockReq, nil
|
|
}
|
|
|
|
// ToBifrostChatResponse converts a Bedrock Converse API response to Bifrost format
|
|
func (response *BedrockConverseResponse) ToBifrostChatResponse(ctx context.Context, model string) (*schemas.BifrostChatResponse, error) {
|
|
if response == nil {
|
|
return nil, fmt.Errorf("bedrock response is nil")
|
|
}
|
|
|
|
// Convert content blocks and tool calls
|
|
var contentStr *string
|
|
var contentBlocks []schemas.ChatContentBlock
|
|
var toolCalls []schemas.ChatAssistantMessageToolCall
|
|
var reasoningDetails []schemas.ChatReasoningDetails
|
|
var reasoningText string
|
|
|
|
if response.Output.Message != nil {
|
|
for _, contentBlock := range response.Output.Message.Content {
|
|
// Handle text content
|
|
if contentBlock.Text != nil && *contentBlock.Text != "" {
|
|
chatContentBlock := schemas.ChatContentBlock{
|
|
Type: schemas.ChatContentBlockTypeText,
|
|
Text: contentBlock.Text,
|
|
}
|
|
contentBlocks = append(contentBlocks, chatContentBlock)
|
|
}
|
|
|
|
if contentBlock.ToolUse != nil {
|
|
// Check if this is the structured output tool
|
|
if structuredOutputToolName, ok := ctx.Value(schemas.BifrostContextKeyStructuredOutputToolName).(string); ok && contentBlock.ToolUse.Name == structuredOutputToolName {
|
|
// This is structured output - set contentStr and skip adding to toolCalls
|
|
if contentBlock.ToolUse.Input != nil {
|
|
jsonStr := string(contentBlock.ToolUse.Input)
|
|
contentStr = &jsonStr
|
|
}
|
|
continue // Skip adding to toolCalls
|
|
}
|
|
|
|
// Regular tool call processing
|
|
var arguments string
|
|
if contentBlock.ToolUse.Input != nil {
|
|
arguments = string(contentBlock.ToolUse.Input)
|
|
} else {
|
|
arguments = "{}"
|
|
}
|
|
|
|
toolUseID := contentBlock.ToolUse.ToolUseID
|
|
toolUseName := contentBlock.ToolUse.Name
|
|
|
|
toolCalls = append(toolCalls, schemas.ChatAssistantMessageToolCall{
|
|
Index: uint16(len(toolCalls)),
|
|
Type: schemas.Ptr("function"),
|
|
ID: &toolUseID,
|
|
Function: schemas.ChatAssistantMessageToolCallFunction{
|
|
Name: &toolUseName,
|
|
Arguments: arguments,
|
|
},
|
|
})
|
|
}
|
|
|
|
// Handle reasoning content
|
|
if contentBlock.ReasoningContent != nil {
|
|
if contentBlock.ReasoningContent.ReasoningText == nil {
|
|
continue
|
|
}
|
|
reasoningDetails = append(reasoningDetails, schemas.ChatReasoningDetails{
|
|
Index: len(reasoningDetails),
|
|
Type: schemas.BifrostReasoningDetailsTypeText,
|
|
Text: contentBlock.ReasoningContent.ReasoningText.Text,
|
|
Signature: contentBlock.ReasoningContent.ReasoningText.Signature,
|
|
})
|
|
if contentBlock.ReasoningContent.ReasoningText.Text != nil {
|
|
reasoningText += *contentBlock.ReasoningContent.ReasoningText.Text + "\n"
|
|
}
|
|
}
|
|
|
|
// Handle document content
|
|
if contentBlock.Document != nil {
|
|
fileBlock := schemas.ChatContentBlock{
|
|
Type: schemas.ChatContentBlockTypeFile,
|
|
File: &schemas.ChatInputFile{},
|
|
}
|
|
|
|
// Set filename from document name
|
|
if contentBlock.Document.Name != "" {
|
|
fileBlock.File.Filename = &contentBlock.Document.Name
|
|
}
|
|
|
|
// Set file type based on format
|
|
if contentBlock.Document.Format != "" {
|
|
var fileType string
|
|
switch contentBlock.Document.Format {
|
|
case "pdf":
|
|
fileType = "application/pdf"
|
|
case "txt":
|
|
fileType = "text/plain"
|
|
case "md":
|
|
fileType = "text/markdown"
|
|
case "html":
|
|
fileType = "text/html"
|
|
case "csv":
|
|
fileType = "text/csv"
|
|
case "doc":
|
|
fileType = "application/msword"
|
|
case "docx":
|
|
fileType = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
|
|
case "xls":
|
|
fileType = "application/vnd.ms-excel"
|
|
case "xlsx":
|
|
fileType = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
|
|
default:
|
|
fileType = "application/pdf"
|
|
}
|
|
fileBlock.File.FileType = &fileType
|
|
}
|
|
|
|
// Convert document source data
|
|
if contentBlock.Document.Source != nil {
|
|
if contentBlock.Document.Source.Bytes != nil {
|
|
fileBlock.File.FileData = contentBlock.Document.Source.Bytes
|
|
} else if contentBlock.Document.Source.Text != nil {
|
|
fileBlock.File.FileData = contentBlock.Document.Source.Text
|
|
}
|
|
}
|
|
|
|
contentBlocks = append(contentBlocks, fileBlock)
|
|
}
|
|
}
|
|
}
|
|
|
|
if len(contentBlocks) == 1 && contentBlocks[0].Type == schemas.ChatContentBlockTypeText {
|
|
contentStr = contentBlocks[0].Text
|
|
contentBlocks = nil
|
|
}
|
|
|
|
// Create the message content
|
|
messageContent := schemas.ChatMessageContent{
|
|
ContentStr: contentStr,
|
|
ContentBlocks: contentBlocks,
|
|
}
|
|
|
|
// Create assistant message if we have tool calls
|
|
var assistantMessage *schemas.ChatAssistantMessage
|
|
if len(toolCalls) > 0 {
|
|
assistantMessage = &schemas.ChatAssistantMessage{
|
|
ToolCalls: toolCalls,
|
|
}
|
|
}
|
|
if len(reasoningDetails) > 0 {
|
|
if assistantMessage == nil {
|
|
assistantMessage = &schemas.ChatAssistantMessage{}
|
|
}
|
|
assistantMessage.ReasoningDetails = reasoningDetails
|
|
if reasoningText != "" {
|
|
assistantMessage.Reasoning = new(reasoningText)
|
|
}
|
|
}
|
|
|
|
// Create the response choice
|
|
choices := []schemas.BifrostResponseChoice{
|
|
{
|
|
Index: 0,
|
|
ChatNonStreamResponseChoice: &schemas.ChatNonStreamResponseChoice{
|
|
Message: &schemas.ChatMessage{
|
|
Role: schemas.ChatMessageRoleAssistant,
|
|
Content: &messageContent,
|
|
ChatAssistantMessage: assistantMessage,
|
|
},
|
|
},
|
|
FinishReason: schemas.Ptr(convertBedrockStopReason(response.StopReason)),
|
|
},
|
|
}
|
|
var usage *schemas.BifrostLLMUsage
|
|
if response.Usage != nil {
|
|
// Convert usage information
|
|
usage = &schemas.BifrostLLMUsage{
|
|
PromptTokens: response.Usage.InputTokens,
|
|
CompletionTokens: response.Usage.OutputTokens,
|
|
TotalTokens: response.Usage.TotalTokens,
|
|
}
|
|
// Handle cached tokens if present
|
|
if response.Usage.CacheReadInputTokens > 0 {
|
|
if usage.PromptTokensDetails == nil {
|
|
usage.PromptTokensDetails = &schemas.ChatPromptTokensDetails{}
|
|
}
|
|
usage.PromptTokensDetails.CachedReadTokens = response.Usage.CacheReadInputTokens
|
|
usage.PromptTokens = usage.PromptTokens + response.Usage.CacheReadInputTokens
|
|
}
|
|
if response.Usage.CacheWriteInputTokens > 0 {
|
|
if usage.PromptTokensDetails == nil {
|
|
usage.PromptTokensDetails = &schemas.ChatPromptTokensDetails{}
|
|
}
|
|
usage.PromptTokensDetails.CachedWriteTokens = response.Usage.CacheWriteInputTokens
|
|
usage.PromptTokens = usage.PromptTokens + response.Usage.CacheWriteInputTokens
|
|
}
|
|
}
|
|
|
|
// Create the final Bifrost response
|
|
bifrostResponse := &schemas.BifrostChatResponse{
|
|
ID: uuid.New().String(),
|
|
Model: model,
|
|
Object: "chat.completion",
|
|
Choices: choices,
|
|
Usage: usage,
|
|
Created: int(time.Now().Unix()),
|
|
ExtraFields: schemas.BifrostResponseExtraFields{
|
|
},
|
|
}
|
|
|
|
if response.ServiceTier != nil && response.ServiceTier.Type != "" {
|
|
bifrostResponse.ServiceTier = &response.ServiceTier.Type
|
|
}
|
|
|
|
return bifrostResponse, nil
|
|
}
|
|
|
|
// BedrockStreamState tracks per-stream tool call index state.
|
|
type BedrockStreamState struct {
|
|
nextToolCallIndex int
|
|
contentBlockToToolCallIdx map[int]int
|
|
}
|
|
|
|
// NewBedrockStreamState returns initialised stream state for one streaming response.
|
|
func NewBedrockStreamState() *BedrockStreamState {
|
|
return &BedrockStreamState{
|
|
contentBlockToToolCallIdx: make(map[int]int),
|
|
}
|
|
}
|
|
|
|
func (chunk *BedrockStreamEvent) ToBifrostChatCompletionStream(state *BedrockStreamState) (*schemas.BifrostChatResponse, *schemas.BifrostError, bool) {
|
|
if state == nil {
|
|
state = NewBedrockStreamState()
|
|
} else if state.contentBlockToToolCallIdx == nil {
|
|
state.contentBlockToToolCallIdx = make(map[int]int)
|
|
}
|
|
|
|
// event with metrics/usage is the last and with stop reason is the second last
|
|
switch {
|
|
case chunk.Role != nil:
|
|
// Send empty response to signal start
|
|
streamResponse := &schemas.BifrostChatResponse{
|
|
Object: "chat.completion.chunk",
|
|
Choices: []schemas.BifrostResponseChoice{
|
|
{
|
|
Index: 0,
|
|
ChatStreamResponseChoice: &schemas.ChatStreamResponseChoice{
|
|
Delta: &schemas.ChatStreamResponseChoiceDelta{
|
|
Role: chunk.Role,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
return streamResponse, nil, false
|
|
|
|
case chunk.Start != nil && chunk.Start.ToolUse != nil:
|
|
toolUseStart := chunk.Start.ToolUse
|
|
|
|
toolCallIdx := 0
|
|
if chunk.ContentBlockIndex != nil {
|
|
toolCallIdx = state.nextToolCallIndex
|
|
state.contentBlockToToolCallIdx[*chunk.ContentBlockIndex] = toolCallIdx
|
|
state.nextToolCallIndex++
|
|
}
|
|
|
|
// Create tool call structure for start event
|
|
var toolCall schemas.ChatAssistantMessageToolCall
|
|
toolCall.Index = uint16(toolCallIdx)
|
|
toolCall.ID = schemas.Ptr(toolUseStart.ToolUseID)
|
|
toolCall.Type = schemas.Ptr("function")
|
|
toolCall.Function.Name = schemas.Ptr(toolUseStart.Name)
|
|
toolCall.Function.Arguments = "" // Start with empty arguments
|
|
|
|
streamResponse := &schemas.BifrostChatResponse{
|
|
Object: "chat.completion.chunk",
|
|
Choices: []schemas.BifrostResponseChoice{
|
|
{
|
|
Index: 0,
|
|
ChatStreamResponseChoice: &schemas.ChatStreamResponseChoice{
|
|
Delta: &schemas.ChatStreamResponseChoiceDelta{
|
|
ToolCalls: []schemas.ChatAssistantMessageToolCall{toolCall},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
return streamResponse, nil, false
|
|
|
|
case chunk.Delta != nil:
|
|
switch {
|
|
case chunk.Delta.Text != nil:
|
|
// Handle text delta
|
|
text := *chunk.Delta.Text
|
|
if text != "" {
|
|
streamResponse := &schemas.BifrostChatResponse{
|
|
Object: "chat.completion.chunk",
|
|
Choices: []schemas.BifrostResponseChoice{
|
|
{
|
|
Index: 0,
|
|
ChatStreamResponseChoice: &schemas.ChatStreamResponseChoice{
|
|
Delta: &schemas.ChatStreamResponseChoiceDelta{
|
|
Content: &text,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
return streamResponse, nil, false
|
|
}
|
|
|
|
case chunk.Delta.ToolUse != nil:
|
|
// Handle tool use delta
|
|
toolUseDelta := chunk.Delta.ToolUse
|
|
|
|
toolCallIdx := 0
|
|
if chunk.ContentBlockIndex != nil {
|
|
toolCallIdx = state.contentBlockToToolCallIdx[*chunk.ContentBlockIndex]
|
|
}
|
|
|
|
// Create tool call structure
|
|
var toolCall schemas.ChatAssistantMessageToolCall
|
|
toolCall.Index = uint16(toolCallIdx)
|
|
toolCall.Type = schemas.Ptr("function")
|
|
|
|
// For streaming, we need to accumulate tool use data
|
|
// This is a simplified approach - in practice, you'd need to track tool calls across chunks
|
|
toolCall.Function.Arguments = toolUseDelta.Input
|
|
|
|
streamResponse := &schemas.BifrostChatResponse{
|
|
Object: "chat.completion.chunk",
|
|
Choices: []schemas.BifrostResponseChoice{
|
|
{
|
|
Index: 0,
|
|
ChatStreamResponseChoice: &schemas.ChatStreamResponseChoice{
|
|
Delta: &schemas.ChatStreamResponseChoiceDelta{
|
|
ToolCalls: []schemas.ChatAssistantMessageToolCall{toolCall},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
return streamResponse, nil, false
|
|
|
|
case chunk.Delta.ReasoningContent != nil:
|
|
// Handle reasoning content delta
|
|
reasoningContentDelta := chunk.Delta.ReasoningContent
|
|
|
|
// Only construct and return a response when either Text or Signature is set
|
|
if (reasoningContentDelta.Text == nil || *reasoningContentDelta.Text == "") && reasoningContentDelta.Signature == nil {
|
|
return nil, nil, false
|
|
}
|
|
|
|
var streamResponse *schemas.BifrostChatResponse
|
|
if reasoningContentDelta.Text != nil && *reasoningContentDelta.Text != "" {
|
|
streamResponse = &schemas.BifrostChatResponse{
|
|
Object: "chat.completion.chunk",
|
|
Choices: []schemas.BifrostResponseChoice{
|
|
{
|
|
Index: 0,
|
|
ChatStreamResponseChoice: &schemas.ChatStreamResponseChoice{
|
|
Delta: &schemas.ChatStreamResponseChoiceDelta{
|
|
Reasoning: reasoningContentDelta.Text,
|
|
ReasoningDetails: []schemas.ChatReasoningDetails{
|
|
{
|
|
Index: 0,
|
|
Type: schemas.BifrostReasoningDetailsTypeText,
|
|
Text: reasoningContentDelta.Text,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
} else if reasoningContentDelta.Signature != nil {
|
|
streamResponse = &schemas.BifrostChatResponse{
|
|
Object: "chat.completion.chunk",
|
|
Choices: []schemas.BifrostResponseChoice{
|
|
{
|
|
Index: 0,
|
|
ChatStreamResponseChoice: &schemas.ChatStreamResponseChoice{
|
|
Delta: &schemas.ChatStreamResponseChoiceDelta{
|
|
ReasoningDetails: []schemas.ChatReasoningDetails{
|
|
{
|
|
Index: 0,
|
|
Type: schemas.BifrostReasoningDetailsTypeText,
|
|
Signature: reasoningContentDelta.Signature,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
}
|
|
|
|
return streamResponse, nil, false
|
|
}
|
|
}
|
|
|
|
return nil, nil, false
|
|
}
|