Merge pull request #2645 from loafoe/feat/bedrock-streaming
feat(bedrock): implement StreamingProvider for real-time token streaming
This commit is contained in:
美電球
GitHub
commit
306f96cfe3
2 changed files with 509 additions and 45 deletions
|
|
@ -135,48 +135,23 @@ func NewProvider(ctx context.Context, opts ...Option) (*Provider, error) {
|
|||
}, nil
|
||||
}
|
||||
|
||||
// Chat sends messages to AWS Bedrock using the Converse API.
|
||||
func (p *Provider) Chat(
|
||||
ctx context.Context,
|
||||
messages []Message,
|
||||
tools []ToolDefinition,
|
||||
model string,
|
||||
options map[string]any,
|
||||
) (*LLMResponse, error) {
|
||||
// Apply request timeout if context doesn't already have a deadline.
|
||||
// Use explicit timeout if set, otherwise fall back to common default.
|
||||
effectiveTimeout := p.requestTimeout
|
||||
if effectiveTimeout <= 0 {
|
||||
effectiveTimeout = common.DefaultRequestTimeout
|
||||
}
|
||||
if _, hasDeadline := ctx.Deadline(); !hasDeadline {
|
||||
var cancel context.CancelFunc
|
||||
ctx, cancel = context.WithTimeout(ctx, effectiveTimeout)
|
||||
defer cancel()
|
||||
// converseParams holds the shared request parameters for Converse and ConverseStream.
|
||||
type converseParams struct {
|
||||
messages []types.Message
|
||||
system []types.SystemContentBlock
|
||||
inferenceConfig *types.InferenceConfiguration
|
||||
toolConfig *types.ToolConfiguration
|
||||
}
|
||||
|
||||
// Build the Converse API input
|
||||
input := &bedrockruntime.ConverseInput{
|
||||
ModelId: aws.String(model),
|
||||
}
|
||||
|
||||
// Convert messages to Bedrock format
|
||||
func buildConverseParams(messages []Message, tools []ToolDefinition, options map[string]any) converseParams {
|
||||
bedrockMessages, systemPrompts := convertMessages(messages)
|
||||
input.Messages = bedrockMessages
|
||||
|
||||
// Set system prompts if any
|
||||
if len(systemPrompts) > 0 {
|
||||
input.System = systemPrompts
|
||||
}
|
||||
|
||||
// Set inference configuration only when options are provided
|
||||
var inferenceConfig *types.InferenceConfiguration
|
||||
|
||||
if maxTokens, ok := common.AsInt(options["max_tokens"]); ok && maxTokens > 0 {
|
||||
if inferenceConfig == nil {
|
||||
inferenceConfig = &types.InferenceConfiguration{}
|
||||
}
|
||||
// Clamp to int32 range to avoid overflow
|
||||
if maxTokens > math.MaxInt32 {
|
||||
maxTokens = math.MaxInt32
|
||||
}
|
||||
|
|
@ -190,23 +165,53 @@ func (p *Provider) Chat(
|
|||
inferenceConfig.Temperature = aws.Float32(float32(temp))
|
||||
}
|
||||
|
||||
if inferenceConfig != nil {
|
||||
input.InferenceConfig = inferenceConfig
|
||||
}
|
||||
|
||||
// Convert tools to Bedrock format
|
||||
// Only set ToolConfig if at least one valid tool was produced
|
||||
var toolConfig *types.ToolConfiguration
|
||||
if len(tools) > 0 {
|
||||
toolConfig := convertTools(tools)
|
||||
if len(toolConfig.Tools) > 0 {
|
||||
input.ToolConfig = toolConfig
|
||||
tc := convertTools(tools)
|
||||
if len(tc.Tools) > 0 {
|
||||
toolConfig = tc
|
||||
}
|
||||
}
|
||||
|
||||
// Call Bedrock Converse API
|
||||
return converseParams{
|
||||
messages: bedrockMessages,
|
||||
system: systemPrompts,
|
||||
inferenceConfig: inferenceConfig,
|
||||
toolConfig: toolConfig,
|
||||
}
|
||||
}
|
||||
|
||||
// Chat sends messages to AWS Bedrock using the Converse API.
|
||||
func (p *Provider) Chat(
|
||||
ctx context.Context,
|
||||
messages []Message,
|
||||
tools []ToolDefinition,
|
||||
model string,
|
||||
options map[string]any,
|
||||
) (*LLMResponse, error) {
|
||||
effectiveTimeout := p.requestTimeout
|
||||
if effectiveTimeout <= 0 {
|
||||
effectiveTimeout = common.DefaultRequestTimeout
|
||||
}
|
||||
if _, hasDeadline := ctx.Deadline(); !hasDeadline {
|
||||
var cancel context.CancelFunc
|
||||
ctx, cancel = context.WithTimeout(ctx, effectiveTimeout)
|
||||
defer cancel()
|
||||
}
|
||||
|
||||
params := buildConverseParams(messages, tools, options)
|
||||
input := &bedrockruntime.ConverseInput{
|
||||
ModelId: aws.String(model),
|
||||
Messages: params.messages,
|
||||
InferenceConfig: params.inferenceConfig,
|
||||
ToolConfig: params.toolConfig,
|
||||
}
|
||||
if len(params.system) > 0 {
|
||||
input.System = params.system
|
||||
}
|
||||
|
||||
output, err := p.client.Converse(ctx, input)
|
||||
if err != nil {
|
||||
// Check for SSO token expiration errors and provide actionable guidance
|
||||
if isSSOTokenError(err) {
|
||||
return nil, fmt.Errorf(
|
||||
"bedrock converse: AWS credentials may have expired. If using AWS SSO, run 'aws sso login' to refresh: %w",
|
||||
|
|
@ -216,10 +221,199 @@ func (p *Provider) Chat(
|
|||
return nil, fmt.Errorf("bedrock converse: %w", err)
|
||||
}
|
||||
|
||||
// Parse the response
|
||||
return parseResponse(output)
|
||||
}
|
||||
|
||||
// ChatStream sends messages to AWS Bedrock using the ConverseStream API.
|
||||
// It streams the accumulated text so far via the onChunk callback and returns the complete response.
|
||||
func (p *Provider) ChatStream(
|
||||
ctx context.Context,
|
||||
messages []Message,
|
||||
tools []ToolDefinition,
|
||||
model string,
|
||||
options map[string]any,
|
||||
onChunk func(accumulated string),
|
||||
) (*LLMResponse, error) {
|
||||
if p.requestTimeout > 0 {
|
||||
if _, hasDeadline := ctx.Deadline(); !hasDeadline {
|
||||
var cancel context.CancelFunc
|
||||
ctx, cancel = context.WithTimeout(ctx, p.requestTimeout)
|
||||
defer cancel()
|
||||
}
|
||||
}
|
||||
|
||||
params := buildConverseParams(messages, tools, options)
|
||||
input := &bedrockruntime.ConverseStreamInput{
|
||||
ModelId: aws.String(model),
|
||||
Messages: params.messages,
|
||||
InferenceConfig: params.inferenceConfig,
|
||||
ToolConfig: params.toolConfig,
|
||||
}
|
||||
if len(params.system) > 0 {
|
||||
input.System = params.system
|
||||
}
|
||||
|
||||
output, err := p.client.ConverseStream(ctx, input)
|
||||
if err != nil {
|
||||
if isSSOTokenError(err) {
|
||||
return nil, fmt.Errorf(
|
||||
"bedrock conversestream: AWS credentials may have expired. If using AWS SSO, run 'aws sso login' to refresh: %w",
|
||||
err,
|
||||
)
|
||||
}
|
||||
return nil, fmt.Errorf("bedrock conversestream: %w", err)
|
||||
}
|
||||
|
||||
return parseStreamResponse(ctx, output.GetStream(), onChunk)
|
||||
}
|
||||
|
||||
// converseStreamReader abstracts the Bedrock event stream so parseStreamResponse
|
||||
// can be unit-tested with a mock event source.
|
||||
type converseStreamReader interface {
|
||||
Events() <-chan types.ConverseStreamOutput
|
||||
Err() error
|
||||
Close() error
|
||||
}
|
||||
|
||||
// parseStreamResponse processes the ConverseStream event stream and accumulates the response.
|
||||
func parseStreamResponse(
|
||||
ctx context.Context,
|
||||
stream converseStreamReader,
|
||||
onChunk func(accumulated string),
|
||||
) (resp *LLMResponse, err error) {
|
||||
if stream == nil {
|
||||
return nil, fmt.Errorf("bedrock conversestream: nil event stream")
|
||||
}
|
||||
defer func() {
|
||||
if closeErr := stream.Close(); closeErr != nil {
|
||||
if err == nil {
|
||||
err = fmt.Errorf("bedrock conversestream: close event stream: %w", closeErr)
|
||||
} else {
|
||||
log.Printf("bedrock conversestream: close event stream: %v", closeErr)
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
var textContent strings.Builder
|
||||
finishReason := "stop"
|
||||
var usage *UsageInfo
|
||||
toolCalls := make([]ToolCall, 0)
|
||||
|
||||
// Track active tool use blocks by index
|
||||
type toolAccum struct {
|
||||
id string
|
||||
name string
|
||||
argsJSON strings.Builder
|
||||
}
|
||||
activeTools := map[int]*toolAccum{}
|
||||
|
||||
events := stream.Events()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
case event, ok := <-events:
|
||||
if !ok {
|
||||
// Stream closed
|
||||
goto done
|
||||
}
|
||||
|
||||
switch e := event.(type) {
|
||||
case *types.ConverseStreamOutputMemberContentBlockStart:
|
||||
// New content block starting
|
||||
if toolUse, ok := e.Value.Start.(*types.ContentBlockStartMemberToolUse); ok {
|
||||
activeTools[int(aws.ToInt32(e.Value.ContentBlockIndex))] = &toolAccum{
|
||||
id: aws.ToString(toolUse.Value.ToolUseId),
|
||||
name: aws.ToString(toolUse.Value.Name),
|
||||
}
|
||||
}
|
||||
|
||||
case *types.ConverseStreamOutputMemberContentBlockDelta:
|
||||
// Content delta
|
||||
switch delta := e.Value.Delta.(type) {
|
||||
case *types.ContentBlockDeltaMemberText:
|
||||
textContent.WriteString(delta.Value)
|
||||
if onChunk != nil {
|
||||
onChunk(textContent.String())
|
||||
}
|
||||
case *types.ContentBlockDeltaMemberToolUse:
|
||||
idx := int(aws.ToInt32(e.Value.ContentBlockIndex))
|
||||
if tool, exists := activeTools[idx]; exists {
|
||||
tool.argsJSON.WriteString(aws.ToString(delta.Value.Input))
|
||||
}
|
||||
}
|
||||
|
||||
case *types.ConverseStreamOutputMemberContentBlockStop:
|
||||
// Content block finished - finalize tool if it was a tool use
|
||||
idx := int(aws.ToInt32(e.Value.ContentBlockIndex))
|
||||
if tool, exists := activeTools[idx]; exists {
|
||||
args := make(map[string]any)
|
||||
argsStr := tool.argsJSON.String()
|
||||
if argsStr != "" {
|
||||
if err := json.Unmarshal([]byte(argsStr), &args); err != nil {
|
||||
log.Printf("bedrock: stream: failed to parse tool arguments for %q: %v", tool.name, err)
|
||||
args = map[string]any{"raw": argsStr}
|
||||
}
|
||||
}
|
||||
funcArgs := argsStr
|
||||
if argsJSON, marshalErr := json.Marshal(args); marshalErr == nil {
|
||||
funcArgs = string(argsJSON)
|
||||
}
|
||||
toolCalls = append(toolCalls, ToolCall{
|
||||
ID: tool.id,
|
||||
Name: tool.name,
|
||||
Arguments: args,
|
||||
Function: &FunctionCall{
|
||||
Name: tool.name,
|
||||
Arguments: funcArgs,
|
||||
},
|
||||
})
|
||||
delete(activeTools, idx)
|
||||
}
|
||||
|
||||
case *types.ConverseStreamOutputMemberMessageStop:
|
||||
// Message complete
|
||||
switch e.Value.StopReason {
|
||||
case types.StopReasonToolUse:
|
||||
finishReason = "tool_calls"
|
||||
case types.StopReasonMaxTokens:
|
||||
finishReason = "length"
|
||||
case types.StopReasonEndTurn:
|
||||
finishReason = "stop"
|
||||
case types.StopReasonStopSequence:
|
||||
finishReason = "stop"
|
||||
case types.StopReasonContentFiltered:
|
||||
finishReason = "content_filter"
|
||||
default:
|
||||
finishReason = "stop"
|
||||
}
|
||||
|
||||
case *types.ConverseStreamOutputMemberMetadata:
|
||||
// Usage metadata
|
||||
if e.Value.Usage != nil {
|
||||
usage = &UsageInfo{
|
||||
PromptTokens: int(aws.ToInt32(e.Value.Usage.InputTokens)),
|
||||
CompletionTokens: int(aws.ToInt32(e.Value.Usage.OutputTokens)),
|
||||
TotalTokens: int(aws.ToInt32(e.Value.Usage.InputTokens)) + int(aws.ToInt32(e.Value.Usage.OutputTokens)),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
done:
|
||||
if err := stream.Err(); err != nil {
|
||||
return nil, fmt.Errorf("bedrock conversestream: %w", err)
|
||||
}
|
||||
|
||||
return &LLMResponse{
|
||||
Content: textContent.String(),
|
||||
ToolCalls: toolCalls,
|
||||
FinishReason: finishReason,
|
||||
Usage: usage,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// GetDefaultModel returns an empty string as Bedrock models are user-configured.
|
||||
func (p *Provider) GetDefaultModel() string {
|
||||
return ""
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@
|
|||
package bedrock
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"testing"
|
||||
|
||||
|
|
@ -605,3 +606,272 @@ func TestIsSSOTokenError(t *testing.T) {
|
|||
})
|
||||
}
|
||||
}
|
||||
|
||||
// mockStreamReader implements bedrockruntime.ConverseStreamOutputReader for testing.
|
||||
type mockStreamReader struct {
|
||||
ch chan types.ConverseStreamOutput
|
||||
err error
|
||||
}
|
||||
|
||||
func (r *mockStreamReader) Events() <-chan types.ConverseStreamOutput { return r.ch }
|
||||
func (r *mockStreamReader) Close() error { return nil }
|
||||
func (r *mockStreamReader) Err() error { return r.err }
|
||||
|
||||
func newMockStream(events []types.ConverseStreamOutput) *bedrockruntime.ConverseStreamEventStream {
|
||||
ch := make(chan types.ConverseStreamOutput, len(events))
|
||||
for _, e := range events {
|
||||
ch <- e
|
||||
}
|
||||
close(ch)
|
||||
|
||||
return bedrockruntime.NewConverseStreamEventStream(func(es *bedrockruntime.ConverseStreamEventStream) {
|
||||
es.Reader = &mockStreamReader{ch: ch}
|
||||
})
|
||||
}
|
||||
|
||||
func TestParseStreamResponse_TextOnly(t *testing.T) {
|
||||
events := []types.ConverseStreamOutput{
|
||||
&types.ConverseStreamOutputMemberContentBlockDelta{
|
||||
Value: types.ContentBlockDeltaEvent{
|
||||
Delta: &types.ContentBlockDeltaMemberText{Value: "Hello "},
|
||||
ContentBlockIndex: aws.Int32(0),
|
||||
},
|
||||
},
|
||||
&types.ConverseStreamOutputMemberContentBlockDelta{
|
||||
Value: types.ContentBlockDeltaEvent{
|
||||
Delta: &types.ContentBlockDeltaMemberText{Value: "World"},
|
||||
ContentBlockIndex: aws.Int32(0),
|
||||
},
|
||||
},
|
||||
&types.ConverseStreamOutputMemberMessageStop{
|
||||
Value: types.MessageStopEvent{StopReason: types.StopReasonEndTurn},
|
||||
},
|
||||
&types.ConverseStreamOutputMemberMetadata{
|
||||
Value: types.ConverseStreamMetadataEvent{
|
||||
Usage: &types.TokenUsage{
|
||||
InputTokens: aws.Int32(10),
|
||||
OutputTokens: aws.Int32(5),
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
var chunks []string
|
||||
stream := newMockStream(events)
|
||||
resp, err := parseStreamResponse(context.Background(), stream, func(accumulated string) {
|
||||
chunks = append(chunks, accumulated)
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "Hello World", resp.Content)
|
||||
assert.Equal(t, "stop", resp.FinishReason)
|
||||
assert.Empty(t, resp.ToolCalls)
|
||||
require.NotNil(t, resp.Usage)
|
||||
assert.Equal(t, 10, resp.Usage.PromptTokens)
|
||||
assert.Equal(t, 5, resp.Usage.CompletionTokens)
|
||||
assert.Equal(t, 15, resp.Usage.TotalTokens)
|
||||
assert.Equal(t, []string{"Hello ", "Hello World"}, chunks)
|
||||
}
|
||||
|
||||
func TestParseStreamResponse_ToolCall(t *testing.T) {
|
||||
events := []types.ConverseStreamOutput{
|
||||
&types.ConverseStreamOutputMemberContentBlockStart{
|
||||
Value: types.ContentBlockStartEvent{
|
||||
ContentBlockIndex: aws.Int32(0),
|
||||
Start: &types.ContentBlockStartMemberToolUse{
|
||||
Value: types.ToolUseBlockStart{
|
||||
ToolUseId: aws.String("call_1"),
|
||||
Name: aws.String("search"),
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
&types.ConverseStreamOutputMemberContentBlockDelta{
|
||||
Value: types.ContentBlockDeltaEvent{
|
||||
ContentBlockIndex: aws.Int32(0),
|
||||
Delta: &types.ContentBlockDeltaMemberToolUse{
|
||||
Value: types.ToolUseBlockDelta{Input: aws.String(`{"q":`)},
|
||||
},
|
||||
},
|
||||
},
|
||||
&types.ConverseStreamOutputMemberContentBlockDelta{
|
||||
Value: types.ContentBlockDeltaEvent{
|
||||
ContentBlockIndex: aws.Int32(0),
|
||||
Delta: &types.ContentBlockDeltaMemberToolUse{
|
||||
Value: types.ToolUseBlockDelta{Input: aws.String(`"test"}`)},
|
||||
},
|
||||
},
|
||||
},
|
||||
&types.ConverseStreamOutputMemberContentBlockStop{
|
||||
Value: types.ContentBlockStopEvent{ContentBlockIndex: aws.Int32(0)},
|
||||
},
|
||||
&types.ConverseStreamOutputMemberMessageStop{
|
||||
Value: types.MessageStopEvent{StopReason: types.StopReasonToolUse},
|
||||
},
|
||||
}
|
||||
|
||||
stream := newMockStream(events)
|
||||
resp, err := parseStreamResponse(context.Background(), stream, nil)
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "tool_calls", resp.FinishReason)
|
||||
require.Len(t, resp.ToolCalls, 1)
|
||||
assert.Equal(t, "call_1", resp.ToolCalls[0].ID)
|
||||
assert.Equal(t, "search", resp.ToolCalls[0].Name)
|
||||
assert.Equal(t, map[string]any{"q": "test"}, resp.ToolCalls[0].Arguments)
|
||||
require.NotNil(t, resp.ToolCalls[0].Function)
|
||||
assert.Equal(t, "search", resp.ToolCalls[0].Function.Name)
|
||||
assert.Equal(t, `{"q":"test"}`, resp.ToolCalls[0].Function.Arguments)
|
||||
}
|
||||
|
||||
func TestParseStreamResponse_TextAndToolCall(t *testing.T) {
|
||||
events := []types.ConverseStreamOutput{
|
||||
&types.ConverseStreamOutputMemberContentBlockDelta{
|
||||
Value: types.ContentBlockDeltaEvent{
|
||||
ContentBlockIndex: aws.Int32(0),
|
||||
Delta: &types.ContentBlockDeltaMemberText{Value: "Let me search that."},
|
||||
},
|
||||
},
|
||||
&types.ConverseStreamOutputMemberContentBlockStart{
|
||||
Value: types.ContentBlockStartEvent{
|
||||
ContentBlockIndex: aws.Int32(1),
|
||||
Start: &types.ContentBlockStartMemberToolUse{
|
||||
Value: types.ToolUseBlockStart{
|
||||
ToolUseId: aws.String("call_2"),
|
||||
Name: aws.String("web"),
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
&types.ConverseStreamOutputMemberContentBlockDelta{
|
||||
Value: types.ContentBlockDeltaEvent{
|
||||
ContentBlockIndex: aws.Int32(1),
|
||||
Delta: &types.ContentBlockDeltaMemberToolUse{
|
||||
Value: types.ToolUseBlockDelta{Input: aws.String(`{"url":"https://example.com"}`)},
|
||||
},
|
||||
},
|
||||
},
|
||||
&types.ConverseStreamOutputMemberContentBlockStop{
|
||||
Value: types.ContentBlockStopEvent{ContentBlockIndex: aws.Int32(1)},
|
||||
},
|
||||
&types.ConverseStreamOutputMemberMessageStop{
|
||||
Value: types.MessageStopEvent{StopReason: types.StopReasonToolUse},
|
||||
},
|
||||
}
|
||||
|
||||
var chunks []string
|
||||
stream := newMockStream(events)
|
||||
resp, err := parseStreamResponse(context.Background(), stream, func(accumulated string) {
|
||||
chunks = append(chunks, accumulated)
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "Let me search that.", resp.Content)
|
||||
assert.Equal(t, "tool_calls", resp.FinishReason)
|
||||
require.Len(t, resp.ToolCalls, 1)
|
||||
assert.Equal(t, "web", resp.ToolCalls[0].Name)
|
||||
assert.Equal(t, []string{"Let me search that."}, chunks)
|
||||
}
|
||||
|
||||
func TestParseStreamResponse_ContextCancelled(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
// Use an unbuffered channel with no events so ctx.Done() is the only ready case.
|
||||
ch := make(chan types.ConverseStreamOutput)
|
||||
|
||||
stream := bedrockruntime.NewConverseStreamEventStream(func(es *bedrockruntime.ConverseStreamEventStream) {
|
||||
es.Reader = &mockStreamReader{ch: ch}
|
||||
})
|
||||
|
||||
_, err := parseStreamResponse(ctx, stream, nil)
|
||||
assert.ErrorIs(t, err, context.Canceled)
|
||||
}
|
||||
|
||||
func TestParseStreamResponse_InvalidToolJSON(t *testing.T) {
|
||||
events := []types.ConverseStreamOutput{
|
||||
&types.ConverseStreamOutputMemberContentBlockStart{
|
||||
Value: types.ContentBlockStartEvent{
|
||||
ContentBlockIndex: aws.Int32(0),
|
||||
Start: &types.ContentBlockStartMemberToolUse{
|
||||
Value: types.ToolUseBlockStart{
|
||||
ToolUseId: aws.String("call_bad"),
|
||||
Name: aws.String("broken"),
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
&types.ConverseStreamOutputMemberContentBlockDelta{
|
||||
Value: types.ContentBlockDeltaEvent{
|
||||
ContentBlockIndex: aws.Int32(0),
|
||||
Delta: &types.ContentBlockDeltaMemberToolUse{
|
||||
Value: types.ToolUseBlockDelta{Input: aws.String(`{not valid json`)},
|
||||
},
|
||||
},
|
||||
},
|
||||
&types.ConverseStreamOutputMemberContentBlockStop{
|
||||
Value: types.ContentBlockStopEvent{ContentBlockIndex: aws.Int32(0)},
|
||||
},
|
||||
&types.ConverseStreamOutputMemberMessageStop{
|
||||
Value: types.MessageStopEvent{StopReason: types.StopReasonToolUse},
|
||||
},
|
||||
}
|
||||
|
||||
stream := newMockStream(events)
|
||||
resp, err := parseStreamResponse(context.Background(), stream, nil)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Len(t, resp.ToolCalls, 1)
|
||||
assert.Equal(t, map[string]any{"raw": `{not valid json`}, resp.ToolCalls[0].Arguments)
|
||||
assert.JSONEq(t, `{"raw":"{not valid json"}`, resp.ToolCalls[0].Function.Arguments)
|
||||
}
|
||||
|
||||
func TestParseStreamResponse_DefaultFinishReason(t *testing.T) {
|
||||
events := []types.ConverseStreamOutput{
|
||||
&types.ConverseStreamOutputMemberContentBlockDelta{
|
||||
Value: types.ContentBlockDeltaEvent{
|
||||
Delta: &types.ContentBlockDeltaMemberText{Value: "partial"},
|
||||
ContentBlockIndex: aws.Int32(0),
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
stream := newMockStream(events)
|
||||
resp, err := parseStreamResponse(context.Background(), stream, nil)
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "stop", resp.FinishReason)
|
||||
}
|
||||
|
||||
func TestParseStreamResponse_NilStream(t *testing.T) {
|
||||
_, err := parseStreamResponse(context.Background(), nil, nil)
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "nil event stream")
|
||||
}
|
||||
|
||||
func TestParseStreamResponse_StopReasons(t *testing.T) {
|
||||
tests := []struct {
|
||||
reason types.StopReason
|
||||
expected string
|
||||
}{
|
||||
{types.StopReasonEndTurn, "stop"},
|
||||
{types.StopReasonMaxTokens, "length"},
|
||||
{types.StopReasonToolUse, "tool_calls"},
|
||||
{types.StopReasonStopSequence, "stop"},
|
||||
{types.StopReasonContentFiltered, "content_filter"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(string(tt.reason), func(t *testing.T) {
|
||||
events := []types.ConverseStreamOutput{
|
||||
&types.ConverseStreamOutputMemberMessageStop{
|
||||
Value: types.MessageStopEvent{StopReason: tt.reason},
|
||||
},
|
||||
}
|
||||
stream := newMockStream(events)
|
||||
resp, err := parseStreamResponse(context.Background(), stream, nil)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, tt.expected, resp.FinishReason)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue