Skip to content

Commit 95bcd45

Browse files
akshaydeoakshaydeo
committed
batch config for each key
1 parent 5dd798c commit 95bcd45

File tree

65 files changed

+3955
-2484
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

65 files changed

+3955
-2484
lines changed

‎Makefile‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -149,11 +149,9 @@ build: build-ui ## Build bifrost-http binary
149149
echo "$(GREEN)╚═══════════════════════════════════════╝$(NC)"; \
150150
fi
151151
@if [ -n "$(DYNAMIC)" ]; then \
152-
echo "$(YELLOW)Note: This will create a dynamically linked build.$(NC)"; \
153-
echo "$(YELLOW)To build with dynamic plugin support.$(NC)"; \
152+
echo "$(YELLOW)Note: This will create a dynamically linked build.$(NC)"; \
154153
else \
155-
echo "$(YELLOW)Note: This will create a statically linked build.$(NC)"; \
156-
echo "$(YELLOW)To build with dynamic plugin support.$(NC)"; \
154+
echo "$(YELLOW)Note: This will create a statically linked build.$(NC)"; \
157155
fi
158156
@mkdir -p ./tmp
159157
@TARGET_OS="$(GOOS)"; \

‎core/bifrost.go‎

Lines changed: 160 additions & 61 deletions
Large diffs are not rendered by default.

‎core/go.mod‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
module github.com/maximhq/bifrost/core
22

3-
go 1.24.3
3+
go 1.25.5
44

