feat: modelcatelog enhancements

Author: Pratham-Mishra04

core/bifrost.go

@@ -1986,8 +1986,8 @@ func executeRequestWithRetries[T any](
 		// Retry if status code or error object indicates rate limiting
 		if (bifrostError.StatusCode != nil && retryableStatusCodes[*bifrostError.StatusCode]) ||
 			(bifrostError.Error != nil &&
-				(isRateLimitError(bifrostError.Error.Message) ||
-					(bifrostError.Error.Type != nil && isRateLimitError(*bifrostError.Error.Type)))) {
+				(IsRateLimitErrorMessage(bifrostError.Error.Message) ||
+					(bifrostError.Error.Type != nil && IsRateLimitErrorMessage(*bifrostError.Error.Type)))) {
 			shouldRetry = true
 			logger.Debug("detected rate limit error in message, will retry: %s", bifrostError.Error.Message)
 		}

core/bifrost_test.go

@@ -353,7 +353,7 @@ func TestCalculateBackoff_MaxBackoffCap(t *testing.T) {
 	}
 }
 
-// Test isRateLimitError - all patterns
+// Test IsRateLimitErrorMessage - all patterns
 func TestIsRateLimitError_AllPatterns(t *testing.T) {
 	// Test all patterns from rateLimitPatterns
 	patterns := []string{
@@ -382,36 +382,36 @@ func TestIsRateLimitError_AllPatterns(t *testing.T) {
 	for _, pattern := range patterns {
 		t.Run(fmt.Sprintf("Pattern_%s", strings.ReplaceAll(pattern, " ", "_")), func(t *testing.T) {
 			// Test exact match
-			if !isRateLimitError(pattern) {
+			if !IsRateLimitErrorMessage(pattern) {
 				t.Errorf("Pattern '%s' should be detected as rate limit error", pattern)
 			}
 
 			// Test case insensitive - uppercase
-			if !isRateLimitError(strings.ToUpper(pattern)) {
+			if !IsRateLimitErrorMessage(strings.ToUpper(pattern)) {
 				t.Errorf("Uppercase pattern '%s' should be detected as rate limit error", strings.ToUpper(pattern))
 			}
 
 			// Test case insensitive - mixed case
-			if !isRateLimitError(strings.Title(pattern)) {
+			if !IsRateLimitErrorMessage(strings.Title(pattern)) {
 				t.Errorf("Title case pattern '%s' should be detected as rate limit error", strings.Title(pattern))
 			}
 
 			// Test as part of larger message
 			message := fmt.Sprintf("Error: %s occurred", pattern)
-			if !isRateLimitError(message) {
+			if !IsRateLimitErrorMessage(message) {
 				t.Errorf("Pattern '%s' in message '%s' should be detected", pattern, message)
 			}
 
 			// Test with prefix and suffix
 			message = fmt.Sprintf("API call failed due to %s - please retry later", pattern)
-			if !isRateLimitError(message) {
+			if !IsRateLimitErrorMessage(message) {
 				t.Errorf("Pattern '%s' in complex message should be detected", pattern)
 			}
 		})
 	}
 }
 
-// Test isRateLimitError - negative cases
+// Test IsRateLimitErrorMessage - negative cases
 func TestIsRateLimitError_NegativeCases(t *testing.T) {
 	negativeCases := []string{
 		"",
@@ -431,31 +431,31 @@ func TestIsRateLimitError_NegativeCases(t *testing.T) {
 
 	for _, testCase := range negativeCases {
 		t.Run(fmt.Sprintf("Negative_%s", strings.ReplaceAll(testCase, " ", "_")), func(t *testing.T) {
-			if isRateLimitError(testCase) {
+			if IsRateLimitErrorMessage(testCase) {
 				t.Errorf("Message '%s' should NOT be detected as rate limit error", testCase)
 			}
 		})
 	}
 }
 
-// Test isRateLimitError - edge cases
+// Test IsRateLimitErrorMessage - edge cases
 func TestIsRateLimitError_EdgeCases(t *testing.T) {
 	t.Run("EmptyString", func(t *testing.T) {
-		if isRateLimitError("") {
+		if IsRateLimitErrorMessage("") {
 			t.Error("Empty string should not be detected as rate limit error")
 		}
 	})
 
 	t.Run("OnlyWhitespace", func(t *testing.T) {
-		if isRateLimitError("   \t\n  ") {
+		if IsRateLimitErrorMessage("   \t\n  ") {
 			t.Error("Whitespace-only string should not be detected as rate limit error")
 		}
 	})
 
 	t.Run("UnicodeCharacters", func(t *testing.T) {
 		// Test with unicode characters that might affect case conversion
 		message := "RATE LIMIT exceeded 🚫"
-		if !isRateLimitError(message) {
+		if !IsRateLimitErrorMessage(message) {
 			t.Error("Message with unicode should still detect rate limit pattern")
 		}
 	})
@@ -544,7 +544,7 @@ func BenchmarkCalculateBackoff(b *testing.B) {
 	}
 }
 
-// Benchmark isRateLimitError performance
+// Benchmark IsRateLimitErrorMessage performance
 func BenchmarkIsRateLimitError(b *testing.B) {
 	messages := []string{
 		"rate limit exceeded",
@@ -559,7 +559,7 @@ func BenchmarkIsRateLimitError(b *testing.B) {
 
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
-		isRateLimitError(messages[i%len(messages)])
+		IsRateLimitErrorMessage(messages[i%len(messages)])
 	}
 }
 

core/utils.go

@@ -92,8 +92,8 @@ func validateRequest(req *schemas.BifrostRequest) *schemas.BifrostError {
 	return nil
 }
 
-// isRateLimitError checks if an error message indicates a rate limit issue
-func isRateLimitError(errorMessage string) bool {
+// IsRateLimitErrorMessage checks if an error message indicates a rate limit issue
+func IsRateLimitErrorMessage(errorMessage string) bool {
 	if errorMessage == "" {
 		return false
 	}

framework/modelcatalog/main.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"fmt"
 	"slices"
+	"strings"
 	"sync"
 	"time"
 
@@ -215,6 +216,51 @@ func (mc *ModelCatalog) GetProvidersForModel(model string) []schemas.ModelProvid
 			providers = append(providers, provider)
 		}
 	}
+
+	// Handler special provider cases
+	// 1. Handler openrouter models
+	if !slices.Contains(providers, schemas.OpenRouter) {
+		for _, provider := range providers {
+			if openRouterModels, ok := mc.modelPool[schemas.OpenRouter]; ok {
+				if slices.Contains(openRouterModels, string(provider)+"/"+model) {
+					providers = append(providers, schemas.OpenRouter)
+				}
+			}
+		}
+	}
+
+	// 2. Handle vertex models
+	if !slices.Contains(providers, schemas.Vertex) {
+		for _, provider := range providers {
+			if vertexModels, ok := mc.modelPool[schemas.Vertex]; ok {
+				if slices.Contains(vertexModels, string(provider)+"/"+model) {
+					providers = append(providers, schemas.Vertex)
+				}
+			}
+		}
+	}
+
+	// 3. Handle openai models for groq
+	if !slices.Contains(providers, schemas.Groq) && strings.Contains(model, "gpt-") {
+		if groqModels, ok := mc.modelPool[schemas.Groq]; ok {
+			if slices.Contains(groqModels, "openai/"+model) {
+				providers = append(providers, schemas.Groq)
+			}
+		}
+	}
+
+	// 4. Handle anthropic models for bedrock
+	if !slices.Contains(providers, schemas.Bedrock) && strings.Contains(model, "claude") {
+		if bedrockModels, ok := mc.modelPool[schemas.Bedrock]; ok {
+			for _, bedrockModel := range bedrockModels {
+				if strings.Contains(bedrockModel, model) {
+					providers = append(providers, schemas.Bedrock)
+					break
+				}
+			}
+		}
+	}
+
 	return providers
 }
 
@@ -235,6 +281,18 @@ func (mc *ModelCatalog) AddModelDataToPool(modelData *schemas.BifrostListModelsR
 	}
 }
 
+// RefineModelForProvider refines the model for a given provider.
+// e.g. "gpt-oss-120b" for groq provider -> "openai/gpt-4o-mini"
+func (mc *ModelCatalog) RefineModelForProvider(provider schemas.ModelProvider, model string) string {
+	switch provider {
+	case schemas.Groq:
+		if model == "gpt-oss-120b" {
+			return "openai/" + model
+		}
+	}
+	return model
+}
+
 // populateModelPool populates the model pool with all available models per provider (thread-safe)
 func (mc *ModelCatalog) populateModelPoolFromPricingData() {
 	// Acquire write lock for the entire rebuild operation

tests/core-providers/README.md

@@ -21,6 +21,29 @@ This directory contains comprehensive tests for all Bifrost AI providers, ensuri
 
 ## 🏃‍♂️ Running Tests
 
+### Parallel Test Execution
+
+All provider tests are configured to run in parallel using Go's `t.Parallel()` function. This allows multiple provider tests to execute concurrently, significantly reducing total test execution time.
+
+**Benefits:**
+- Faster test execution when testing multiple providers
+- Better resource utilization
+- Isolated test execution (each test creates its own client instance)
+
+**Usage:**
+```bash
+# Default: Tests run in parallel (up to GOMAXPROCS concurrent tests)
+go test -v ./tests/core-providers/
+
+# Explicitly set number of parallel tests
+go test -v ./tests/core-providers/ -parallel 10
+
+# Run tests sequentially (disable parallel execution)
+go test -v ./tests/core-providers/ -parallel 1
+```
+
+**Note:** Each test function creates its own isolated Bifrost client instance via `config.SetupTest()`, ensuring no shared state between parallel test executions.
+
 ### Development with Local Bifrost Core
 
 To test changes with a forked or local version of bifrost-core:
@@ -91,10 +114,16 @@ export GEMINI_API_KEY="your-gemini-key"
 # Run all tests with verbose output (recommended)
 go test -v ./tests/core-providers/
 
+# Run all tests in parallel (faster execution)
+# Tests are configured to run in parallel by default
+go test -v ./tests/core-providers/ -parallel 10
+
 # Run with debug logs
 go test -v ./tests/core-providers/ -debug
 ```
 
+**Note**: All provider tests are configured to run in parallel using `t.Parallel()`. This means multiple provider tests can execute concurrently, significantly reducing total test execution time. The number of parallel tests can be controlled using the `-parallel` flag (default is the number of CPUs).
+
 ### Run Specific Provider Tests
 
 ```bash

tests/core-providers/anthropic_test.go

@@ -10,6 +10,7 @@ import (
 )
 
 func TestAnthropic(t *testing.T) {
+	t.Parallel()
 	if os.Getenv("ANTHROPIC_API_KEY") == "" {
 		t.Skip("Skipping Anthropic tests because ANTHROPIC_API_KEY is not set")
 	}

tests/core-providers/azure_test.go

@@ -10,6 +10,7 @@ import (
 )
 
 func TestAzure(t *testing.T) {
+	t.Parallel()
 	t.Skip("Skipping Azure tests because Azure.")
 
 	if os.Getenv("AZURE_API_KEY") == "" {

tests/core-providers/bedrock_test.go

@@ -10,6 +10,7 @@ import (
 )
 
 func TestBedrock(t *testing.T) {
+	t.Parallel()
 	if os.Getenv("AWS_ACCESS_KEY_ID") == "" || os.Getenv("AWS_SECRET_ACCESS_KEY") == "" {
 		t.Skip("Skipping Bedrock embedding: AWS credentials not set")
 	}

tests/core-providers/cerebras_test.go

@@ -10,6 +10,7 @@ import (
 )
 
 func TestCerebras(t *testing.T) {
+	t.Parallel()
 	if os.Getenv("CEREBRAS_API_KEY") == "" {
 		t.Skip("Skipping Cerebras tests because CEREBRAS_API_KEY is not set")
 	}

tests/core-providers/cohere_test.go

@@ -10,6 +10,7 @@ import (
 )
 
 func TestCohere(t *testing.T) {
+	t.Parallel()
 	if os.Getenv("COHERE_API_KEY") == "" {
 		t.Skip("Skipping Cohere tests because COHERE_API_KEY is not set")
 	}