gemini-sdk-batch-support

Author: akshaydeo

core/bifrost.go

@@ -1260,6 +1260,84 @@ func (bifrost *Bifrost) BatchResultsRequest(ctx context.Context, req *schemas.Bi
 	return response, nil
 }
 
+// BatchDeleteRequest deletes a batch job.
+func (bifrost *Bifrost) BatchDeleteRequest(ctx context.Context, req *schemas.BifrostBatchDeleteRequest) (*schemas.BifrostBatchDeleteResponse, *schemas.BifrostError) {
+	if req == nil {
+		return nil, &schemas.BifrostError{
+			IsBifrostError: false,
+			Error: &schemas.ErrorField{
+				Message: "batch delete request is nil",
+			},
+		}
+	}
+	if req.Provider == "" {
+		return nil, &schemas.BifrostError{
+			IsBifrostError: false,
+			Error: &schemas.ErrorField{
+				Message: "provider is required for batch delete request",
+			},
+		}
+	}
+	if req.BatchID == "" {
+		return nil, &schemas.BifrostError{
+			IsBifrostError: false,
+			Error: &schemas.ErrorField{
+				Message: "batch_id is required for batch delete request",
+			},
+		}
+	}
+	if ctx == nil {
+		ctx = bifrost.ctx
+	}
+
+	provider := bifrost.getProviderByKey(req.Provider)
+	if provider == nil {
+		return nil, &schemas.BifrostError{
+			IsBifrostError: false,
+			Error: &schemas.ErrorField{
+				Message: "provider not found for batch delete request",
+			},
+		}
+	}
+
+	config, err := bifrost.account.GetConfigForProvider(req.Provider)
+	if err != nil {
+		return nil, newBifrostErrorFromMsg(fmt.Sprintf("failed to get config for provider %s: %v", req.Provider, err.Error()))
+	}
+	if config == nil {
+		return nil, newBifrostErrorFromMsg(fmt.Sprintf("config is nil for provider %s", req.Provider))
+	}
+
+	// Determine the base provider type for key requirement checks
+	baseProvider := req.Provider
+	if config.CustomProviderConfig != nil && config.CustomProviderConfig.BaseProviderType != "" {
+		baseProvider = config.CustomProviderConfig.BaseProviderType
+	}
+
+	var key schemas.Key
+	if providerRequiresKey(baseProvider, config.CustomProviderConfig) {
+		keys, keyErr := bifrost.getAllSupportedKeys(&ctx, req.Provider, baseProvider)
+		if keyErr != nil {
+			return nil, newBifrostError(keyErr)
+		}
+		if len(keys) > 0 {
+			key = keys[0]
+		}
+	}
+
+	response, bifrostErr := executeRequestWithRetries(&ctx, config, func() (*schemas.BifrostBatchDeleteResponse, *schemas.BifrostError) {
+		return provider.BatchDelete(ctx, key, req)
+	}, schemas.BatchDeleteRequest, req.Provider, "")
+	if bifrostErr != nil {
+		bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{
+			RequestType: schemas.BatchDeleteRequest,
+			Provider:    req.Provider,
+		}
+		return nil, bifrostErr
+	}
+	return response, nil
+}
+
 // FileUploadRequest uploads a file to the specified provider.
 func (bifrost *Bifrost) FileUploadRequest(ctx context.Context, req *schemas.BifrostFileUploadRequest) (*schemas.BifrostFileUploadResponse, *schemas.BifrostError) {
 	if req == nil {

core/providers/anthropic/batch.go

@@ -377,3 +377,8 @@ func formatAnthropicTimestamp(unixTime int64) string {
 	}
 	return time.Unix(unixTime, 0).UTC().Format(time.RFC3339)
 }
+
+// BatchDelete is not supported by Anthropic provider.
+func (provider *AnthropicProvider) BatchDelete(ctx context.Context, key schemas.Key, request *schemas.BifrostBatchDeleteRequest) (*schemas.BifrostBatchDeleteResponse, *schemas.BifrostError) {
+	return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchDeleteRequest, provider.GetProviderKey())
+}

core/providers/cerebras/cerebras.go

@@ -260,3 +260,8 @@ func (provider *CerebrasProvider) BatchCancel(_ context.Context, _ schemas.Key,
 func (provider *CerebrasProvider) BatchResults(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchResultsRequest) (*schemas.BifrostBatchResultsResponse, *schemas.BifrostError) {
 	return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchResultsRequest, provider.GetProviderKey())
 }
+
+// BatchDelete is not supported by Cerebras provider.
+func (provider *CerebrasProvider) BatchDelete(ctx context.Context, key schemas.Key, request *schemas.BifrostBatchDeleteRequest) (*schemas.BifrostBatchDeleteResponse, *schemas.BifrostError) {
+	return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchDeleteRequest, provider.GetProviderKey())
+}

core/providers/cohere/cohere.go