55
require (
66
github.com/aws/aws-sdk-go-v2 v1.41.0

‎core/internal/testutil/account.go‎

Lines changed: 79 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -66,12 +66,12 @@ type ComprehensiveTestConfig struct {
6666
TranscriptionModel string
6767
SpeechSynthesisModel string
6868
Scenarios TestScenarios
69-
Fallbacks []schemas.Fallback // for chat, responses, image and reasoning tests
70-
TextCompletionFallbacks []schemas.Fallback // for text completion tests
71-
TranscriptionFallbacks []schemas.Fallback // for transcription tests
72-
SpeechSynthesisFallbacks []schemas.Fallback // for speech synthesis tests
73-
EmbeddingFallbacks []schemas.Fallback // for embedding tests
74-
SkipReason string // Reason to skip certain tests
69+
Fallbacks []schemas.Fallback // for chat, responses, image and reasoning tests
70+
TextCompletionFallbacks []schemas.Fallback // for text completion tests
71+
TranscriptionFallbacks []schemas.Fallback // for transcription tests
72+
SpeechSynthesisFallbacks []schemas.Fallback // for speech synthesis tests
73+
EmbeddingFallbacks []schemas.Fallback // for embedding tests
74+
SkipReason string // Reason to skip certain tests
7575
BatchExtraParams map[string]interface{} // Extra params for batch operations (e.g., role_arn, output_s3_uri for Bedrock)
7676
FileExtraParams map[string]interface{} // Extra params for file operations (e.g., s3_bucket for Bedrock)
7777
}
@@ -117,25 +117,28 @@ func (account *ComprehensiveTestAccount) GetKeysForProvider(ctx *context.Context
117117
case schemas.OpenAI:
118118
return []schemas.Key{
119119
{
120-
Value: os.Getenv("OPENAI_API_KEY"),
121-
Models: []string{},
122-
Weight: 1.0,
120+
Value: os.Getenv("OPENAI_API_KEY"),
121+
Models: []string{},
122+
Weight: 1.0,
123+
UseForBatchAPI: bifrost.Ptr(true),
123124
},
124125
}, nil
125126
case ProviderOpenAICustom:
126127
return []schemas.Key{
127128
{
128-
Value: os.Getenv("OPENAI_API_KEY"), // Use GROQ API key for OpenAI-compatible endpoint
129-
Models: []string{},
130-
Weight: 1.0,
129+
Value: os.Getenv("OPENAI_API_KEY"), // Use GROQ API key for OpenAI-compatible endpoint
130+
Models: []string{},
131+
Weight: 1.0,
132+
UseForBatchAPI: bifrost.Ptr(true),
131133
},
132134
}, nil
133135
case schemas.Anthropic:
134136
return []schemas.Key{
135137
{
136-
Value: os.Getenv("ANTHROPIC_API_KEY"),
137-
Models: []string{},
138-
Weight: 1.0,
138+
Value: os.Getenv("ANTHROPIC_API_KEY"),
139+
Models: []string{},
140+
Weight: 1.0,
141+
UseForBatchAPI: bifrost.Ptr(true),
139142
},
140143
}, nil
141144
case schemas.Bedrock:
@@ -157,6 +160,25 @@ func (account *ComprehensiveTestAccount) GetKeysForProvider(ctx *context.Context
157160
},
158161
},
159162
},
163+
{
164+
Models: []string{},
165+
Weight: 1.0,
166+
BedrockKeyConfig: &schemas.BedrockKeyConfig{
167+
AccessKey: os.Getenv("AWS_ACCESS_KEY_ID"),
168+
SecretKey: os.Getenv("AWS_SECRET_ACCESS_KEY"),
169+
SessionToken: bifrost.Ptr(os.Getenv("AWS_SESSION_TOKEN")),
170+
Region: bifrost.Ptr(getEnvWithDefault("AWS_REGION", "us-east-1")),
171+
ARN: bifrost.Ptr(os.Getenv("AWS_BEDROCK_ARN")),
172+
Deployments: map[string]string{
173+
"claude-3.5-sonnet": "anthropic.claude-3-5-sonnet-20240620-v1:0",
174+
"claude-3.7-sonnet": "us.anthropic.claude-3-7-sonnet-20250219-v1:0",
175+
"claude-4-sonnet": "global.anthropic.claude-sonnet-4-20250514-v1:0",
176+
"claude-4.5-sonnet": "global.anthropic.claude-sonnet-4-5-20250929-v1:0",
177+
"claude-4.5-haiku": "global.anthropic.claude-haiku-4-5-20251001-v1:0",
178+
},
179+
},
180+
UseForBatchAPI: bifrost.Ptr(true),
181+
},
160182
{
161183
Models: []string{"cohere.embed-v4:0"},
162184
Weight: 1.0,
@@ -171,9 +193,10 @@ func (account *ComprehensiveTestAccount) GetKeysForProvider(ctx *context.Context
171193
case schemas.Cohere:
172194
return []schemas.Key{
173195
{
174-
Value: os.Getenv("COHERE_API_KEY"),
175-
Models: []string{},
176-
Weight: 1.0,
196+
Value: os.Getenv("COHERE_API_KEY"),
197+
Models: []string{},
198+
Weight: 1.0,
199+
UseForBatchAPI: bifrost.Ptr(true),
177200
},
178201
}, nil
179202
case schemas.Azure:
@@ -192,6 +215,7 @@ func (account *ComprehensiveTestAccount) GetKeysForProvider(ctx *context.Context
192215
"text-embedding-ada-002": "text-embedding-ada-002",
193216
},
194217
},
218+
UseForBatchAPI: bifrost.Ptr(true),
195219
},
196220
}, nil
197221
case schemas.Vertex:
@@ -205,78 +229,88 @@ func (account *ComprehensiveTestAccount) GetKeysForProvider(ctx *context.Context
205229
Region: getEnvWithDefault("VERTEX_REGION", "us-central1"),
206230
AuthCredentials: os.Getenv("VERTEX_CREDENTIALS"),
207231
},
232+
UseForBatchAPI: bifrost.Ptr(true),
208233
},
209234
}, nil
210235
case schemas.Mistral:
211236
return []schemas.Key{
212237
{
213-
Value: os.Getenv("MISTRAL_API_KEY"),
214-
Models: []string{},
215-
Weight: 1.0,
238+
Value: os.Getenv("MISTRAL_API_KEY"),
239+
Models: []string{},
240+
Weight: 1.0,
241+
UseForBatchAPI: bifrost.Ptr(true),
216242
},
217243
}, nil
218244
case schemas.Groq:
219245
return []schemas.Key{
220246
{
221-
Value: os.Getenv("GROQ_API_KEY"),
222-
Models: []string{},
223-
Weight: 1.0,
247+
Value: os.Getenv("GROQ_API_KEY"),
248+
Models: []string{},
249+
Weight: 1.0,
250+
UseForBatchAPI: bifrost.Ptr(true),
224251
},
225252
}, nil
226253
case schemas.Parasail:
227254
return []schemas.Key{
228255
{
229-
Value: os.Getenv("PARASAIL_API_KEY"),
230-
Models: []string{},
231-
Weight: 1.0,
256+
Value: os.Getenv("PARASAIL_API_KEY"),
257+
Models: []string{},
258+
Weight: 1.0,
259+
UseForBatchAPI: bifrost.Ptr(true),
232260
},
233261
}, nil
234262
case schemas.Elevenlabs:
235263
return []schemas.Key{
236264
{
237-
Value: os.Getenv("ELEVENLABS_API_KEY"),
238-
Models: []string{},
239-
Weight: 1.0,
265+
Value: os.Getenv("ELEVENLABS_API_KEY"),
266+
Models: []string{},
267+
Weight: 1.0,
268+
UseForBatchAPI: bifrost.Ptr(true),
240269
},
241270
}, nil
242271
case schemas.Perplexity:
243272
return []schemas.Key{
244273
{
245-
Value: os.Getenv("PERPLEXITY_API_KEY"),
246-
Models: []string{},
247-
Weight: 1.0,
274+
Value: os.Getenv("PERPLEXITY_API_KEY"),
275+
Models: []string{},
276+
Weight: 1.0,
277+
UseForBatchAPI: bifrost.Ptr(true),
248278
},
249279
}, nil
250280
case schemas.Cerebras:
251281
return []schemas.Key{
252282
{
253-
Value: os.Getenv("CEREBRAS_API_KEY"),
254-
Models: []string{},
255-
Weight: 1.0,
283+
Value: os.Getenv("CEREBRAS_API_KEY"),
284+
Models: []string{},
285+
Weight: 1.0,
286+
UseForBatchAPI: bifrost.Ptr(true),
256287
},
257288
}, nil
258289
case schemas.Gemini:
259290
return []schemas.Key{
260291
{
261-
Value: os.Getenv("GEMINI_API_KEY"),
262-
Models: []string{},
263-
Weight: 1.0,
292+
Value: os.Getenv("GEMINI_API_KEY"),
293+
Models: []string{},
294+
Weight: 1.0,
295+
UseForBatchAPI: bifrost.Ptr(true),
264296
},
265297
}, nil
266298
case schemas.OpenRouter:
267299
return []schemas.Key{
268300
{
269-
Value: os.Getenv("OPENROUTER_API_KEY"),
270-
Models: []string{},
271-
Weight: 1.0,
301+
Value: os.Getenv("OPENROUTER_API_KEY"),
302+
Models: []string{},
303+
Weight: 1.0,
304+
UseForBatchAPI: bifrost.Ptr(true),
272305
},
273306
}, nil
274307
case schemas.Nebius:
275308
return []schemas.Key{
276309
{
277-
Value: os.Getenv("NEBIUS_API_KEY"),
278-
Models: []string{},
279-
Weight: 1.0,
310+
Value: os.Getenv("NEBIUS_API_KEY"),
311+
Models: []string{},
312+
Weight: 1.0,
313+
UseForBatchAPI: bifrost.Ptr(true),
280314
},
281315
}, nil
282316
default:

