Enhance Gemini to OpenAI response conversion

Refactor response handling to support multiple candidates and improve parameter management.

internal/translator/gemini/openai/chat-completions/gemini_openai_response.go

@@ -21,7 +21,8 @@ import (
 // convertGeminiResponseToOpenAIChatParams holds parameters for response conversion.
 type convertGeminiResponseToOpenAIChatParams struct {
 	UnixTimestamp int64
-	FunctionIndex int
+	// FunctionIndex tracks tool call indices per candidate index to support multiple candidates.
+	FunctionIndex map[int]int
 }
 
 // functionCallIDCounter provides a process-wide unique counter for function call identifiers.
@@ -42,13 +43,20 @@ var functionCallIDCounter uint64
 // Returns:
 //   - []string: A slice of strings, each containing an OpenAI-compatible JSON response
 func ConvertGeminiResponseToOpenAI(_ context.Context, _ string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, param *any) []string {
+	// Initialize parameters if nil.
 	if *param == nil {
 		*param = &convertGeminiResponseToOpenAIChatParams{
 			UnixTimestamp: 0,
-			FunctionIndex: 0,
+			FunctionIndex: make(map[int]int),
 		}
 	}
 	
+	// Ensure the Map is initialized (handling cases where param might be reused from older context).
+	p := (*param).(*convertGeminiResponseToOpenAIChatParams)
+	if p.FunctionIndex == nil {
+		p.FunctionIndex = make(map[int]int)
+	}
+
 	if bytes.HasPrefix(rawJSON, []byte("data:")) {
 		rawJSON = bytes.TrimSpace(rawJSON[5:])
 	}
@@ -57,64 +65,79 @@ func ConvertGeminiResponseToOpenAI(_ context.Context, _ string, originalRequestR
 		return []string{}
 	}
 
-	// Initialize the OpenAI SSE template.
+	// Initialize the OpenAI SSE base template.
-	template := `{"id":"","object":"chat.completion.chunk","created":12345,"model":"model","choices":[{"index":0,"delta":{"role":null,"content":null,"reasoning_content":null,"tool_calls":null},"finish_reason":null,"native_finish_reason":null}]}`
+	// We use a base template and clone it for each candidate to support multiple candidates.
+	baseTemplate := `{"id":"","object":"chat.completion.chunk","created":12345,"model":"model","choices":[{"index":0,"delta":{"role":null,"content":null,"reasoning_content":null,"tool_calls":null},"finish_reason":null,"native_finish_reason":null}]}`
 
 	// Extract and set the model version.
 	if modelVersionResult := gjson.GetBytes(rawJSON, "modelVersion"); modelVersionResult.Exists() {
-		template, _ = sjson.Set(template, "model", modelVersionResult.String())
+		baseTemplate, _ = sjson.Set(baseTemplate, "model", modelVersionResult.String())
 	}
 
 	// Extract and set the creation timestamp.
 	if createTimeResult := gjson.GetBytes(rawJSON, "createTime"); createTimeResult.Exists() {
 		t, err := time.Parse(time.RFC3339Nano, createTimeResult.String())
 		if err == nil {
-			(*param).(*convertGeminiResponseToOpenAIChatParams).UnixTimestamp = t.Unix()
+			p.UnixTimestamp = t.Unix()
 		}
-		template, _ = sjson.Set(template, "created", (*param).(*convertGeminiResponseToOpenAIChatParams).UnixTimestamp)
+		baseTemplate, _ = sjson.Set(baseTemplate, "created", p.UnixTimestamp)
 	} else {
-		template, _ = sjson.Set(template, "created", (*param).(*convertGeminiResponseToOpenAIChatParams).UnixTimestamp)
+		baseTemplate, _ = sjson.Set(baseTemplate, "created", p.UnixTimestamp)
 	}
 
 	// Extract and set the response ID.
 	if responseIDResult := gjson.GetBytes(rawJSON, "responseId"); responseIDResult.Exists() {
-		template, _ = sjson.Set(template, "id", responseIDResult.String())
+		baseTemplate, _ = sjson.Set(baseTemplate, "id", responseIDResult.String())
-	}
-
-	// Extract and set the finish reason.
-	if finishReasonResult := gjson.GetBytes(rawJSON, "candidates.0.finishReason"); finishReasonResult.Exists() {
-		template, _ = sjson.Set(template, "choices.0.finish_reason", strings.ToLower(finishReasonResult.String()))
-		template, _ = sjson.Set(template, "choices.0.native_finish_reason", strings.ToLower(finishReasonResult.String()))
 	}
 
 	// Extract and set usage metadata (token counts).
+	// Usage is applied to the base template so it appears in the chunks.
 	if usageResult := gjson.GetBytes(rawJSON, "usageMetadata"); usageResult.Exists() {
 		cachedTokenCount := usageResult.Get("cachedContentTokenCount").Int()
 		if candidatesTokenCountResult := usageResult.Get("candidatesTokenCount"); candidatesTokenCountResult.Exists() {
-			template, _ = sjson.Set(template, "usage.completion_tokens", candidatesTokenCountResult.Int())
+			baseTemplate, _ = sjson.Set(baseTemplate, "usage.completion_tokens", candidatesTokenCountResult.Int())
 		}
 		if totalTokenCountResult := usageResult.Get("totalTokenCount"); totalTokenCountResult.Exists() {
-			template, _ = sjson.Set(template, "usage.total_tokens", totalTokenCountResult.Int())
+			baseTemplate, _ = sjson.Set(baseTemplate, "usage.total_tokens", totalTokenCountResult.Int())
 		}
 		promptTokenCount := usageResult.Get("promptTokenCount").Int() - cachedTokenCount
 		thoughtsTokenCount := usageResult.Get("thoughtsTokenCount").Int()
-		template, _ = sjson.Set(template, "usage.prompt_tokens", promptTokenCount+thoughtsTokenCount)
+		baseTemplate, _ = sjson.Set(baseTemplate, "usage.prompt_tokens", promptTokenCount+thoughtsTokenCount)
 		if thoughtsTokenCount > 0 {
-			template, _ = sjson.Set(template, "usage.completion_tokens_details.reasoning_tokens", thoughtsTokenCount)
+			baseTemplate, _ = sjson.Set(baseTemplate, "usage.completion_tokens_details.reasoning_tokens", thoughtsTokenCount)
 		}
 		// Include cached token count if present (indicates prompt caching is working)
 		if cachedTokenCount > 0 {
 			var err error
-			template, err = sjson.Set(template, "usage.prompt_tokens_details.cached_tokens", cachedTokenCount)
+			baseTemplate, err = sjson.Set(baseTemplate, "usage.prompt_tokens_details.cached_tokens", cachedTokenCount)
 			if err != nil {
 				log.Warnf("gemini openai response: failed to set cached_tokens in streaming: %v", err)
 			}
 		}
 	}
 
-	// Process the main content part of the response.
+	var responseStrings []string
-	partsResult := gjson.GetBytes(rawJSON, "candidates.0.content.parts")
+	candidates := gjson.GetBytes(rawJSON, "candidates")
+
+	// Iterate over all candidates to support candidate_count > 1.
+	if candidates.IsArray() {
+		candidates.ForEach(func(_, candidate gjson.Result) bool {
+			// Clone the template for the current candidate.
+			template := baseTemplate
+
+			// Set the specific index for this candidate.
+			candidateIndex := int(candidate.Get("index").Int())
+			template, _ = sjson.Set(template, "choices.0.index", candidateIndex)
+
+			// Extract and set the finish reason.
+			if finishReasonResult := candidate.Get("finishReason"); finishReasonResult.Exists() {
+				template, _ = sjson.Set(template, "choices.0.finish_reason", strings.ToLower(finishReasonResult.String()))
+				template, _ = sjson.Set(template, "choices.0.native_finish_reason", strings.ToLower(finishReasonResult.String()))
+			}
+
+			partsResult := candidate.Get("content.parts")
 			hasFunctionCall := false
+
 			if partsResult.IsArray() {
 				partResults := partsResult.Array()
 				for i := 0; i < len(partResults); i++ {
@@ -151,8 +174,11 @@ func ConvertGeminiResponseToOpenAI(_ context.Context, _ string, originalRequestR
 						// Handle function call content.
 						hasFunctionCall = true
 						toolCallsResult := gjson.Get(template, "choices.0.delta.tool_calls")
-				functionCallIndex := (*param).(*convertGeminiResponseToOpenAIChatParams).FunctionIndex
+						
-				(*param).(*convertGeminiResponseToOpenAIChatParams).FunctionIndex++
+						// Retrieve the function index for this specific candidate.
+						functionCallIndex := p.FunctionIndex[candidateIndex]
+						p.FunctionIndex[candidateIndex]++
+
 						if toolCallsResult.Exists() && toolCallsResult.IsArray() {
 							functionCallIndex = len(toolCallsResult.Array())
 						} else {
@@ -201,7 +227,17 @@ func ConvertGeminiResponseToOpenAI(_ context.Context, _ string, originalRequestR
 				template, _ = sjson.Set(template, "choices.0.native_finish_reason", "tool_calls")
 			}
 
-	return []string{template}
+			responseStrings = append(responseStrings, template)
+			return true // continue loop
+		})
+	} else {
+		// If there are no candidates (e.g., a pure usageMetadata chunk), return the usage chunk if present.
+		if gjson.GetBytes(rawJSON, "usageMetadata").Exists() && len(responseStrings) == 0 {
+			responseStrings = append(responseStrings, baseTemplate)
+		}
+	}
+
+	return responseStrings
 }
 
 // ConvertGeminiResponseToOpenAINonStream converts a non-streaming Gemini response to a non-streaming OpenAI response.
@@ -219,7 +255,9 @@ func ConvertGeminiResponseToOpenAI(_ context.Context, _ string, originalRequestR
 //   - string: An OpenAI-compatible JSON response containing all message content and metadata
 func ConvertGeminiResponseToOpenAINonStream(_ context.Context, _ string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, _ *any) string {
 	var unixTimestamp int64
-	template := `{"id":"","object":"chat.completion","created":123456,"model":"model","choices":[{"index":0,"message":{"role":"assistant","content":null,"reasoning_content":null,"tool_calls":null},"finish_reason":null,"native_finish_reason":null}]}`
+	// Initialize template with an empty choices array to support multiple candidates.
+	template := `{"id":"","object":"chat.completion","created":123456,"model":"model","choices":[]}`
+
 	if modelVersionResult := gjson.GetBytes(rawJSON, "modelVersion"); modelVersionResult.Exists() {
 		template, _ = sjson.Set(template, "model", modelVersionResult.String())
 	}
@@ -238,11 +276,6 @@ func ConvertGeminiResponseToOpenAINonStream(_ context.Context, _ string, origina
 		template, _ = sjson.Set(template, "id", responseIDResult.String())
 	}
 
-	if finishReasonResult := gjson.GetBytes(rawJSON, "candidates.0.finishReason"); finishReasonResult.Exists() {
-		template, _ = sjson.Set(template, "choices.0.finish_reason", strings.ToLower(finishReasonResult.String()))
-		template, _ = sjson.Set(template, "choices.0.native_finish_reason", strings.ToLower(finishReasonResult.String()))
-	}
-
 	if usageResult := gjson.GetBytes(rawJSON, "usageMetadata"); usageResult.Exists() {
 		if candidatesTokenCountResult := usageResult.Get("candidatesTokenCount"); candidatesTokenCountResult.Exists() {
 			template, _ = sjson.Set(template, "usage.completion_tokens", candidatesTokenCountResult.Int())
@@ -267,8 +300,23 @@ func ConvertGeminiResponseToOpenAINonStream(_ context.Context, _ string, origina
 		}
 	}
 
-	// Process the main content part of the response.
+	// Process the main content part of the response for all candidates.
-	partsResult := gjson.GetBytes(rawJSON, "candidates.0.content.parts")
+	candidates := gjson.GetBytes(rawJSON, "candidates")
+	if candidates.IsArray() {
+		candidates.ForEach(func(_, candidate gjson.Result) bool {
+			// Construct a single Choice object.
+			choiceTemplate := `{"index":0,"message":{"role":"assistant","content":null,"reasoning_content":null,"tool_calls":null},"finish_reason":null,"native_finish_reason":null}`
+
+			// Set the index for this choice.
+			choiceTemplate, _ = sjson.Set(choiceTemplate, "index", candidate.Get("index").Int())
+
+			// Set finish reason.
+			if finishReasonResult := candidate.Get("finishReason"); finishReasonResult.Exists() {
+				choiceTemplate, _ = sjson.Set(choiceTemplate, "finish_reason", strings.ToLower(finishReasonResult.String()))
+				choiceTemplate, _ = sjson.Set(choiceTemplate, "native_finish_reason", strings.ToLower(finishReasonResult.String()))
+			}
+
+			partsResult := candidate.Get("content.parts")
 			hasFunctionCall := false
 			if partsResult.IsArray() {
 				partsResults := partsResult.Array()
@@ -284,17 +332,19 @@ func ConvertGeminiResponseToOpenAINonStream(_ context.Context, _ string, origina
 					if partTextResult.Exists() {
 						// Append text content, distinguishing between regular content and reasoning.
 						if partResult.Get("thought").Bool() {
-					template, _ = sjson.Set(template, "choices.0.message.reasoning_content", partTextResult.String())
+							oldVal := gjson.Get(choiceTemplate, "message.reasoning_content").String()
+							choiceTemplate, _ = sjson.Set(choiceTemplate, "message.reasoning_content", oldVal+partTextResult.String())
 						} else {
-					template, _ = sjson.Set(template, "choices.0.message.content", partTextResult.String())
+							oldVal := gjson.Get(choiceTemplate, "message.content").String()
+							choiceTemplate, _ = sjson.Set(choiceTemplate, "message.content", oldVal+partTextResult.String())
 						}
-				template, _ = sjson.Set(template, "choices.0.message.role", "assistant")
+						choiceTemplate, _ = sjson.Set(choiceTemplate, "message.role", "assistant")
 					} else if functionCallResult.Exists() {
 						// Append function call content to the tool_calls array.
 						hasFunctionCall = true
-				toolCallsResult := gjson.Get(template, "choices.0.message.tool_calls")
+						toolCallsResult := gjson.Get(choiceTemplate, "message.tool_calls")
 						if !toolCallsResult.Exists() || !toolCallsResult.IsArray() {
-					template, _ = sjson.SetRaw(template, "choices.0.message.tool_calls", `[]`)
+							choiceTemplate, _ = sjson.SetRaw(choiceTemplate, "message.tool_calls", `[]`)
 						}
 						functionCallItemTemplate := `{"id": "","type": "function","function": {"name": "","arguments": ""}}`
 						fcName := functionCallResult.Get("name").String()
@@ -303,13 +353,11 @@ func ConvertGeminiResponseToOpenAINonStream(_ context.Context, _ string, origina
 						if fcArgsResult := functionCallResult.Get("args"); fcArgsResult.Exists() {
 							functionCallItemTemplate, _ = sjson.Set(functionCallItemTemplate, "function.arguments", fcArgsResult.Raw)
 						}
-				template, _ = sjson.Set(template, "choices.0.message.role", "assistant")
+						choiceTemplate, _ = sjson.Set(choiceTemplate, "message.role", "assistant")
-				template, _ = sjson.SetRaw(template, "choices.0.message.tool_calls.-1", functionCallItemTemplate)
+						choiceTemplate, _ = sjson.SetRaw(choiceTemplate, "message.tool_calls.-1", functionCallItemTemplate)
 					} else if inlineDataResult.Exists() {
 						data := inlineDataResult.Get("data").String()
-				if data == "" {
+						if data != "" {
-					continue
-				}
 							mimeType := inlineDataResult.Get("mimeType").String()
 							if mimeType == "" {
 								mimeType = inlineDataResult.Get("mime_type").String()
@@ -318,23 +366,30 @@ func ConvertGeminiResponseToOpenAINonStream(_ context.Context, _ string, origina
 								mimeType = "image/png"
 							}
 							imageURL := fmt.Sprintf("data:%s;base64,%s", mimeType, data)
-				imagesResult := gjson.Get(template, "choices.0.message.images")
+							imagesResult := gjson.Get(choiceTemplate, "message.images")
 							if !imagesResult.Exists() || !imagesResult.IsArray() {
-					template, _ = sjson.SetRaw(template, "choices.0.message.images", `[]`)
+								choiceTemplate, _ = sjson.SetRaw(choiceTemplate, "message.images", `[]`)
 							}
-				imageIndex := len(gjson.Get(template, "choices.0.message.images").Array())
+							imageIndex := len(gjson.Get(choiceTemplate, "message.images").Array())
 							imagePayload := `{"type":"image_url","image_url":{"url":""}}`
 							imagePayload, _ = sjson.Set(imagePayload, "index", imageIndex)
 							imagePayload, _ = sjson.Set(imagePayload, "image_url.url", imageURL)
-				template, _ = sjson.Set(template, "choices.0.message.role", "assistant")
+							choiceTemplate, _ = sjson.Set(choiceTemplate, "message.role", "assistant")
-				template, _ = sjson.SetRaw(template, "choices.0.message.images.-1", imagePayload)
+							choiceTemplate, _ = sjson.SetRaw(choiceTemplate, "message.images.-1", imagePayload)
+						}
 					}
 				}
 			}
 
 			if hasFunctionCall {
-		template, _ = sjson.Set(template, "choices.0.finish_reason", "tool_calls")
+				choiceTemplate, _ = sjson.Set(choiceTemplate, "finish_reason", "tool_calls")
-		template, _ = sjson.Set(template, "choices.0.native_finish_reason", "tool_calls")
+				choiceTemplate, _ = sjson.Set(choiceTemplate, "native_finish_reason", "tool_calls")
+			}
+
+			// Append the constructed choice to the main choices array.
+			template, _ = sjson.SetRaw(template, "choices.-1", choiceTemplate)
+			return true
+		})
 	}
 
 	return template