@@ -866,6 +866,12 @@ func (provider *CohereProvider) BatchResults(_ context.Context, _ schemas.Key, _
 	return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchResultsRequest, provider.GetProviderKey())
 }
 
+// BatchDelete is not supported by Cohere provider.
+func (provider *CohereProvider) BatchDelete(ctx context.Context, key schemas.Key, request *schemas.BifrostBatchDeleteRequest) (*schemas.BifrostBatchDeleteResponse, *schemas.BifrostError) {
+	return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchDeleteRequest, provider.GetProviderKey())
+}
+
+
 // FileUpload is not supported by Cohere provider.
 func (provider *CohereProvider) FileUpload(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileUploadRequest) (*schemas.BifrostFileUploadResponse, *schemas.BifrostError) {
 	return nil, providerUtils.NewUnsupportedOperationError(schemas.FileUploadRequest, provider.GetProviderKey())

core/providers/gemini/batch.go

@@ -240,3 +240,146 @@ func extractGeminiUsageMetadata(geminiResponse *GenerateContentResponse) (int, i
 	}
 	return inputTokens, outputTokens, totalTokens
 }
+
+// ==================== SDK RESPONSE CONVERTERS ====================
+// These functions convert Bifrost batch responses to Google GenAI SDK format.
+
+// ToGeminiJobState converts Bifrost batch status to Gemini SDK job state.
+func ToGeminiJobState(status schemas.BatchStatus) string {
+	switch status {
+	case schemas.BatchStatusValidating:
+		return GeminiJobStatePending
+	case schemas.BatchStatusInProgress:
+		return GeminiJobStateRunning
+	case schemas.BatchStatusFinalizing:
+		return GeminiJobStateRunning
+	case schemas.BatchStatusCompleted:
+		return GeminiJobStateSucceeded
+	case schemas.BatchStatusFailed:
+		return GeminiJobStateFailed
+	case schemas.BatchStatusCancelling:
+		return GeminiJobStateCancelling
+	case schemas.BatchStatusCancelled:
+		return GeminiJobStateCancelled
+	case schemas.BatchStatusExpired:
+		return GeminiJobStateFailed
+	default:
+		return GeminiJobStatePending
+	}
+}
+
+// ToGeminiBatchJobResponse converts a BifrostBatchCreateResponse to Gemini SDK format.
+func ToGeminiBatchJobResponse(resp *schemas.BifrostBatchCreateResponse) *GeminiBatchJobResponseSDK {
+	if resp == nil {
+		return nil
+	}
+
+	result := &GeminiBatchJobResponseSDK{
+		Name:  resp.ID,
+		State: ToGeminiJobState(resp.Status),
+	}
+
+	// Add metadata if available
+	if resp.CreatedAt > 0 {
+		result.Metadata = &GeminiBatchMetadata{
+			Name:       resp.ID,
+			State:      ToGeminiJobState(resp.Status),
+			CreateTime: time.Unix(resp.CreatedAt, 0).Format(time.RFC3339),
+			BatchStats: &GeminiBatchStats{
+				RequestCount:           resp.RequestCounts.Total,
+				PendingRequestCount:    resp.RequestCounts.Total - resp.RequestCounts.Completed,
+				SuccessfulRequestCount: resp.RequestCounts.Completed - resp.RequestCounts.Failed,
+			},
+		}
+	}
+
+	return result
+}
+
+// ToGeminiBatchRetrieveResponse converts a BifrostBatchRetrieveResponse to Gemini SDK format.
+func ToGeminiBatchRetrieveResponse(resp *schemas.BifrostBatchRetrieveResponse) *GeminiBatchJobResponseSDK {
+	if resp == nil {
+		return nil
+	}
+
+	result := &GeminiBatchJobResponseSDK{
+		Name:  resp.ID,
+		State: ToGeminiJobState(resp.Status),
+	}
+
+	// Add metadata
+	result.Metadata = &GeminiBatchMetadata{
+		Name:       resp.ID,
+		State:      ToGeminiJobState(resp.Status),
+		CreateTime: time.Unix(resp.CreatedAt, 0).Format(time.RFC3339),
+		BatchStats: &GeminiBatchStats{
+			RequestCount:           resp.RequestCounts.Total,
+			PendingRequestCount:    resp.RequestCounts.Total - resp.RequestCounts.Completed,
+			SuccessfulRequestCount: resp.RequestCounts.Completed - resp.RequestCounts.Failed,
+		},
+	}
+
+	if resp.CompletedAt != nil {
+		result.Metadata.EndTime = time.Unix(*resp.CompletedAt, 0).Format(time.RFC3339)
+	}
+
+	// Add output file info if available
+	if resp.OutputFileID != nil {
+		result.Dest = &GeminiBatchDest{
+			FileName: *resp.OutputFileID,
+		}
+	}
+
+	return result
+}
+
+// ToGeminiBatchListResponse converts a BifrostBatchListResponse to Gemini SDK format.
+func ToGeminiBatchListResponse(resp *schemas.BifrostBatchListResponse) *GeminiBatchListResponseSDK {
+	if resp == nil {
+		return nil
+	}
+
+	jobs := make([]GeminiBatchJobResponseSDK, 0, len(resp.Data))
+	for _, batch := range resp.Data {
+		job := GeminiBatchJobResponseSDK{
+			Name:  batch.ID,
+			State: ToGeminiJobState(batch.Status),
+		}
+
+		// Add metadata
+		job.Metadata = &GeminiBatchMetadata{
+			Name:       batch.ID,
+			State:      ToGeminiJobState(batch.Status),
+			CreateTime: time.Unix(batch.CreatedAt, 0).Format(time.RFC3339),
+			BatchStats: &GeminiBatchStats{
+				RequestCount:           batch.RequestCounts.Total,
+				PendingRequestCount:    batch.RequestCounts.Total - batch.RequestCounts.Completed,
+				SuccessfulRequestCount: batch.RequestCounts.Completed - batch.RequestCounts.Failed,
+			},
+		}
+
+		jobs = append(jobs, job)
+	}
+
+	result := &GeminiBatchListResponseSDK{
+		BatchJobs: jobs,
+	}
+
+	if resp.NextCursor != nil {
+		result.NextPageToken = *resp.NextCursor
+	}
+
+	return result
+}
+
+// ToGeminiBatchCancelResponse converts a BifrostBatchCancelResponse to Gemini SDK format.
+func ToGeminiBatchCancelResponse(resp *schemas.BifrostBatchCancelResponse) *GeminiBatchJobResponseSDK {
+	if resp == nil {
+		return nil
+	}
+
+	return &GeminiBatchJobResponseSDK{
+		Name:  resp.ID,
+		State: ToGeminiJobState(resp.Status),
+	}
+}

