History

Merge pull request #874 from MohammadErfan-Jabbari/fix/streaming-finish-reason-tool-calls

fix(antigravity): preserve finish_reason tool_calls across streaming chunks

Luis Pater committed 9 months ago

Commit 73db4e62 parents 69ca0a8 fe6043a2 files · +154 −10Browse files

Unified · Split

2 files
  1. internal/translator/antigravity/openai/chat-completions/antigravity_openai_response.go +26 −10
  2. internal/translator/antigravity/openai/chat-completions/antigravity_openai_response_test.go +128 −0
internal/translator/antigravity/openai/chat-completions/antigravity_openai_response.go+26−10
@@ -22,8 +22,10 @@ import ( · 21 unmodified lines
22 22
23// convertCliResponseToOpenAIChatParams holds parameters for response conversion.23// convertCliResponseToOpenAIChatParams holds parameters for response conversion.
24type convertCliResponseToOpenAIChatParams struct {24type convertCliResponseToOpenAIChatParams struct {
25 UnixTimestamp int6425 UnixTimestamp int64
26 FunctionIndex int26 FunctionIndex int
27 SawToolCall bool // Tracks if any tool call was seen in the entire stream
28 UpstreamFinishReason string // Caches the upstream finish reason for final chunk
27}29}
28 30
29// functionCallIDCounter provides a process-wide unique counter for function call identifiers.31// functionCallIDCounter provides a process-wide unique counter for function call identifiers.
@@ -79,10 +81,9 @@ func ConvertAntigravityResponseToOpenAI(_ context.Context, _ string, originalReq · 49 unmodified lines
79 template, _ = sjson.Set(template, "id", responseIDResult.String())81 template, _ = sjson.Set(template, "id", responseIDResult.String())
80 }82 }
81 83
82 // Extract and set the finish reason.84 // Cache the finish reason - do NOT set it in output yet (will be set on final chunk)
83 if finishReasonResult := gjson.GetBytes(rawJSON, "response.candidates.0.finishReason"); finishReasonResult.Exists() {85 if finishReasonResult := gjson.GetBytes(rawJSON, "response.candidates.0.finishReason"); finishReasonResult.Exists() {
84 template, _ = sjson.Set(template, "choices.0.finish_reason", strings.ToLower(finishReasonResult.String()))86 (*param).(*convertCliResponseToOpenAIChatParams).UpstreamFinishReason = strings.ToUpper(finishReasonResult.String())
85 template, _ = sjson.Set(template, "choices.0.native_finish_reason", strings.ToLower(finishReasonResult.String()))
86 }87 }
87 88
88 // Extract and set usage metadata (token counts).89 // Extract and set usage metadata (token counts).
@@ -112,7 +113,6 @@ func ConvertAntigravityResponseToOpenAI(_ context.Context, _ string, originalReq · 23 unmodified lines
112 113
113 // Process the main content part of the response.114 // Process the main content part of the response.
114 partsResult := gjson.GetBytes(rawJSON, "response.candidates.0.content.parts")115 partsResult := gjson.GetBytes(rawJSON, "response.candidates.0.content.parts")
115 hasFunctionCall := false
116 if partsResult.IsArray() {116 if partsResult.IsArray() {
117 partResults := partsResult.Array()117 partResults := partsResult.Array()
118 for i := 0; i < len(partResults); i++ {118 for i := 0; i < len(partResults); i++ {
@@ -148,7 +148,7 @@ func ConvertAntigravityResponseToOpenAI(_ context.Context, _ string, originalReq · 29 unmodified lines
148 template, _ = sjson.Set(template, "choices.0.delta.role", "assistant")148 template, _ = sjson.Set(template, "choices.0.delta.role", "assistant")
149 } else if functionCallResult.Exists() {149 } else if functionCallResult.Exists() {
150 // Handle function call content.150 // Handle function call content.
151 hasFunctionCall = true151 (*param).(*convertCliResponseToOpenAIChatParams).SawToolCall = true // Persist across chunks
152 toolCallsResult := gjson.Get(template, "choices.0.delta.tool_calls")152 toolCallsResult := gjson.Get(template, "choices.0.delta.tool_calls")
153 functionCallIndex := (*param).(*convertCliResponseToOpenAIChatParams).FunctionIndex153 functionCallIndex := (*param).(*convertCliResponseToOpenAIChatParams).FunctionIndex
154 (*param).(*convertCliResponseToOpenAIChatParams).FunctionIndex++154 (*param).(*convertCliResponseToOpenAIChatParams).FunctionIndex++
@@ -195,9 +195,25 @@ func ConvertAntigravityResponseToOpenAI(_ context.Context, _ string, originalReq · 40 unmodified lines
195 }195 }
196 }196 }
197 197
198 if hasFunctionCall {198 // Determine finish_reason only on the final chunk (has both finishReason and usage metadata)
199 template, _ = sjson.Set(template, "choices.0.finish_reason", "tool_calls")199 params := (*param).(*convertCliResponseToOpenAIChatParams)
200 template, _ = sjson.Set(template, "choices.0.native_finish_reason", "tool_calls")200 upstreamFinishReason := params.UpstreamFinishReason
201 sawToolCall := params.SawToolCall
202
203 usageExists := gjson.GetBytes(rawJSON, "response.usageMetadata").Exists()
204 isFinalChunk := upstreamFinishReason != "" && usageExists
205
206 if isFinalChunk {
207 var finishReason string
208 if sawToolCall {
209 finishReason = "tool_calls"
210 } else if upstreamFinishReason == "MAX_TOKENS" {
211 finishReason = "max_tokens"
212 } else {
213 finishReason = "stop"
214 }
215 template, _ = sjson.Set(template, "choices.0.finish_reason", finishReason)
216 template, _ = sjson.Set(template, "choices.0.native_finish_reason", strings.ToLower(upstreamFinishReason))
201 }217 }
202 218
203 return []string{template}219 return []string{template}
internal/translator/antigravity/openai/chat-completions/antigravity_openai_response_test.go+128−0
@@ -0,0 +1,128 @@
1package chat_completions
2
3import (
4 "context"
5 "testing"
6
7 "github.com/tidwall/gjson"
8)
9
10func TestFinishReasonToolCallsNotOverwritten(t *testing.T) {
11 ctx := context.Background()
12 var param any
13
14 // Chunk 1: Contains functionCall - should set SawToolCall = true
15 chunk1 := []byte(`{"response":{"candidates":[{"content":{"parts":[{"functionCall":{"name":"list_files","args":{"path":"."}}}]}}]}}`)
16 result1 := ConvertAntigravityResponseToOpenAI(ctx, "model", nil, nil, chunk1, &param)
17
18 // Verify chunk1 has no finish_reason (null)
19 if len(result1) != 1 {
20 t.Fatalf("Expected 1 result from chunk1, got %d", len(result1))
21 }
22 fr1 := gjson.Get(result1[0], "choices.0.finish_reason")
23 if fr1.Exists() && fr1.String() != "" && fr1.Type.String() != "Null" {
24 t.Errorf("Expected finish_reason to be null in chunk1, got: %v", fr1.String())
25 }
26
27 // Chunk 2: Contains finishReason STOP + usage (final chunk, no functionCall)
28 // This simulates what the upstream sends AFTER the tool call chunk
29 chunk2 := []byte(`{"response":{"candidates":[{"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":20,"totalTokenCount":30}}}`)
30 result2 := ConvertAntigravityResponseToOpenAI(ctx, "model", nil, nil, chunk2, &param)
31
32 // Verify chunk2 has finish_reason: "tool_calls" (not "stop")
33 if len(result2) != 1 {
34 t.Fatalf("Expected 1 result from chunk2, got %d", len(result2))
35 }
36 fr2 := gjson.Get(result2[0], "choices.0.finish_reason").String()
37 if fr2 != "tool_calls" {
38 t.Errorf("Expected finish_reason 'tool_calls', got: %s", fr2)
39 }
40
41 // Verify native_finish_reason is lowercase upstream value
42 nfr2 := gjson.Get(result2[0], "choices.0.native_finish_reason").String()
43 if nfr2 != "stop" {
44 t.Errorf("Expected native_finish_reason 'stop', got: %s", nfr2)
45 }
46}
47
48func TestFinishReasonStopForNormalText(t *testing.T) {
49 ctx := context.Background()
50 var param any
51
52 // Chunk 1: Text content only
53 chunk1 := []byte(`{"response":{"candidates":[{"content":{"parts":[{"text":"Hello world"}]}}]}}`)
54 ConvertAntigravityResponseToOpenAI(ctx, "model", nil, nil, chunk1, &param)
55
56 // Chunk 2: Final chunk with STOP
57 chunk2 := []byte(`{"response":{"candidates":[{"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":5,"totalTokenCount":15}}}`)
58 result2 := ConvertAntigravityResponseToOpenAI(ctx, "model", nil, nil, chunk2, &param)
59
60 // Verify finish_reason is "stop" (no tool calls were made)
61 fr := gjson.Get(result2[0], "choices.0.finish_reason").String()
62 if fr != "stop" {
63 t.Errorf("Expected finish_reason 'stop', got: %s", fr)
64 }
65}
66
67func TestFinishReasonMaxTokens(t *testing.T) {
68 ctx := context.Background()
69 var param any
70
71 // Chunk 1: Text content
72 chunk1 := []byte(`{"response":{"candidates":[{"content":{"parts":[{"text":"Hello"}]}}]}}`)
73 ConvertAntigravityResponseToOpenAI(ctx, "model", nil, nil, chunk1, &param)
74
75 // Chunk 2: Final chunk with MAX_TOKENS
76 chunk2 := []byte(`{"response":{"candidates":[{"finishReason":"MAX_TOKENS"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":100,"totalTokenCount":110}}}`)
77 result2 := ConvertAntigravityResponseToOpenAI(ctx, "model", nil, nil, chunk2, &param)
78
79 // Verify finish_reason is "max_tokens"
80 fr := gjson.Get(result2[0], "choices.0.finish_reason").String()
81 if fr != "max_tokens" {
82 t.Errorf("Expected finish_reason 'max_tokens', got: %s", fr)
83 }
84}
85
86func TestToolCallTakesPriorityOverMaxTokens(t *testing.T) {
87 ctx := context.Background()
88 var param any
89
90 // Chunk 1: Contains functionCall
91 chunk1 := []byte(`{"response":{"candidates":[{"content":{"parts":[{"functionCall":{"name":"test","args":{}}}]}}]}}`)
92 ConvertAntigravityResponseToOpenAI(ctx, "model", nil, nil, chunk1, &param)
93
94 // Chunk 2: Final chunk with MAX_TOKENS (but we had a tool call, so tool_calls should win)
95 chunk2 := []byte(`{"response":{"candidates":[{"finishReason":"MAX_TOKENS"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":100,"totalTokenCount":110}}}`)
96 result2 := ConvertAntigravityResponseToOpenAI(ctx, "model", nil, nil, chunk2, &param)
97
98 // Verify finish_reason is "tool_calls" (takes priority over max_tokens)
99 fr := gjson.Get(result2[0], "choices.0.finish_reason").String()
100 if fr != "tool_calls" {
101 t.Errorf("Expected finish_reason 'tool_calls', got: %s", fr)
102 }
103}
104
105func TestNoFinishReasonOnIntermediateChunks(t *testing.T) {
106 ctx := context.Background()
107 var param any
108
109 // Chunk 1: Text content (no finish reason, no usage)
110 chunk1 := []byte(`{"response":{"candidates":[{"content":{"parts":[{"text":"Hello"}]}}]}}`)
111 result1 := ConvertAntigravityResponseToOpenAI(ctx, "model", nil, nil, chunk1, &param)
112
113 // Verify no finish_reason on intermediate chunk
114 fr1 := gjson.Get(result1[0], "choices.0.finish_reason")
115 if fr1.Exists() && fr1.String() != "" && fr1.Type.String() != "Null" {
116 t.Errorf("Expected no finish_reason on intermediate chunk, got: %v", fr1)
117 }
118
119 // Chunk 2: More text (no finish reason, no usage)
120 chunk2 := []byte(`{"response":{"candidates":[{"content":{"parts":[{"text":" world"}]}}]}}`)
121 result2 := ConvertAntigravityResponseToOpenAI(ctx, "model", nil, nil, chunk2, &param)
122
123 // Verify no finish_reason on intermediate chunk
124 fr2 := gjson.Get(result2[0], "choices.0.finish_reason")
125 if fr2.Exists() && fr2.String() != "" && fr2.Type.String() != "Null" {
126 t.Errorf("Expected no finish_reason on intermediate chunk, got: %v", fr2)
127 }
128}