fix: integration test cases fixes

core/changelog.md

@@ -1,4 +1,6 @@
 feat: send back raw request in extra fields
 feat: added support for reasoning in chat completions
 feat: enhanced reasoning support in responses api
-enhancement: improved internal inter provider conversions for integrations
+enhancement: improved internal inter provider conversions for integrations
+feat: switched to gemini native api
+feat: fallback to supported request type for custom models used in integration

core/providers/cohere/chat.go

@@ -690,5 +690,11 @@ func (cm *CohereMessage) ToBifrostChatMessage() *schemas.ChatMessage {
 		Content:              messageContent,
 		ChatAssistantMessage: assistantMessage,
 	}
+
+	if cm.Role == "tool" {
+		bifrostMessage.ChatToolMessage = &schemas.ChatToolMessage{
+			ToolCallID: cm.ToolCallID,
+		}
+	}
 	return bifrostMessage
 }

core/providers/gemini/chat.go

@@ -187,8 +187,8 @@ func (response *GenerateContentResponse) ToBifrostChatCompletionStream() (*schem
 
 	candidate := response.Candidates[0]
 
-	// Determine if this is the last chunk based on finish reason or usage metadata
-	isLastChunk := candidate.FinishReason != "" && candidate.FinishReason != FinishReasonUnspecified
+	// Determine if this is the last chunk based on finish reason and usage metadata
+	isLastChunk := candidate.FinishReason != "" && response.UsageMetadata != nil
 
 	// Create the streaming response
 	streamResponse := &schemas.BifrostChatResponse{
@@ -256,15 +256,6 @@ func (response *GenerateContentResponse) ToBifrostChatCompletionStream() (*schem
 					},
 				}
 
-				// Preserve thought signature if present (required for Gemini 3 Pro)
-				if len(part.ThoughtSignature) > 0 {
-					toolCall.ExtraContent = map[string]interface{}{
-						"google": map[string]interface{}{
-							"thought_signature": string(part.ThoughtSignature),
-						},
-					}
-				}
-
 				toolCalls = append(toolCalls, toolCall)
 
 			case part.FunctionResponse != nil:

core/providers/gemini/gemini.go

@@ -243,6 +243,12 @@ func (provider *GeminiProvider) ChatCompletion(ctx context.Context, key schemas.
 	bifrostResponse.ExtraFields.ModelRequested = request.Model
 	bifrostResponse.ExtraFields.Latency = latency.Milliseconds()
 
+	// Set raw request if enabled
+	if providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest) {
+		providerUtils.ParseAndSetRawRequest(&bifrostResponse.ExtraFields, jsonData)
+	}
+
+	// Set raw response if enabled
 	if providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse) {
 		bifrostResponse.ExtraFields.RawResponse = rawResponse
 	}
@@ -291,6 +297,7 @@ func (provider *GeminiProvider) ChatCompletionStream(ctx context.Context, postHo
 		jsonData,
 		headers,
 		provider.networkConfig.ExtraHeaders,
+		providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest),
 		providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse),
 		provider.GetProviderKey(),
 		request.Model,
@@ -308,6 +315,7 @@ func HandleGeminiChatCompletionStream(
 	jsonBody []byte,
 	headers map[string]string,
 	extraHeaders map[string]string,
+	sendBackRawRequest bool,
 	sendBackRawResponse bool,
 	providerName schemas.ModelProvider,
 	model string,
@@ -378,6 +386,11 @@ func HandleGeminiChatCompletionStream(
 		var modelName string
 
 		for scanner.Scan() {
+			select {
+			case <-ctx.Done():
+				return
+			default:
+			}
 			line := scanner.Text()
 
 			// Skip empty lines and comments
@@ -473,6 +486,9 @@ func HandleGeminiChatCompletionStream(
 				chunkIndex++
 
 				if isLastChunk {
+					if sendBackRawRequest {
+						providerUtils.ParseAndSetRawRequest(&response.ExtraFields, jsonBody)
+					}
 					response.ExtraFields.Latency = time.Since(startTime).Milliseconds()
 					ctx = context.WithValue(ctx, schemas.BifrostContextKeyStreamEndIndicator, true)
 					providerUtils.ProcessAndSendResponse(ctx, postHookRunner, providerUtils.GetBifrostResponseForStreamResponse(nil, response, nil, nil, nil), responseChan)
@@ -584,6 +600,7 @@ func (provider *GeminiProvider) ResponsesStream(ctx context.Context, postHookRun
 		jsonData,
 		headers,
 		provider.networkConfig.ExtraHeaders,
+		providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest),
 		providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse),
 		provider.GetProviderKey(),
 		request.Model,
@@ -601,6 +618,7 @@ func HandleGeminiResponsesStream(
 	jsonBody []byte,
 	headers map[string]string,
 	extraHeaders map[string]string,
+	sendBackRawRequest bool,
 	sendBackRawResponse bool,
 	providerName schemas.ModelProvider,
 	model string,
@@ -675,6 +693,11 @@ func HandleGeminiResponsesStream(
 		var lastUsageMetadata *GenerateContentResponseUsageMetadata
 
 		for scanner.Scan() {
+			select {
+			case <-ctx.Done():
+				return
+			default:
+			}
 			line := scanner.Text()
 
 			// Skip empty lines and comments
@@ -771,6 +794,9 @@ func HandleGeminiResponsesStream(
 					}
 
 					if isLastChunk {
+						if sendBackRawRequest {
+							providerUtils.ParseAndSetRawRequest(&response.ExtraFields, jsonBody)
+						}
 						response.ExtraFields.Latency = time.Since(startTime).Milliseconds()
 						ctx = context.WithValue(ctx, schemas.BifrostContextKeyStreamEndIndicator, true)
 						providerUtils.ProcessAndSendResponse(ctx, postHookRunner, providerUtils.GetBifrostResponseForStreamResponse(nil, nil, response, nil, nil), responseChan)
@@ -985,6 +1011,11 @@ func (provider *GeminiProvider) SpeechStream(ctx context.Context, postHookRunner
 		lastChunkTime := startTime
 
 		for scanner.Scan() {
+			select {
+			case <-ctx.Done():
+				return
+			default:
+			}
 			line := scanner.Text()
 
 			// Skip empty lines
@@ -1243,6 +1274,11 @@ func (provider *GeminiProvider) TranscriptionStream(ctx context.Context, postHoo
 		var fullTranscriptionText string
 
 		for scanner.Scan() {
+			select {
+			case <-ctx.Done():
+				return
+			default:
+			}
 			line := scanner.Text()
 
 			// Skip empty lines