Skip to content

Commit 5a7e7a2

Browse files
akshaydeoakshaydeo
committed
gemini-sdk-batch-support
1 parent 6a51a9f commit 5a7e7a2

File tree

25 files changed

+1485
-55
lines changed

25 files changed

+1485
-55
lines changed

‎core/bifrost.go‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1130,6 +1130,84 @@ func (bifrost *Bifrost) BatchResultsRequest(ctx context.Context, req *schemas.Bi
11301130
return response, nil
11311131
}
11321132

1133+
// BatchDeleteRequest deletes a batch job.
1134+
func (bifrost *Bifrost) BatchDeleteRequest(ctx context.Context, req *schemas.BifrostBatchDeleteRequest) (*schemas.BifrostBatchDeleteResponse, *schemas.BifrostError) {
1135+
if req == nil {
1136+
return nil, &schemas.BifrostError{
1137+
IsBifrostError: false,
1138+
Error: &schemas.ErrorField{
1139+
Message: "batch delete request is nil",
1140+
},
1141+
}
1142+
}
1143+
if req.Provider == "" {
1144+
return nil, &schemas.BifrostError{
1145+
IsBifrostError: false,
1146+
Error: &schemas.ErrorField{
1147+
Message: "provider is required for batch delete request",
1148+
},
1149+
}
1150+
}
1151+
if req.BatchID == "" {
1152+
return nil, &schemas.BifrostError{
1153+
IsBifrostError: false,
1154+
Error: &schemas.ErrorField{
1155+
Message: "batch_id is required for batch delete request",
1156+
},
1157+
}
1158+
}
1159+
if ctx == nil {
1160+
ctx = bifrost.ctx
1161+
}
1162+
1163+
provider := bifrost.getProviderByKey(req.Provider)
1164+
if provider == nil {
1165+
return nil, &schemas.BifrostError{
1166+
IsBifrostError: false,
1167+
Error: &schemas.ErrorField{
1168+
Message: "provider not found for batch delete request",
1169+
},
1170+
}
1171+
}
1172+
1173+
config, err := bifrost.account.GetConfigForProvider(req.Provider)
1174+
if err != nil {
1175+
return nil, newBifrostErrorFromMsg(fmt.Sprintf("failed to get config for provider %s: %v", req.Provider, err.Error()))
1176+
}
1177+
if config == nil {
1178+
return nil, newBifrostErrorFromMsg(fmt.Sprintf("config is nil for provider %s", req.Provider))
1179+
}
1180+
1181+
// Determine the base provider type for key requirement checks
1182+
baseProvider := req.Provider
1183+
if config.CustomProviderConfig != nil && config.CustomProviderConfig.BaseProviderType != "" {
1184+
baseProvider = config.CustomProviderConfig.BaseProviderType
1185+
}
1186+
1187+
var key schemas.Key
1188+
if providerRequiresKey(baseProvider, config.CustomProviderConfig) {
1189+
keys, keyErr := bifrost.getAllSupportedKeys(&ctx, req.Provider, baseProvider)
1190+
if keyErr != nil {
1191+
return nil, newBifrostError(keyErr)
1192+
}
1193+
if len(keys) > 0 {
1194+
key = keys[0]
1195+
}
1196+
}
1197+
1198+
response, bifrostErr := executeRequestWithRetries(&ctx, config, func() (*schemas.BifrostBatchDeleteResponse, *schemas.BifrostError) {
1199+
return provider.BatchDelete(ctx, key, req)
1200+
}, schemas.BatchDeleteRequest, req.Provider, "")
1201+
if bifrostErr != nil {
1202+
bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{
1203+
RequestType: schemas.BatchDeleteRequest,
1204+
Provider: req.Provider,
1205+
}
1206+
return nil, bifrostErr
1207+
}
1208+
return response, nil
1209+
}
1210+
11331211
// FileUploadRequest uploads a file to the specified provider.
11341212
func (bifrost *Bifrost) FileUploadRequest(ctx context.Context, req *schemas.BifrostFileUploadRequest) (*schemas.BifrostFileUploadResponse, *schemas.BifrostError) {
11351213
if req == nil {

‎core/providers/anthropic/batch.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -779,3 +779,8 @@ func formatAnthropicTimestamp(unixTime int64) string {
779779
}
780780
return time.Unix(unixTime, 0).UTC().Format(time.RFC3339)
781781
}
782+
783+
// BatchDelete is not supported by Anthropic provider.
784+
func (provider *AnthropicProvider) BatchDelete(ctx context.Context, key schemas.Key, request *schemas.BifrostBatchDeleteRequest) (*schemas.BifrostBatchDeleteResponse, *schemas.BifrostError) {
785+
return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchDeleteRequest, provider.GetProviderKey())
786+
}

