Skip to content

Commit 307ae76

Browse files
luispaterluispater
committed
refactor: streamline ConvertCodexResponseToGeminiNonStream by removing unnecessary buffer and improving response handling
1 parent 735b213 commit 307ae76

File tree

1 file changed

+113
-127
lines changed

1 file changed

+113
-127
lines changed

‎internal/translator/codex/gemini/codex_gemini_response.go‎

Lines changed: 113 additions & 127 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@
55
package gemini
66

77
import (
8-
"bufio"
98
"bytes"
109
"context"
1110
"encoding/json"
@@ -152,159 +151,146 @@ func ConvertCodexResponseToGemini(_ context.Context, modelName string, originalR
152151
// Returns:
153152
// - string: A Gemini-compatible JSON response containing all message content and metadata
154153
func ConvertCodexResponseToGeminiNonStream(_ context.Context, modelName string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, _ *any) string {
155-
scanner := bufio.NewScanner(bytes.NewReader(rawJSON))
156-
buffer := make([]byte, 20_971_520)
157-
scanner.Buffer(buffer, 20_971_520)
158-
for scanner.Scan() {
159-
line := scanner.Bytes()
160-
// log.Debug(string(line))
161-
if !bytes.HasPrefix(line, dataTag) {
162-
continue
163-
}
164-
rawJSON = bytes.TrimSpace(rawJSON[5:])
165-
166-
rootResult := gjson.ParseBytes(rawJSON)
154+
rootResult := gjson.ParseBytes(rawJSON)
167155

168-
// Verify this is a response.completed event
169-
if rootResult.Get("type").String() != "response.completed" {
170-
continue
171-
}
156+
// Verify this is a response.completed event
157+
if rootResult.Get("type").String() != "response.completed" {
158+
return ""
159+
}
172160

173-
// Base Gemini response template for non-streaming
174-
template := `{"candidates":[{"content":{"role":"model","parts":[]},"finishReason":"STOP"}],"usageMetadata":{"trafficType":"PROVISIONED_THROUGHPUT"},"modelVersion":"","createTime":"","responseId":""}`
161+
// Base Gemini response template for non-streaming
162+
template := `{"candidates":[{"content":{"role":"model","parts":[]},"finishReason":"STOP"}],"usageMetadata":{"trafficType":"PROVISIONED_THROUGHPUT"},"modelVersion":"","createTime":"","responseId":""}`
175163

176-
// Set model version
177-
template, _ = sjson.Set(template, "modelVersion", modelName)
164+
// Set model version
165+
template, _ = sjson.Set(template, "modelVersion", modelName)
178166

179-
// Set response metadata from the completed response
180-
responseData := rootResult.Get("response")
181-
if responseData.Exists() {
182-
// Set response ID
183-
if responseId := responseData.Get("id"); responseId.Exists() {
184-
template, _ = sjson.Set(template, "responseId", responseId.String())
185-
}
167+
// Set response metadata from the completed response
168+
responseData := rootResult.Get("response")
169+
if responseData.Exists() {
170+
// Set response ID
171+
if responseId := responseData.Get("id"); responseId.Exists() {
172+
template, _ = sjson.Set(template, "responseId", responseId.String())
173+
}
186174

187-
// Set creation time
188-
if createdAt := responseData.Get("created_at"); createdAt.Exists() {
189-
template, _ = sjson.Set(template, "createTime", time.Unix(createdAt.Int(), 0).Format(time.RFC3339Nano))
190-
}
175+
// Set creation time
176+
if createdAt := responseData.Get("created_at"); createdAt.Exists() {
177+
template, _ = sjson.Set(template, "createTime", time.Unix(createdAt.Int(), 0).Format(time.RFC3339Nano))
178+
}
191179

192-
// Set usage metadata
193-
if usage := responseData.Get("usage"); usage.Exists() {
194-
inputTokens := usage.Get("input_tokens").Int()
195-
outputTokens := usage.Get("output_tokens").Int()
196-
totalTokens := inputTokens + outputTokens
180+
// Set usage metadata
181+
if usage := responseData.Get("usage"); usage.Exists() {
182+
inputTokens := usage.Get("input_tokens").Int()
183+
outputTokens := usage.Get("output_tokens").Int()
184+
totalTokens := inputTokens + outputTokens
197185

198-
template, _ = sjson.Set(template, "usageMetadata.promptTokenCount", inputTokens)
199-
template, _ = sjson.Set(template, "usageMetadata.candidatesTokenCount", outputTokens)
200-
template, _ = sjson.Set(template, "usageMetadata.totalTokenCount", totalTokens)
201-
}
186+
template, _ = sjson.Set(template, "usageMetadata.promptTokenCount", inputTokens)
187+
template, _ = sjson.Set(template, "usageMetadata.candidatesTokenCount", outputTokens)
188+
template, _ = sjson.Set(template, "usageMetadata.totalTokenCount", totalTokens)
189+
}
202190

203-
// Process output content to build parts array
204-
var parts []interface{}
205-
hasToolCall := false
206-
var pendingFunctionCalls []interface{}
207-
208-
flushPendingFunctionCalls := func() {
209-
if len(pendingFunctionCalls) > 0 {
210-
// Add all pending function calls as individual parts
211-
// This maintains the original Gemini API format while ensuring consecutive calls are grouped together
212-
for _, fc := range pendingFunctionCalls {
213-
parts = append(parts, fc)
214-
}
215-
pendingFunctionCalls = nil
191+
// Process output content to build parts array
192+
var parts []interface{}
193+
hasToolCall := false
194+
var pendingFunctionCalls []interface{}
195+
196+
flushPendingFunctionCalls := func() {
197+
if len(pendingFunctionCalls) > 0 {
198+
// Add all pending function calls as individual parts
199+
// This maintains the original Gemini API format while ensuring consecutive calls are grouped together
200+
for _, fc := range pendingFunctionCalls {
201+
parts = append(parts, fc)
216202
}
203+
pendingFunctionCalls = nil
217204
}
205+
}
218206

219-
if output := responseData.Get("output"); output.Exists() && output.IsArray() {
220-
output.ForEach(func(key, value gjson.Result) bool {
221-
itemType := value.Get("type").String()
207+
if output := responseData.Get("output"); output.Exists() && output.IsArray() {
208+
output.ForEach(func(key, value gjson.Result) bool {
209+
itemType := value.Get("type").String()
222210

223-
switch itemType {
224-
case "reasoning":
225-
// Flush any pending function calls before adding non-function content
226-
flushPendingFunctionCalls()
211+
switch itemType {
212+
case "reasoning":
213+
// Flush any pending function calls before adding non-function content
214+
flushPendingFunctionCalls()
227215

228-
// Add thinking content
229-
if content := value.Get("content"); content.Exists() {
230-
part := map[string]interface{}{
231-
"thought": true,
232-
"text": content.String(),
233-
}
234-
parts = append(parts, part)
216+
// Add thinking content
217+
if content := value.Get("content"); content.Exists() {
218+
part := map[string]interface{}{
219+
"thought": true,
220+
"text": content.String(),
235221
}
222+
parts = append(parts, part)
223+
}
236224

237-
case "message":
238-
// Flush any pending function calls before adding non-function content
239-
flushPendingFunctionCalls()
240-
241-
// Add regular text content
242-
if content := value.Get("content"); content.Exists() && content.IsArray() {
243-
content.ForEach(func(_, contentItem gjson.Result) bool {
244-
if contentItem.Get("type").String() == "output_text" {
245-
if text := contentItem.Get("text"); text.Exists() {
246-
part := map[string]interface{}{
247-
"text": text.String(),
248-
}
249-
parts = append(parts, part)
225+
case "message":
226+
// Flush any pending function calls before adding non-function content
227+
flushPendingFunctionCalls()
228+
229+
// Add regular text content
230+
if content := value.Get("content"); content.Exists() && content.IsArray() {
231+
content.ForEach(func(_, contentItem gjson.Result) bool {
232+
if contentItem.Get("type").String() == "output_text" {
233+
if text := contentItem.Get("text"); text.Exists() {
234+
part := map[string]interface{}{
235+
"text": text.String(),
250236
}
237+
parts = append(parts, part)
251238
}
252-
return true
253-
})
254-
}
255-
256-
case "function_call":
257-
// Collect function call for potential merging with consecutive ones
258-
hasToolCall = true
259-
functionCall := map[string]interface{}{
260-
"functionCall": map[string]interface{}{
261-
"name": func() string {
262-
n := value.Get("name").String()
263-
rev := buildReverseMapFromGeminiOriginal(originalRequestRawJSON)
264-
if orig, ok := rev[n]; ok {
265-
return orig
266-
}
267-
return n
268-
}(),
269-
"args": map[string]interface{}{},
270-
},
271-
}
239+
}
240+
return true
241+
})
242+
}
272243

273-
// Parse and set arguments
274-
if argsStr := value.Get("arguments").String(); argsStr != "" {
275-
argsResult := gjson.Parse(argsStr)
276-
if argsResult.IsObject() {
277-
var args map[string]interface{}
278-
if err := json.Unmarshal([]byte(argsStr), &args); err == nil {
279-
functionCall["functionCall"].(map[string]interface{})["args"] = args
244+
case "function_call":
245+
// Collect function call for potential merging with consecutive ones
246+
hasToolCall = true
247+
functionCall := map[string]interface{}{
248+
"functionCall": map[string]interface{}{
249+
"name": func() string {
250+
n := value.Get("name").String()
251+
rev := buildReverseMapFromGeminiOriginal(originalRequestRawJSON)
252+
if orig, ok := rev[n]; ok {
253+
return orig
280254
}
255+
return n
256+
}(),
257+
"args": map[string]interface{}{},
258+
},
259+
}
260+
261+
// Parse and set arguments
262+
if argsStr := value.Get("arguments").String(); argsStr != "" {
263+
argsResult := gjson.Parse(argsStr)
264+
if argsResult.IsObject() {
265+
var args map[string]interface{}
266+
if err := json.Unmarshal([]byte(argsStr), &args); err == nil {
267+
functionCall["functionCall"].(map[string]interface{})["args"] = args
281268
}
282269
}
283-
284-
pendingFunctionCalls = append(pendingFunctionCalls, functionCall)
285270
}
286-
return true
287-
})
288271

289-
// Handle any remaining pending function calls at the end
290-
flushPendingFunctionCalls()
291-
}
272+
pendingFunctionCalls = append(pendingFunctionCalls, functionCall)
273+
}
274+
return true
275+
})
292276

293-
// Set the parts array
294-
if len(parts) > 0 {
295-
template, _ = sjson.SetRaw(template, "candidates.0.content.parts", mustMarshalJSON(parts))
296-
}
277+
// Handle any remaining pending function calls at the end
278+
flushPendingFunctionCalls()
279+
}
297280

298-
// Set finish reason based on whether there were tool calls
299-
if hasToolCall {
300-
template, _ = sjson.Set(template, "candidates.0.finishReason", "STOP")
301-
} else {
302-
template, _ = sjson.Set(template, "candidates.0.finishReason", "STOP")
303-
}
281+
// Set the parts array
282+
if len(parts) > 0 {
283+
template, _ = sjson.SetRaw(template, "candidates.0.content.parts", mustMarshalJSON(parts))
284+
}
285+
286+
// Set finish reason based on whether there were tool calls
287+
if hasToolCall {
288+
template, _ = sjson.Set(template, "candidates.0.finishReason", "STOP")
289+
} else {
290+
template, _ = sjson.Set(template, "candidates.0.finishReason", "STOP")
304291
}
305-
return template
306292
}
307-
return ""
293+
return template
308294
}
309295

310296
// buildReverseMapFromGeminiOriginal builds a map[short]original from original Gemini request tools.

0 commit comments

Comments
 (0)