From 941e75c66262e84e8225a9d45097558893dddda7 Mon Sep 17 00:00:00 2001 From: tejas ghatte Date: Tue, 12 May 2026 14:34:38 +0530 Subject: [PATCH] feat: bedrock system tools --- core/providers/anthropic/types.go | 4 +- core/providers/bedrock/responses.go | 1074 ++++++++++++++--- core/providers/bedrock/types.go | 49 +- core/schemas/utils.go | 4 + .../integrations/python/tests/test_bedrock.py | 414 +++++++ 5 files changed, 1359 insertions(+), 186 deletions(-) diff --git a/core/providers/anthropic/types.go b/core/providers/anthropic/types.go index 08eb9a71118..4b4870c2828 100644 --- a/core/providers/anthropic/types.go +++ b/core/providers/anthropic/types.go @@ -183,7 +183,9 @@ var ProviderFeatures = map[schemas.ModelProvider]ProviderFeatureSupport{ // WebSearch, CodeExecution, FastMode, TaskBudgets, AdvisorTool, // InferenceGeo, RedactThinking, AdvancedToolUse (full), PromptCachingScope. schemas.Bedrock: { - ComputerUse: true, Bash: true, Memory: true, TextEditor: true, ToolSearch: true, + WebSearch: true, + CodeExecution: true, + ComputerUse: true, Bash: true, Memory: true, TextEditor: true, ToolSearch: true, ContainerBasic: true, // StructuredOutputs: kept true to match pre-existing behavior and the // provider_feature_support_test.go assertion, but NEITHER B-header diff --git a/core/providers/bedrock/responses.go b/core/providers/bedrock/responses.go index 6743cb68e47..f9bc4a800a1 100644 --- a/core/providers/bedrock/responses.go +++ b/core/providers/bedrock/responses.go @@ -19,20 +19,24 @@ import ( // BedrockResponsesStreamState tracks state during streaming conversion for responses API type BedrockResponsesStreamState struct { - ContentIndexToOutputIndex map[int]int // Maps Bedrock contentBlockIndex to OpenAI output_index - ToolArgumentBuffers map[int]string // Maps output_index to accumulated tool argument JSON - ItemIDs map[int]string // Maps output_index to item ID for stable IDs - ToolCallIDs map[int]string // Maps output_index to tool call ID (callID) - ToolCallNames map[int]string // Maps output_index to tool call name - ReasoningContentIndices map[int]bool // Tracks which content indices are reasoning blocks - CompletedOutputIndices map[int]bool // Tracks which output indices have been completed - CurrentOutputIndex int // Current output index counter - MessageID *string // Message ID (generated) - Model *string // Model name - StopReason *string // Stop reason for the message - CreatedAt int // Timestamp for created_at consistency - HasEmittedCreated bool // Whether we've emitted response.created - HasEmittedInProgress bool // Whether we've emitted response.in_progress + ContentIndexToOutputIndex map[int]int // Maps Bedrock contentBlockIndex to OpenAI output_index + ToolArgumentBuffers map[int]string // Maps output_index to accumulated tool argument JSON + ItemIDs map[int]string // Maps output_index to item ID for stable IDs + ToolCallIDs map[int]string // Maps output_index to tool call ID (callID) + ToolCallNames map[int]string // Maps output_index to tool call name + ReasoningContentIndices map[int]bool // Tracks which content indices are reasoning blocks + CodeInterpreterIndices map[int]bool // Tracks which output indices are nova_code_interpreter calls + NovaGroundingIndices map[int]bool // Tracks which output indices are nova_grounding (web_search_call) blocks + NovaGroundingCitations map[int][]schemas.ResponsesWebSearchToolCallActionSearchSource // Collected citation sources per nova_grounding output index + CompletedOutputIndices map[int]bool // Tracks which output indices have been completed + AnnotationIndices map[int]int // Maps output_index to next annotation index for sequential citation numbering + CurrentOutputIndex int // Current output index counter + MessageID *string // Message ID (generated) + Model *string // Model name + StopReason *string // Stop reason for the message + CreatedAt int // Timestamp for created_at consistency + HasEmittedCreated bool // Whether we've emitted response.created + HasEmittedInProgress bool // Whether we've emitted response.in_progress } // bedrockResponsesStreamStatePool provides a pool for Bedrock responses stream state objects. @@ -45,7 +49,11 @@ var bedrockResponsesStreamStatePool = sync.Pool{ ToolCallIDs: make(map[int]string), ToolCallNames: make(map[int]string), ReasoningContentIndices: make(map[int]bool), + CodeInterpreterIndices: make(map[int]bool), + NovaGroundingIndices: make(map[int]bool), + NovaGroundingCitations: make(map[int][]schemas.ResponsesWebSearchToolCallActionSearchSource), CompletedOutputIndices: make(map[int]bool), + AnnotationIndices: make(map[int]int), CurrentOutputIndex: 0, CreatedAt: int(time.Now().Unix()), HasEmittedCreated: false, @@ -89,11 +97,31 @@ func acquireBedrockResponsesStreamState() *BedrockResponsesStreamState { } else { clear(state.ReasoningContentIndices) } + if state.CodeInterpreterIndices == nil { + state.CodeInterpreterIndices = make(map[int]bool) + } else { + clear(state.CodeInterpreterIndices) + } + if state.NovaGroundingIndices == nil { + state.NovaGroundingIndices = make(map[int]bool) + } else { + clear(state.NovaGroundingIndices) + } + if state.NovaGroundingCitations == nil { + state.NovaGroundingCitations = make(map[int][]schemas.ResponsesWebSearchToolCallActionSearchSource) + } else { + clear(state.NovaGroundingCitations) + } if state.CompletedOutputIndices == nil { state.CompletedOutputIndices = make(map[int]bool) } else { clear(state.CompletedOutputIndices) } + if state.AnnotationIndices == nil { + state.AnnotationIndices = make(map[int]int) + } else { + clear(state.AnnotationIndices) + } // Reset other fields state.CurrentOutputIndex = 0 state.MessageID = nil @@ -145,11 +173,31 @@ func (state *BedrockResponsesStreamState) flush() { } else { clear(state.ReasoningContentIndices) } + if state.CodeInterpreterIndices == nil { + state.CodeInterpreterIndices = make(map[int]bool) + } else { + clear(state.CodeInterpreterIndices) + } + if state.NovaGroundingIndices == nil { + state.NovaGroundingIndices = make(map[int]bool) + } else { + clear(state.NovaGroundingIndices) + } + if state.NovaGroundingCitations == nil { + state.NovaGroundingCitations = make(map[int][]schemas.ResponsesWebSearchToolCallActionSearchSource) + } else { + clear(state.NovaGroundingCitations) + } if state.CompletedOutputIndices == nil { state.CompletedOutputIndices = make(map[int]bool) } else { clear(state.CompletedOutputIndices) } + if state.AnnotationIndices == nil { + state.AnnotationIndices = make(map[int]int) + } else { + clear(state.AnnotationIndices) + } state.CurrentOutputIndex = 0 state.MessageID = nil state.Model = nil @@ -404,78 +452,83 @@ func (chunk *BedrockStreamEvent) ToBifrostResponsesStream(sequenceNumber int, st prevItemID := state.ItemIDs[prevOutputIndex] prevToolName := state.ToolCallNames[prevOutputIndex] accumulatedArgs := state.ToolArgumentBuffers[prevOutputIndex] + statusCompleted := "completed" - // Emit content_part.done for tool call - emptyText := "" - part := &schemas.ResponsesMessageContentBlock{ - Type: schemas.ResponsesOutputMessageContentTypeText, - Text: &emptyText, - ResponsesOutputMessageContentText: &schemas.ResponsesOutputMessageContentText{ - LogProbs: []schemas.ResponsesOutputMessageContentTextLogProb{}, - Annotations: []schemas.ResponsesOutputMessageContentTextAnnotation{}, - }, - } - responses = append(responses, &schemas.BifrostResponsesStreamResponse{ - Type: schemas.ResponsesStreamResponseTypeContentPartDone, - SequenceNumber: sequenceNumber + len(responses), - OutputIndex: schemas.Ptr(prevOutputIndex), - ContentIndex: schemas.Ptr(prevContentIndex), - ItemID: &prevItemID, - Part: part, - }) + if state.CodeInterpreterIndices[prevOutputIndex] { + ciEvents := emitCodeInterpreterDoneEvents(prevOutputIndex, prevContentIndex, prevItemID, prevToolCallID, accumulatedArgs, sequenceNumber+len(responses)) + responses = append(responses, ciEvents...) + } else if state.NovaGroundingIndices[prevOutputIndex] { + citations := state.NovaGroundingCitations[prevOutputIndex] + wsEvents := emitNovaGroundingDoneEvents(prevOutputIndex, prevContentIndex, prevItemID, citations, accumulatedArgs, sequenceNumber+len(responses)) + responses = append(responses, wsEvents...) + } else { + // Close a regular function_call block + emptyText := "" + part := &schemas.ResponsesMessageContentBlock{ + Type: schemas.ResponsesOutputMessageContentTypeText, + Text: &emptyText, + ResponsesOutputMessageContentText: &schemas.ResponsesOutputMessageContentText{ + LogProbs: []schemas.ResponsesOutputMessageContentTextLogProb{}, + Annotations: []schemas.ResponsesOutputMessageContentTextAnnotation{}, + }, + } + responses = append(responses, &schemas.BifrostResponsesStreamResponse{ + Type: schemas.ResponsesStreamResponseTypeContentPartDone, + SequenceNumber: sequenceNumber + len(responses), + OutputIndex: schemas.Ptr(prevOutputIndex), + ContentIndex: schemas.Ptr(prevContentIndex), + ItemID: &prevItemID, + Part: part, + }) - // Emit function_call_arguments.done with full arguments - if accumulatedArgs != "" { - var doneItem *schemas.ResponsesMessage - if prevToolCallID != "" || prevToolName != "" { - doneItem = &schemas.ResponsesMessage{ - ResponsesToolMessage: &schemas.ResponsesToolMessage{}, + if accumulatedArgs != "" { + var doneItem *schemas.ResponsesMessage + if prevToolCallID != "" || prevToolName != "" { + doneItem = &schemas.ResponsesMessage{ + ResponsesToolMessage: &schemas.ResponsesToolMessage{}, + } + if prevToolCallID != "" { + doneItem.ResponsesToolMessage.CallID = &prevToolCallID + } + if prevToolName != "" { + doneItem.ResponsesToolMessage.Name = &prevToolName + } + } + argsDoneResponse := &schemas.BifrostResponsesStreamResponse{ + Type: schemas.ResponsesStreamResponseTypeFunctionCallArgumentsDone, + SequenceNumber: sequenceNumber + len(responses), + OutputIndex: schemas.Ptr(prevOutputIndex), + Arguments: &accumulatedArgs, } - if prevToolCallID != "" { - doneItem.ResponsesToolMessage.CallID = &prevToolCallID + if prevItemID != "" { + argsDoneResponse.ItemID = &prevItemID } - if prevToolName != "" { - doneItem.ResponsesToolMessage.Name = &prevToolName + if doneItem != nil { + argsDoneResponse.Item = doneItem } + responses = append(responses, argsDoneResponse) } - argsDoneResponse := &schemas.BifrostResponsesStreamResponse{ - Type: schemas.ResponsesStreamResponseTypeFunctionCallArgumentsDone, + toolDoneItem := &schemas.ResponsesMessage{ + ID: &prevItemID, + Type: schemas.Ptr(schemas.ResponsesMessageTypeFunctionCall), + Status: &statusCompleted, + ResponsesToolMessage: &schemas.ResponsesToolMessage{ + CallID: &prevToolCallID, + Name: &prevToolName, + Arguments: &accumulatedArgs, + }, + } + responses = append(responses, &schemas.BifrostResponsesStreamResponse{ + Type: schemas.ResponsesStreamResponseTypeOutputItemDone, SequenceNumber: sequenceNumber + len(responses), OutputIndex: schemas.Ptr(prevOutputIndex), - Arguments: &accumulatedArgs, - } - if prevItemID != "" { - argsDoneResponse.ItemID = &prevItemID - } - if doneItem != nil { - argsDoneResponse.Item = doneItem - } - responses = append(responses, argsDoneResponse) - } - - // Emit output_item.done for tool call - statusCompleted := "completed" - toolDoneItem := &schemas.ResponsesMessage{ - ID: &prevItemID, - Type: schemas.Ptr(schemas.ResponsesMessageTypeFunctionCall), - Status: &statusCompleted, - ResponsesToolMessage: &schemas.ResponsesToolMessage{ - CallID: &prevToolCallID, - Name: &prevToolName, - Arguments: &accumulatedArgs, - }, + ContentIndex: schemas.Ptr(prevContentIndex), + ItemID: &prevItemID, + Item: toolDoneItem, + }) } - responses = append(responses, &schemas.BifrostResponsesStreamResponse{ - Type: schemas.ResponsesStreamResponseTypeOutputItemDone, - SequenceNumber: sequenceNumber + len(responses), - OutputIndex: schemas.Ptr(prevOutputIndex), - ContentIndex: schemas.Ptr(prevContentIndex), - ItemID: &prevItemID, - Item: toolDoneItem, - }) - // Mark this output index as completed state.CompletedOutputIndices[prevOutputIndex] = true } @@ -483,37 +536,101 @@ func (chunk *BedrockStreamEvent) ToBifrostResponsesStream(sequenceNumber int, st // Create new output index for this tool use outputIndex := state.CurrentOutputIndex state.ContentIndexToOutputIndex[contentBlockIndex] = outputIndex - state.CurrentOutputIndex++ // Increment for next use + state.CurrentOutputIndex++ - // Store tool use ID as item ID and call ID toolUseID := chunk.Start.ToolUse.ToolUseID toolName := chunk.Start.ToolUse.Name state.ItemIDs[outputIndex] = toolUseID state.ToolCallIDs[outputIndex] = toolUseID state.ToolCallNames[outputIndex] = toolName - statusInProgress := "in_progress" - item := &schemas.ResponsesMessage{ - ID: &toolUseID, - Type: schemas.Ptr(schemas.ResponsesMessageTypeFunctionCall), - Status: &statusInProgress, - ResponsesToolMessage: &schemas.ResponsesToolMessage{ - CallID: &toolUseID, - Name: &toolName, - Arguments: schemas.Ptr(""), // Arguments will be filled by deltas - }, - } - - // Initialize argument buffer for this tool call + // Initialize argument buffer state.ToolArgumentBuffers[outputIndex] = "" - responses = append(responses, &schemas.BifrostResponsesStreamResponse{ - Type: schemas.ResponsesStreamResponseTypeOutputItemAdded, - SequenceNumber: sequenceNumber + len(responses), - OutputIndex: schemas.Ptr(outputIndex), - ContentIndex: schemas.Ptr(contentBlockIndex), - Item: item, - }) + statusInProgress := "in_progress" + + if toolName == "nova_code_interpreter" { + // Emit output_item.added then code_interpreter_call.in_progress + state.CodeInterpreterIndices[outputIndex] = true + item := &schemas.ResponsesMessage{ + ID: &toolUseID, + Type: schemas.Ptr(schemas.ResponsesMessageTypeCodeInterpreterCall), + Status: &statusInProgress, + ResponsesToolMessage: &schemas.ResponsesToolMessage{ + ResponsesCodeInterpreterToolCall: &schemas.ResponsesCodeInterpreterToolCall{ + ContainerID: toolUseID, + Outputs: []schemas.ResponsesCodeInterpreterOutput{}, + }, + }, + } + responses = append(responses, &schemas.BifrostResponsesStreamResponse{ + Type: schemas.ResponsesStreamResponseTypeOutputItemAdded, + SequenceNumber: sequenceNumber + len(responses), + OutputIndex: schemas.Ptr(outputIndex), + ContentIndex: schemas.Ptr(contentBlockIndex), + Item: item, + }) + responses = append(responses, &schemas.BifrostResponsesStreamResponse{ + Type: schemas.ResponsesStreamResponseTypeCodeInterpreterCallInProgress, + SequenceNumber: sequenceNumber + len(responses), + OutputIndex: schemas.Ptr(outputIndex), + ContentIndex: schemas.Ptr(contentBlockIndex), + Item: item, + }) + } else if toolName == string(BedrockSystemToolNovaGrounding) { + state.NovaGroundingIndices[outputIndex] = true + state.NovaGroundingCitations[outputIndex] = nil + item := &schemas.ResponsesMessage{ + ID: &toolUseID, + Type: schemas.Ptr(schemas.ResponsesMessageTypeWebSearchCall), + Status: &statusInProgress, + ResponsesToolMessage: &schemas.ResponsesToolMessage{ + CallID: &toolUseID, + Action: &schemas.ResponsesToolMessageActionStruct{ + ResponsesWebSearchToolCallAction: &schemas.ResponsesWebSearchToolCallAction{ + Type: "search", + }, + }, + }, + } + responses = append(responses, &schemas.BifrostResponsesStreamResponse{ + Type: schemas.ResponsesStreamResponseTypeOutputItemAdded, + SequenceNumber: sequenceNumber + len(responses), + OutputIndex: schemas.Ptr(outputIndex), + ContentIndex: schemas.Ptr(contentBlockIndex), + Item: item, + }) + responses = append(responses, &schemas.BifrostResponsesStreamResponse{ + Type: schemas.ResponsesStreamResponseTypeWebSearchCallInProgress, + SequenceNumber: sequenceNumber + len(responses), + OutputIndex: schemas.Ptr(outputIndex), + ItemID: &toolUseID, + }) + responses = append(responses, &schemas.BifrostResponsesStreamResponse{ + Type: schemas.ResponsesStreamResponseTypeWebSearchCallSearching, + SequenceNumber: sequenceNumber + len(responses), + OutputIndex: schemas.Ptr(outputIndex), + ItemID: &toolUseID, + }) + } else { + item := &schemas.ResponsesMessage{ + ID: &toolUseID, + Type: schemas.Ptr(schemas.ResponsesMessageTypeFunctionCall), + Status: &statusInProgress, + ResponsesToolMessage: &schemas.ResponsesToolMessage{ + CallID: &toolUseID, + Name: &toolName, + Arguments: schemas.Ptr(""), + }, + } + responses = append(responses, &schemas.BifrostResponsesStreamResponse{ + Type: schemas.ResponsesStreamResponseTypeOutputItemAdded, + SequenceNumber: sequenceNumber + len(responses), + OutputIndex: schemas.Ptr(outputIndex), + ContentIndex: schemas.Ptr(contentBlockIndex), + Item: item, + }) + } return responses, nil, false } @@ -705,21 +822,71 @@ func (chunk *BedrockStreamEvent) ToBifrostResponsesStream(sequenceNumber int, st return []*schemas.BifrostResponsesStreamResponse{response}, nil, false } + case chunk.Delta.Citation != nil: + citation := chunk.Delta.Citation + if citation.Location.Web != nil { + if state.NovaGroundingIndices[outputIndex] { + domain := citation.Location.Web.Domain + state.NovaGroundingCitations[outputIndex] = append( + state.NovaGroundingCitations[outputIndex], + schemas.ResponsesWebSearchToolCallActionSearchSource{ + Type: "url", + URL: citation.Location.Web.URL, + Title: &domain, + }, + ) + } + // Emit as url_citation annotation (covers both nova_grounding and text blocks). + itemID := state.ItemIDs[outputIndex] + annotationIndex := state.AnnotationIndices[outputIndex] + state.AnnotationIndices[outputIndex]++ + annotation := &schemas.ResponsesOutputMessageContentTextAnnotation{ + Type: "url_citation", + URL: schemas.Ptr(citation.Location.Web.URL), + Title: schemas.Ptr(citation.Location.Web.Domain), + } + response := &schemas.BifrostResponsesStreamResponse{ + Type: schemas.ResponsesStreamResponseTypeOutputTextAnnotationAdded, + SequenceNumber: sequenceNumber, + OutputIndex: schemas.Ptr(outputIndex), + ContentIndex: &contentBlockIndex, + AnnotationIndex: &annotationIndex, + Annotation: annotation, + } + if itemID != "" { + response.ItemID = &itemID + } + return []*schemas.BifrostResponsesStreamResponse{response}, nil, false + } + case chunk.Delta.ToolUse != nil: - // Handle tool use delta - function call arguments + // Handle tool use delta - function call arguments or code interpreter code toolUseDelta := chunk.Delta.ToolUse if toolUseDelta.Input != "" { - // Accumulate argument deltas state.ToolArgumentBuffers[outputIndex] += toolUseDelta.Input itemID := state.ItemIDs[outputIndex] - response := &schemas.BifrostResponsesStreamResponse{ - Type: schemas.ResponsesStreamResponseTypeFunctionCallArgumentsDelta, - SequenceNumber: sequenceNumber, - OutputIndex: schemas.Ptr(outputIndex), - ContentIndex: &contentBlockIndex, - Delta: &toolUseDelta.Input, + + var response *schemas.BifrostResponsesStreamResponse + if state.CodeInterpreterIndices[outputIndex] { + // Each nova_code_interpreter delta is a complete JSON object {"snippet":"..."}. + codeDelta := providerUtils.GetJSONField([]byte(toolUseDelta.Input), "snippet").String() + response = &schemas.BifrostResponsesStreamResponse{ + Type: schemas.ResponsesStreamResponseTypeCodeInterpreterCallCodeDelta, + SequenceNumber: sequenceNumber, + OutputIndex: schemas.Ptr(outputIndex), + ContentIndex: &contentBlockIndex, + Delta: &codeDelta, + } + } else { + response = &schemas.BifrostResponsesStreamResponse{ + Type: schemas.ResponsesStreamResponseTypeFunctionCallArgumentsDelta, + SequenceNumber: sequenceNumber, + OutputIndex: schemas.Ptr(outputIndex), + ContentIndex: &contentBlockIndex, + Delta: &toolUseDelta.Input, + } } if itemID != "" { response.ItemID = &itemID @@ -853,6 +1020,92 @@ func (chunk *BedrockStreamEvent) ToBifrostResponsesStream(sequenceNumber int, st return nil, nil, false } +// emitCodeInterpreterDoneEvents extracts the code from accumulated JSON args and emits +// code_interpreter_call.code.done + code_interpreter_call.completed + output_item.done in sequence. +func emitCodeInterpreterDoneEvents(outputIndex, contentIndex int, itemID, containerID, accumulatedArgs string, baseSequenceNumber int) []*schemas.BifrostResponsesStreamResponse { + code := providerUtils.GetJSONField([]byte(accumulatedArgs), "snippet").String() + statusCompleted := "completed" + codeDone := &schemas.BifrostResponsesStreamResponse{ + Type: schemas.ResponsesStreamResponseTypeCodeInterpreterCallCodeDone, + SequenceNumber: baseSequenceNumber, + OutputIndex: schemas.Ptr(outputIndex), + ContentIndex: &contentIndex, + ItemID: &itemID, + Delta: &code, + } + doneItem := &schemas.ResponsesMessage{ + ID: &itemID, + Type: schemas.Ptr(schemas.ResponsesMessageTypeCodeInterpreterCall), + Status: &statusCompleted, + ResponsesToolMessage: &schemas.ResponsesToolMessage{ + ResponsesCodeInterpreterToolCall: &schemas.ResponsesCodeInterpreterToolCall{ + Code: &code, + ContainerID: containerID, + Outputs: []schemas.ResponsesCodeInterpreterOutput{}, + }, + }, + } + completed := &schemas.BifrostResponsesStreamResponse{ + Type: schemas.ResponsesStreamResponseTypeCodeInterpreterCallCompleted, + SequenceNumber: baseSequenceNumber + 1, + OutputIndex: schemas.Ptr(outputIndex), + ContentIndex: &contentIndex, + ItemID: &itemID, + Item: doneItem, + } + outputDone := &schemas.BifrostResponsesStreamResponse{ + Type: schemas.ResponsesStreamResponseTypeOutputItemDone, + SequenceNumber: baseSequenceNumber + 2, + OutputIndex: schemas.Ptr(outputIndex), + ContentIndex: &contentIndex, + ItemID: &itemID, + Item: doneItem, + } + return []*schemas.BifrostResponsesStreamResponse{codeDone, completed, outputDone} +} + +// emitNovaGroundingDoneEvents emits web_search_call.completed + output_item.done for a nova_grounding block. +// accumulatedArgs holds the raw toolUse input JSON (e.g. `{"query":"..."}`) from the block's deltas. +func emitNovaGroundingDoneEvents(outputIndex, contentIndex int, itemID string, citations []schemas.ResponsesWebSearchToolCallActionSearchSource, accumulatedArgs string, baseSequenceNumber int) []*schemas.BifrostResponsesStreamResponse { + statusCompleted := "completed" + action := &schemas.ResponsesWebSearchToolCallAction{ + Type: "search", + Sources: citations, + } + // Extract the search query from the accumulated toolUse input. + if q := providerUtils.GetJSONField([]byte(accumulatedArgs), "query").String(); q != "" { + action.Query = &q + action.Queries = []string{q} + } + doneItem := &schemas.ResponsesMessage{ + ID: &itemID, + Type: schemas.Ptr(schemas.ResponsesMessageTypeWebSearchCall), + Status: &statusCompleted, + ResponsesToolMessage: &schemas.ResponsesToolMessage{ + CallID: &itemID, + Action: &schemas.ResponsesToolMessageActionStruct{ + ResponsesWebSearchToolCallAction: action, + }, + }, + } + return []*schemas.BifrostResponsesStreamResponse{ + { + Type: schemas.ResponsesStreamResponseTypeWebSearchCallCompleted, + SequenceNumber: baseSequenceNumber, + OutputIndex: schemas.Ptr(outputIndex), + ItemID: &itemID, + }, + { + Type: schemas.ResponsesStreamResponseTypeOutputItemDone, + SequenceNumber: baseSequenceNumber + 1, + OutputIndex: schemas.Ptr(outputIndex), + ContentIndex: &contentIndex, + ItemID: &itemID, + Item: doneItem, + }, + } +} + // FinalizeBedrockStream finalizes the stream by closing any open items and emitting completed event func FinalizeBedrockStream(state *BedrockResponsesStreamState, sequenceNumber int, usage *schemas.ResponsesResponseUsage) []*schemas.BifrostResponsesStreamResponse { var responses []*schemas.BifrostResponsesStreamResponse @@ -917,80 +1170,84 @@ func FinalizeBedrockStream(state *BedrockResponsesStreamState, sequenceNumber in isToolCall := toolCallID != "" if isToolCall { - // This is a tool call that needs to be closed - - // Emit content_part.done for tool call - emptyText := "" - part := &schemas.ResponsesMessageContentBlock{ - Type: schemas.ResponsesOutputMessageContentTypeText, - Text: &emptyText, - ResponsesOutputMessageContentText: &schemas.ResponsesOutputMessageContentText{ - LogProbs: []schemas.ResponsesOutputMessageContentTextLogProb{}, - Annotations: []schemas.ResponsesOutputMessageContentTextAnnotation{}, - }, - } - responses = append(responses, &schemas.BifrostResponsesStreamResponse{ - Type: schemas.ResponsesStreamResponseTypeContentPartDone, - SequenceNumber: sequenceNumber + len(responses), - OutputIndex: schemas.Ptr(outputIndex), - ContentIndex: &contentIndex, - ItemID: &itemID, - Part: part, - }) - - // Emit function_call_arguments.done with full arguments toolName := state.ToolCallNames[outputIndex] accumulatedArgs := state.ToolArgumentBuffers[outputIndex] - if accumulatedArgs != "" { - var doneItem *schemas.ResponsesMessage - if toolCallID != "" || toolName != "" { - doneItem = &schemas.ResponsesMessage{ - ResponsesToolMessage: &schemas.ResponsesToolMessage{}, + statusCompleted := "completed" + + if state.CodeInterpreterIndices[outputIndex] { + ciEvents := emitCodeInterpreterDoneEvents(outputIndex, contentIndex, itemID, toolCallID, accumulatedArgs, sequenceNumber+len(responses)) + responses = append(responses, ciEvents...) + } else if state.NovaGroundingIndices[outputIndex] { + citations := state.NovaGroundingCitations[outputIndex] + wsEvents := emitNovaGroundingDoneEvents(outputIndex, contentIndex, itemID, citations, accumulatedArgs, sequenceNumber+len(responses)) + responses = append(responses, wsEvents...) + } else { + // Close a regular function_call + emptyText := "" + part := &schemas.ResponsesMessageContentBlock{ + Type: schemas.ResponsesOutputMessageContentTypeText, + Text: &emptyText, + ResponsesOutputMessageContentText: &schemas.ResponsesOutputMessageContentText{ + LogProbs: []schemas.ResponsesOutputMessageContentTextLogProb{}, + Annotations: []schemas.ResponsesOutputMessageContentTextAnnotation{}, + }, + } + responses = append(responses, &schemas.BifrostResponsesStreamResponse{ + Type: schemas.ResponsesStreamResponseTypeContentPartDone, + SequenceNumber: sequenceNumber + len(responses), + OutputIndex: schemas.Ptr(outputIndex), + ContentIndex: &contentIndex, + ItemID: &itemID, + Part: part, + }) + + if accumulatedArgs != "" { + var doneItem *schemas.ResponsesMessage + if toolCallID != "" || toolName != "" { + doneItem = &schemas.ResponsesMessage{ + ResponsesToolMessage: &schemas.ResponsesToolMessage{}, + } + if toolCallID != "" { + doneItem.ResponsesToolMessage.CallID = &toolCallID + } + if toolName != "" { + doneItem.ResponsesToolMessage.Name = &toolName + } + } + response := &schemas.BifrostResponsesStreamResponse{ + Type: schemas.ResponsesStreamResponseTypeFunctionCallArgumentsDone, + SequenceNumber: sequenceNumber + len(responses), + OutputIndex: schemas.Ptr(outputIndex), + Arguments: &accumulatedArgs, } - if toolCallID != "" { - doneItem.ResponsesToolMessage.CallID = &toolCallID + if itemID != "" { + response.ItemID = &itemID } - if toolName != "" { - doneItem.ResponsesToolMessage.Name = &toolName + if doneItem != nil { + response.Item = doneItem } + responses = append(responses, response) } - response := &schemas.BifrostResponsesStreamResponse{ - Type: schemas.ResponsesStreamResponseTypeFunctionCallArgumentsDone, + doneItem := &schemas.ResponsesMessage{ + ID: &itemID, + Type: schemas.Ptr(schemas.ResponsesMessageTypeFunctionCall), + Status: &statusCompleted, + ResponsesToolMessage: &schemas.ResponsesToolMessage{ + CallID: &toolCallID, + Name: &toolName, + Arguments: &accumulatedArgs, + }, + } + responses = append(responses, &schemas.BifrostResponsesStreamResponse{ + Type: schemas.ResponsesStreamResponseTypeOutputItemDone, SequenceNumber: sequenceNumber + len(responses), OutputIndex: schemas.Ptr(outputIndex), - Arguments: &accumulatedArgs, - } - if itemID != "" { - response.ItemID = &itemID - } - if doneItem != nil { - response.Item = doneItem - } - responses = append(responses, response) - } - - // Emit output_item.done for tool call - statusCompleted := "completed" - doneItem := &schemas.ResponsesMessage{ - ID: &itemID, - Type: schemas.Ptr(schemas.ResponsesMessageTypeFunctionCall), - Status: &statusCompleted, - ResponsesToolMessage: &schemas.ResponsesToolMessage{ - CallID: &toolCallID, - Name: &toolName, - Arguments: &accumulatedArgs, - }, - } - - responses = append(responses, &schemas.BifrostResponsesStreamResponse{ - Type: schemas.ResponsesStreamResponseTypeOutputItemDone, - SequenceNumber: sequenceNumber + len(responses), - OutputIndex: schemas.Ptr(outputIndex), - ContentIndex: &contentIndex, - ItemID: &itemID, - Item: doneItem, - }) + ContentIndex: &contentIndex, + ItemID: &itemID, + Item: doneItem, + }) + } // end else (regular function call) } else { // This is likely a text item that needs to be closed @@ -1193,14 +1450,24 @@ func ToBedrockConverseStreamResponse(bifrostResp *schemas.BifrostResponsesStream return nil, nil case schemas.ResponsesStreamResponseTypeOutputItemAdded: - // Content block start + // Content block start — handles nova_grounding (web_search_call), function calls, and text items. if bifrostResp.Item != nil && bifrostResp.Item.ResponsesToolMessage != nil { - // Tool use start - if bifrostResp.Item.ResponsesToolMessage.Name != nil && bifrostResp.Item.ResponsesToolMessage.CallID != nil { - contentBlockIndex := 0 - if bifrostResp.ContentIndex != nil { - contentBlockIndex = *bifrostResp.ContentIndex + contentBlockIndex := 0 + if bifrostResp.ContentIndex != nil { + contentBlockIndex = *bifrostResp.ContentIndex + } + // web_search_call (nova_grounding): CallID is set, Name is nil + if bifrostResp.Item.Type != nil && *bifrostResp.Item.Type == schemas.ResponsesMessageTypeWebSearchCall && + bifrostResp.Item.ResponsesToolMessage.CallID != nil { + event.ContentBlockIndex = &contentBlockIndex + event.Start = &BedrockContentBlockStart{ + ToolUse: &BedrockToolUseStart{ + ToolUseID: *bifrostResp.Item.ResponsesToolMessage.CallID, + Name: string(BedrockSystemToolNovaGrounding), + }, } + } else if bifrostResp.Item.ResponsesToolMessage.Name != nil && bifrostResp.Item.ResponsesToolMessage.CallID != nil { + // Regular function call event.ContentBlockIndex = &contentBlockIndex event.Start = &BedrockContentBlockStart{ ToolUse: &BedrockToolUseStart{ @@ -1208,15 +1475,97 @@ func ToBedrockConverseStreamResponse(bifrostResp *schemas.BifrostResponsesStream Name: *bifrostResp.Item.ResponsesToolMessage.Name, }, } + } else { + return nil, nil } } else if bifrostResp.Item != nil { // Text item added - Bedrock doesn't have an explicit text start event, so we skip it - // Check if it's a text message (has content blocks or is a message type) if bifrostResp.Item.Content != nil || (bifrostResp.Item.Type != nil && *bifrostResp.Item.Type == schemas.ResponsesMessageTypeMessage) { return nil, nil } } + case schemas.ResponsesStreamResponseTypeOutputTextAnnotationAdded: + // url_citation annotation → contentBlockDelta.citation + if bifrostResp.Annotation != nil && bifrostResp.Annotation.URL != nil { + contentBlockIndex := 0 + if bifrostResp.ContentIndex != nil { + contentBlockIndex = *bifrostResp.ContentIndex + } + domain := "" + if bifrostResp.Annotation.Title != nil { + domain = *bifrostResp.Annotation.Title + } + event.ContentBlockIndex = &contentBlockIndex + event.Delta = &BedrockContentBlockDelta{ + Citation: &BedrockCitation{ + Location: BedrockCitationLocation{ + Web: &BedrockWebCitationLocation{ + URL: *bifrostResp.Annotation.URL, + Domain: domain, + }, + }, + }, + } + } else { + return nil, nil + } + + case schemas.ResponsesStreamResponseTypeWebSearchCallInProgress, + schemas.ResponsesStreamResponseTypeWebSearchCallSearching, + schemas.ResponsesStreamResponseTypeWebSearchCallCompleted, + schemas.ResponsesStreamResponseTypeWebSearchCallResultsAdded, + schemas.ResponsesStreamResponseTypeWebSearchCallResultsCompleted: + // No Bedrock equivalent for these status events — skip. + return nil, nil + + case schemas.ResponsesStreamResponseTypeCodeInterpreterCallInProgress: + // nova_code_interpreter → contentBlockStart + if bifrostResp.Item != nil && bifrostResp.Item.ResponsesToolMessage != nil && + bifrostResp.Item.ResponsesToolMessage.ResponsesCodeInterpreterToolCall != nil { + toolUseID := bifrostResp.Item.ResponsesToolMessage.ResponsesCodeInterpreterToolCall.ContainerID + if toolUseID == "" && bifrostResp.Item.ID != nil { + toolUseID = *bifrostResp.Item.ID + } + contentBlockIndex := 0 + if bifrostResp.ContentIndex != nil { + contentBlockIndex = *bifrostResp.ContentIndex + } + event.ContentBlockIndex = &contentBlockIndex + event.Start = &BedrockContentBlockStart{ + ToolUse: &BedrockToolUseStart{ + ToolUseID: toolUseID, + Name: string(BedrockSystemToolNovaCodeInterpreter), + }, + } + } else { + return nil, nil + } + + case schemas.ResponsesStreamResponseTypeCodeInterpreterCallCodeDelta: + // nova_code_interpreter toolUse delta — wrap snippet back into {"snippet":"..."} JSON + if bifrostResp.Delta != nil && *bifrostResp.Delta != "" { + contentBlockIndex := 0 + if bifrostResp.ContentIndex != nil { + contentBlockIndex = *bifrostResp.ContentIndex + } + inputJSON, _ := json.Marshal(map[string]string{"snippet": *bifrostResp.Delta}) + event.ContentBlockIndex = &contentBlockIndex + event.Delta = &BedrockContentBlockDelta{ + ToolUse: &BedrockToolUseDelta{ + Input: string(inputJSON), + }, + } + } else { + return nil, nil + } + + case schemas.ResponsesStreamResponseTypeCodeInterpreterCallCodeDone, + schemas.ResponsesStreamResponseTypeCodeInterpreterCallCompleted, + schemas.ResponsesStreamResponseTypeCodeInterpreterCallInterpreting: + // No Bedrock equivalent — skip. + return nil, nil + case schemas.ResponsesStreamResponseTypeOutputTextDelta: // Text delta if bifrostResp.Delta != nil && *bifrostResp.Delta != "" { @@ -1464,6 +1813,18 @@ func (request *BedrockConverseRequest) ToBifrostResponsesRequest(ctx *schemas.Bi } bifrostReq.Params.Tools = append(bifrostReq.Params.Tools, bifrostTool) + } else if tool.SystemTool != nil { + // Nova system tools: nova_grounding → web_search, nova_code_interpreter → code_interpreter + var toolType schemas.ResponsesToolType + switch tool.SystemTool.Name { + case BedrockSystemToolNovaGrounding: + toolType = schemas.ResponsesToolTypeWebSearch + case BedrockSystemToolNovaCodeInterpreter: + toolType = schemas.ResponsesToolTypeCodeInterpreter + default: + continue + } + bifrostReq.Params.Tools = append(bifrostReq.Params.Tools, schemas.ResponsesTool{Type: toolType}) } else if tool.CachePoint != nil && !schemas.IsNovaModel(bifrostReq.Model) { // add cache control to last tool in tools array if len(bifrostReq.Params.Tools) > 0 { @@ -1959,7 +2320,24 @@ func ToBedrockResponsesRequest(ctx *schemas.BifrostContext, bifrostReq *schemas. // Convert tools if bifrostReq.Params != nil && bifrostReq.Params.Tools != nil { var bedrockTools []BedrockTool + isNova2 := schemas.IsNova2Model(bifrostReq.Model) for _, tool := range bifrostReq.Params.Tools { + if tool.Type == schemas.ResponsesToolTypeWebSearch || tool.Type == schemas.ResponsesToolTypeCodeInterpreter { + if !isNova2 { + return nil, fmt.Errorf("tool type %q is only supported on Nova 2 models in Bedrock; got model %q", tool.Type, bifrostReq.Model) + } + var systemToolName BedrockSystemToolType + switch tool.Type { + case schemas.ResponsesToolTypeWebSearch: + systemToolName = BedrockSystemToolNovaGrounding + case schemas.ResponsesToolTypeCodeInterpreter: + systemToolName = BedrockSystemToolNovaCodeInterpreter + } + bedrockTools = append(bedrockTools, BedrockTool{ + SystemTool: &BedrockSystemTool{Name: systemToolName}, + }) + continue + } if tool.ResponsesToolFunction != nil { // Create the complete schema object that Bedrock expects var schemaObject interface{} @@ -2167,9 +2545,19 @@ func ToBedrockConverseResponse(bifrostResp *schemas.BifrostResponsesResponse) (* message.Content = append(message.Content, bedrockMsg.Content...) } - // Check for tool use in the content blocks + // Check for tool use in the content blocks. Server-managed tools + // (nova_grounding, nova_code_interpreter) return both toolUse and + // toolResult in the same message — their stop reason is "end_turn", + // not "tool_use". Only flag hasToolUse when there is an unmatched + // toolUse (i.e. the model is waiting for a client-side tool result). + resolvedToolUseIDs := make(map[string]bool) + for _, block := range message.Content { + if block.ToolResult != nil { + resolvedToolUseIDs[block.ToolResult.ToolUseID] = true + } + } for _, block := range message.Content { - if block.ToolUse != nil { + if block.ToolUse != nil && !resolvedToolUseIDs[block.ToolUse.ToolUseID] { hasToolUse = true break } @@ -2239,16 +2627,17 @@ func ensureResponsesToolConfigForConversation(bifrostReq *schemas.BifrostRespons return // Already has tool config } - hasToolContent, tools := extractToolsFromResponsesConversationHistory(bifrostReq.Input) + hasToolContent, tools := extractToolsFromResponsesConversationHistory(bifrostReq.Input, bifrostReq.Model) if hasToolContent && len(tools) > 0 { bedrockReq.ToolConfig = &BedrockToolConfig{Tools: tools} } } // extractToolsFromResponsesConversationHistory extracts tools from Responses conversation history -func extractToolsFromResponsesConversationHistory(messages []schemas.ResponsesMessage) (bool, []BedrockTool) { +func extractToolsFromResponsesConversationHistory(messages []schemas.ResponsesMessage, model string) (bool, []BedrockTool) { var hasToolContent bool toolMap := make(map[string]*schemas.ResponsesTool) // Use map to deduplicate by name + var hasNovaGrounding, hasNovaCodeInterpreter bool for _, msg := range messages { // Check if message contains tool use or tool result @@ -2273,11 +2662,17 @@ func extractToolsFromResponsesConversationHistory(messages []schemas.ResponsesMe } } } + case schemas.ResponsesMessageTypeWebSearchCall: + hasToolContent = true + hasNovaGrounding = true + case schemas.ResponsesMessageTypeCodeInterpreterCall: + hasToolContent = true + hasNovaCodeInterpreter = true } } } - // Convert map to slice + // Convert function tool map to BedrockTool slice var tools []BedrockTool for _, tool := range toolMap { if tool.Name != nil && tool.ResponsesToolFunction != nil { @@ -2308,6 +2703,16 @@ func extractToolsFromResponsesConversationHistory(messages []schemas.ResponsesMe } } + // Append system tools found in history — only valid on Nova 2 models + if schemas.IsNova2Model(model) { + if hasNovaGrounding { + tools = append(tools, BedrockTool{SystemTool: &BedrockSystemTool{Name: BedrockSystemToolNovaGrounding}}) + } + if hasNovaCodeInterpreter { + tools = append(tools, BedrockTool{SystemTool: &BedrockSystemTool{Name: BedrockSystemToolNovaCodeInterpreter}}) + } + } + return hasToolContent, tools } @@ -2572,6 +2977,9 @@ func ConvertBifrostMessagesToBedrockMessages(ctx context.Context, bifrostMessage var bedrockMessages []BedrockMessage var systemMessages []BedrockSystemMessage var pendingReasoningContentBlocks []BedrockContentBlock + // pendingServerToolBlocks accumulates nova_grounding / nova_code_interpreter toolUse+toolResult + // blocks that must be prepended to the next assistant text message (same-turn server-managed tools). + var pendingServerToolBlocks []BedrockContentBlock // Initialize the state manager for tracking tool calls and results stateManager := NewToolCallStateManager() @@ -2910,6 +3318,12 @@ func ConvertBifrostMessagesToBedrockMessages(ctx context.Context, bifrostMessage // Convert user/assistant text message bedrockMsg := convertBifrostMessageToBedrockMessage(ctx, &msg) if bedrockMsg != nil { + // Prepend buffered server-managed tool blocks (nova_grounding / nova_code_interpreter) + // to the assistant message they belong to — they're part of the same turn. + if bedrockMsg.Role == BedrockMessageRoleAssistant && len(pendingServerToolBlocks) > 0 { + bedrockMsg.Content = append(pendingServerToolBlocks, bedrockMsg.Content...) + pendingServerToolBlocks = nil + } bedrockMessages = append(bedrockMessages, *bedrockMsg) } } @@ -2921,9 +3335,110 @@ func ConvertBifrostMessagesToBedrockMessages(ctx context.Context, bifrostMessage if len(reasoningBlocks) > 0 { pendingReasoningContentBlocks = append(pendingReasoningContentBlocks, reasoningBlocks...) } + + case schemas.ResponsesMessageTypeWebSearchCall: + // Convert web_search_call → nova_grounding toolUse + toolResult. + if msg.ResponsesToolMessage == nil || msg.ResponsesToolMessage.CallID == nil { + continue + } + callID := *msg.ResponsesToolMessage.CallID + // Build toolUse input from the search query (matches original Bedrock format). + inputMap := map[string]string{} + if msg.ResponsesToolMessage.Action != nil && + msg.ResponsesToolMessage.Action.ResponsesWebSearchToolCallAction != nil { + action := msg.ResponsesToolMessage.Action.ResponsesWebSearchToolCallAction + if action.Query != nil { + inputMap["query"] = *action.Query + } + } + inputBytes, _ := json.Marshal(inputMap) + toolUseBlock := BedrockContentBlock{ + ToolUse: &BedrockToolUse{ + ToolUseID: callID, + Name: string(BedrockSystemToolNovaGrounding), + Input: json.RawMessage(inputBytes), + Type: "server_tool_use", + }, + } + // Serialize sources as JSON for the toolResult content; preserve type and status. + sourcesText := "[]" + if msg.ResponsesToolMessage.Action != nil && + msg.ResponsesToolMessage.Action.ResponsesWebSearchToolCallAction != nil { + action := msg.ResponsesToolMessage.Action.ResponsesWebSearchToolCallAction + if len(action.Sources) > 0 { + if b, err := json.Marshal(action.Sources); err == nil { + sourcesText = string(b) + } + } + } + resultType := BedrockNovaGroundingResultType + toolResultBlock := BedrockContentBlock{ + ToolResult: &BedrockToolResult{ + ToolUseID: callID, + Type: &resultType, + Status: schemas.Ptr("success"), + Content: []BedrockContentBlock{{Text: &sourcesText}}, + }, + } + pendingServerToolBlocks = append(pendingServerToolBlocks, toolUseBlock, toolResultBlock) + + case schemas.ResponsesMessageTypeCodeInterpreterCall: + // Convert code_interpreter_call → nova_code_interpreter toolUse + toolResult. + // Both blocks are buffered and prepended to the next assistant message. + if msg.ResponsesToolMessage == nil || msg.ResponsesToolMessage.ResponsesCodeInterpreterToolCall == nil { + continue + } + ci := msg.ResponsesToolMessage.ResponsesCodeInterpreterToolCall + toolUseID := ci.ContainerID + if toolUseID == "" && msg.ID != nil { + toolUseID = *msg.ID + } + code := "" + if ci.Code != nil { + code = *ci.Code + } + inputBytes, _ := json.Marshal(map[string]string{"snippet": code}) + toolUseBlock := BedrockContentBlock{ + ToolUse: &BedrockToolUse{ + ToolUseID: toolUseID, + Name: string(BedrockSystemToolNovaCodeInterpreter), + Input: json.RawMessage(inputBytes), + Type: "server_tool_use", + }, + } + // Build toolResult from outputs (stdout/stderr). + var stdOut, stdErr string + for _, output := range ci.Outputs { + if output.ResponsesCodeInterpreterOutputLogs != nil { + stdOut += output.ResponsesCodeInterpreterOutputLogs.Logs + } + } + execResultBytes, _ := json.Marshal(struct { + StdOut string `json:"stdOut"` + StdErr string `json:"stdErr"` + }{StdOut: stdOut, StdErr: stdErr}) + execResultStr := string(execResultBytes) + resultType := BedrockNovaCodeInterpreterResultType + toolResultBlock := BedrockContentBlock{ + ToolResult: &BedrockToolResult{ + ToolUseID: toolUseID, + Type: &resultType, + Content: []BedrockContentBlock{{Text: &execResultStr}}, + }, + } + pendingServerToolBlocks = append(pendingServerToolBlocks, toolUseBlock, toolResultBlock) } } + // Flush any remaining server-managed tool blocks (no following assistant message). + if len(pendingServerToolBlocks) > 0 { + bedrockMessages = append(bedrockMessages, BedrockMessage{ + Role: BedrockMessageRoleAssistant, + Content: pendingServerToolBlocks, + }) + pendingServerToolBlocks = nil + } + // Flush any remaining pending tool calls flushPendingToolCalls() @@ -3130,7 +3645,55 @@ func convertSingleBedrockMessageToBifrostMessages(ctx *schemas.BifrostContext, m } } + // Pre-scan: build toolUseId → toolResult map for nova_code_interpreter_result blocks + // so we can attach execution output when we encounter the matching toolUse block. + novaCodeResults := make(map[string]*BedrockToolResult) + for i := range msg.Content { + r := msg.Content[i].ToolResult + if r != nil && r.Type != nil && *r.Type == BedrockNovaCodeInterpreterResultType { + novaCodeResults[r.ToolUseID] = r + } + } + + // Pre-scan: collect nova_grounding toolUseIDs and citation sources from citationsContent. + // nova_grounding toolResults (paired with the toolUse) are skipped in the main loop; + // citation URLs from text blocks are surfaced as sources on the web_search_call item. + novaGroundingToolUseIDs := make(map[string]bool) + var novaGroundingSources []schemas.ResponsesWebSearchToolCallActionSearchSource + seenCitationURLs := make(map[string]bool) + for i := range msg.Content { + if msg.Content[i].ToolUse != nil && msg.Content[i].ToolUse.Name == string(BedrockSystemToolNovaGrounding) { + novaGroundingToolUseIDs[msg.Content[i].ToolUse.ToolUseID] = true + } + if msg.Content[i].CitationsContent != nil { + for _, citation := range msg.Content[i].CitationsContent.Citations { + if citation.Location.Web != nil && !seenCitationURLs[citation.Location.Web.URL] { + seenCitationURLs[citation.Location.Web.URL] = true + domain := citation.Location.Web.Domain + novaGroundingSources = append(novaGroundingSources, schemas.ResponsesWebSearchToolCallActionSearchSource{ + Type: "url", + URL: citation.Location.Web.URL, + Title: &domain, + }) + } + } + } + } + + // lastTextOutputIdx tracks the index into outputMessages of the most recently appended + // text message, so standalone citationsContent blocks can be attached to it as annotations. + lastTextOutputIdx := -1 + for _, block := range msg.Content { + // Skip nova_code_interpreter_result tool results — they are consumed via novaCodeResults above. + if block.ToolResult != nil && block.ToolResult.Type != nil && *block.ToolResult.Type == BedrockNovaCodeInterpreterResultType { + continue + } + // Skip nova_grounding tool results — server-managed, consumed by the pre-scan above. + if block.ToolResult != nil && novaGroundingToolUseIDs[block.ToolResult.ToolUseID] { + continue + } + if block.Text != nil { // Text content role := convertBedrockRoleToBifrostRole(msg.Role) @@ -3147,6 +3710,37 @@ func convertSingleBedrockMessageToBifrostMessages(ctx *schemas.BifrostContext, m bifrostMsg.ID = schemas.Ptr("msg_" + fmt.Sprintf("%d", time.Now().UnixNano())) } outputMessages = append(outputMessages, bifrostMsg) + // Track this message so standalone citationsContent blocks can be attached to it. + lastTextOutputIdx = len(outputMessages) - 1 + + } else if block.CitationsContent != nil { + // Standalone citationsContent block — attach citations as url_citation annotations + // to the most recently created text message (interleaved in the Bedrock format). + if lastTextOutputIdx >= 0 { + lastMsg := &outputMessages[lastTextOutputIdx] + if lastMsg.Content != nil && len(lastMsg.Content.ContentBlocks) > 0 { + cb := &lastMsg.Content.ContentBlocks[0] + if cb.ResponsesOutputMessageContentText == nil { + cb.ResponsesOutputMessageContentText = &schemas.ResponsesOutputMessageContentText{ + LogProbs: []schemas.ResponsesOutputMessageContentTextLogProb{}, + Annotations: []schemas.ResponsesOutputMessageContentTextAnnotation{}, + } + } + for _, citation := range block.CitationsContent.Citations { + if citation.Location.Web == nil { + continue + } + cb.ResponsesOutputMessageContentText.Annotations = append( + cb.ResponsesOutputMessageContentText.Annotations, + schemas.ResponsesOutputMessageContentTextAnnotation{ + Type: "url_citation", + URL: schemas.Ptr(citation.Location.Web.URL), + Title: schemas.Ptr(citation.Location.Web.Domain), + }, + ) + } + } + } } else if block.ReasoningContent != nil { // Reasoning content - collect to create a single reasoning message @@ -3181,6 +3775,96 @@ func convertSingleBedrockMessageToBifrostMessages(ctx *schemas.BifrostContext, m bifrostMsg.ID = schemas.Ptr("msg_" + fmt.Sprintf("%d", time.Now().UnixNano())) } outputMessages = append(outputMessages, bifrostMsg) + } else if toolUseName == "nova_code_interpreter" { + // Nova code interpreter: build a code_interpreter_call message. + // Bedrock returns the code under the "snippet" key in toolUse.input. + var snippetInput []byte + if block.ToolUse.Input != nil { + snippetInput = block.ToolUse.Input + } + codeSnippet := providerUtils.GetJSONField(snippetInput, "snippet").String() + + // Build outputs from the paired toolResult (pre-scanned above). + var ciOutputs []schemas.ResponsesCodeInterpreterOutput + if result, ok := novaCodeResults[toolUseID]; ok { + // Extract the JSON payload: {"stdOut":"...","stdErr":"...","exitCode":0,"isError":false} + var execResult struct { + StdOut string `json:"stdOut"` + StdErr string `json:"stdErr"` + } + for _, c := range result.Content { + if c.Text != nil { + _ = json.Unmarshal([]byte(*c.Text), &execResult) + break + } + } + if execResult.StdOut != "" { + ciOutputs = append(ciOutputs, schemas.ResponsesCodeInterpreterOutput{ + ResponsesCodeInterpreterOutputLogs: &schemas.ResponsesCodeInterpreterOutputLogs{ + Type: "logs", + Logs: execResult.StdOut, + }, + }) + } + if execResult.StdErr != "" { + ciOutputs = append(ciOutputs, schemas.ResponsesCodeInterpreterOutput{ + ResponsesCodeInterpreterOutputLogs: &schemas.ResponsesCodeInterpreterOutputLogs{ + Type: "logs", + Logs: execResult.StdErr, + }, + }) + } + } + if ciOutputs == nil { + ciOutputs = []schemas.ResponsesCodeInterpreterOutput{} + } + + ciMsg := schemas.ResponsesMessage{ + Type: schemas.Ptr(schemas.ResponsesMessageTypeCodeInterpreterCall), + Status: schemas.Ptr("completed"), + ResponsesToolMessage: &schemas.ResponsesToolMessage{ + ResponsesCodeInterpreterToolCall: &schemas.ResponsesCodeInterpreterToolCall{ + Code: &codeSnippet, + ContainerID: toolUseID, + Outputs: ciOutputs, + }, + }, + } + if isOutputMessage { + ciMsg.ID = schemas.Ptr("msg_" + fmt.Sprintf("%d", time.Now().UnixNano())) + role := schemas.ResponsesInputMessageRoleAssistant + ciMsg.Role = &role + } + outputMessages = append(outputMessages, ciMsg) + + } else if toolUseName == string(BedrockSystemToolNovaGrounding) { + // nova_grounding → web_search_call with query from toolUse.input and citations from text blocks. + wsAction := &schemas.ResponsesWebSearchToolCallAction{ + Type: "search", + Sources: novaGroundingSources, + } + if block.ToolUse.Input != nil { + if q := providerUtils.GetJSONField(block.ToolUse.Input, "query").String(); q != "" { + wsAction.Query = &q + wsAction.Queries = []string{q} + } + } + wsMsg := schemas.ResponsesMessage{ + Type: schemas.Ptr(schemas.ResponsesMessageTypeWebSearchCall), + Status: schemas.Ptr("completed"), + ResponsesToolMessage: &schemas.ResponsesToolMessage{ + CallID: &toolUseID, + Action: &schemas.ResponsesToolMessageActionStruct{ + ResponsesWebSearchToolCallAction: wsAction, + }, + }, + } + if isOutputMessage { + wsMsg.ID = schemas.Ptr("msg_" + fmt.Sprintf("%d", time.Now().UnixNano())) + role := schemas.ResponsesInputMessageRoleAssistant + wsMsg.Role = &role + } + outputMessages = append(outputMessages, wsMsg) } else { // Normal tool call message arguments := "{}" @@ -3509,6 +4193,32 @@ func convertBifrostResponsesMessageContentBlocksToBedrockContentBlocks(ctx conte blocks = append(blocks, bedrockBlock) } + // For text blocks: emit a citationsContent block per url_citation annotation, + // reconstructing the interleaved text+citation structure Bedrock uses. + if bedrockBlock.Text != nil && block.ResponsesOutputMessageContentText != nil { + for _, annotation := range block.ResponsesOutputMessageContentText.Annotations { + if annotation.Type != "url_citation" || annotation.URL == nil { + continue + } + domain := "" + if annotation.Title != nil { + domain = *annotation.Title + } + blocks = append(blocks, BedrockContentBlock{ + CitationsContent: &BedrockCitationsContent{ + Citations: []BedrockCitation{{ + Location: BedrockCitationLocation{ + Web: &BedrockWebCitationLocation{ + URL: *annotation.URL, + Domain: domain, + }, + }, + }}, + }, + }) + } + } + if block.CacheControl != nil { blocks = append(blocks, BedrockContentBlock{ CachePoint: &BedrockCachePoint{ diff --git a/core/providers/bedrock/types.go b/core/providers/bedrock/types.go index 9faec96ad29..4c5367c3253 100644 --- a/core/providers/bedrock/types.go +++ b/core/providers/bedrock/types.go @@ -210,6 +210,9 @@ type BedrockContentBlock struct { // Cache point for the content block CachePoint *BedrockCachePoint `json:"cachePoint,omitempty"` + + // Citations from nova_grounding — co-located with a text block in the same content block + CitationsContent *BedrockCitationsContent `json:"citationsContent,omitempty"` } type BedrockCachePointType string @@ -249,9 +252,10 @@ type BedrockDocumentSourceData struct { // BedrockToolUse represents a tool use request type BedrockToolUse struct { - ToolUseID string `json:"toolUseId"` // Required: Unique identifier for this tool use - Name string `json:"name"` // Required: Name of the tool to use - Input json.RawMessage `json:"input"` // Required: Input parameters for the tool (json.RawMessage preserves key ordering for prompt caching) + ToolUseID string `json:"toolUseId"` // Required: Unique identifier for this tool use + Name string `json:"name"` // Required: Name of the tool to use + Input json.RawMessage `json:"input"` // Required: Input parameters for the tool (json.RawMessage preserves key ordering for prompt caching) + Type string `json:"type,omitempty"` // Optional: "server_tool_use" for Nova system tools } // BedrockToolResult represents the result of a tool use @@ -259,6 +263,7 @@ type BedrockToolResult struct { ToolUseID string `json:"toolUseId"` // Required: ID of the tool use this result corresponds to Content []BedrockContentBlock `json:"content"` // Required: Content of the tool result Status *string `json:"status,omitempty"` // Optional: Status of tool execution ("success" or "error") + Type *string `json:"type,omitempty"` // Optional: result type e.g. "nova_code_interpreter_result" } // BedrockGuardContent represents guard content for guardrails @@ -308,6 +313,22 @@ type BedrockToolConfig struct { type BedrockTool struct { ToolSpec *BedrockToolSpec `json:"toolSpec,omitempty"` // Tool specification CachePoint *BedrockCachePoint `json:"cachePoint,omitempty"` // Cache point for the tool + SystemTool *BedrockSystemTool `json:"systemTool,omitempty"` // Nova system tool (nova_grounding, nova_code_interpreter) +} + +type BedrockSystemToolType string + +const ( + BedrockSystemToolNovaGrounding BedrockSystemToolType = "nova_grounding" + BedrockSystemToolNovaCodeInterpreter BedrockSystemToolType = "nova_code_interpreter" +) + +const BedrockNovaCodeInterpreterResultType = "nova_code_interpreter_result" +const BedrockNovaGroundingResultType = "nova_grounding_result" + +// BedrockSystemTool represents a Nova-managed system tool +type BedrockSystemTool struct { + Name BedrockSystemToolType `json:"name"` // "nova_grounding" | "nova_code_interpreter" } // BedrockToolSpec represents the specification of a tool @@ -647,6 +668,28 @@ type BedrockContentBlockDelta struct { Text *string `json:"text,omitempty"` // Text content delta ReasoningContent *BedrockReasoningContentText `json:"reasoningContent,omitempty"` // Reasoning content delta ToolUse *BedrockToolUseDelta `json:"toolUse,omitempty"` // Tool use delta + Citation *BedrockCitation `json:"citation,omitempty"` // nova_grounding citation delta +} + +// BedrockWebCitationLocation represents the web location of a citation +type BedrockWebCitationLocation struct { + URL string `json:"url"` + Domain string `json:"domain"` +} + +// BedrockCitationLocation represents the location of a citation (union type) +type BedrockCitationLocation struct { + Web *BedrockWebCitationLocation `json:"web,omitempty"` +} + +// BedrockCitation represents a single citation returned by nova_grounding +type BedrockCitation struct { + Location BedrockCitationLocation `json:"location"` +} + +// BedrockCitationsContent represents the citations block embedded in a text content block +type BedrockCitationsContent struct { + Citations []BedrockCitation `json:"citations"` } // BedrockToolUseDelta represents incremental tool use content diff --git a/core/schemas/utils.go b/core/schemas/utils.go index 169fb1238e8..189bd72e15e 100644 --- a/core/schemas/utils.go +++ b/core/schemas/utils.go @@ -1263,6 +1263,10 @@ func IsNovaModel(model string) bool { return strings.Contains(model, "nova") } +func IsNova2Model(model string) bool { + return strings.Contains(model, "nova-2") && (strings.Contains(model, "lite") || strings.Contains(model, "sonic")) +} + // IsAnthropicModel checks if the model is an Anthropic model. func IsAnthropicModel(model string) bool { return strings.Contains(model, "anthropic.") || strings.Contains(model, "claude") diff --git a/tests/integrations/python/tests/test_bedrock.py b/tests/integrations/python/tests/test_bedrock.py index 8640390e333..fbe0a1f0550 100644 --- a/tests/integrations/python/tests/test_bedrock.py +++ b/tests/integrations/python/tests/test_bedrock.py @@ -46,6 +46,12 @@ 27. Count tokens from long text - Cross-provider 28. Count tokens from multi-turn conversation - Cross-provider +Nova System Tools Tests (TestNovaSystemTools): +50. nova_grounding non-streaming (converse) +51. nova_grounding streaming (converse-stream) +52. nova_code_interpreter non-streaming (converse) +53. nova_code_interpreter streaming (converse-stream) + Invoke Endpoint — Image Generation Tests (TestBedrockInvokeEndpoint): 29. Titan image generation via invoke (taskType=TEXT_IMAGE) 30. Titan embeddings via invoke (inputText) @@ -2893,3 +2899,411 @@ def test_49_invoke_stream_anthropic_messages(self, bedrock_client): full_text = "".join(text_parts) assert full_text, f"Expected non-empty streamed text, got: {full_text!r}" print(f" ✓ event_types={event_types}, text={full_text[:60]!r}") + + +# --------------------------------------------------------------------------- +# Nova System Tools Tests (nova_grounding and nova_code_interpreter) +# --------------------------------------------------------------------------- +# These tests exercise Bedrock Nova system tools through the Bifrost converse +# and converse-stream paths. Nova system tools are AWS-managed: the model +# invokes them automatically (no client-side tool execution required). +# +# nova_grounding → maps to web_search in Bifrost neutral schema +# nova_code_interpreter → maps to code_interpreter in Bifrost neutral schema +# --------------------------------------------------------------------------- + + +class TestNovaSystemTools: + """ + Tests for Amazon Nova system tools via Bedrock Converse and Converse-Stream. + + Both tools are server-managed by AWS — the model calls them and AWS executes + them automatically in the same response. No client-side tool loop is needed. + + 50. nova_grounding non-streaming + 51. nova_grounding streaming + 52. nova_code_interpreter non-streaming + 53. nova_code_interpreter streaming + """ + + NOVA_MODEL = "us.amazon.nova-2-lite-v1:0" + + # ------------------------------------------------------------------ # + # 50. nova_grounding — non-streaming # + # ------------------------------------------------------------------ # + @skip_if_no_api_key("bedrock") + def test_50_nova_grounding_non_streaming(self, bedrock_client): + """Test Case 50: nova_grounding system tool via Bedrock Converse (non-streaming). + + Sends a converse request with systemTool nova_grounding enabled. The model + automatically searches the web and returns a grounded text response. Bifrost + maps nova_grounding → web_search in the neutral schema and converts back. + """ + print("\n=== Test 50: nova_grounding via converse (non-streaming) ===") + + tool_config = { + "tools": [ + {"systemTool": {"name": "nova_grounding"}} + ] + } + + try: + response = bedrock_client.converse( + modelId=self.NOVA_MODEL, + messages=[ + { + "role": "user", + "content": [ + { + "text": ( + "Use web search to find a brief description of the Eiffel Tower " + "and tell me when it was built." + ) + } + ], + } + ], + toolConfig=tool_config, + inferenceConfig={"maxTokens": 500}, + ) + except Exception as e: + err_str = str(e).lower() + if "validation" in err_str or "unknown" in err_str or "not supported" in err_str: + pytest.skip(f"nova_grounding not available or schema rejected: {e}") + raise + + assert "output" in response, f"Expected 'output' in response, got: {list(response.keys())}" + msg = response["output"].get("message", {}) + assert msg.get("role") == "assistant", f"Expected role='assistant', got: {msg.get('role')}" + + content_blocks = msg.get("content", []) + assert isinstance(content_blocks, list) and len(content_blocks) > 0, ( + f"Expected non-empty content blocks, got: {content_blocks}" + ) + + # nova_grounding returns multiple content blocks: empty text, toolUse, + # toolResult, then the actual grounded text (possibly split across blocks). + # Collect all non-empty text across every block. + full_text = " ".join(b["text"] for b in content_blocks if b.get("text", "").strip()) + assert full_text, ( + f"Expected non-empty text in grounding response, got: {content_blocks}" + ) + + # nova_grounding should produce text about the Eiffel Tower + assert any(kw in full_text.lower() for kw in ["eiffel", "paris", "tower", "france", "1889"]), ( + f"Expected Eiffel Tower info in response, got: {full_text[:200]}" + ) + + stop_reason = response.get("stopReason", "") + print(stop_reason) + assert stop_reason in ("end_turn", "max_tokens"), ( + f"Unexpected stopReason: {stop_reason}" + ) + print(f" ✓ stopReason={stop_reason!r}, text={full_text[:80]!r}") + + # ------------------------------------------------------------------ # + # 51. nova_grounding — streaming # + # ------------------------------------------------------------------ # + @skip_if_no_api_key("bedrock") + def test_51_nova_grounding_streaming(self, bedrock_client): + """Test Case 51: nova_grounding system tool via Bedrock Converse-Stream. + + Per AWS docs, nova_grounding streaming produces citation deltas inline within + the text stream (no separate contentBlockStart for the tool block): + messageStart + contentBlockStart (text block) + contentBlockDelta { delta: { citation: { location: { web: { url, domain } } } } } (0-N) + contentBlockDelta { delta: { text: "..." } } (1-N) + contentBlockStop + messageStop + + Bifrost must reproduce these citation deltas as contentBlockDelta.citation events. + The query asks for real-time information to ensure the model uses grounding. + """ + print("\n=== Test 51: nova_grounding via converse-stream (streaming) ===") + + tool_config = { + "tools": [ + {"systemTool": {"name": "nova_grounding"}} + ] + } + + try: + response_stream = bedrock_client.converse_stream( + modelId=self.NOVA_MODEL, + messages=[ + { + "role": "user", + "content": [ + { + # Use a real-time query so the model actually invokes grounding + "text": ( + "Search the web and tell me today's date and one current headline. " + "You must use web search." + ) + } + ], + } + ], + toolConfig=tool_config, + inferenceConfig={"maxTokens": 500}, + ) + except AttributeError: + pytest.skip("converse_stream not available in this boto3 version") + except Exception as e: + err_str = str(e).lower() + if "validation" in err_str or "unknown" in err_str or "not supported" in err_str: + pytest.skip(f"nova_grounding streaming not available: {e}") + raise + + stream = response_stream.get("stream") + if stream is None: + stream = response_stream.get("eventStream") + assert stream is not None, "Response missing 'stream' or 'eventStream'" + + citation_urls = [] # contentBlockDelta.citation events + text_parts = [] # contentBlockDelta.text events + got_message_stop = False + start_time = time.time() + timeout = 60 + + for event in stream: + print(event) + if time.time() - start_time > timeout: + pytest.fail(f"Streaming timed out after {timeout}s") + + if "contentBlockDelta" in event: + delta = event["contentBlockDelta"].get("delta", {}) + if "text" in delta and delta["text"]: + text_parts.append(delta["text"]) + elif "citation" in delta: + # Citation delta produced by nova_grounding: { citation: { location: { web: { url, domain } } } } + web = delta["citation"].get("location", {}).get("web", {}) + if web.get("url"): + citation_urls.append(web["url"]) + + elif "messageStop" in event: + got_message_stop = True + + assert got_message_stop, "Expected 'messageStop' event" + + full_text = "".join(text_parts) + assert full_text, "Expected non-empty streamed text from nova_grounding response" + + # Grounding must produce citation deltas alongside the text + assert len(citation_urls) > 0, ( + f"Expected at least one contentBlockDelta.citation event — " + f"nova_grounding must emit citation deltas that Bifrost preserves as " + f"contentBlockDelta.citation on the converse-stream route. " + f"text_parts={len(text_parts)}, text={full_text[:100]!r}" + ) + + print( + f" ✓ {len(citation_urls)} citation(s), {len(text_parts)} text delta(s), " + f"text={full_text[:80]!r}" + ) + + # ------------------------------------------------------------------ # + # 52. nova_code_interpreter — non-streaming # + # ------------------------------------------------------------------ # + @skip_if_no_api_key("bedrock") + def test_52_nova_code_interpreter_non_streaming(self, bedrock_client): + """Test Case 52: nova_code_interpreter system tool via Bedrock Converse (non-streaming). + + AWS Bedrock executes the generated code automatically and returns both the + toolUse (code) and toolResult (stdout/stderr) in the same assistant message. + Bifrost merges these into a code_interpreter_call output item and converts + back to Bedrock format, producing a text explanation of the result. + """ + print("\n=== Test 52: nova_code_interpreter via converse (non-streaming) ===") + + tool_config = { + "tools": [ + {"systemTool": {"name": "nova_code_interpreter"}} + ] + } + + try: + response = bedrock_client.converse( + modelId=self.NOVA_MODEL, + messages=[ + { + "role": "user", + "content": [ + { + "text": ( + "Write and execute Python code to calculate the factorial of 10 " + "and print the result." + ) + } + ], + } + ], + toolConfig=tool_config, + inferenceConfig={"maxTokens": 500}, + ) + except Exception as e: + err_str = str(e).lower() + if "validation" in err_str or "unknown" in err_str or "not supported" in err_str: + pytest.skip(f"nova_code_interpreter not available or schema rejected: {e}") + raise + + assert "output" in response, f"Expected 'output' in response, got: {list(response.keys())}" + msg = response["output"].get("message", {}) + assert msg.get("role") == "assistant", f"Expected role='assistant', got: {msg.get('role')}" + + content_blocks = msg.get("content", []) + assert isinstance(content_blocks, list) and len(content_blocks) > 0, ( + f"Expected non-empty content blocks, got: {content_blocks}" + ) + + # nova_code_interpreter returns: empty text, toolUse (code), toolResult + # (stdout), then the model's explanation — possibly split across blocks. + # Collect all non-empty text and verify the factorial result appears. + has_tool_use = any("toolUse" in b for b in content_blocks) + has_tool_result = any("toolResult" in b for b in content_blocks) + full_text = " ".join(b["text"] for b in content_blocks if b.get("text", "").strip()) + + assert has_tool_use, f"Expected toolUse block, got: {content_blocks}" + assert full_text or has_tool_result, ( + f"Expected execution result (toolResult) or explanatory text, got: {content_blocks}" + ) + + # The combined text should mention the factorial result or the computation + if full_text: + assert any(kw in full_text.lower() for kw in ["3628800", "factorial", "result", "10", "code"]), ( + f"Expected factorial-related text, got: {full_text[:200]}" + ) + + stop_reason = response.get("stopReason", "") + assert stop_reason in ("end_turn", "max_tokens"), ( + f"Unexpected stopReason: {stop_reason}" + ) + print(f" ✓ stopReason={stop_reason!r}, content_blocks={len(content_blocks)}") + if full_text: + print(f" ✓ text={full_text[:80]!r}") + + # ------------------------------------------------------------------ # + # 53. nova_code_interpreter — streaming # + # ------------------------------------------------------------------ # + @skip_if_no_api_key("bedrock") + def test_53_nova_code_interpreter_streaming(self, bedrock_client): + """Test Case 53: nova_code_interpreter system tool via Bedrock Converse-Stream. + + Bedrock streaming for nova_code_interpreter produces (per AWS docs): + messageStart + contentBlockStart { start: { toolUse: { name: "nova_code_interpreter", ... } } } + contentBlockDelta { delta: { toolUse: { input: '{"snippet":"..."}' } } } (1-N) + contentBlockStop + contentBlockStart (text block) + contentBlockDelta { delta: { text: "..." } } (1-N) + contentBlockStop + messageStop + + Each toolUse delta is a complete JSON object {"snippet":""}. + Bifrost must reproduce this exact shape on the converse-stream route. + """ + print("\n=== Test 53: nova_code_interpreter via converse-stream (streaming) ===") + + tool_config = { + "tools": [ + {"systemTool": {"name": "nova_code_interpreter"}} + ] + } + + try: + response_stream = bedrock_client.converse_stream( + modelId=self.NOVA_MODEL, + messages=[ + { + "role": "user", + "content": [ + { + "text": ( + "Write and run Python code to compute 2 raised to the power of 20." + ) + } + ], + } + ], + toolConfig=tool_config, + inferenceConfig={"maxTokens": 500}, + ) + except AttributeError: + pytest.skip("converse_stream not available in this boto3 version") + except Exception as e: + err_str = str(e).lower() + if "validation" in err_str or "unknown" in err_str or "not supported" in err_str: + pytest.skip(f"nova_code_interpreter streaming not available: {e}") + raise + + stream = response_stream.get("stream") + if stream is None: + stream = response_stream.get("eventStream") + assert stream is not None, "Response missing 'stream' or 'eventStream'" + + has_code_interpreter_block_start = False # contentBlockStart with nova_code_interpreter + code_snippets = [] # parsed snippet values from toolUse deltas + text_parts = [] + got_message_stop = False + current_block_start = None + start_time = time.time() + timeout = 60 + + for event in stream: + if time.time() - start_time > timeout: + pytest.fail(f"Streaming timed out after {timeout}s") + + if "messageStart" in event: + pass + + elif "contentBlockStart" in event: + current_block_start = event["contentBlockStart"].get("start", {}) + tool_use = current_block_start.get("toolUse", {}) + if tool_use.get("name") == "nova_code_interpreter": + has_code_interpreter_block_start = True + + elif "contentBlockStop" in event: + current_block_start = None + + elif "contentBlockDelta" in event: + delta = event["contentBlockDelta"].get("delta", {}) + if "text" in delta and delta["text"]: + text_parts.append(delta["text"]) + elif "toolUse" in delta: + raw_input = delta["toolUse"].get("input", "") + if raw_input: + # Each delta is a complete JSON object: {"snippet": ""} + try: + parsed = json.loads(raw_input) + snippet = parsed.get("snippet", "") + if snippet: + code_snippets.append(snippet) + except json.JSONDecodeError: + pass # Unexpected — deltas should be complete JSON + + elif "messageStop" in event: + got_message_stop = True + + assert got_message_stop, "Expected 'messageStop' event" + + # Bedrock contract: a contentBlockStart for nova_code_interpreter MUST appear + assert has_code_interpreter_block_start, ( + "Expected contentBlockStart with toolUse.name='nova_code_interpreter' — " + "Bifrost must emit this event for the nova_code_interpreter block" + ) + + # nova_code_interpreter must stream at least one toolUse delta with a snippet + assert len(code_snippets) > 0, ( + "Expected at least one contentBlockDelta.toolUse.input with a 'snippet' field " + "from nova_code_interpreter — Bifrost must emit these code deltas" + ) + + full_code = "".join(code_snippets) + assert full_code.strip(), "Expected non-empty code snippet from nova_code_interpreter" + + full_text = "".join(text_parts) + print( + f" ✓ code_interpreter block started, {len(code_snippets)} code delta(s), " + f"code={full_code[:60]!r}, text={full_text[:60]!r}" + )