‎core/providers/azure/batch.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -434,3 +434,8 @@ func splitJSONL(data []byte) [][]byte {
434434
}
435435
return lines
436436
}
437+
438+
// BatchDelete is not supported by Azure provider.
439+
func (provider *AzureProvider) BatchDelete(ctx context.Context, key schemas.Key, request *schemas.BifrostBatchDeleteRequest) (*schemas.BifrostBatchDeleteResponse, *schemas.BifrostError) {
440+
return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchDeleteRequest, provider.GetProviderKey())
441+
}

‎core/providers/bedrock/batch.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -927,3 +927,8 @@ func splitJSONL(data []byte) [][]byte {
927927
}
928928
return lines
929929
}
930+
931+
// BatchDelete is not supported by Bedrock provider.
932+
func (provider *BedrockProvider) BatchDelete(ctx context.Context, key schemas.Key, request *schemas.BifrostBatchDeleteRequest) (*schemas.BifrostBatchDeleteResponse, *schemas.BifrostError) {
933+
return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchDeleteRequest, provider.GetProviderKey())
934+
}

‎core/providers/cerebras/batch.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,3 +32,8 @@ func (provider *CerebrasProvider) BatchResults(ctx context.Context, key schemas.
3232
return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchResultsRequest, provider.GetProviderKey())
3333
}
3434

35+
// BatchDelete is not supported by Cerebras provider.
36+
func (provider *CerebrasProvider) BatchDelete(ctx context.Context, key schemas.Key, request *schemas.BifrostBatchDeleteRequest) (*schemas.BifrostBatchDeleteResponse, *schemas.BifrostError) {
37+
return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchDeleteRequest, provider.GetProviderKey())
38+
}
39+

‎core/providers/cohere/batch.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,3 +32,8 @@ func (provider *CohereProvider) BatchResults(ctx context.Context, key schemas.Ke
3232
return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchResultsRequest, provider.GetProviderKey())
3333
}
3434

35+
// BatchDelete is not supported by Cohere provider.
36+
func (provider *CohereProvider) BatchDelete(ctx context.Context, key schemas.Key, request *schemas.BifrostBatchDeleteRequest) (*schemas.BifrostBatchDeleteResponse, *schemas.BifrostError) {
37+
return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchDeleteRequest, provider.GetProviderKey())
38+
}
39+

‎core/providers/elevenlabs/batch.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,3 +32,8 @@ func (provider *ElevenlabsProvider) BatchResults(ctx context.Context, key schema
3232
return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchResultsRequest, provider.GetProviderKey())
3333
}
3434

35+
// BatchDelete is not supported by Elevenlabs provider.
36+
func (provider *ElevenlabsProvider) BatchDelete(ctx context.Context, key schemas.Key, request *schemas.BifrostBatchDeleteRequest) (*schemas.BifrostBatchDeleteResponse, *schemas.BifrostError) {
37+
return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchDeleteRequest, provider.GetProviderKey())
38+
}
39+

‎core/providers/gemini/batch.go‎