core/providers/gemini/gemini.go

@@ -2101,6 +2101,65 @@ func (provider *GeminiProvider) BatchResults(ctx context.Context, key schemas.Ke
 	return batchResultsResp, nil
 }
 
+// BatchDelete deletes a batch job for Gemini.
+func (provider *GeminiProvider) BatchDelete(ctx context.Context, key schemas.Key, request *schemas.BifrostBatchDeleteRequest) (*schemas.BifrostBatchDeleteResponse, *schemas.BifrostError) {
+	if err := providerUtils.CheckOperationAllowed(schemas.Gemini, provider.customProviderConfig, schemas.BatchDeleteRequest); err != nil {
+		return nil, err
+	}
+
+	providerName := provider.GetProviderKey()
+
+	if request.BatchID == "" {
+		return nil, providerUtils.NewBifrostOperationError("batch_id is required", nil, providerName)
+	}
+
+	// Create HTTP request
+	req := fasthttp.AcquireRequest()
+	resp := fasthttp.AcquireResponse()
+	defer fasthttp.ReleaseRequest(req)
+	defer fasthttp.ReleaseResponse(resp)
+
+	// Build URL for delete operation
+	batchID := request.BatchID
+	var url string
+	if strings.HasPrefix(batchID, "batches/") {
+		url = fmt.Sprintf("%s/%s", provider.networkConfig.BaseURL, batchID)
+	} else {
+		url = fmt.Sprintf("%s/batches/%s", provider.networkConfig.BaseURL, batchID)
+	}
+
+	provider.logger.Debug("gemini batch delete url: " + url)
+	providerUtils.SetExtraHeaders(ctx, req, provider.networkConfig.ExtraHeaders, nil)
+	req.SetRequestURI(url)
+	req.Header.SetMethod(http.MethodDelete)
+	if key.Value != "" {
+		req.Header.Set("x-goog-api-key", key.Value)
+	}
+	req.Header.SetContentType("application/json")
+
+	// Make request
+	latency, bifrostErr := providerUtils.MakeRequestWithContext(ctx, provider.client, req, resp)
+	if bifrostErr != nil {
+		return nil, bifrostErr
+	}
+
+	// Handle response
+	if resp.StatusCode() != fasthttp.StatusOK && resp.StatusCode() != fasthttp.StatusNoContent {
+		return nil, parseGeminiError(resp)
+	}
+
+	return &schemas.BifrostBatchDeleteResponse{
+		ID:      request.BatchID,
+		Object:  "batch",
+		Deleted: true,
+		ExtraFields: schemas.BifrostResponseExtraFields{
+			RequestType: schemas.BatchDeleteRequest,
+			Provider:    providerName,
+			Latency:     latency.Milliseconds(),
+		},
+	}, nil
+}
+
 // FileUpload uploads a file to Gemini.
 func (provider *GeminiProvider) FileUpload(ctx context.Context, key schemas.Key, request *schemas.BifrostFileUploadRequest) (*schemas.BifrostFileUploadResponse, *schemas.BifrostError) {
 	if err := providerUtils.CheckOperationAllowed(schemas.Gemini, provider.customProviderConfig, schemas.FileUploadRequest); err != nil {