[DO NOT MERGE] gemini-sdk-batch-support

core/bifrost.go

@@ -1130,6 +1130,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

@@ -789,3 +789,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/azure/batch.go

@@ -434,3 +434,8 @@ func splitJSONL(data []byte) [][]byte {
 	}
 	return lines
 }
+
+// BatchDelete is not supported by Azure provider.
+func (provider *AzureProvider) BatchDelete(ctx context.Context, key schemas.Key, request *schemas.BifrostBatchDeleteRequest) (*schemas.BifrostBatchDeleteResponse, *schemas.BifrostError) {
+	return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchDeleteRequest, provider.GetProviderKey())
+}

core/providers/bedrock/batch.go

@@ -1060,3 +1060,8 @@ func splitJSONL(data []byte) [][]byte {
 	}
 	return lines
 }
+
+// BatchDelete is not supported by Bedrock provider.
+func (provider *BedrockProvider) 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/batch.go

@@ -32,3 +32,8 @@ func (provider *CerebrasProvider) BatchResults(ctx context.Context, key schemas.
 	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/batch.go

@@ -32,3 +32,8 @@ func (provider *CohereProvider) BatchResults(ctx context.Context, key schemas.Ke
 	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())
+}
+

core/providers/elevenlabs/batch.go

@@ -32,3 +32,8 @@ func (provider *ElevenlabsProvider) BatchResults(ctx context.Context, key schema
 	return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchResultsRequest, provider.GetProviderKey())
 }
 
+// BatchDelete is not supported by Elevenlabs provider.
+func (provider *ElevenlabsProvider) BatchDelete(ctx context.Context, key schemas.Key, request *schemas.BifrostBatchDeleteRequest) (*schemas.BifrostBatchDeleteResponse, *schemas.BifrostError) {
+	return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchDeleteRequest, provider.GetProviderKey())
+}
+

core/providers/gemini/batch.go

@@ -839,3 +839,205 @@ func (provider *GeminiProvider) BatchResults(ctx context.Context, key schemas.Ke
 		},
 	}, nil
 }
+
+// ==================== 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),
+	}
+}
+
+// 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(providerName, resp)
+	}
+
+	return &schemas.BifrostBatchDeleteResponse{
+		ID:      request.BatchID,
+		Object:  "batch",
+		Deleted: true,
+		ExtraFields: schemas.BifrostResponseExtraFields{
+			RequestType: schemas.BatchDeleteRequest,
+			Provider:    providerName,
+			Latency:     latency.Milliseconds(),
+		},
+	}, nil
+}

core/providers/gemini/files.go

@@ -463,17 +463,23 @@ func (provider *GeminiProvider) FileContent(ctx context.Context, key schemas.Key
 }
 
 // ToGeminiFileUploadResponse converts a Bifrost file upload response to Gemini format.
-func ToGeminiFileUploadResponse(resp *schemas.BifrostFileUploadResponse) map[string]interface{} {
-	return map[string]interface{}{
-		"file": map[string]interface{}{
-			"name":           resp.ID,
-			"displayName":    resp.Filename,
-			"mimeType":       "application/octet-stream",
-			"sizeBytes":      fmt.Sprintf("%d", resp.Bytes),
-			"createTime":     formatGeminiTimestamp(resp.CreatedAt),
-			"state":          toGeminiFileState(resp.Status),
-			"uri":            resp.StorageURI,
-			"expirationTime": formatGeminiTimestamp(safeDerefInt64(resp.ExpiresAt)),
+// Uses snake_case field names to match Google's API format.
+// GeminiFileUploadResponseWrapper is a wrapper that contains the file response for the upload API.
+type GeminiFileUploadResponseWrapper struct {
+	File GeminiFileResponse `json:"file"`
+}
+
+func ToGeminiFileUploadResponse(resp *schemas.BifrostFileUploadResponse) *GeminiFileUploadResponseWrapper {
+	return &GeminiFileUploadResponseWrapper{
+		File: GeminiFileResponse{
+			Name:           resp.ID,
+			DisplayName:    resp.Filename,
+			MimeType:       "application/octet-stream",
+			SizeBytes:      fmt.Sprintf("%d", resp.Bytes),
+			CreateTime:     formatGeminiTimestamp(resp.CreatedAt),
+			State:          toGeminiFileState(resp.Status),
+			URI:            resp.StorageURI,
+			ExpirationTime: formatGeminiTimestamp(safeDerefInt64(resp.ExpiresAt)),
 		},
 	}
 }