From 7e11e2f3205a17e710237bedfbc3df34174f37b1 Mon Sep 17 00:00:00 2001 From: DENG <33118163+XinhhD@users.noreply.github.com> Date: Wed, 17 Jun 2026 11:40:38 +0800 Subject: [PATCH] fix(ai-proxy): preserve Vertex thoughtSignature in OpenAI tool calls (#3973) Signed-off-by: DENG <33118163+XinhhD@users.noreply.github.com> --- .../extensions/ai-proxy/provider/model.go | 47 ++- .../extensions/ai-proxy/provider/vertex.go | 10 +- .../ai-proxy/provider/vertex_test.go | 362 ++++++++++++++++++ 3 files changed, 413 insertions(+), 6 deletions(-) diff --git a/plugins/wasm-go/extensions/ai-proxy/provider/model.go b/plugins/wasm-go/extensions/ai-proxy/provider/model.go index 2e072e442..8475d3688 100644 --- a/plugins/wasm-go/extensions/ai-proxy/provider/model.go +++ b/plugins/wasm-go/extensions/ai-proxy/provider/model.go @@ -492,10 +492,49 @@ func (m *chatMessage) ParseContent() []chatMessageContent { } type toolCall struct { - Index int `json:"index"` - Id string `json:"id,omitempty"` - Type string `json:"type"` - Function functionCall `json:"function"` + Index int `json:"index"` + Id string `json:"id,omitempty"` + Type string `json:"type"` + Function functionCall `json:"function"` + ThoughtSignature string `json:"thought_signature,omitempty"` + ExtraContent map[string]any `json:"extra_content,omitempty"` +} + +func (t *toolCall) getThoughtSignature() string { + if t == nil { + return "" + } + if t.ThoughtSignature != "" { + return t.ThoughtSignature + } + return getNestedString(t.ExtraContent, "google", "thought_signature") +} + +func buildGoogleThoughtSignatureExtraContent(signature string) map[string]any { + if signature == "" { + return nil + } + return map[string]any{ + "google": map[string]any{ + "thought_signature": signature, + }, + } +} + +func getNestedString(data map[string]any, path ...string) string { + var current any = data + for _, key := range path { + currentMap, ok := current.(map[string]any) + if !ok { + return "" + } + current, ok = currentMap[key] + if !ok { + return "" + } + } + value, _ := current.(string) + return value } type functionCall struct { diff --git a/plugins/wasm-go/extensions/ai-proxy/provider/vertex.go b/plugins/wasm-go/extensions/ai-proxy/provider/vertex.go index 44c9b9077..7c7769cad 100644 --- a/plugins/wasm-go/extensions/ai-proxy/provider/vertex.go +++ b/plugins/wasm-go/extensions/ai-proxy/provider/vertex.go @@ -871,7 +871,9 @@ func (v *vertexProvider) buildChatCompletionResponse(ctx wrapper.HttpContext, re args, _ := json.Marshal(part.FunctionCall.Args) choice.Message.ToolCalls = []toolCall{ { - Type: "function", + Type: "function", + ThoughtSignature: part.ThoughtSignature, + ExtraContent: buildGoogleThoughtSignatureExtraContent(part.ThoughtSignature), Function: functionCall{ Name: part.FunctionCall.Name, Arguments: string(args), @@ -979,7 +981,9 @@ func (v *vertexProvider) buildChatCompletionStreamResponse(ctx wrapper.HttpConte choice.Delta = &chatMessage{ ToolCalls: []toolCall{ { - Type: "function", + Type: "function", + ThoughtSignature: part.ThoughtSignature, + ExtraContent: buildGoogleThoughtSignatureExtraContent(part.ThoughtSignature), Function: functionCall{ Name: part.FunctionCall.Name, Arguments: string(args), @@ -1143,6 +1147,7 @@ func (v *vertexProvider) buildVertexChatRequest(request *chatCompletionRequest) Name: lastFunctionName, Args: args, }, + ThoughtSignature: message.ToolCalls[0].getThoughtSignature(), }) } else { for _, part := range message.ParseContent() { @@ -1359,6 +1364,7 @@ type vertexPart struct { FunctionCall *vertexFunctionCall `json:"functionCall,omitempty"` FunctionResponse *vertexFunctionResponse `json:"functionResponse,omitempty"` Thounght *bool `json:"thought,omitempty"` + ThoughtSignature string `json:"thoughtSignature,omitempty"` } type blob struct { diff --git a/plugins/wasm-go/extensions/ai-proxy/provider/vertex_test.go b/plugins/wasm-go/extensions/ai-proxy/provider/vertex_test.go index cdf264089..d7a88b333 100644 --- a/plugins/wasm-go/extensions/ai-proxy/provider/vertex_test.go +++ b/plugins/wasm-go/extensions/ai-proxy/provider/vertex_test.go @@ -578,3 +578,365 @@ func TestVertexAnthropicPassthrough_MaxTokensDefault(t *testing.T) { "client-supplied max_tokens must not be overwritten by the default") }) } + +func TestVertexProviderPreservesFunctionCallThoughtSignature(t *testing.T) { + v := &vertexProvider{} + ctx := newMockMultipartHttpContext() + ctx.SetContext(ctxKeyFinalRequestModel, "gemini-3.1-pro-preview") + + response := v.buildChatCompletionResponse(ctx, &vertexChatResponse{ + ResponseId: "vertex-response-id", + Candidates: []vertexChatCandidate{ + { + Index: 0, + Content: vertexChatContent{ + Role: "model", + Parts: []vertexPart{ + { + FunctionCall: &vertexFunctionCall{ + Name: "Skill", + Args: map[string]interface{}{"query": "intelligentization"}, + }, + ThoughtSignature: "thought-signature-from-vertex", + }, + }, + }, + FinishReason: "STOP", + }, + }, + }) + + require.Len(t, response.Choices, 1) + require.NotNil(t, response.Choices[0].Message) + require.Len(t, response.Choices[0].Message.ToolCalls, 1) + assert.Equal(t, "thought-signature-from-vertex", response.Choices[0].Message.ToolCalls[0].ThoughtSignature) + assert.Equal( + t, + "thought-signature-from-vertex", + getNestedString(response.Choices[0].Message.ToolCalls[0].ExtraContent, "google", "thought_signature"), + ) +} + +func TestVertexProviderRestoresFunctionCallThoughtSignature(t *testing.T) { + v := &vertexProvider{} + req := &chatCompletionRequest{ + Model: "gemini-3.1-pro-preview", + Messages: []chatMessage{ + {Role: roleUser, Content: "search docs"}, + { + Role: roleAssistant, + ToolCalls: []toolCall{ + { + Type: "function", + ThoughtSignature: "thought-signature-from-client", + Function: functionCall{ + Name: "Skill", + Arguments: `{"query":"intelligentization"}`, + }, + }, + }, + }, + {Role: roleTool, Content: "tool result"}, + }, + } + + vertexReq, err := v.buildVertexChatRequest(req) + require.NoError(t, err) + require.NotNil(t, vertexReq) + require.Len(t, vertexReq.Contents, 3) + require.Len(t, vertexReq.Contents[1].Parts, 1) + require.NotNil(t, vertexReq.Contents[1].Parts[0].FunctionCall) + assert.Equal(t, "thought-signature-from-client", vertexReq.Contents[1].Parts[0].ThoughtSignature) +} + +func TestVertexProviderRestoresFunctionCallThoughtSignatureFromGoogleExtraContent(t *testing.T) { + v := &vertexProvider{} + req := &chatCompletionRequest{ + Model: "gemini-3.1-pro-preview", + Messages: []chatMessage{ + {Role: roleUser, Content: "search docs"}, + { + Role: roleAssistant, + ToolCalls: []toolCall{ + { + Type: "function", + ExtraContent: map[string]any{ + "google": map[string]any{ + "thought_signature": "thought-signature-from-extra-content", + }, + }, + Function: functionCall{ + Name: "Skill", + Arguments: `{"query":"intelligentization"}`, + }, + }, + }, + }, + {Role: roleTool, Content: "tool result"}, + }, + } + + vertexReq, err := v.buildVertexChatRequest(req) + require.NoError(t, err) + require.NotNil(t, vertexReq) + require.Len(t, vertexReq.Contents, 3) + require.Len(t, vertexReq.Contents[1].Parts, 1) + require.NotNil(t, vertexReq.Contents[1].Parts[0].FunctionCall) + assert.Equal(t, "thought-signature-from-extra-content", vertexReq.Contents[1].Parts[0].ThoughtSignature) +} + +func TestVertexProviderRestoresFunctionCallThoughtSignatureInvalidArguments(t *testing.T) { + v := &vertexProvider{} + req := &chatCompletionRequest{ + Model: "gemini-3.1-pro-preview", + Messages: []chatMessage{ + { + Role: roleAssistant, + ToolCalls: []toolCall{ + { + Type: "function", + ThoughtSignature: "thought-signature-from-client", + Function: functionCall{ + Name: "Skill", + Arguments: `invalid-json`, + }, + }, + }, + }, + }, + } + + vertexReq, err := v.buildVertexChatRequest(req) + require.NoError(t, err) + require.NotNil(t, vertexReq) + require.Len(t, vertexReq.Contents, 1) + require.Len(t, vertexReq.Contents[0].Parts, 1) + require.NotNil(t, vertexReq.Contents[0].Parts[0].FunctionCall) + assert.Equal(t, "thought-signature-from-client", vertexReq.Contents[0].Parts[0].ThoughtSignature) +} + +func TestVertexProviderStreamPreservesFunctionCallThoughtSignature(t *testing.T) { + v := &vertexProvider{} + ctx := newMockMultipartHttpContext() + ctx.SetContext(ctxKeyFinalRequestModel, "gemini-3.1-pro-preview") + + response := v.buildChatCompletionStreamResponse(ctx, &vertexChatResponse{ + ResponseId: "vertex-response-id", + Candidates: []vertexChatCandidate{ + { + Index: 0, + Content: vertexChatContent{ + Role: "model", + Parts: []vertexPart{ + { + FunctionCall: &vertexFunctionCall{ + Name: "Skill", + Args: map[string]interface{}{"query": "intelligentization"}, + }, + ThoughtSignature: "thought-signature-from-vertex-stream", + }, + }, + }, + FinishReason: "STOP", + }, + }, + }) + + require.NotNil(t, response) + require.NotNil(t, response.Choices) + require.Len(t, response.Choices, 1) + require.NotNil(t, response.Choices[0].Delta) + require.Len(t, response.Choices[0].Delta.ToolCalls, 1) + assert.Equal(t, "thought-signature-from-vertex-stream", response.Choices[0].Delta.ToolCalls[0].ThoughtSignature) + assert.Equal( + t, + "thought-signature-from-vertex-stream", + getNestedString(response.Choices[0].Delta.ToolCalls[0].ExtraContent, "google", "thought_signature"), + ) +} + +func TestToolCallGetThoughtSignatureAllPaths(t *testing.T) { + // 1. nil toolCall pointer + var nilTc *toolCall + assert.Equal(t, "", nilTc.getThoughtSignature()) + + // 2. ThoughtSignature is directly set + tc1 := &toolCall{ + ThoughtSignature: "direct-sig", + } + assert.Equal(t, "direct-sig", tc1.getThoughtSignature()) + + // 3. ExtraContent is nil + tc2 := &toolCall{} + assert.Equal(t, "", tc2.getThoughtSignature()) + + // 4. ExtraContent does not contain "google" + tc3 := &toolCall{ + ExtraContent: map[string]any{ + "other": "val", + }, + } + assert.Equal(t, "", tc3.getThoughtSignature()) + + // 5. ExtraContent contains "google" but not as map[string]any + tc4 := &toolCall{ + ExtraContent: map[string]any{ + "google": "not-a-map", + }, + } + assert.Equal(t, "", tc4.getThoughtSignature()) + + // 6. ExtraContent contains "google" map but not "thought_signature" + tc5 := &toolCall{ + ExtraContent: map[string]any{ + "google": map[string]any{ + "other": "val", + }, + }, + } + assert.Equal(t, "", tc5.getThoughtSignature()) + + // 7. ExtraContent contains "google" map with "thought_signature" but not string + tc6 := &toolCall{ + ExtraContent: map[string]any{ + "google": map[string]any{ + "thought_signature": 12345, + }, + }, + } + assert.Equal(t, "", tc6.getThoughtSignature()) + + // 8. ExtraContent contains "google" map with "thought_signature" as string + tc7 := &toolCall{ + ExtraContent: map[string]any{ + "google": map[string]any{ + "thought_signature": "google-extra-sig", + }, + }, + } + assert.Equal(t, "google-extra-sig", tc7.getThoughtSignature()) +} + +func TestBuildGoogleThoughtSignatureExtraContentEmpty(t *testing.T) { + assert.Nil(t, buildGoogleThoughtSignatureExtraContent("")) +} + +func TestVertexProviderStreamThoughtAndText(t *testing.T) { + v := &vertexProvider{} + + t.Run("thinking start and continue", func(t *testing.T) { + ctx := newMockMultipartHttpContext() + + // 1. Thinking start + resp1 := v.buildChatCompletionStreamResponse(ctx, &vertexChatResponse{ + Candidates: []vertexChatCandidate{ + { + Content: vertexChatContent{ + Parts: []vertexPart{ + { + Text: "thinking process", + Thounght: util.Ptr(true), + }, + }, + }, + }, + }, + }) + require.NotNil(t, resp1) + assert.Equal(t, reasoningStartTag+"thinking process", resp1.Choices[0].Delta.Content) + assert.Equal(t, true, ctx.GetContext("thinking_start")) + + // 2. Thinking continue + resp2 := v.buildChatCompletionStreamResponse(ctx, &vertexChatResponse{ + Candidates: []vertexChatCandidate{ + { + Content: vertexChatContent{ + Parts: []vertexPart{ + { + Text: " more thinking", + Thounght: util.Ptr(true), + }, + }, + }, + }, + }, + }) + require.NotNil(t, resp2) + assert.Equal(t, " more thinking", resp2.Choices[0].Delta.Content) + }) + + t.Run("thinking end and text continue", func(t *testing.T) { + ctx := newMockMultipartHttpContext() + ctx.SetContext("thinking_start", true) + + // 1. Thinking end + resp1 := v.buildChatCompletionStreamResponse(ctx, &vertexChatResponse{ + Candidates: []vertexChatCandidate{ + { + Content: vertexChatContent{ + Parts: []vertexPart{ + { + Text: "final answer", + }, + }, + }, + }, + }, + }) + require.NotNil(t, resp1) + assert.Equal(t, reasoningEndTag+"final answer", resp1.Choices[0].Delta.Content) + assert.Equal(t, true, ctx.GetContext("thinking_end")) + + // 2. Text continue + resp2 := v.buildChatCompletionStreamResponse(ctx, &vertexChatResponse{ + Candidates: []vertexChatCandidate{ + { + Content: vertexChatContent{ + Parts: []vertexPart{ + { + Text: " suffix", + }, + }, + }, + }, + }, + }) + require.NotNil(t, resp2) + assert.Equal(t, " suffix", resp2.Choices[0].Delta.Content) + }) +} + +func TestGetNestedString(t *testing.T) { + // 1. empty path + assert.Equal(t, "", getNestedString(map[string]any{"a": 1})) + + // 2. nested string exists + data := map[string]any{ + "google": map[string]any{ + "thought_signature": "sig", + }, + } + assert.Equal(t, "sig", getNestedString(data, "google", "thought_signature")) + + // 3. nested key not a map + data2 := map[string]any{ + "google": "not-a-map", + } + assert.Equal(t, "", getNestedString(data2, "google", "thought_signature")) + + // 4. nested key is missing + data3 := map[string]any{ + "google": map[string]any{ + "other": "val", + }, + } + assert.Equal(t, "", getNestedString(data3, "google", "thought_signature")) + + // 5. nested value is not a string + data4 := map[string]any{ + "google": map[string]any{ + "thought_signature": 1234, + }, + } + assert.Equal(t, "", getNestedString(data4, "google", "thought_signature")) +}