From e30687c368e2b86ad86f253ca052089065091acf Mon Sep 17 00:00:00 2001 From: helloxjade <52953016+helloxjade@users.noreply.github.com> Date: Thu, 15 Oct 2026 14:40:47 +0800 Subject: [PATCH 1/2] fix(adk): parse Mistral reasoning content chunks Decode typed content arrays from the raw OpenAI response in streaming and non-streaming completions. Keep thinking separate from answer text and preserve native thinking chunks for history replay. Fixes #3126 Signed-off-by: helloxjade <52953016+helloxjade@users.noreply.github.com> --- go/adk/pkg/models/openai_adk.go | 49 +++-- go/adk/pkg/models/openai_content.go | 95 +++++++++ go/adk/pkg/models/openai_content_test.go | 237 +++++++++++++++++++++++ 3 files changed, 363 insertions(+), 18 deletions(-) create mode 100644 go/adk/pkg/models/openai_content.go create mode 100644 go/adk/pkg/models/openai_content_test.go diff --git a/go/adk/pkg/models/openai_adk.go b/go/adk/pkg/models/openai_adk.go index 11bab883d0..492690ee95 100644 --- a/go/adk/pkg/models/openai_adk.go +++ b/go/adk/pkg/models/openai_adk.go @@ -262,7 +262,7 @@ func genaiContentsToOpenAIMessages(contents []*genai.Content, config *genai.Gene if part == nil { continue } - if part.Text != "" { + if part.Text != "" && !part.Thought { textParts = append(textParts, part.Text) } else if part.FunctionCall != nil { functionCalls = append(functionCalls, part.FunctionCall) @@ -273,7 +273,8 @@ func genaiContentsToOpenAIMessages(contents []*genai.Content, config *genai.Gene } } - if len(functionCalls) > 0 && (role == openAIRoleModel || role == openAIRoleAssistant) { + mistralContent := mistralAssistantContent(content.Parts) + if (len(functionCalls) > 0 || len(mistralContent) > 0) && (role == openAIRoleModel || role == openAIRoleAssistant) { toolCalls := make([]openai.ChatCompletionMessageToolCallUnionParam, 0, len(functionCalls)) var toolResponseMessages []openai.ChatCompletionMessageParamUnion for _, fc := range functionCalls { @@ -313,7 +314,9 @@ func genaiContentsToOpenAIMessages(contents []*genai.Content, config *genai.Gene Role: constant.Assistant("assistant"), ToolCalls: toolCalls, } - if len(textParts) > 0 { + if len(mistralContent) > 0 { + asst.Content.OfArrayOfContentParts = mistralContent + } else if len(textParts) > 0 { asst.Content.OfString = param.NewOpt(textContent) } messages = append(messages, openai.ChatCompletionMessageParamUnion{OfAssistant: &asst}) @@ -383,6 +386,7 @@ func runStreaming(ctx context.Context, m *OpenAIModel, params openai.ChatComplet stream := m.Client.Chat.Completions.NewStreaming(ctx, params, openAIPassthroughOpts(ctx, m)...) defer stream.Close() + var aggregatedParts []*genai.Part var aggregatedText strings.Builder toolCallsAcc := make(map[int64]map[string]any) var finishReason string @@ -400,12 +404,27 @@ func runStreaming(ctx context.Context, m *OpenAIModel, params openai.ChatComplet } choice := chunk.Choices[0] delta := choice.Delta - if delta.Content != "" { - aggregatedText.WriteString(delta.Content) + parts, err := openAIContentToParts(delta.Content, delta.JSON.Content.Raw()) + if err != nil { + yield(&model.LLMResponse{ErrorCode: openAIContentError, ErrorMessage: err.Error()}, nil) + return + } + for _, part := range parts { + if part.Thought { + if aggregatedText.Len() > 0 { + aggregatedParts = append(aggregatedParts, &genai.Part{Text: aggregatedText.String()}) + aggregatedText.Reset() + } + aggregatedParts = append(aggregatedParts, part) + } else { + aggregatedText.WriteString(part.Text) + } + } + if len(parts) > 0 { if !yield(&model.LLMResponse{ Partial: true, TurnComplete: choice.FinishReason != "", - Content: &genai.Content{Role: string(genai.RoleModel), Parts: []*genai.Part{{Text: delta.Content}}}, + Content: &genai.Content{Role: string(genai.RoleModel), Parts: parts}, }, nil) { return } @@ -449,10 +468,9 @@ func runStreaming(ctx context.Context, m *OpenAIModel, params openai.ChatComplet indices = append(indices, k) } slices.Sort(indices) - finalParts := make([]*genai.Part, 0, 1+nToolCalls) - text := aggregatedText.String() - if text != "" { - finalParts = append(finalParts, &genai.Part{Text: text}) + finalParts := aggregatedParts + if aggregatedText.Len() > 0 { + finalParts = append(finalParts, &genai.Part{Text: aggregatedText.String()}) } for _, idx := range indices { tc := toolCallsAcc[idx] @@ -505,14 +523,9 @@ func runNonStreaming(ctx context.Context, m *OpenAIModel, params openai.ChatComp func chatCompletionToLLMResponse(completion *openai.ChatCompletion) *model.LLMResponse { choice := completion.Choices[0] msg := choice.Message - nParts := 0 - if msg.Content != "" { - nParts++ - } - nParts += len(msg.ToolCalls) - parts := make([]*genai.Part, 0, nParts) - if msg.Content != "" { - parts = append(parts, &genai.Part{Text: msg.Content}) + parts, err := openAIContentToParts(msg.Content, msg.JSON.Content.Raw()) + if err != nil { + return &model.LLMResponse{ErrorCode: openAIContentError, ErrorMessage: err.Error()} } for _, tc := range msg.ToolCalls { if tc.Type == openAIToolTypeFunction && tc.Function.Name != "" { diff --git a/go/adk/pkg/models/openai_content.go b/go/adk/pkg/models/openai_content.go new file mode 100644 index 0000000000..212b3693bb --- /dev/null +++ b/go/adk/pkg/models/openai_content.go @@ -0,0 +1,95 @@ +package models + +import ( + "encoding/json" + "fmt" + "strings" + + "github.com/openai/openai-go/v3" + "github.com/openai/openai-go/v3/packages/param" + "google.golang.org/genai" +) + +const ( + openAIContentError = "CONTENT_ERROR" + openAIChunkTypeText = "text" + openAIChunkTypeThinking = "thinking" + mistralThinkingChunkKey = "mistral_thinking_chunk" +) + +type openAITextChunk struct { + Type string `json:"type"` + Text string `json:"text"` +} + +// The SDK models content as a string, but Mistral returns typed chunk arrays. +// Inspect the wire JSON so a string containing JSON remains ordinary answer text. +func openAIContentToParts(text, raw string) ([]*genai.Part, error) { + if !strings.HasPrefix(strings.TrimSpace(raw), "[") { + if text == "" { + return nil, nil + } + return []*genai.Part{{Text: text}}, nil + } + var chunks []json.RawMessage + if err := json.Unmarshal([]byte(raw), &chunks); err != nil { + return nil, fmt.Errorf("failed to decode OpenAI content chunks: %w", err) + } + parts := make([]*genai.Part, 0, len(chunks)) + for _, rawChunk := range chunks { + var chunk struct { + openAITextChunk + Thinking []openAITextChunk `json:"thinking"` + } + if err := json.Unmarshal(rawChunk, &chunk); err != nil { + return nil, fmt.Errorf("failed to decode OpenAI content chunk: %w", err) + } + switch chunk.Type { + case openAIChunkTypeText: + if chunk.Text != "" { + parts = append(parts, &genai.Part{Text: chunk.Text}) + } + case openAIChunkTypeThinking: + var thought strings.Builder + for _, inner := range chunk.Thinking { + if inner.Type != openAIChunkTypeText { + return nil, fmt.Errorf("unsupported OpenAI thinking chunk type %q", inner.Type) + } + thought.WriteString(inner.Text) + } + parts = append(parts, &genai.Part{ + Text: thought.String(), Thought: true, + // Preserve the native block, including closed, across persistence. + PartMetadata: map[string]any{mistralThinkingChunkKey: string(rawChunk)}, + }) + default: + return nil, fmt.Errorf("unsupported OpenAI content chunk type %q", chunk.Type) + } + } + return parts, nil +} + +// Only native Mistral thoughts opt a message into Mistral's chunk-list format. +func mistralAssistantContent(parts []*genai.Part) []openai.ChatCompletionAssistantMessageParamContentArrayOfContentPartUnion { + var chunks []openai.ChatCompletionAssistantMessageParamContentArrayOfContentPartUnion + var hasThinking bool + for _, part := range parts { + if part == nil { + continue + } + if part.Thought { + if raw, ok := part.PartMetadata[mistralThinkingChunkKey].(string); ok { + hasThinking = true + chunks = append(chunks, param.Override[openai.ChatCompletionAssistantMessageParamContentArrayOfContentPartUnion](json.RawMessage(raw))) + } + } else if part.Text != "" { + chunks = append(chunks, openai.ChatCompletionAssistantMessageParamContentArrayOfContentPartUnion{ + OfText: &openai.ChatCompletionContentPartTextParam{Text: part.Text}, + }) + } + } + if !hasThinking { + return nil + } + return chunks +} diff --git a/go/adk/pkg/models/openai_content_test.go b/go/adk/pkg/models/openai_content_test.go new file mode 100644 index 0000000000..f43a03b9e4 --- /dev/null +++ b/go/adk/pkg/models/openai_content_test.go @@ -0,0 +1,237 @@ +package models + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/openai/openai-go/v3" + "github.com/openai/openai-go/v3/option" + "github.com/stretchr/testify/require" + "google.golang.org/adk/v2/model" + "google.golang.org/genai" +) + +func TestOpenAIMistralContent(t *testing.T) { + for _, stream := range []bool{false, true} { + t.Run(fmt.Sprintf("stream=%t", stream), func(t *testing.T) { + thinking := `{"type":"thinking","thinking":[{"type":"text","text":"Let me think."}],"closed":true}` + answer := `{"type":"text","text":"{\"answer\":"}` + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if stream { + w.Header().Set("Content-Type", "text/event-stream") + for _, content := range []string{ + `[{"type":"thinking","thinking":[{"type":"text","text":"Let me "}],"closed":false}]`, + `[{"type":"thinking","thinking":[{"type":"text","text":"think."}],"closed":true},` + answer + `]`, + `"42}"`, + } { + fmt.Fprintf(w, "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":%s}}]}\n\n", content) + } + fmt.Fprint(w, "data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":3,\"completion_tokens\":5,\"total_tokens\":8}}\n\ndata: [DONE]\n\n") + } else { + w.Header().Set("Content-Type", "application/json") + fmt.Fprintf(w, `{"choices":[{"message":{"role":"assistant","content":[%s,{"type":"text","text":"{\"answer\":42}"}]},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":5,"total_tokens":8}}`, thinking) + } + })) + t.Cleanup(server.Close) + m := &OpenAIModel{ + Config: &OpenAIConfig{Model: "mistral-test"}, + Client: openai.NewClient(option.WithBaseURL(server.URL), option.WithAPIKey("test")), + } + var final *model.LLMResponse + var partialThought, partialAnswer strings.Builder + for resp, err := range m.GenerateContent(context.Background(), &model.LLMRequest{ + Contents: []*genai.Content{genai.NewContentFromText("hello", genai.RoleUser)}, + }, stream) { + require.NoError(t, err) + require.Empty(t, resp.ErrorCode, resp.ErrorMessage) + if !resp.Partial { + final = resp + continue + } + for _, part := range resp.Content.Parts { + if part.Thought { + partialThought.WriteString(part.Text) + } else { + partialAnswer.WriteString(part.Text) + } + } + } + require.NotNil(t, final) + require.True(t, final.TurnComplete) + require.Equal(t, genai.FinishReasonStop, final.FinishReason) + require.Equal(t, int32(8), final.UsageMetadata.TotalTokenCount) + var thoughtText, answerText strings.Builder + for _, part := range final.Content.Parts { + if part.Thought { + thoughtText.WriteString(part.Text) + } else { + answerText.WriteString(part.Text) + } + } + require.JSONEq(t, `{"answer":42}`, answerText.String()) + require.Equal(t, "Let me think.", thoughtText.String()) + if stream { + require.Equal(t, thoughtText.String(), partialThought.String()) + require.Equal(t, answerText.String(), partialAnswer.String()) + } + + // Persisted ADK content must retain the native chunks for the next turn. + stored, err := json.Marshal(final.Content) + require.NoError(t, err) + var history genai.Content + require.NoError(t, json.Unmarshal(stored, &history)) + messages, _ := genaiContentsToOpenAIMessages([]*genai.Content{&history}, nil) + require.Len(t, messages, 1) + require.NotNil(t, messages[0].OfAssistant) + wire, err := json.Marshal(messages[0]) + require.NoError(t, err) + var replay struct { + Role string `json:"role"` + Content []struct { + Type string `json:"type"` + Text string `json:"text"` + Thinking []struct { + Text string `json:"text"` + } `json:"thinking"` + Closed *bool `json:"closed"` + } `json:"content"` + } + require.NoError(t, json.Unmarshal(wire, &replay)) + require.Equal(t, "assistant", replay.Role) + if stream { + require.Len(t, replay.Content, 3) + require.NotNil(t, replay.Content[0].Closed) + require.False(t, *replay.Content[0].Closed) + require.NotNil(t, replay.Content[1].Closed) + require.True(t, *replay.Content[1].Closed) + } else { + require.Len(t, replay.Content, 2) + require.NotNil(t, replay.Content[0].Closed) + require.True(t, *replay.Content[0].Closed) + } + var replayThought, replayAnswer strings.Builder + for _, chunk := range replay.Content { + switch chunk.Type { + case "thinking": + require.NotNil(t, chunk.Closed) + for _, inner := range chunk.Thinking { + replayThought.WriteString(inner.Text) + } + case "text": + replayAnswer.WriteString(chunk.Text) + default: + t.Fatalf("unexpected content chunk: %s", chunk.Type) + } + } + require.Equal(t, thoughtText.String(), replayThought.String()) + require.Equal(t, answerText.String(), replayAnswer.String()) + }) + } +} + +func TestOpenAIStreamingContentVariants(t *testing.T) { + for _, tt := range []struct { + name string + content []string + want string + wantErr bool + }{ + {name: "plain text and empty deltas", content: []string{`null`, `"hello"`, `[]`, `" world"`, `""`}, want: "hello world"}, + {name: "literal JSON string", content: []string{`"[{\"type\":\"text\",\"text\":\"literal\"}]"`}, want: `[{"type":"text","text":"literal"}]`}, + {name: "invalid thinking", content: []string{`[{"type":"thinking","thinking":"invalid"}]`}, wantErr: true}, + } { + t.Run(tt.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + for _, content := range tt.content { + fmt.Fprintf(w, "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":%s}}]}\n\n", content) + } + fmt.Fprint(w, "data: [DONE]\n\n") + })) + t.Cleanup(server.Close) + m := &OpenAIModel{Client: openai.NewClient(option.WithBaseURL(server.URL), option.WithAPIKey("test"))} + var responses []*model.LLMResponse + runStreaming(context.Background(), m, openai.ChatCompletionNewParams{}, func(resp *model.LLMResponse, err error) bool { + require.NoError(t, err) + responses = append(responses, resp) + return true + }) + require.NotEmpty(t, responses) + final := responses[len(responses)-1] + if tt.wantErr { + require.NotEmpty(t, final.ErrorCode) + require.NotEmpty(t, final.ErrorMessage) + require.Nil(t, final.Content) + } else { + require.Empty(t, final.ErrorCode) + require.False(t, final.Partial) + require.Len(t, final.Content.Parts, 1) + require.Equal(t, tt.want, final.Content.Parts[0].Text) + require.False(t, final.Content.Parts[0].Thought) + } + }) + } +} + +func TestChatCompletionContentStrings(t *testing.T) { + for _, text := range []string{"hello", `[{"type":"text","text":"literal JSON answer"}]`, ""} { + t.Run(text, func(t *testing.T) { + content, err := json.Marshal(text) + require.NoError(t, err) + var completion openai.ChatCompletion + require.NoError(t, json.Unmarshal([]byte(`{"choices":[{"message":{"role":"assistant","content":`+string(content)+`}}]}`), &completion)) + resp := chatCompletionToLLMResponse(&completion) + if text == "" { + require.Empty(t, resp.Content.Parts) + } else { + require.Len(t, resp.Content.Parts, 1) + require.Equal(t, text, resp.Content.Parts[0].Text) + require.False(t, resp.Content.Parts[0].Thought) + } + }) + } +} + +func TestMistralContentWithToolCalls(t *testing.T) { + content := `[{"type":"thinking","thinking":[{"type":"text","text":"Use a tool."}],"closed":true}]` + var completion openai.ChatCompletion + require.NoError(t, json.Unmarshal([]byte(`{"choices":[{"message":{"role":"assistant","content":`+content+`,"tool_calls":[{"type":"function","id":"call_1","function":{"name":"lookup","arguments":"{}"}}]},"finish_reason":"tool_calls"}]}`), &completion)) + resp := chatCompletionToLLMResponse(&completion) + require.Len(t, resp.Content.Parts, 2) + require.True(t, resp.Content.Parts[0].Thought) + require.Equal(t, "lookup", resp.Content.Parts[1].FunctionCall.Name) + result := genai.NewPartFromFunctionResponse("lookup", map[string]any{"result": "found"}) + result.FunctionResponse.ID = "call_1" + messages, _ := genaiContentsToOpenAIMessages([]*genai.Content{ + resp.Content, + {Role: "user", Parts: []*genai.Part{result}}, + }, nil) + require.Len(t, messages, 2) + wire, err := json.Marshal(messages[0]) + require.NoError(t, err) + require.JSONEq(t, `{"role":"assistant","content":`+content+`,"tool_calls":[{"type":"function","id":"call_1","function":{"name":"lookup","arguments":"{}"}}]}`, string(wire)) + require.Equal(t, "call_1", messages[1].OfTool.ToolCallID) + require.Equal(t, "found", messages[1].OfTool.Content.OfString.Value) +} + +func TestChatCompletionInvalidContentChunks(t *testing.T) { + for _, content := range []string{ + `[{"type":"thinking","thinking":"invalid"}]`, + `[{"type":"text","text":42}]`, + `[{"type":"unknown","text":"must not become an answer"}]`, + } { + t.Run(content, func(t *testing.T) { + var completion openai.ChatCompletion + require.NoError(t, json.Unmarshal([]byte(`{"choices":[{"message":{"content":`+content+`}}]}`), &completion)) + resp := chatCompletionToLLMResponse(&completion) + require.NotEmpty(t, resp.ErrorCode) + require.NotEmpty(t, resp.ErrorMessage) + require.Nil(t, resp.Content) + }) + } +} From 64262e127e938c6d170a42e346109bb787d0b642 Mon Sep 17 00:00:00 2001 From: helloxjade <52953016+helloxjade@users.noreply.github.com> Date: Sat, 17 Oct 2026 17:52:49 +0800 Subject: [PATCH 2/2] fix(adk): aggregate streamed thoughts and tolerate extensions Persist streamed Mistral thinking as one complete closed block before history replay. Skip unsupported non-text chunks without aborting valid answers, while still rejecting malformed known chunks. Signed-off-by: helloxjade <52953016+helloxjade@users.noreply.github.com> --- go/adk/pkg/models/openai_adk.go | 7 ++ go/adk/pkg/models/openai_content.go | 81 +++++++++++++++++++++--- go/adk/pkg/models/openai_content_test.go | 41 +++++++----- 3 files changed, 105 insertions(+), 24 deletions(-) diff --git a/go/adk/pkg/models/openai_adk.go b/go/adk/pkg/models/openai_adk.go index 492690ee95..b3717e7b3e 100644 --- a/go/adk/pkg/models/openai_adk.go +++ b/go/adk/pkg/models/openai_adk.go @@ -420,6 +420,8 @@ func runStreaming(ctx context.Context, m *OpenAIModel, params openai.ChatComplet aggregatedText.WriteString(part.Text) } } + // Closing-only thinking deltas carry replay state, not displayable content. + parts = slices.DeleteFunc(parts, func(part *genai.Part) bool { return part.Text == "" }) if len(parts) > 0 { if !yield(&model.LLMResponse{ Partial: true, @@ -472,6 +474,11 @@ func runStreaming(ctx context.Context, m *OpenAIModel, params openai.ChatComplet if aggregatedText.Len() > 0 { finalParts = append(finalParts, &genai.Part{Text: aggregatedText.String()}) } + finalParts, err := mergeMistralThinkingParts(finalParts) + if err != nil { + yield(&model.LLMResponse{ErrorCode: openAIContentError, ErrorMessage: err.Error()}, nil) + return + } for _, idx := range indices { tc := toolCallsAcc[idx] argsStr, _ := tc["arguments"].(string) diff --git a/go/adk/pkg/models/openai_content.go b/go/adk/pkg/models/openai_content.go index 212b3693bb..f6c2dac3ae 100644 --- a/go/adk/pkg/models/openai_content.go +++ b/go/adk/pkg/models/openai_content.go @@ -22,6 +22,12 @@ type openAITextChunk struct { Text string `json:"text"` } +type mistralThinkingChunk struct { + Type string `json:"type"` + Thinking []json.RawMessage `json:"thinking"` + Closed bool `json:"closed"` +} + // The SDK models content as a string, but Mistral returns typed chunk arrays. // Inspect the wire JSON so a string containing JSON remains ordinary answer text. func openAIContentToParts(text, raw string) ([]*genai.Part, error) { @@ -38,37 +44,94 @@ func openAIContentToParts(text, raw string) ([]*genai.Part, error) { parts := make([]*genai.Part, 0, len(chunks)) for _, rawChunk := range chunks { var chunk struct { - openAITextChunk - Thinking []openAITextChunk `json:"thinking"` + Type string `json:"type"` } if err := json.Unmarshal(rawChunk, &chunk); err != nil { return nil, fmt.Errorf("failed to decode OpenAI content chunk: %w", err) } switch chunk.Type { case openAIChunkTypeText: - if chunk.Text != "" { - parts = append(parts, &genai.Part{Text: chunk.Text}) + var textChunk openAITextChunk + if err := json.Unmarshal(rawChunk, &textChunk); err != nil { + return nil, fmt.Errorf("failed to decode OpenAI text chunk: %w", err) + } + if textChunk.Text != "" { + parts = append(parts, &genai.Part{Text: textChunk.Text}) } case openAIChunkTypeThinking: + var thinkingChunk mistralThinkingChunk + if err := json.Unmarshal(rawChunk, &thinkingChunk); err != nil { + return nil, fmt.Errorf("failed to decode OpenAI thinking chunk: %w", err) + } var thought strings.Builder - for _, inner := range chunk.Thinking { - if inner.Type != openAIChunkTypeText { - return nil, fmt.Errorf("unsupported OpenAI thinking chunk type %q", inner.Type) + for _, rawInner := range thinkingChunk.Thinking { + var inner struct { + Type string `json:"type"` + } + if err := json.Unmarshal(rawInner, &inner); err != nil { + return nil, fmt.Errorf("failed to decode OpenAI thinking content: %w", err) + } + if inner.Type == "" { + return nil, fmt.Errorf("OpenAI thinking content is missing its type") + } + if inner.Type == openAIChunkTypeText { + var textChunk openAITextChunk + if err := json.Unmarshal(rawInner, &textChunk); err != nil { + return nil, fmt.Errorf("failed to decode OpenAI thinking text: %w", err) + } + thought.WriteString(textChunk.Text) } - thought.WriteString(inner.Text) } parts = append(parts, &genai.Part{ Text: thought.String(), Thought: true, // Preserve the native block, including closed, across persistence. PartMetadata: map[string]any{mistralThinkingChunkKey: string(rawChunk)}, }) + case "": + return nil, fmt.Errorf("OpenAI content chunk is missing its type") default: - return nil, fmt.Errorf("unsupported OpenAI content chunk type %q", chunk.Type) + // Unsupported non-text extensions must not abort the answer or become text. + continue } } return parts, nil } +// Complete a streamed thinking block before persisting or replaying the turn. +func mergeMistralThinkingParts(parts []*genai.Part) ([]*genai.Part, error) { + var merged *genai.Part + var text strings.Builder + var thinking []json.RawMessage + var result []*genai.Part + for _, part := range parts { + if !part.Thought { + result = append(result, part) + continue + } + var chunk mistralThinkingChunk + raw, _ := part.PartMetadata[mistralThinkingChunkKey].(string) + if err := json.Unmarshal([]byte(raw), &chunk); err != nil { + return nil, fmt.Errorf("failed to decode streamed Mistral thinking chunk: %w", err) + } + thinking = append(thinking, chunk.Thinking...) + text.WriteString(part.Text) + if merged == nil { + merged = &genai.Part{Thought: true} + result = append(result, merged) + } + } + if merged != nil { + chunk := mistralThinkingChunk{Type: openAIChunkTypeThinking, Thinking: thinking, Closed: true} + raw, err := json.Marshal(chunk) + if err != nil { + return nil, fmt.Errorf("failed to encode complete Mistral thinking chunk: %w", err) + } + merged.Text = text.String() + merged.PartMetadata = map[string]any{mistralThinkingChunkKey: string(raw)} + } + return result, nil +} + // Only native Mistral thoughts opt a message into Mistral's chunk-list format. func mistralAssistantContent(parts []*genai.Part) []openai.ChatCompletionAssistantMessageParamContentArrayOfContentPartUnion { var chunks []openai.ChatCompletionAssistantMessageParamContentArrayOfContentPartUnion diff --git a/go/adk/pkg/models/openai_content_test.go b/go/adk/pkg/models/openai_content_test.go index f43a03b9e4..bf67b7ae40 100644 --- a/go/adk/pkg/models/openai_content_test.go +++ b/go/adk/pkg/models/openai_content_test.go @@ -21,20 +21,23 @@ func TestOpenAIMistralContent(t *testing.T) { t.Run(fmt.Sprintf("stream=%t", stream), func(t *testing.T) { thinking := `{"type":"thinking","thinking":[{"type":"text","text":"Let me think."}],"closed":true}` answer := `{"type":"text","text":"{\"answer\":"}` + reference := `{"type":"reference","reference_ids":[1]}` server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if stream { w.Header().Set("Content-Type", "text/event-stream") for _, content := range []string{ `[{"type":"thinking","thinking":[{"type":"text","text":"Let me "}],"closed":false}]`, - `[{"type":"thinking","thinking":[{"type":"text","text":"think."}],"closed":true},` + answer + `]`, + `[{"type":"thinking","thinking":[{"type":"text","text":"think."}],"closed":false}]`, + `[{"type":"thinking","thinking":[],"closed":true},` + answer + `]`, `"42}"`, + `[` + reference + `]`, } { fmt.Fprintf(w, "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":%s}}]}\n\n", content) } fmt.Fprint(w, "data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":3,\"completion_tokens\":5,\"total_tokens\":8}}\n\ndata: [DONE]\n\n") } else { w.Header().Set("Content-Type", "application/json") - fmt.Fprintf(w, `{"choices":[{"message":{"role":"assistant","content":[%s,{"type":"text","text":"{\"answer\":42}"}]},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":5,"total_tokens":8}}`, thinking) + fmt.Fprintf(w, `{"choices":[{"message":{"role":"assistant","content":[%s,{"type":"text","text":"{\"answer\":42}"},%s]},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":5,"total_tokens":8}}`, thinking, reference) } })) t.Cleanup(server.Close) @@ -54,6 +57,7 @@ func TestOpenAIMistralContent(t *testing.T) { continue } for _, part := range resp.Content.Parts { + require.NotEmpty(t, part.Text) if part.Thought { partialThought.WriteString(part.Text) } else { @@ -62,6 +66,7 @@ func TestOpenAIMistralContent(t *testing.T) { } } require.NotNil(t, final) + require.Len(t, final.Content.Parts, 2) require.True(t, final.TurnComplete) require.Equal(t, genai.FinishReasonStop, final.FinishReason) require.Equal(t, int32(8), final.UsageMetadata.TotalTokenCount) @@ -90,6 +95,11 @@ func TestOpenAIMistralContent(t *testing.T) { require.NotNil(t, messages[0].OfAssistant) wire, err := json.Marshal(messages[0]) require.NoError(t, err) + wantThinking := thinking + if stream { + wantThinking = `{"type":"thinking","thinking":[{"type":"text","text":"Let me "},{"type":"text","text":"think."}],"closed":true}` + } + require.JSONEq(t, `{"role":"assistant","content":[`+wantThinking+`,{"type":"text","text":"{\"answer\":42}"}],"tool_calls":[]}`, string(wire)) var replay struct { Role string `json:"role"` Content []struct { @@ -103,17 +113,9 @@ func TestOpenAIMistralContent(t *testing.T) { } require.NoError(t, json.Unmarshal(wire, &replay)) require.Equal(t, "assistant", replay.Role) - if stream { - require.Len(t, replay.Content, 3) - require.NotNil(t, replay.Content[0].Closed) - require.False(t, *replay.Content[0].Closed) - require.NotNil(t, replay.Content[1].Closed) - require.True(t, *replay.Content[1].Closed) - } else { - require.Len(t, replay.Content, 2) - require.NotNil(t, replay.Content[0].Closed) - require.True(t, *replay.Content[0].Closed) - } + require.Len(t, replay.Content, 2) + require.NotNil(t, replay.Content[0].Closed) + require.True(t, *replay.Content[0].Closed) var replayThought, replayAnswer strings.Builder for _, chunk := range replay.Content { switch chunk.Type { @@ -140,10 +142,17 @@ func TestOpenAIStreamingContentVariants(t *testing.T) { content []string want string wantErr bool + thought bool }{ {name: "plain text and empty deltas", content: []string{`null`, `"hello"`, `[]`, `" world"`, `""`}, want: "hello world"}, {name: "literal JSON string", content: []string{`"[{\"type\":\"text\",\"text\":\"literal\"}]"`}, want: `[{"type":"text","text":"literal"}]`}, {name: "invalid thinking", content: []string{`[{"type":"thinking","thinking":"invalid"}]`}, wantErr: true}, + {name: "unknown extension", content: []string{`[{"type":"future","text":{"value":42},"thinking":false},{"type":"text","text":"answer"}]`}, want: "answer"}, + {name: "thinking fragments marked closed", content: []string{ + `[{"type":"thinking","thinking":[{"type":"text","text":"a "}],"closed":true}]`, + `[{"type":"thinking","thinking":[{"type":"text","text":"plan"}],"closed":true}]`, + }, want: "a plan", thought: true}, + {name: "thinking extension", content: []string{`[{"type":"thinking","thinking":[{"type":"future","text":{"value":42}},{"type":"text","text":"a plan"}],"closed":true}]`}, want: "a plan", thought: true}, } { t.Run(tt.name, func(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -172,7 +181,7 @@ func TestOpenAIStreamingContentVariants(t *testing.T) { require.False(t, final.Partial) require.Len(t, final.Content.Parts, 1) require.Equal(t, tt.want, final.Content.Parts[0].Text) - require.False(t, final.Content.Parts[0].Thought) + require.Equal(t, tt.thought, final.Content.Parts[0].Thought) } }) } @@ -223,7 +232,9 @@ func TestChatCompletionInvalidContentChunks(t *testing.T) { for _, content := range []string{ `[{"type":"thinking","thinking":"invalid"}]`, `[{"type":"text","text":42}]`, - `[{"type":"unknown","text":"must not become an answer"}]`, + `[{"type":"thinking","thinking":[{"type":"text","text":42}]}]`, + `[{"type":"thinking","thinking":[{"text":"missing type"}]}]`, + `[{"text":"missing type"}]`, } { t.Run(content, func(t *testing.T) { var completion openai.ChatCompletion