Lines changed: 202 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -839,3 +839,205 @@ func (provider *GeminiProvider) BatchResults(ctx context.Context, key schemas.Ke
839839
},
840840
}, nil
841841
}
842+
843+
// ==================== SDK RESPONSE CONVERTERS ====================
844+
// These functions convert Bifrost batch responses to Google GenAI SDK format.
845+
846+
// ToGeminiJobState converts Bifrost batch status to Gemini SDK job state.
847+
func ToGeminiJobState(status schemas.BatchStatus) string {
848+
switch status {
849+
case schemas.BatchStatusValidating:
850+
return GeminiJobStatePending
851+
case schemas.BatchStatusInProgress:
852+
return GeminiJobStateRunning
853+
case schemas.BatchStatusFinalizing:
854+
return GeminiJobStateRunning
855+
case schemas.BatchStatusCompleted:
856+
return GeminiJobStateSucceeded
857+
case schemas.BatchStatusFailed:
858+
return GeminiJobStateFailed
859+
case schemas.BatchStatusCancelling:
860+
return GeminiJobStateCancelling
861+
case schemas.BatchStatusCancelled:
862+
return GeminiJobStateCancelled
863+
case schemas.BatchStatusExpired:
864+
return GeminiJobStateFailed
865+
default:
866+
return GeminiJobStatePending
867+
}
868+
}
869+
870+
// ToGeminiBatchJobResponse converts a BifrostBatchCreateResponse to Gemini SDK format.
871+
func ToGeminiBatchJobResponse(resp *schemas.BifrostBatchCreateResponse) *GeminiBatchJobResponseSDK {
872+
if resp == nil {
873+
return nil
874+
}
875+
876+
result := &GeminiBatchJobResponseSDK{
877+
Name: resp.ID,
878+
State: ToGeminiJobState(resp.Status),
879+
}
880+
881+
// Add metadata if available
882+
if resp.CreatedAt > 0 {
883+
result.Metadata = &GeminiBatchMetadata{
884+
Name: resp.ID,
885+
State: ToGeminiJobState(resp.Status),
886+
CreateTime: time.Unix(resp.CreatedAt, 0).Format(time.RFC3339),
887+
BatchStats: &GeminiBatchStats{
888+
RequestCount: resp.RequestCounts.Total,
889+
PendingRequestCount: resp.RequestCounts.Total - resp.RequestCounts.Completed,
890+
SuccessfulRequestCount: resp.RequestCounts.Completed - resp.RequestCounts.Failed,
891+
},
892+
}
893+
}
894+
895+
return result
896+
}
897+
898+
// ToGeminiBatchRetrieveResponse converts a BifrostBatchRetrieveResponse to Gemini SDK format.
899+
func ToGeminiBatchRetrieveResponse(resp *schemas.BifrostBatchRetrieveResponse) *GeminiBatchJobResponseSDK {
900+
if resp == nil {
901+
return nil
902+
}
903+
904+
result := &GeminiBatchJobResponseSDK{
905+
Name: resp.ID,
906+
State: ToGeminiJobState(resp.Status),
907+
}
908+
909+
// Add metadata
910+
result.Metadata = &GeminiBatchMetadata{
911+
Name: resp.ID,
912+
State: ToGeminiJobState(resp.Status),
913+
CreateTime: time.Unix(resp.CreatedAt, 0).Format(time.RFC3339),
914+
BatchStats: &GeminiBatchStats{
915+
RequestCount: resp.RequestCounts.Total,
916+
PendingRequestCount: resp.RequestCounts.Total - resp.RequestCounts.Completed,
917+
SuccessfulRequestCount: resp.RequestCounts.Completed - resp.RequestCounts.Failed,
918+
},
919+
}
920+
921+
if resp.CompletedAt != nil {
922+
result.Metadata.EndTime = time.Unix(*resp.CompletedAt, 0).Format(time.RFC3339)
923+
}
924+
925+
// Add output file info if available
926+
if resp.OutputFileID != nil {
927+
result.Dest = &GeminiBatchDest{
928+
FileName: *resp.OutputFileID,
929+
}
930+
}
931+
932+
return result
933+
}
934+
935+
// ToGeminiBatchListResponse converts a BifrostBatchListResponse to Gemini SDK format.
936+
func ToGeminiBatchListResponse(resp *schemas.BifrostBatchListResponse) *GeminiBatchListResponseSDK {
937+
if resp == nil {
938+
return nil
939+
}
940+
941+
jobs := make([]GeminiBatchJobResponseSDK, 0, len(resp.Data))
942+
for _, batch := range resp.Data {
943+
job := GeminiBatchJobResponseSDK{
944+
Name: batch.ID,
945+
State: ToGeminiJobState(batch.Status),
946+
}
947+
948+
// Add metadata
949+
job.Metadata = &GeminiBatchMetadata{
950+
Name: batch.ID,
951+
State: ToGeminiJobState(batch.Status),
952+
CreateTime: time.Unix(batch.CreatedAt, 0).Format(time.RFC3339),
953+
BatchStats: &GeminiBatchStats{
954+
RequestCount: batch.RequestCounts.Total,
955+
PendingRequestCount: batch.RequestCounts.Total - batch.RequestCounts.Completed,
956+
SuccessfulRequestCount: batch.RequestCounts.Completed - batch.RequestCounts.Failed,
957+
},
958+
}
959+
960+
jobs = append(jobs, job)
961+
}
962+
963+
result := &GeminiBatchListResponseSDK{
964+
BatchJobs: jobs,
965+
}
966+
967+
if resp.NextCursor != nil {
968+
result.NextPageToken = *resp.NextCursor
969+
}
970+
971+
return result
972+
}
973+
974+
// ToGeminiBatchCancelResponse converts a BifrostBatchCancelResponse to Gemini SDK format.
975+
func ToGeminiBatchCancelResponse(resp *schemas.BifrostBatchCancelResponse) *GeminiBatchJobResponseSDK {
976+
if resp == nil {
977+
return nil
978+
}
979+
980+
return &GeminiBatchJobResponseSDK{
981+
Name: resp.ID,
982+
State: ToGeminiJobState(resp.Status),
983+
}
984+
}
985+
986+
// BatchDelete deletes a batch job for Gemini.
987+
func (provider *GeminiProvider) BatchDelete(ctx context.Context, key schemas.Key, request *schemas.BifrostBatchDeleteRequest) (*schemas.BifrostBatchDeleteResponse, *schemas.BifrostError) {
988+
if err := providerUtils.CheckOperationAllowed(schemas.Gemini, provider.customProviderConfig, schemas.BatchDeleteRequest); err != nil {
989+
return nil, err
990+
}
991+
992+
providerName := provider.GetProviderKey()
993+
994+
if request.BatchID == "" {
995+
return nil, providerUtils.NewBifrostOperationError("batch_id is required", nil, providerName)
996+
}
997+
998+
// Create HTTP request
999+
req := fasthttp.AcquireRequest()
1000+
resp := fasthttp.AcquireResponse()
1001+
defer fasthttp.ReleaseRequest(req)
1002+
defer fasthttp.ReleaseResponse(resp)
1003+
1004+
// Build URL for delete operation
1005+
batchID := request.BatchID
1006+
var url string
1007+
if strings.HasPrefix(batchID, "batches/") {
1008+
url = fmt.Sprintf("%s/%s", provider.networkConfig.BaseURL, batchID)
1009+
} else {
1010+
url = fmt.Sprintf("%s/batches/%s", provider.networkConfig.BaseURL, batchID)
1011+
}
1012+
1013+
provider.logger.Debug("gemini batch delete url: " + url)
1014+
providerUtils.SetExtraHeaders(ctx, req, provider.networkConfig.ExtraHeaders, nil)
1015+
req.SetRequestURI(url)
1016+
req.Header.SetMethod(http.MethodDelete)
1017+
if key.Value != "" {
1018+
req.Header.Set("x-goog-api-key", key.Value)
1019+
}
1020+
req.Header.SetContentType("application/json")
1021+
1022+
// Make request
1023+
latency, bifrostErr := providerUtils.MakeRequestWithContext(ctx, provider.client, req, resp)
1024+
if bifrostErr != nil {
1025+
return nil, bifrostErr
1026+
}
1027+
1028+
// Handle response
1029+
if resp.StatusCode() != fasthttp.StatusOK && resp.StatusCode() != fasthttp.StatusNoContent {
1030+
return nil, parseGeminiError(providerName, resp)
1031+
}
1032+
1033+
return &schemas.BifrostBatchDeleteResponse{
1034+
ID: request.BatchID,
1035+
Object: "batch",
1036+
Deleted: true,
1037+
ExtraFields: schemas.BifrostResponseExtraFields{
1038+
RequestType: schemas.BatchDeleteRequest,
1039+
Provider: providerName,
1040+
Latency: latency.Milliseconds(),
1041+
},
1042+
}, nil
1043+
}

