chore: tests fixes

Author: Pratham-Mishra04

.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,8 @@
 <!-- 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

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/test_helpers.go

@@ -0,0 +1,40 @@
+package providers
+
+import (
+	"context"
+
+	"github.com/maximhq/bifrost/core/schemas"
+)
+
+// newTestLogger creates a simple test logger
+func newTestLogger() schemas.Logger {
+	return &testLogger{}
+}
+
+type testLogger struct{}
+
+func (l *testLogger) Debug(msg string, args ...interface{})                  {}
+func (l *testLogger) Info(msg string, args ...interface{})                   {}
+func (l *testLogger) Warn(msg string, args ...interface{})                   {}
+func (l *testLogger) Error(msg string, args ...interface{})                  {}
+func (l *testLogger) Fatal(msg string, args ...interface{})                  {}
+func (l *testLogger) SetLevel(level schemas.LogLevel)                        {}
+func (l *testLogger) SetOutputType(outputType schemas.LoggerOutputType)      {}
+
+// Helper functions for tests
+func stringPtr(s string) *string {
+	return &s
+}
+
+func float64Ptr(f float64) *float64 {
+	return &f
+}
+
+func intPtr(i int) *int {
+	return &i
+}
+
+// mockPostHookRunner is a mock implementation of PostHookRunner for testing
+func mockPostHookRunner(ctx *context.Context, result *schemas.BifrostResponse, err *schemas.BifrostError) (*schemas.BifrostResponse, *schemas.BifrostError) {
+	return result, err
+}

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 {