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
2222
2323 // convertCliResponseToOpenAIChatParams holds parameters for response conversion.
2424 type convertCliResponseToOpenAIChatParams struct {
25− UnixTimestamp int64
26− FunctionIndex int
25+ UnixTimestamp int64
26+ 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
2729 }
2830
2931 // functionCallIDCounter provides a process-wide unique counter for function call identifiers.
@@ -79,10 +81,9 @@ func ConvertAntigravityResponseToOpenAI(_ context.Context, _ string, originalReq · 49 unmodified lines
7981 template, _ = sjson.Set(template, "id", responseIDResult.String())
8082 }
8183
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)
8385 if finishReasonResult := gjson.GetBytes(rawJSON, "response.candidates.0.finishReason"); finishReasonResult.Exists() {
84− template, _ = sjson.Set(template, "choices.0.finish_reason", strings.ToLower(finishReasonResult.String()))
85− template, _ = sjson.Set(template, "choices.0.native_finish_reason", strings.ToLower(finishReasonResult.String()))
86+ (*param).(*convertCliResponseToOpenAIChatParams).UpstreamFinishReason = strings.ToUpper(finishReasonResult.String())
8687 }
8788
8889 // Extract and set usage metadata (token counts).
@@ -112,7 +113,6 @@ func ConvertAntigravityResponseToOpenAI(_ context.Context, _ string, originalReq · 23 unmodified lines
112113
113114 // Process the main content part of the response.
114115 partsResult := gjson.GetBytes(rawJSON, "response.candidates.0.content.parts")
115− hasFunctionCall := false
116116 if partsResult.IsArray() {
117117 partResults := partsResult.Array()
118118 for i := 0; i < len(partResults); i++ {
@@ -148,7 +148,7 @@ func ConvertAntigravityResponseToOpenAI(_ context.Context, _ string, originalReq · 29 unmodified lines
148148 template, _ = sjson.Set(template, "choices.0.delta.role", "assistant")
149149 } else if functionCallResult.Exists() {
150150 // Handle function call content.
151− hasFunctionCall = true
151+ (*param).(*convertCliResponseToOpenAIChatParams).SawToolCall = true // Persist across chunks
152152 toolCallsResult := gjson.Get(template, "choices.0.delta.tool_calls")
153153 functionCallIndex := (*param).(*convertCliResponseToOpenAIChatParams).FunctionIndex
154154 (*param).(*convertCliResponseToOpenAIChatParams).FunctionIndex++
@@ -195,9 +195,25 @@ func ConvertAntigravityResponseToOpenAI(_ context.Context, _ string, originalReq · 40 unmodified lines
195195 }
196196 }
197197
198− if hasFunctionCall {
199− template, _ = sjson.Set(template, "choices.0.finish_reason", "tool_calls")
200− template, _ = sjson.Set(template, "choices.0.native_finish_reason", "tool_calls")
198+ // Determine finish_reason only on the final chunk (has both finishReason and usage metadata)
199+ params := (*param).(*convertCliResponseToOpenAIChatParams)
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))
201217 }
202218
203219 return []string{template}
internal/translator/antigravity/openai/chat-completions/antigravity_openai_response_test.go+128−0
@@ -0,0 +1,128 @@
1+package chat_completions
2+
3+import (
4+ "context"
5+ "testing"
6+
7+ "github.com/tidwall/gjson"
8+)
9+
10+func 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+
48+func 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+
67+func 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+
86+func 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+
105+func 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+}