‎core/providers/gemini/files.go‎

Lines changed: 17 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -463,17 +463,23 @@ func (provider *GeminiProvider) FileContent(ctx context.Context, key schemas.Key
463463
}
464464

465465
// ToGeminiFileUploadResponse converts a Bifrost file upload response to Gemini format.
466-
func ToGeminiFileUploadResponse(resp *schemas.BifrostFileUploadResponse) map[string]interface{} {
467-
return map[string]interface{}{
468-
"file": map[string]interface{}{
469-
"name": resp.ID,
470-
"displayName": resp.Filename,
471-
"mimeType": "application/octet-stream",
472-
"sizeBytes": fmt.Sprintf("%d", resp.Bytes),
473-
"createTime": formatGeminiTimestamp(resp.CreatedAt),
474-
"state": toGeminiFileState(resp.Status),
475-
"uri": resp.StorageURI,
476-
"expirationTime": formatGeminiTimestamp(safeDerefInt64(resp.ExpiresAt)),
466+
// Uses snake_case field names to match Google's API format.
467+
// GeminiFileUploadResponseWrapper is a wrapper that contains the file response for the upload API.
468+
type GeminiFileUploadResponseWrapper struct {
469+
File GeminiFileResponse `json:"file"`
470+
}
471+
472+
func ToGeminiFileUploadResponse(resp *schemas.BifrostFileUploadResponse) *GeminiFileUploadResponseWrapper {
473+
return &GeminiFileUploadResponseWrapper{
474+
File: GeminiFileResponse{
475+
Name: resp.ID,
476+
DisplayName: resp.Filename,
477+
MimeType: "application/octet-stream",
478+
SizeBytes: fmt.Sprintf("%d", resp.Bytes),
479+
CreateTime: formatGeminiTimestamp(resp.CreatedAt),
480+
State: toGeminiFileState(resp.Status),
481+
URI: resp.StorageURI,
482+
ExpirationTime: formatGeminiTimestamp(safeDerefInt64(resp.ExpiresAt)),
477483
},
478484
}
479485
}

0 commit comments

Comments
 (0)