‎core/internal/testutil/batch.go‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ func RunBatchCreateTest(t *testing.T, client *bifrost.Bifrost, ctx context.Conte
2222
// Create a batch request with a simple chat completion
2323
request := &schemas.BifrostBatchCreateRequest{
2424
Provider: testConfig.Provider,
25-
Model: testConfig.ChatModel,
25+
Model: schemas.Ptr(testConfig.ChatModel),
2626
Endpoint: schemas.BatchEndpointChatCompletions,
2727
Requests: []schemas.BatchRequestItem{
2828
{
@@ -113,7 +113,7 @@ func RunBatchRetrieveTest(t *testing.T, client *bifrost.Bifrost, ctx context.Con
113113
// Note: In real tests, you might want to use an existing batch ID
114114
createRequest := &schemas.BifrostBatchCreateRequest{
115115
Provider: testConfig.Provider,
116-
Model: testConfig.ChatModel,
116+
Model: schemas.Ptr(testConfig.ChatModel),
117117
Endpoint: schemas.BatchEndpointChatCompletions,
118118
Requests: []schemas.BatchRequestItem{
119119
{
@@ -185,7 +185,7 @@ func RunBatchCancelTest(t *testing.T, client *bifrost.Bifrost, ctx context.Conte
185185
// First, create a batch to cancel
186186
createRequest := &schemas.BifrostBatchCreateRequest{
187187
Provider: testConfig.Provider,
188-
Model: testConfig.ChatModel,
188+
Model: schemas.Ptr(testConfig.ChatModel),
189189
Endpoint: schemas.BatchEndpointChatCompletions,
190190
Requests: []schemas.BatchRequestItem{
191191
{
@@ -295,7 +295,7 @@ func RunBatchUnsupportedTest(t *testing.T, client *bifrost.Bifrost, ctx context.
295295
// Try to create a batch - should fail with unsupported error
296296
request := &schemas.BifrostBatchCreateRequest{
297297
Provider: testConfig.Provider,
298-
Model: testConfig.ChatModel,
298+
Model: schemas.Ptr(testConfig.ChatModel),
299299
Endpoint: schemas.BatchEndpointChatCompletions,
300300
Requests: []schemas.BatchRequestItem{
301301
{
@@ -691,7 +691,7 @@ func RunFileAndBatchIntegrationTest(t *testing.T, client *bifrost.Bifrost, ctx c
691691
// Step 2: Create a batch using the uploaded file
692692
batchRequest := &schemas.BifrostBatchCreateRequest{
693693
Provider: testConfig.Provider,
694-
Model: testConfig.ChatModel,
694+
Model: schemas.Ptr(testConfig.ChatModel),
695695
InputFileID: uploadResponse.ID,
696696
Endpoint: schemas.BatchEndpointChatCompletions,
697697
CompletionWindow: "24h",

0 commit comments

Comments
 (0)