feat: tests fixes and improvements

.github/workflows/scripts/release-core.sh

@@ -32,10 +32,14 @@ fi
 # Building core
 go mod download
 go build ./...
-go test ./...
 cd ..
 echo "✅ Core build validation successful"
 
+# Run core provider tests
+echo "🔧 Running core provider tests..."
+cd tests/core-providers
+go test -v ./...
+cd ../..
 
 # Capturing changelog
 CHANGELOG_BODY=$(cat core/changelog.md)

.gitignore

@@ -33,4 +33,6 @@ go.work.sum
 # Sqlite DBs
 *.db
 *.db-shm
-*.db-wal
+*.db-wal
+
+.claude

core/bifrost.go

@@ -1229,7 +1229,11 @@ func (bifrost *Bifrost) handleRequest(ctx context.Context, req *schemas.BifrostR
 	primaryResult, primaryErr := bifrost.tryRequest(ctx, req)
 
 	if primaryErr != nil {
-		bifrost.logger.Debug(fmt.Sprintf("Primary provider %s with model %s returned error: %v", provider, model, primaryErr))
+		if primaryErr.Error != nil {
+			bifrost.logger.Debug(fmt.Sprintf("Primary provider %s with model %s returned error: %s", provider, model, primaryErr.Error.Message))
+		} else {
+			bifrost.logger.Debug(fmt.Sprintf("Primary provider %s with model %s returned error: %v", provider, model, primaryErr))
+		}
 		if len(fallbacks) > 0 {
 			bifrost.logger.Debug(fmt.Sprintf("Check if we should try %d fallbacks", len(fallbacks)))
 		}
@@ -1629,7 +1633,7 @@ func (bifrost *Bifrost) requestWorker(provider schemas.Provider, config *schemas
 				time.Sleep(backoff)
 			}
 
-			bifrost.logger.Debug("attempting request for provider %s", provider.GetProviderKey())
+			bifrost.logger.Debug("attempting %s request for provider %s", req.RequestType, provider.GetProviderKey())
 
 			// Attempt the request
 			if IsStreamRequestType(req.RequestType) {
@@ -1644,7 +1648,7 @@ func (bifrost *Bifrost) requestWorker(provider schemas.Provider, config *schemas
 				}
 			}
 
-			bifrost.logger.Debug("request for provider %s completed", provider.GetProviderKey())
+			bifrost.logger.Debug("request %s for provider %s completed", req.RequestType, provider.GetProviderKey())
 
 			// Check if successful or if we should retry
 			if bifrostError == nil ||

core/changelog.md

@@ -1,6 +1,9 @@
 <!-- The pattern we follow here is to keep the changelog for the latest version -->
 <!-- Old changelogs are automatically attached to the GitHub releases -->
 
-- fix: openai specific parameters filtered for openai compatibile providers
-- fix: error response unmarshalling for gemini provider
-- BREAKING FIX: json_schema field correctly renamed to schema; ResponsesTextConfigFormatJSONSchema restructured
+- bug: fixed embedding request not being handled in `GetExtraFields()` method of `BifrostResponse`
+- fix: added latency calculation for vertex native requests
+- feat: added cached tokens and reasoning tokens to the usage metadata for chat completions
+- feat: added global region support for vertex API
+- fix: added filter for extra fields in chat completions request for Mistral provider
+- fix: fixed ResponsesComputerToolCallPendingSafetyCheck code field

core/providers/gemini.go

@@ -381,8 +381,8 @@ func (provider *GeminiProvider) SpeechStream(ctx context.Context, postHookRunner
 
 		scanner := bufio.NewScanner(resp.Body)
 		// Increase buffer size to handle large chunks (especially for audio data)
-		buf := make([]byte, 0, 256*1024) // 256KB buffer
-		scanner.Buffer(buf, 1024*1024)   // Allow up to 1MB tokens
+		buf := make([]byte, 0, 1024*1024) // 1MB initial buffer
+		scanner.Buffer(buf, 10*1024*1024) // Allow up to 10MB tokens
 		chunkIndex := -1
 		usage := &schemas.SpeechUsage{}
 		startTime := time.Now()
@@ -658,6 +658,9 @@ func (provider *GeminiProvider) TranscriptionStream(ctx context.Context, postHoo
 		defer resp.Body.Close()
 
 		scanner := bufio.NewScanner(resp.Body)
+		// Increase buffer size to handle large chunks (especially for audio data)
+		buf := make([]byte, 0, 1024*1024) // 1MB initial buffer
+		scanner.Buffer(buf, 10*1024*1024) // Allow up to 10MB tokens
 		chunkIndex := -1
 		usage := &schemas.TranscriptionUsage{}
 		startTime := time.Now()
@@ -674,8 +677,8 @@ func (provider *GeminiProvider) TranscriptionStream(ctx context.Context, postHoo
 			}
 			var jsonData string
 			// Parse SSE data
-			if strings.HasPrefix(line, "data: ") {
-				jsonData = strings.TrimPrefix(line, "data: ")
+			if after, ok := strings.CutPrefix(line, "data: "); ok {
+				jsonData = after
 			} else {
 				// Handle raw JSON errors (without "data: " prefix)
 				jsonData = line

core/providers/groq.go

@@ -124,13 +124,18 @@ func (provider *GroqProvider) TextCompletionStream(ctx context.Context, postHook
 				responseChan <- response
 				continue
 			}
-			response.ToTextCompletionResponse()
-			if response.BifrostTextCompletionResponse != nil {
-				response.BifrostTextCompletionResponse.ExtraFields.RequestType = schemas.TextCompletionRequest
-				response.BifrostTextCompletionResponse.ExtraFields.Provider = provider.GetProviderKey()
-				response.BifrostTextCompletionResponse.ExtraFields.ModelRequested = request.Model
+			if response.BifrostChatResponse != nil {
+				textCompletionResponse := response.BifrostChatResponse.ToTextCompletionResponse()
+				if textCompletionResponse != nil {
+					textCompletionResponse.ExtraFields.RequestType = schemas.TextCompletionRequest
+					textCompletionResponse.ExtraFields.Provider = provider.GetProviderKey()
+					textCompletionResponse.ExtraFields.ModelRequested = request.Model
+
+					responseChan <- &schemas.BifrostStream{
+						BifrostTextCompletionResponse: textCompletionResponse,
+					}
+				}
 			}
-			responseChan <- response
 		}
 	}()
 	return responseChan, nil

core/providers/vertex.go

@@ -203,10 +203,19 @@ func (provider *VertexProvider) ChatCompletion(ctx context.Context, key schemas.
 		return nil, newConfigurationError("region is not set in key config", schemas.Vertex)
 	}
 
-	url := fmt.Sprintf("https://%s-aiplatform.googleapis.com/v1beta1/projects/%s/locations/%s/endpoints/openapi/chat/completions", region, projectID, region)
-
+	var url string
 	if strings.Contains(request.Model, "claude") {
-		url = fmt.Sprintf("https://%s-aiplatform.googleapis.com/v1/projects/%s/locations/%s/publishers/anthropic/models/%s:rawPredict", region, projectID, region, request.Model)
+		if region == "global" {
+			url = fmt.Sprintf("https://aiplatform.googleapis.com/v1/projects/%s/locations/global/publishers/anthropic/models/%s:rawPredict", projectID, request.Model)
+		} else {
+			url = fmt.Sprintf("https://%s-aiplatform.googleapis.com/v1/projects/%s/locations/%s/publishers/anthropic/models/%s:rawPredict", region, projectID, region, request.Model)
+		}
+	} else {
+		if region == "global" {
+			url = fmt.Sprintf("https://aiplatform.googleapis.com/v1beta1/projects/%s/locations/global/endpoints/openapi/chat/completions", projectID)
+		} else {
+			url = fmt.Sprintf("https://%s-aiplatform.googleapis.com/v1beta1/projects/%s/locations/%s/endpoints/openapi/chat/completions", region, projectID, region)
+		}
 	}
 
 	// Create request
@@ -286,12 +295,19 @@ func (provider *VertexProvider) ChatCompletion(ctx context.Context, key schemas.
 		}
 
 		var openAIErr schemas.BifrostError
-		var vertexErr []VertexError
 
+		var vertexErr []VertexError
 		if err := sonic.Unmarshal(body, &openAIErr); err != nil {
 			// Try Vertex error format if OpenAI format fails
 			if err := sonic.Unmarshal(body, &vertexErr); err != nil {
-				return nil, newBifrostOperationError(schemas.ErrProviderResponseUnmarshal, err, schemas.Vertex)
+
+				//try with single Vertex error format
+				var vertexErr VertexError
+				if err := sonic.Unmarshal(body, &vertexErr); err != nil {
+					return nil, newBifrostOperationError(schemas.ErrProviderResponseUnmarshal, err, schemas.Vertex)
+				}
+
+				return nil, newProviderAPIError(vertexErr.Error.Message, nil, resp.StatusCode, schemas.Vertex, nil, nil)
 			}
 
 			if len(vertexErr) > 0 {
@@ -395,7 +411,12 @@ func (provider *VertexProvider) ChatCompletionStream(ctx context.Context, postHo
 		delete(requestBody, "model")
 		delete(requestBody, "region")
 
-		url := fmt.Sprintf("https://%s-aiplatform.googleapis.com/v1/projects/%s/locations/%s/publishers/anthropic/models/%s:streamRawPredict", region, projectID, region, request.Model)
+		var url string
+		if region == "global" {
+			url = fmt.Sprintf("https://aiplatform.googleapis.com/v1/projects/%s/locations/global/publishers/anthropic/models/%s:streamRawPredict", projectID, request.Model)
+		} else {
+			url = fmt.Sprintf("https://%s-aiplatform.googleapis.com/v1/projects/%s/locations/%s/publishers/anthropic/models/%s:streamRawPredict", region, projectID, region, request.Model)
+		}
 
 		// Prepare headers for Vertex Anthropic
 		headers := map[string]string{
@@ -418,7 +439,12 @@ func (provider *VertexProvider) ChatCompletionStream(ctx context.Context, postHo
 			provider.logger,
 		)
 	} else {
-		url := fmt.Sprintf("https://%s-aiplatform.googleapis.com/v1beta1/projects/%s/locations/%s/endpoints/openapi/chat/completions", region, projectID, region)
+		var url string
+		if region == "global" {
+			url = fmt.Sprintf("https://aiplatform.googleapis.com/v1beta1/projects/%s/locations/global/endpoints/openapi/chat/completions", projectID)
+		} else {
+			url = fmt.Sprintf("https://%s-aiplatform.googleapis.com/v1beta1/projects/%s/locations/%s/endpoints/openapi/chat/completions", region, projectID, region)
+		}
 		authHeader := map[string]string{}
 		if key.Value != "" {
 			authHeader["Authorization"] = "Bearer " + key.Value

core/schemas/bifrost.go

@@ -232,6 +232,8 @@ func (r *BifrostResponse) GetExtraFields() *BifrostResponseExtraFields {
 		return &r.ResponsesResponse.ExtraFields
 	case r.ResponsesStreamResponse != nil:
 		return &r.ResponsesStreamResponse.ExtraFields
+	case r.EmbeddingResponse != nil:
+		return &r.EmbeddingResponse.ExtraFields
 	case r.SpeechResponse != nil:
 		return &r.SpeechResponse.ExtraFields
 	case r.SpeechStreamResponse != nil:

core/schemas/providers/gemini/chat.go

@@ -31,10 +31,10 @@ func (request *GeminiGenerationRequest) ToBifrostChatRequest() *schemas.BifrostC
 
 	allGenAiMessages := []Content{}
 	if request.SystemInstruction != nil {
-		allGenAiMessages = append(allGenAiMessages, request.SystemInstruction.ToGenAIContent())
+		allGenAiMessages = append(allGenAiMessages, *request.SystemInstruction)
 	}
 	for _, content := range request.Contents {
-		allGenAiMessages = append(allGenAiMessages, content.ToGenAIContent())
+		allGenAiMessages = append(allGenAiMessages, content)
 	}
 
 	for _, content := range allGenAiMessages {

core/schemas/providers/gemini/embedding.go

@@ -26,8 +26,8 @@ func ToGeminiEmbeddingRequest(bifrostReq *schemas.BifrostEmbeddingRequest) *Gemi
 	// Create the Gemini embedding request
 	request := &GeminiEmbeddingRequest{
 		Model: bifrostReq.Model,
-		Content: &CustomContent{
-			Parts: []*CustomPart{
+		Content: &Content{
+			Parts: []*Part{
 				{
 					Text: text,
 				},

core/schemas/providers/gemini/responses.go

@@ -482,21 +482,21 @@ func convertPropertyToGeminiSchema(prop interface{}) *Schema {
 }
 
 // convertResponsesMessagesToGeminiContents converts Responses messages to Gemini contents
-func convertResponsesMessagesToGeminiContents(messages []schemas.ResponsesMessage) ([]CustomContent, *CustomContent, error) {
-	var contents []CustomContent
-	var systemInstruction *CustomContent
+func convertResponsesMessagesToGeminiContents(messages []schemas.ResponsesMessage) ([]Content, *Content, error) {
+	var contents []Content
+	var systemInstruction *Content
 
 	for _, msg := range messages {
 		// Handle system messages separately
 		if msg.Role != nil && *msg.Role == schemas.ResponsesInputMessageRoleSystem {
 			if systemInstruction == nil {
-				systemInstruction = &CustomContent{}
+				systemInstruction = &Content{}
 			}
 
 			// Convert system message content
 			if msg.Content != nil {
 				if msg.Content.ContentStr != nil {
-					systemInstruction.Parts = append(systemInstruction.Parts, &CustomPart{
+					systemInstruction.Parts = append(systemInstruction.Parts, &Part{
 						Text: *msg.Content.ContentStr,
 					})
 				}
@@ -517,7 +517,7 @@ func convertResponsesMessagesToGeminiContents(messages []schemas.ResponsesMessag
 		}
 
 		// Handle regular messages
-		content := CustomContent{}
+		content := Content{}
 
 		if msg.Role != nil {
 			content.Role = string(*msg.Role)
@@ -528,7 +528,7 @@ func convertResponsesMessagesToGeminiContents(messages []schemas.ResponsesMessag
 		// Convert message content
 		if msg.Content != nil {
 			if msg.Content.ContentStr != nil {
-				content.Parts = append(content.Parts, &CustomPart{
+				content.Parts = append(content.Parts, &Part{
 					Text: *msg.Content.ContentStr,
 				})
 			}
@@ -559,7 +559,7 @@ func convertResponsesMessagesToGeminiContents(messages []schemas.ResponsesMessag
 						}
 					}
 
-					part := &CustomPart{
+					part := &Part{
 						FunctionCall: &FunctionCall{
 							Name: *msg.ResponsesToolMessage.Name,
 							Args: argsMap,
@@ -586,7 +586,7 @@ func convertResponsesMessagesToGeminiContents(messages []schemas.ResponsesMessag
 						funcName = *msg.ResponsesToolMessage.CallID
 					}
 
-					part := &CustomPart{
+					part := &Part{
 						FunctionResponse: &FunctionResponse{
 							Name:     funcName,
 							Response: responseMap,
@@ -608,11 +608,11 @@ func convertResponsesMessagesToGeminiContents(messages []schemas.ResponsesMessag
 }
 
 // convertContentBlockToGeminiPart converts a content block to Gemini part
-func convertContentBlockToGeminiPart(block schemas.ResponsesMessageContentBlock) (*CustomPart, error) {
+func convertContentBlockToGeminiPart(block schemas.ResponsesMessageContentBlock) (*Part, error) {
 	switch block.Type {
 	case schemas.ResponsesInputMessageContentBlockTypeText:
 		if block.Text != nil {
-			return &CustomPart{
+			return &Part{
 				Text: *block.Text,
 			}, nil
 		}
@@ -645,14 +645,14 @@ func convertContentBlockToGeminiPart(block schemas.ResponsesMessageContentBlock)
 					return nil, fmt.Errorf("failed to decode base64 image data: %w", err)
 				}
 
-				return &CustomPart{
-					InlineData: &CustomBlob{
+				return &Part{
+					InlineData: &Blob{
 						MIMEType: mimeType,
 						Data:     decodedData,
 					},
 				}, nil
 			} else {
-				return &CustomPart{
+				return &Part{
 					FileData: &FileData{
 						MIMEType: mimeType,
 						FileURI:  sanitizedURL,
@@ -669,8 +669,8 @@ func convertContentBlockToGeminiPart(block schemas.ResponsesMessageContentBlock)
 				return nil, fmt.Errorf("failed to decode base64 audio data: %w", err)
 			}
 
-			return &CustomPart{
-				InlineData: &CustomBlob{
+			return &Part{
+				InlineData: &Blob{
 					MIMEType: func() string {
 						f := strings.ToLower(strings.TrimSpace(block.Audio.Format))
 						if f == "" {
@@ -689,15 +689,15 @@ func convertContentBlockToGeminiPart(block schemas.ResponsesMessageContentBlock)
 	case schemas.ResponsesInputMessageContentBlockTypeFile:
 		if block.ResponsesInputMessageContentBlockFile != nil {
 			if block.ResponsesInputMessageContentBlockFile.FileURL != nil {
-				return &CustomPart{
+				return &Part{
 					FileData: &FileData{
 						MIMEType: "application/octet-stream", // default
 						FileURI:  *block.ResponsesInputMessageContentBlockFile.FileURL,
 					},
 				}, nil
 			} else if block.ResponsesInputMessageContentBlockFile.FileData != nil {
-				return &CustomPart{
-					InlineData: &CustomBlob{
+				return &Part{
+					InlineData: &Blob{
 						MIMEType: "application/octet-stream", // default
 						Data:     []byte(*block.ResponsesInputMessageContentBlockFile.FileData),
 					},

core/schemas/providers/gemini/speech.go

@@ -27,9 +27,9 @@ func ToGeminiSpeechRequest(bifrostReq *schemas.BifrostSpeechRequest, responseMod
 
 	// Convert speech input to Gemini format
 	if bifrostReq.Input.Input != "" {
-		geminiReq.Contents = []CustomContent{
+		geminiReq.Contents = []Content{
 			{
-				Parts: []*CustomPart{
+				Parts: []*Part{
 					{
 						Text: bifrostReq.Input.Input,
 					},

core/schemas/providers/gemini/transcription.go

@@ -47,23 +47,23 @@ func ToGeminiTranscriptionRequest(bifrostReq *schemas.BifrostTranscriptionReques
 	}
 
 	// Create parts for the transcription request
-	parts := []*CustomPart{
+	parts := []*Part{
 		{
 			Text: prompt,
 		},
 	}
 
 	// Add audio file if present
 	if len(bifrostReq.Input.File) > 0 {
-		parts = append(parts, &CustomPart{
-			InlineData: &CustomBlob{
+		parts = append(parts, &Part{
+			InlineData: &Blob{
 				MIMEType: detectAudioMimeType(bifrostReq.Input.File),
 				Data:     bifrostReq.Input.File,
 			},
 		})
 	}
 
-	geminiReq.Contents = []CustomContent{
+	geminiReq.Contents = []Content{
 		{
 			Parts: parts,
 		},