diff --git a/core/internal/llmtests/account.go b/core/internal/llmtests/account.go index 6a2b0e5d8d..3dfc797b3b 100644 --- a/core/internal/llmtests/account.go +++ b/core/internal/llmtests/account.go @@ -131,8 +131,10 @@ type ComprehensiveTestConfig struct { VideoGenerationModel string // Model for video generation ExternalTTSProvider schemas.ModelProvider // External TTS provider to use for testing ExternalTTSModel string // External TTS model to use for testing - BatchExtraParams map[string]interface{} // Extra params for batch operations (e.g., role_arn, output_s3_uri for Bedrock) - FileExtraParams map[string]interface{} // Extra params for file operations (e.g., s3_bucket for Bedrock) + BatchExtraParams map[string]interface{} // Extra params for batch operations (e.g., role_arn, output_s3_uri for Bedrock) + BatchOutputFolder *schemas.BatchOutputFolder // Typed batch output location (e.g., GCS gs:// prefix for Vertex) + FileExtraParams map[string]interface{} // Extra params for file operations (e.g., s3_bucket for Bedrock) + FileStorageConfig *schemas.FileStorageConfig // Typed storage config for file operations (e.g., GCS bucket for Vertex) DisableParallelFor []string // Test scenarios to disable parallel execution for (e.g., "Transcription" for rate-limited APIs) ExpectRawRequestResponse bool // When true, validate rawRequest/rawResponse in ExtraFields PassthroughModel string // Model for passthrough API tests; defaults to ChatModel when empty diff --git a/core/internal/llmtests/batch.go b/core/internal/llmtests/batch.go index a6717ebe37..6c118f1fcc 100644 --- a/core/internal/llmtests/batch.go +++ b/core/internal/llmtests/batch.go @@ -84,6 +84,7 @@ func RunBatchCreateTest(t *testing.T, client *bifrost.Bifrost, ctx context.Conte }, CompletionWindow: "24h", ExtraParams: testConfig.BatchExtraParams, + OutputFolder: testConfig.BatchOutputFolder, } bfCtx := schemas.NewBifrostContext(ctx, schemas.NoDeadline) return client.BatchCreateRequest(bfCtx, request) @@ -231,6 +232,7 @@ func RunBatchRetrieveTest(t *testing.T, client *bifrost.Bifrost, ctx context.Con }, CompletionWindow: "24h", ExtraParams: testConfig.BatchExtraParams, + OutputFolder: testConfig.BatchOutputFolder, } bfCtx := schemas.NewBifrostContext(ctx, schemas.NoDeadline) return client.BatchCreateRequest(bfCtx, createRequest) @@ -358,6 +360,7 @@ func RunBatchCancelTest(t *testing.T, client *bifrost.Bifrost, ctx context.Conte }, CompletionWindow: "24h", ExtraParams: testConfig.BatchExtraParams, + OutputFolder: testConfig.BatchOutputFolder, } bfCtx := schemas.NewBifrostContext(ctx, schemas.NoDeadline) return client.BatchCreateRequest(bfCtx, createRequest) @@ -603,9 +606,10 @@ func RunFileUploadTest(t *testing.T, client *bifrost.Bifrost, ctx context.Contex request := &schemas.BifrostFileUploadRequest{ Provider: testConfig.Provider, File: fileContent, - Filename: "test_batch.jsonl", - Purpose: "batch", - ExtraParams: testConfig.FileExtraParams, + Filename: "test_batch.jsonl", + Purpose: "batch", + ExtraParams: testConfig.FileExtraParams, + StorageConfig: testConfig.FileStorageConfig, } bfCtx := schemas.NewBifrostContext(ctx, schemas.NoDeadline) return client.FileUploadRequest(bfCtx, request) @@ -672,9 +676,10 @@ func RunFileListTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context, response, err := WithFileListTestRetry(t, fileListRetryConfig, retryContext, expectations, "FileList", func() (*schemas.BifrostFileListResponse, *schemas.BifrostError) { request := &schemas.BifrostFileListRequest{ - Provider: testConfig.Provider, - Limit: 10, - ExtraParams: testConfig.FileExtraParams, + Provider: testConfig.Provider, + Limit: 10, + ExtraParams: testConfig.FileExtraParams, + StorageConfig: testConfig.FileStorageConfig, } bfCtx := schemas.NewBifrostContext(ctx, schemas.NoDeadline) return client.FileListRequest(bfCtx, request) @@ -742,9 +747,10 @@ func RunFileRetrieveTest(t *testing.T, client *bifrost.Bifrost, ctx context.Cont uploadRequest := &schemas.BifrostFileUploadRequest{ Provider: testConfig.Provider, File: fileContent, - Filename: "test_retrieve.jsonl", - Purpose: "batch", - ExtraParams: testConfig.FileExtraParams, + Filename: "test_retrieve.jsonl", + Purpose: "batch", + ExtraParams: testConfig.FileExtraParams, + StorageConfig: testConfig.FileStorageConfig, } bfCtx := schemas.NewBifrostContext(ctx, schemas.NoDeadline) return client.FileUploadRequest(bfCtx, uploadRequest) @@ -860,9 +866,10 @@ func RunFileDeleteTest(t *testing.T, client *bifrost.Bifrost, ctx context.Contex uploadRequest := &schemas.BifrostFileUploadRequest{ Provider: testConfig.Provider, File: fileContent, - Filename: "test_delete.jsonl", - Purpose: "batch", - ExtraParams: testConfig.FileExtraParams, + Filename: "test_delete.jsonl", + Purpose: "batch", + ExtraParams: testConfig.FileExtraParams, + StorageConfig: testConfig.FileStorageConfig, } bfCtx := schemas.NewBifrostContext(ctx, schemas.NoDeadline) return client.FileUploadRequest(bfCtx, uploadRequest) @@ -978,9 +985,10 @@ func RunFileContentTest(t *testing.T, client *bifrost.Bifrost, ctx context.Conte uploadRequest := &schemas.BifrostFileUploadRequest{ Provider: testConfig.Provider, File: originalContent, - Filename: "test_content.jsonl", - Purpose: "batch", - ExtraParams: testConfig.FileExtraParams, + Filename: "test_content.jsonl", + Purpose: "batch", + ExtraParams: testConfig.FileExtraParams, + StorageConfig: testConfig.FileStorageConfig, } bfCtx := schemas.NewBifrostContext(ctx, schemas.NoDeadline) return client.FileUploadRequest(bfCtx, uploadRequest) @@ -1117,9 +1125,10 @@ func RunFileAndBatchIntegrationTest(t *testing.T, client *bifrost.Bifrost, ctx c uploadRequest := &schemas.BifrostFileUploadRequest{ Provider: testConfig.Provider, File: fileContent, - Filename: "integration_test_batch.jsonl", - Purpose: "batch", - ExtraParams: testConfig.FileExtraParams, + Filename: "integration_test_batch.jsonl", + Purpose: "batch", + ExtraParams: testConfig.FileExtraParams, + StorageConfig: testConfig.FileStorageConfig, } bfCtx := schemas.NewBifrostContext(ctx, schemas.NoDeadline) @@ -1148,6 +1157,7 @@ func RunFileAndBatchIntegrationTest(t *testing.T, client *bifrost.Bifrost, ctx c Endpoint: schemas.BatchEndpointChatCompletions, CompletionWindow: "24h", ExtraParams: testConfig.BatchExtraParams, + OutputFolder: testConfig.BatchOutputFolder, } bfCtx2 := schemas.NewBifrostContext(ctx, schemas.NoDeadline) diff --git a/core/providers/vertex/batch.go b/core/providers/vertex/batch.go new file mode 100644 index 0000000000..48623e8a03 --- /dev/null +++ b/core/providers/vertex/batch.go @@ -0,0 +1,458 @@ +package vertex + +import ( + "bytes" + "fmt" + "strconv" + "strings" + "time" + + "github.com/bytedance/sonic" + providerUtils "github.com/maximhq/bifrost/core/providers/utils" + "github.com/maximhq/bifrost/core/schemas" +) + +// vertexBatchCustomIDLabel is the request label used to carry the Bifrost custom_id +// through a batch prediction job (Vertex JSONL has no native custom_id field; the +// request — labels included — is echoed back in each output line). +const vertexBatchCustomIDLabel = "bifrost_custom_id" + +// vertexJobStateToBatchStatus maps Vertex JOB_STATE_* values to Bifrost batch statuses. +func vertexJobStateToBatchStatus(state string) schemas.BatchStatus { + switch state { + case "JOB_STATE_QUEUED", "JOB_STATE_PENDING": + return schemas.BatchStatusValidating + case "JOB_STATE_RUNNING", "JOB_STATE_UPDATING": + return schemas.BatchStatusInProgress + case "JOB_STATE_SUCCEEDED", "JOB_STATE_PARTIALLY_SUCCEEDED": + return schemas.BatchStatusCompleted + case "JOB_STATE_FAILED": + return schemas.BatchStatusFailed + case "JOB_STATE_CANCELLING": + return schemas.BatchStatusCancelling + case "JOB_STATE_CANCELLED": + return schemas.BatchStatusCancelled + case "JOB_STATE_EXPIRED": + return schemas.BatchStatusExpired + default: + return schemas.BatchStatus(state) + } +} + +// vertexBatchJobsBaseURL returns ".../v1/projects/{project}/locations/{region}" for the +// key's configured project and region. Batch prediction requires a regional endpoint. +func vertexBatchJobsBaseURL(key schemas.Key) (string, *schemas.BifrostError) { + if key.VertexKeyConfig == nil { + return "", providerUtils.NewConfigurationError("vertex key config is not set") + } + projectID := strings.TrimSpace(key.VertexKeyConfig.ProjectID.GetValue()) + if projectID == "" { + return "", providerUtils.NewConfigurationError("project ID is not set") + } + region := key.VertexKeyConfig.Region.GetValue() + if region == "" || region == "global" { + return "", providerUtils.NewConfigurationError("a regional vertex key (e.g. us-central1) is required for batch prediction; global is not supported") + } + return getVertexProjectLocationURL(region, "v1", projectID), nil +} + +// vertexBatchJobURL resolves a Bifrost batch ID (bare job ID or full resource name) +// to the job's REST URL. +func vertexBatchJobURL(key schemas.Key, batchID string) (string, *schemas.BifrostError) { + if strings.HasPrefix(batchID, "projects/") { + // Full resource name: projects/{p}/locations/{r}/batchPredictionJobs/{id} + parts := strings.Split(batchID, "/") + if len(parts) >= 6 && parts[2] == "locations" { + return getVertexAPIBaseURL(parts[3], "v1") + "/" + batchID, nil + } + return "", providerUtils.NewBifrostOperationError(fmt.Sprintf("invalid Vertex batch ID %q", batchID), nil) + } + base, cfgErr := vertexBatchJobsBaseURL(key) + if cfgErr != nil { + return "", cfgErr + } + return base + "/batchPredictionJobs/" + batchID, nil +} + +// vertexBatchJobIDFromName extracts the bare job ID from a full resource name. +func vertexBatchJobIDFromName(name string) string { + if idx := strings.LastIndexByte(name, '/'); idx >= 0 { + return name[idx+1:] + } + return name +} + +// vertexBatchJobToBifrost maps a BatchPredictionJob resource to the Bifrost retrieve response. +func vertexBatchJobToBifrost(job *VertexBatchPredictionJob) schemas.BifrostBatchRetrieveResponse { + status := vertexJobStateToBatchStatus(job.State) + resp := schemas.BifrostBatchRetrieveResponse{ + ID: job.Name, + Object: "batch", + Status: status, + CreatedAt: gcsParseTime(job.CreateTime), + } + if job.DisplayName != "" { + resp.DisplayName = schemas.Ptr(job.DisplayName) + } + if job.InputConfig.GcsSource != nil && len(job.InputConfig.GcsSource.Uris) > 0 { + resp.InputFileID = job.InputConfig.GcsSource.Uris[0] + } + if job.OutputInfo != nil && job.OutputInfo.GcsOutputDirectory != "" { + resp.OutputFileID = schemas.Ptr(job.OutputInfo.GcsOutputDirectory) + } + if job.StartTime != "" { + resp.InProgressAt = schemas.Ptr(gcsParseTime(job.StartTime)) + } + if job.EndTime != "" { + endTime := gcsParseTime(job.EndTime) + switch status { + case schemas.BatchStatusCompleted: + resp.CompletedAt = &endTime + case schemas.BatchStatusFailed: + resp.FailedAt = &endTime + case schemas.BatchStatusCancelled: + resp.CancelledAt = &endTime + case schemas.BatchStatusExpired: + resp.ExpiredAt = &endTime + } + } + if job.CompletionStats != nil { + succeeded := gcsParseSize(job.CompletionStats.SuccessfulCount) + failed := gcsParseSize(job.CompletionStats.FailedCount) + incomplete := gcsParseSize(job.CompletionStats.IncompleteCount) + resp.RequestCounts = schemas.BatchRequestCounts{ + Total: int(succeeded + failed + incomplete), + Completed: int(succeeded), + Failed: int(failed), + } + } + if job.Error != nil && job.Error.Message != "" { + resp.Errors = &schemas.BatchErrors{ + Data: []schemas.BatchError{{Code: fmt.Sprintf("%d", job.Error.Code), Message: job.Error.Message}}, + } + } + return resp +} + +// parseVertexJobAPIError parses a Vertex AI error response (same envelope as GCS). +func parseVertexJobAPIError(body []byte, statusCode int, op string) *schemas.BifrostError { + var apiErr gcsErrorBody + _ = sonic.Unmarshal(body, &apiErr) + msg := apiErr.Error.Message + if msg == "" { + msg = fmt.Sprintf("Vertex %s failed with HTTP %d", op, statusCode) + } + return providerUtils.NewProviderAPIError(msg, nil, statusCode, nil, nil) +} + +// vertexBatchControlParams are extra_params consumed by BatchCreate for Bifrost-side +// storage/routing; they are stripped before the rest of extra_params is passed through to +// the Vertex BatchPredictionJob body. +var vertexBatchControlParams = []string{"output_uri", "gcs_bucket", "gcs_prefix"} + +// ToVertexBatchCreateRequest maps a Bifrost batch create request to a Vertex +// BatchPredictionJob request. The model, display name and input/output GCS config are +// mapped explicitly; every other field is taken from extra_params (e.g. modelParameters, +// labels, modelVersionId, encryptionSpec) and merged verbatim into the job body. +func ToVertexBatchCreateRequest(request *schemas.BifrostBatchCreateRequest, displayName, inputURI, outputURI string) *VertexBatchCreateRequest { + model := "" + if request.Model != nil { + model = *request.Model + } + if model != "" && !strings.Contains(model, "/") { + model = "publishers/google/models/" + model + } + + req := &VertexBatchCreateRequest{ + DisplayName: displayName, + Model: model, + InputConfig: VertexBatchInputConfig{ + InstancesFormat: "jsonl", + GcsSource: &VertexGcsSource{Uris: []string{inputURI}}, + }, + OutputConfig: VertexBatchOutputConfig{ + PredictionsFormat: "jsonl", + GcsDestination: &VertexGcsDestination{OutputUriPrefix: outputURI}, + }, + ExtraParams: request.ExtraParams, + } + + // Strip Bifrost control keys so they are not forwarded to Vertex as unknown fields. + if len(req.ExtraParams) > 0 { + stripped := make(map[string]interface{}, len(req.ExtraParams)) + for k, v := range req.ExtraParams { + stripped[k] = v + } + for _, k := range vertexBatchControlParams { + delete(stripped, k) + } + req.ExtraParams = stripped + } + + return req +} + +// vertexConvertRequestsToJSONL converts inline batch request items to Vertex batch JSONL. +// Bodies are passed through as-is (callers provide Gemini-native request bodies, mirroring +// the Anthropic/Bedrock providers); each custom_id is carried in request labels. +func vertexConvertRequestsToJSONL(requests []schemas.BatchRequestItem) ([]byte, error) { + var buf bytes.Buffer + for i, item := range requests { + body := item.Body + if body == nil { + body = item.Params + } + if body == nil { + return nil, fmt.Errorf("batch request item %d (custom_id %q) has no body", i, item.CustomID) + } + if item.CustomID != "" { + // Shallow-copy before injecting labels so the caller's map is not mutated. + withLabels := make(map[string]interface{}, len(body)+1) + for k, v := range body { + withLabels[k] = v + } + labels := map[string]interface{}{} + if existing, ok := withLabels["labels"].(map[string]interface{}); ok { + for k, v := range existing { + labels[k] = v + } + } + labels[vertexBatchCustomIDLabel] = item.CustomID + withLabels["labels"] = labels + body = withLabels + } + line, err := providerUtils.MarshalSorted(map[string]interface{}{"request": body}) + if err != nil { + return nil, fmt.Errorf("failed to marshal batch request item %d (custom_id %q): %w", i, item.CustomID, err) + } + buf.Write(line) + buf.WriteByte('\n') + } + return buf.Bytes(), nil +} + +// ============================ Integration Converters ============================ +// Convert between the native Vertex BatchPredictionJob wire shape (used by the aiplatform +// JobServiceClient) and Bifrost's neutral batch types, for the genai HTTP integration. +// Key/project selection happens in Bifrost from the vertex key config, so the project and +// location in the inbound request path are placeholders — only the job body is converted. + +// batchStatusToVertexJobState is the inverse of vertexJobStateToBatchStatus. +func batchStatusToVertexJobState(status schemas.BatchStatus) string { + switch status { + case schemas.BatchStatusValidating: + return "JOB_STATE_PENDING" + case schemas.BatchStatusInProgress, schemas.BatchStatusFinalizing: + return "JOB_STATE_RUNNING" + case schemas.BatchStatusCompleted, schemas.BatchStatusEnded: + return "JOB_STATE_SUCCEEDED" + case schemas.BatchStatusFailed: + return "JOB_STATE_FAILED" + case schemas.BatchStatusCancelling: + return "JOB_STATE_CANCELLING" + case schemas.BatchStatusCancelled: + return "JOB_STATE_CANCELLED" + case schemas.BatchStatusExpired: + return "JOB_STATE_EXPIRED" + default: + return "JOB_STATE_UNSPECIFIED" + } +} + +// formatVertexBatchTime renders a Unix timestamp as an RFC3339 string, empty when zero. +func formatVertexBatchTime(unix int64) string { + if unix <= 0 { + return "" + } + return time.Unix(unix, 0).UTC().Format(time.RFC3339) +} + +// vertexCompletionStatsFromCounts maps Bifrost request counts to Vertex completion stats. +func vertexCompletionStatsFromCounts(c schemas.BatchRequestCounts) *VertexBatchCompletionStats { + if c.Total == 0 && c.Completed == 0 && c.Failed == 0 { + return nil + } + incomplete := c.Total - c.Completed - c.Failed + if incomplete < 0 { + incomplete = 0 + } + return &VertexBatchCompletionStats{ + SuccessfulCount: strconv.Itoa(c.Completed), + FailedCount: strconv.Itoa(c.Failed), + IncompleteCount: strconv.Itoa(incomplete), + } +} + +// ToBifrostBatchCreateRequest maps an inbound native Vertex BatchPredictionJob (as sent by +// the aiplatform JobServiceClient) to a Bifrost batch create request. The model, GCS input +// URI and display name are mapped to typed Bifrost fields; the GCS output prefix and every +// other Vertex-native create-input field (modelParameters, labels, modelVersionId, +// encryptionSpec, instanceConfig, ...) are carried through ExtraParams keyed by their Vertex +// JSON names, so ToVertexBatchCreateRequest can merge them back into the job body verbatim +// for a lossless round trip. Server-populated, output-only fields (state, outputInfo, error, +// timestamps, completionStats, partialFailures, satisfiesPz*, ...) are intentionally omitted. +func ToBifrostBatchCreateRequest(job *VertexBatchPredictionJob) *schemas.BifrostBatchCreateRequest { + req := &schemas.BifrostBatchCreateRequest{Provider: schemas.Vertex} + if job == nil { + return req + } + if job.Model != "" { + req.Model = schemas.Ptr(job.Model) + } + if job.InputConfig.GcsSource != nil && len(job.InputConfig.GcsSource.Uris) > 0 { + req.InputFileID = job.InputConfig.GcsSource.Uris[0] + } + // Display name maps to the typed DisplayName field (read back by BatchCreate to set + // the outbound Vertex displayName) so it survives a full round trip. + if job.DisplayName != "" { + req.DisplayName = schemas.Ptr(job.DisplayName) + } + + // Output destination maps to the typed OutputFolder (read back by BatchCreate as the + // gs:// output prefix), so it survives a full round trip without going through extra_params. + if job.OutputConfig.GcsDestination != nil && job.OutputConfig.GcsDestination.OutputUriPrefix != "" { + req.OutputFolder = &schemas.BatchOutputFolder{URL: job.OutputConfig.GcsDestination.OutputUriPrefix} + } + + // Remaining create-input fields → ExtraParams, keyed by their Vertex JSON names so they + // merge cleanly into the outbound BatchPredictionJob body. Each is guarded so zero values + // are not forwarded (mirroring the native struct's omitempty tags). + extra := map[string]interface{}{} + if job.ModelVersionID != "" { + extra["modelVersionId"] = job.ModelVersionID + } + if job.UnmanagedContainerModel != nil { + extra["unmanagedContainerModel"] = job.UnmanagedContainerModel + } + if job.InstanceConfig != nil { + extra["instanceConfig"] = job.InstanceConfig + } + if job.ModelParameters != nil { + extra["modelParameters"] = job.ModelParameters + } + if job.DedicatedResources != nil { + extra["dedicatedResources"] = job.DedicatedResources + } + if job.ServiceAccount != "" { + extra["serviceAccount"] = job.ServiceAccount + } + if job.ManualBatchTuningParameters != nil { + extra["manualBatchTuningParameters"] = job.ManualBatchTuningParameters + } + if job.GenerateExplanation { + extra["generateExplanation"] = job.GenerateExplanation + } + if len(job.ExplanationSpec) > 0 { + extra["explanationSpec"] = job.ExplanationSpec + } + if len(job.Labels) > 0 { + extra["labels"] = job.Labels + } + if job.EncryptionSpec != nil { + extra["encryptionSpec"] = job.EncryptionSpec + } + if len(job.ModelMonitoringConfig) > 0 { + extra["modelMonitoringConfig"] = job.ModelMonitoringConfig + } + if job.DisableContainerLogging { + extra["disableContainerLogging"] = job.DisableContainerLogging + } + if len(extra) > 0 { + req.ExtraParams = extra + } + return req +} + +// vertexBatchJobShell builds the BatchPredictionJob fields shared by the create and retrieve +// response converters. name is whatever Bifrost returns (bare id or full resource name); +// displayName is the human-readable job name, kept distinct from name. +func vertexBatchJobShell(name, displayName string, status schemas.BatchStatus, createdAt int64, inputFileID string, outputFileID *string) *VertexBatchPredictionJob { + job := &VertexBatchPredictionJob{ + Name: name, + DisplayName: displayName, + State: batchStatusToVertexJobState(status), + CreateTime: formatVertexBatchTime(createdAt), + } + if inputFileID != "" { + job.InputConfig = VertexBatchInputConfig{ + InstancesFormat: "jsonl", + GcsSource: &VertexGcsSource{Uris: []string{inputFileID}}, + } + } + if outputFileID != nil && *outputFileID != "" { + job.OutputConfig = VertexBatchOutputConfig{ + PredictionsFormat: "jsonl", + GcsDestination: &VertexGcsDestination{OutputUriPrefix: *outputFileID}, + } + job.OutputInfo = &VertexBatchOutputInfo{GcsOutputDirectory: *outputFileID} + } + return job +} + +// ToVertexBatchCreateResponse maps a Bifrost batch create response to a native Vertex +// BatchPredictionJob. +func ToVertexBatchCreateResponse(resp *schemas.BifrostBatchCreateResponse) *VertexBatchPredictionJob { + if resp == nil { + return nil + } + displayName := "" + if resp.DisplayName != nil { + displayName = *resp.DisplayName + } + job := vertexBatchJobShell(resp.ID, displayName, resp.Status, resp.CreatedAt, resp.InputFileID, resp.OutputFileID) + job.CompletionStats = vertexCompletionStatsFromCounts(resp.RequestCounts) + return job +} + +// ToVertexBatchRetrieveResponse maps a Bifrost batch retrieve response to a native Vertex +// BatchPredictionJob, including timestamps, completion stats and any terminal error. +func ToVertexBatchRetrieveResponse(resp *schemas.BifrostBatchRetrieveResponse) *VertexBatchPredictionJob { + if resp == nil { + return nil + } + displayName := "" + if resp.DisplayName != nil { + displayName = *resp.DisplayName + } + job := vertexBatchJobShell(resp.ID, displayName, resp.Status, resp.CreatedAt, resp.InputFileID, resp.OutputFileID) + if resp.InProgressAt != nil { + job.StartTime = formatVertexBatchTime(*resp.InProgressAt) + } + switch { + case resp.CompletedAt != nil: + job.EndTime = formatVertexBatchTime(*resp.CompletedAt) + case resp.FailedAt != nil: + job.EndTime = formatVertexBatchTime(*resp.FailedAt) + case resp.CancelledAt != nil: + job.EndTime = formatVertexBatchTime(*resp.CancelledAt) + case resp.ExpiredAt != nil: + job.EndTime = formatVertexBatchTime(*resp.ExpiredAt) + } + job.CompletionStats = vertexCompletionStatsFromCounts(resp.RequestCounts) + if resp.Errors != nil && len(resp.Errors.Data) > 0 { + code := 0 + if c, err := strconv.Atoi(resp.Errors.Data[0].Code); err == nil { + code = c + } + job.Error = &VertexBatchJobError{Code: code, Message: resp.Errors.Data[0].Message} + } + return job +} + +// ToVertexBatchListResponse maps a Bifrost batch list response to the native Vertex +// batchPredictionJobs.list response envelope. +func ToVertexBatchListResponse(resp *schemas.BifrostBatchListResponse) *VertexBatchJobListResponse { + out := &VertexBatchJobListResponse{} + if resp == nil { + return out + } + for i := range resp.Data { + if job := ToVertexBatchRetrieveResponse(&resp.Data[i]); job != nil { + out.BatchPredictionJobs = append(out.BatchPredictionJobs, *job) + } + } + if resp.NextCursor != nil { + out.NextPageToken = *resp.NextCursor + } + return out +} diff --git a/core/providers/vertex/types.go b/core/providers/vertex/types.go index d9db81374f..f155426bde 100644 --- a/core/providers/vertex/types.go +++ b/core/providers/vertex/types.go @@ -247,13 +247,208 @@ type VertexCountTokensResponse struct { CachedContentTokenCount int32 `json:"cachedContentTokenCount,omitempty"` } +// ================================ Batch Prediction API Types ================================ + +// VertexGcsSource is the GCS input source for a batch prediction job. +type VertexGcsSource struct { + Uris []string `json:"uris"` +} + +// VertexBigQuerySource is the BigQuery input source for a batch prediction job. +type VertexBigQuerySource struct { + InputUri string `json:"inputUri"` +} + +// VertexBatchInputConfig is the input configuration for a batch prediction job. +type VertexBatchInputConfig struct { + InstancesFormat string `json:"instancesFormat"` + GcsSource *VertexGcsSource `json:"gcsSource,omitempty"` + BigquerySource *VertexBigQuerySource `json:"bigquerySource,omitempty"` +} + +// VertexBatchInstanceConfig controls how input instances are converted to prediction instances. +type VertexBatchInstanceConfig struct { + InstanceType string `json:"instanceType,omitempty"` + KeyField string `json:"keyField,omitempty"` + IncludedFields []string `json:"includedFields,omitempty"` + ExcludedFields []string `json:"excludedFields,omitempty"` +} + +// VertexGcsDestination is the GCS output destination for a batch prediction job. +type VertexGcsDestination struct { + OutputUriPrefix string `json:"outputUriPrefix"` +} + +// VertexBigQueryDestination is the BigQuery output destination for a batch prediction job. +type VertexBigQueryDestination struct { + OutputUri string `json:"outputUri"` +} + +// VertexBatchOutputConfig is the output configuration for a batch prediction job. +type VertexBatchOutputConfig struct { + PredictionsFormat string `json:"predictionsFormat"` + GcsDestination *VertexGcsDestination `json:"gcsDestination,omitempty"` + BigqueryDestination *VertexBigQueryDestination `json:"bigqueryDestination,omitempty"` +} + +// VertexBatchOutputInfo describes where a finished job wrote its output. +type VertexBatchOutputInfo struct { + GcsOutputDirectory string `json:"gcsOutputDirectory,omitempty"` + BigqueryOutputDataset string `json:"bigqueryOutputDataset,omitempty"` + BigqueryOutputTable string `json:"bigqueryOutputTable,omitempty"` +} + +// VertexBatchCompletionStats tracks per-request completion counts of a job. +type VertexBatchCompletionStats struct { + SuccessfulCount string `json:"successfulCount"` // int64 serialised as string + FailedCount string `json:"failedCount"` // int64 serialised as string + IncompleteCount string `json:"incompleteCount"` // int64 serialised as string + SuccessfulForecastPointCount string `json:"successfulForecastPointCount,omitempty"` // int64 serialised as string +} + +// VertexResourcesConsumed reports resources consumed by a batch prediction job. +type VertexResourcesConsumed struct { + ReplicaHours float64 `json:"replicaHours,omitempty"` +} + +// VertexManualBatchTuningParameters configures batch behaviour (only with dedicatedResources). +type VertexManualBatchTuningParameters struct { + BatchSize int `json:"batchSize,omitempty"` +} + +// VertexReservationAffinity configures the reservation a MachineSpec draws resources from. +type VertexReservationAffinity struct { + ReservationAffinityType string `json:"reservationAffinityType,omitempty"` + Key string `json:"key,omitempty"` + Values []string `json:"values,omitempty"` +} + +// VertexMachineSpec is the compute machine configuration for dedicated resources. +type VertexMachineSpec struct { + MachineType string `json:"machineType,omitempty"` + AcceleratorType string `json:"acceleratorType,omitempty"` + AcceleratorCount int `json:"acceleratorCount,omitempty"` + TpuTopology string `json:"tpuTopology,omitempty"` + ReservationAffinity *VertexReservationAffinity `json:"reservationAffinity,omitempty"` +} + +// VertexBatchDedicatedResources is the dedicated compute config used during batch prediction. +type VertexBatchDedicatedResources struct { + MachineSpec *VertexMachineSpec `json:"machineSpec,omitempty"` + StartingReplicaCount int `json:"startingReplicaCount,omitempty"` + MaxReplicaCount int `json:"maxReplicaCount,omitempty"` +} + +// VertexEncryptionSpec is the customer-managed encryption key configuration. +type VertexEncryptionSpec struct { + KmsKeyName string `json:"kmsKeyName,omitempty"` +} + +// VertexPredictSchemata describes the instance/parameter/prediction schemas of a model. +type VertexPredictSchemata struct { + InstanceSchemaUri string `json:"instanceSchemaUri,omitempty"` + ParametersSchemaUri string `json:"parametersSchemaUri,omitempty"` + PredictionSchemaUri string `json:"predictionSchemaUri,omitempty"` +} + +// VertexUnmanagedContainerModel describes a model used without registry upload. +type VertexUnmanagedContainerModel struct { + ArtifactUri string `json:"artifactUri,omitempty"` + PredictSchemata *VertexPredictSchemata `json:"predictSchemata,omitempty"` + // ContainerSpec (ModelContainerSpec) is deeply nested; kept generic for passthrough. + ContainerSpec map[string]interface{} `json:"containerSpec,omitempty"` +} + +// VertexBatchJobError mirrors google.rpc.Status. Used for the job's terminal error as well +// as partialFailures and modelMonitoringStatus. The details array holds google.protobuf.Any +// entries with no fixed schema, so it is kept generic. +type VertexBatchJobError struct { + Code int `json:"code"` + Message string `json:"message"` + Details []map[string]interface{} `json:"details,omitempty"` +} + +// VertexBatchPredictionJob is the BatchPredictionJob resource returned by the Vertex AI API. +// Fields Bifrost interprets are typed; the deeply-nested explanation/monitoring config trees +// (rarely used for Gemini batch) are kept generic for lossless passthrough. +type VertexBatchPredictionJob struct { + Name string `json:"name,omitempty"` + DisplayName string `json:"displayName"` + Model string `json:"model,omitempty"` + ModelVersionID string `json:"modelVersionId,omitempty"` + UnmanagedContainerModel *VertexUnmanagedContainerModel `json:"unmanagedContainerModel,omitempty"` + InputConfig VertexBatchInputConfig `json:"inputConfig"` + InstanceConfig *VertexBatchInstanceConfig `json:"instanceConfig,omitempty"` + ModelParameters interface{} `json:"modelParameters,omitempty"` + OutputConfig VertexBatchOutputConfig `json:"outputConfig"` + DedicatedResources *VertexBatchDedicatedResources `json:"dedicatedResources,omitempty"` + ServiceAccount string `json:"serviceAccount,omitempty"` + ManualBatchTuningParameters *VertexManualBatchTuningParameters `json:"manualBatchTuningParameters,omitempty"` + GenerateExplanation bool `json:"generateExplanation,omitempty"` + // ExplanationSpec is deeply nested; kept generic for passthrough. + ExplanationSpec map[string]interface{} `json:"explanationSpec,omitempty"` + OutputInfo *VertexBatchOutputInfo `json:"outputInfo,omitempty"` + State string `json:"state,omitempty"` + Error *VertexBatchJobError `json:"error,omitempty"` + PartialFailures []VertexBatchJobError `json:"partialFailures,omitempty"` + ResourcesConsumed *VertexResourcesConsumed `json:"resourcesConsumed,omitempty"` + CompletionStats *VertexBatchCompletionStats `json:"completionStats,omitempty"` + CreateTime string `json:"createTime,omitempty"` // RFC3339 + StartTime string `json:"startTime,omitempty"` // RFC3339 + EndTime string `json:"endTime,omitempty"` // RFC3339 + UpdateTime string `json:"updateTime,omitempty"` // RFC3339 + Labels map[string]string `json:"labels,omitempty"` + EncryptionSpec *VertexEncryptionSpec `json:"encryptionSpec,omitempty"` + // ModelMonitoringConfig / ModelMonitoringStatsAnomalies are deeply nested; kept generic. + ModelMonitoringConfig map[string]interface{} `json:"modelMonitoringConfig,omitempty"` + ModelMonitoringStatsAnomalies []map[string]interface{} `json:"modelMonitoringStatsAnomalies,omitempty"` + ModelMonitoringStatus *VertexBatchJobError `json:"modelMonitoringStatus,omitempty"` + DisableContainerLogging bool `json:"disableContainerLogging,omitempty"` + SatisfiesPzs bool `json:"satisfiesPzs,omitempty"` + SatisfiesPzi bool `json:"satisfiesPzi,omitempty"` +} + +// VertexBatchCreateRequest is the request body for creating a BatchPredictionJob. Only the +// fields Bifrost maps directly are typed; any other Vertex-native field (modelParameters, +// labels, modelVersionId, encryptionSpec, instanceConfig, ...) is passed through ExtraParams +// and merged into the body by CheckContextAndGetRequestBody. +type VertexBatchCreateRequest struct { + DisplayName string `json:"displayName"` + Model string `json:"model"` + InputConfig VertexBatchInputConfig `json:"inputConfig"` + OutputConfig VertexBatchOutputConfig `json:"outputConfig"` + + ExtraParams map[string]interface{} `json:"-"` +} + +// GetExtraParams implements the providerUtils.RequestBodyWithExtraParams interface. +func (r *VertexBatchCreateRequest) GetExtraParams() map[string]interface{} { + return r.ExtraParams +} + +// VertexBatchJobListResponse is the batchPredictionJobs.list response envelope. +type VertexBatchJobListResponse struct { + BatchPredictionJobs []VertexBatchPredictionJob `json:"batchPredictionJobs"` + NextPageToken string `json:"nextPageToken"` +} + +// VertexBatchOutputLine is one line of a predictions-*.jsonl batch output file. +// The original request is echoed back; labels carry the Bifrost custom_id. +type VertexBatchOutputLine struct { + Status string `json:"status,omitempty"` // error string for failed records, empty on success + Request struct { + Labels map[string]string `json:"labels"` + } `json:"request"` + Response map[string]interface{} `json:"response,omitempty"` +} + // ================================ GCS File API Types ================================ // gcsObjectMetadata represents GCS object metadata as returned by the JSON API. type gcsObjectMetadata struct { Name string `json:"name"` Bucket string `json:"bucket"` - Size string `json:"size"` // int64 serialised as string by GCS + Size string `json:"size"` // int64 serialised as string by GCS ContentType string `json:"contentType"` TimeCreated string `json:"timeCreated"` // RFC3339 Updated string `json:"updated"` // RFC3339 diff --git a/core/providers/vertex/vertex.go b/core/providers/vertex/vertex.go index 58c9d309ca..fafc443a7c 100644 --- a/core/providers/vertex/vertex.go +++ b/core/providers/vertex/vertex.go @@ -2623,34 +2623,680 @@ func stripVertexGeminiUnsupportedFieldsRaw(jsonBody []byte) []byte { return out } -// BatchCreate is not supported by Vertex AI provider. +// BatchCreate creates a Vertex AI batch prediction job. +// +// Input modes (mutually exclusive, mirroring the Gemini provider): +// - InputFileID: a gs:// URI of an existing Vertex-format JSONL file. +// - Requests: inline items converted to JSONL and uploaded to GCS via FileUpload. +// +// The output destination is taken from the typed output_folder.url (a gs:// prefix). func (provider *VertexProvider) BatchCreate(ctx *schemas.BifrostContext, key schemas.Key, request *schemas.BifrostBatchCreateRequest) (*schemas.BifrostBatchCreateResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchCreateRequest, provider.GetProviderKey()) + if request.Model == nil || *request.Model == "" { + return nil, providerUtils.NewBifrostOperationError("model is required for Vertex batch API", nil) + } + hasFileInput := request.InputFileID != "" + hasInlineRequests := len(request.Requests) > 0 + if hasFileInput && hasInlineRequests { + return nil, providerUtils.NewBifrostOperationError("cannot specify both input_file_id and requests", nil) + } + if !hasFileInput && !hasInlineRequests { + return nil, providerUtils.NewBifrostOperationError("either input_file_id (gs:// JSONL URI) or requests is required for Vertex batch API", nil) + } + + baseURL, cfgErr := vertexBatchJobsBaseURL(key) + if cfgErr != nil { + return nil, cfgErr + } + + // Output destination is the typed output_folder.url (a gs:// prefix). Vertex writes + // results into its own subdirectory under this prefix. + outputURI := "" + if request.OutputFolder != nil { + outputURI = strings.TrimSpace(request.OutputFolder.URL) + } + if outputURI == "" { + return nil, providerUtils.NewBifrostOperationError("output_folder.url (gs:// prefix) is required for Vertex batch API", nil) + } + + jobName := fmt.Sprintf("bifrost-batch-%d", time.Now().Unix()) + if request.DisplayName != nil && *request.DisplayName != "" { + jobName = *request.DisplayName + } else if request.Metadata != nil { + // Back-compat: OpenAI-compatible clients may pass the job name via metadata. + if name, ok := request.Metadata["job_name"]; ok && name != "" { + jobName = name + } + } + + // Inline mode: convert to JSONL and upload next to the output location (Bedrock pattern). + inputFileID := request.InputFileID + if inputFileID == "" { + jsonlData, err := vertexConvertRequestsToJSONL(request.Requests) + if err != nil { + return nil, providerUtils.NewBifrostOperationError("failed to convert requests to Vertex JSONL", err) + } + outBucket, outKey, parseErr := parseGCSURI(outputURI) + if parseErr != nil { + return nil, providerUtils.NewBifrostOperationError(parseErr.Error(), nil) + } + // Place the input alongside the output directory (sibling, not child) so the + // generated JSONL does not live inside the directory Vertex writes results to. + inputPrefix := "vertex-batches-input" + if trimmed := strings.Trim(outKey, "/"); trimmed != "" { + if idx := strings.LastIndexByte(trimmed, '/'); idx >= 0 { + inputPrefix = trimmed[:idx] + "/input" + } else { + inputPrefix = "input" + } + } + uploadResp, uploadErr := provider.FileUpload(ctx, key, &schemas.BifrostFileUploadRequest{ + Provider: schemas.Vertex, + File: jsonlData, + Filename: jobName + "-input.jsonl", + Purpose: schemas.FilePurposeBatch, + ContentType: schemas.Ptr("application/jsonl"), + StorageConfig: &schemas.FileStorageConfig{ + GCS: &schemas.GCSStorageConfig{Bucket: outBucket, Prefix: inputPrefix}, + }, + }) + if uploadErr != nil { + return nil, uploadErr + } + inputFileID = uploadResp.ID + } + + jsonData, bodyErr := providerUtils.CheckContextAndGetRequestBody( + ctx, + request, + func() (providerUtils.RequestBodyWithExtraParams, error) { + return ToVertexBatchCreateRequest(request, jobName, inputFileID, outputURI), nil + }, + ) + if bodyErr != nil { + return nil, bodyErr + } + + authHeader, authErr := gcsGetAuthHeader(key) + if authErr != nil { + return nil, providerUtils.NewBifrostOperationError(authErr.Error(), nil) + } + + req := fasthttp.AcquireRequest() + resp := fasthttp.AcquireResponse() + defer fasthttp.ReleaseRequest(req) + defer fasthttp.ReleaseResponse(resp) + + req.SetRequestURI(baseURL + "/batchPredictionJobs") + req.Header.SetMethod(http.MethodPost) + req.Header.SetContentType("application/json") + providerUtils.SetExtraHeaders(ctx, req, provider.networkConfig.ExtraHeaders, nil) + req.Header.Set("Authorization", authHeader) + req.SetBody(jsonData) + + sendBackRawRequest := providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest) + sendBackRawResponse := providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse) + + startTime := time.Now() + _, bifrostErr, wait := providerUtils.MakeRequestWithContext(ctx, provider.client, req, resp) + defer wait() + if bifrostErr != nil { + return nil, providerUtils.EnrichError(ctx, bifrostErr, jsonData, nil, provider.sendBackRawRequest, provider.sendBackRawResponse) + } + + if resp.StatusCode() != fasthttp.StatusOK { + if resp.StatusCode() == fasthttp.StatusUnauthorized || resp.StatusCode() == fasthttp.StatusForbidden { + removeVertexClient(key.VertexKeyConfig.AuthCredentials.GetValue()) + } + return nil, providerUtils.EnrichError(ctx, parseVertexJobAPIError(resp.Body(), resp.StatusCode(), "batch create"), jsonData, resp.Body(), provider.sendBackRawRequest, provider.sendBackRawResponse) + } + + var created VertexBatchPredictionJob + rawRequest, rawResponse, parseErr := providerUtils.HandleProviderResponse(resp.Body(), &created, jsonData, sendBackRawRequest, sendBackRawResponse) + if parseErr != nil { + return nil, providerUtils.EnrichError(ctx, parseErr, jsonData, resp.Body(), provider.sendBackRawRequest, provider.sendBackRawResponse) + } + + result := &schemas.BifrostBatchCreateResponse{ + ID: created.Name, + Object: "batch", + InputFileID: inputFileID, + Status: vertexJobStateToBatchStatus(created.State), + CreatedAt: gcsParseTime(created.CreateTime), + Metadata: request.Metadata, + ExtraFields: schemas.BifrostResponseExtraFields{ + Latency: time.Since(startTime).Milliseconds(), + }, + } + if created.DisplayName != "" { + result.DisplayName = schemas.Ptr(created.DisplayName) + } + if sendBackRawRequest { + result.ExtraFields.RawRequest = rawRequest + } + if sendBackRawResponse { + result.ExtraFields.RawResponse = rawResponse + } + return result, nil } -// BatchList is not supported by Vertex AI provider. -func (provider *VertexProvider) BatchList(_ *schemas.BifrostContext, _ []schemas.Key, _ *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchListRequest, provider.GetProviderKey()) +// BatchList lists Vertex AI batch prediction jobs across all keys, paginating one key +// at a time. Each Vertex key carries its own project/region, and batch jobs are scoped to +// that project/region, so the serial helper walks every key (exhausting all of its pages +// before advancing) to avoid hiding jobs created under any key but the first. +func (provider *VertexProvider) BatchList(ctx *schemas.BifrostContext, keys []schemas.Key, request *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { + if len(keys) == 0 { + return nil, providerUtils.NewBifrostOperationError("no keys provided for Vertex BatchList", nil) + } + + // The OpenAI-compatible /v1/batches route feeds the cursor back via After. + helper, err := providerUtils.NewSerialListHelper(keys, request.After, provider.logger, true) + if err != nil { + return nil, providerUtils.NewBifrostOperationError("invalid pagination cursor", err) + } + + key, nativeCursor, ok := helper.GetCurrentKey() + if !ok { + // All keys exhausted. + return &schemas.BifrostBatchListResponse{ + Object: "list", + Data: []schemas.BifrostBatchRetrieveResponse{}, + }, nil + } + + // Query the current key with its native Vertex page token. + modifiedRequest := *request + if nativeCursor != "" { + modifiedRequest.PageToken = &nativeCursor + } else { + modifiedRequest.PageToken = nil + } + + resp, latency, bifrostErr := provider.batchListByKey(ctx, key, &modifiedRequest) + if bifrostErr != nil { + return nil, bifrostErr + } + + nativeNextCursor := "" + if resp.NextCursor != nil { + nativeNextCursor = *resp.NextCursor + } + nextCursor, hasMore := helper.BuildNextCursor(resp.HasMore, nativeNextCursor) + + resp.HasMore = hasMore + if nextCursor != "" { + resp.NextCursor = &nextCursor + } else { + resp.NextCursor = nil + } + resp.ExtraFields.Latency = latency.Milliseconds() + return resp, nil } -// BatchRetrieve is not supported by Vertex AI provider. -func (provider *VertexProvider) BatchRetrieve(_ *schemas.BifrostContext, _ []schemas.Key, _ *schemas.BifrostBatchRetrieveRequest) (*schemas.BifrostBatchRetrieveResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchRetrieveRequest, provider.GetProviderKey()) +// batchListByKey lists batch prediction jobs for a single Vertex key/project/region. +// The native Vertex page token (if any) is taken from request.PageToken; the returned +// NextCursor carries Vertex's nextPageToken verbatim for the caller to re-encode. +func (provider *VertexProvider) batchListByKey(ctx *schemas.BifrostContext, key schemas.Key, request *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, time.Duration, *schemas.BifrostError) { + baseURL, cfgErr := vertexBatchJobsBaseURL(key) + if cfgErr != nil { + return nil, 0, cfgErr + } + + params := url.Values{} + pageSize := request.PageSize + if pageSize <= 0 { + pageSize = request.Limit + } + if pageSize <= 0 { + pageSize = 20 + } + params.Set("pageSize", fmt.Sprintf("%d", pageSize)) + if request.PageToken != nil && *request.PageToken != "" { + params.Set("pageToken", *request.PageToken) + } + + authHeader, authErr := gcsGetAuthHeader(key) + if authErr != nil { + return nil, 0, providerUtils.NewBifrostOperationError(authErr.Error(), nil) + } + + req := fasthttp.AcquireRequest() + resp := fasthttp.AcquireResponse() + defer fasthttp.ReleaseRequest(req) + defer fasthttp.ReleaseResponse(resp) + + req.SetRequestURI(baseURL + "/batchPredictionJobs?" + params.Encode()) + req.Header.SetMethod(http.MethodGet) + providerUtils.SetExtraHeaders(ctx, req, provider.networkConfig.ExtraHeaders, nil) + req.Header.Set("Authorization", authHeader) + + sendBackRawResponse := providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse) + + startTime := time.Now() + _, bifrostErr, wait := providerUtils.MakeRequestWithContext(ctx, provider.client, req, resp) + defer wait() + if bifrostErr != nil { + return nil, 0, providerUtils.EnrichError(ctx, bifrostErr, nil, nil, provider.sendBackRawRequest, provider.sendBackRawResponse) + } + + if resp.StatusCode() != fasthttp.StatusOK { + if resp.StatusCode() == fasthttp.StatusUnauthorized || resp.StatusCode() == fasthttp.StatusForbidden { + removeVertexClient(key.VertexKeyConfig.AuthCredentials.GetValue()) + } + return nil, 0, providerUtils.EnrichError(ctx, parseVertexJobAPIError(resp.Body(), resp.StatusCode(), "batch list"), nil, resp.Body(), provider.sendBackRawRequest, provider.sendBackRawResponse) + } + + // GET request: no request body, so raw request capture is skipped by HandleProviderResponse. + var listResp VertexBatchJobListResponse + _, rawResponse, parseErr := providerUtils.HandleProviderResponse(resp.Body(), &listResp, nil, false, sendBackRawResponse) + if parseErr != nil { + return nil, 0, providerUtils.EnrichError(ctx, parseErr, nil, resp.Body(), provider.sendBackRawRequest, provider.sendBackRawResponse) + } + + data := make([]schemas.BifrostBatchRetrieveResponse, 0, len(listResp.BatchPredictionJobs)) + for i := range listResp.BatchPredictionJobs { + data = append(data, vertexBatchJobToBifrost(&listResp.BatchPredictionJobs[i])) + } + + var nextCursor *string + if listResp.NextPageToken != "" { + nextCursor = &listResp.NextPageToken + } + + result := &schemas.BifrostBatchListResponse{ + Object: "list", + Data: data, + HasMore: listResp.NextPageToken != "", + NextCursor: nextCursor, + } + if sendBackRawResponse { + result.ExtraFields.RawResponse = rawResponse + } + return result, time.Since(startTime), nil } -// BatchCancel is not supported by Vertex AI provider. -func (provider *VertexProvider) BatchCancel(_ *schemas.BifrostContext, _ []schemas.Key, _ *schemas.BifrostBatchCancelRequest) (*schemas.BifrostBatchCancelResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchCancelRequest, provider.GetProviderKey()) +// BatchRetrieve fetches a Vertex AI batch prediction job by ID (bare or full resource name). +func (provider *VertexProvider) BatchRetrieve(ctx *schemas.BifrostContext, keys []schemas.Key, request *schemas.BifrostBatchRetrieveRequest) (*schemas.BifrostBatchRetrieveResponse, *schemas.BifrostError) { + if len(keys) == 0 { + return nil, providerUtils.NewBifrostOperationError("no keys provided for Vertex BatchRetrieve", nil) + } + + // A job ID is scoped to the project/region of the key that created it, so try each key + // until one resolves the job; return the last error only if all keys fail. + var lastErr *schemas.BifrostError + for _, key := range keys { + startTime := time.Now() + job, rawResponse, bifrostErr := provider.vertexGetBatchJob(ctx, key, request.BatchID) + if bifrostErr != nil { + lastErr = bifrostErr + continue + } + + result := vertexBatchJobToBifrost(job) + result.ExtraFields = schemas.BifrostResponseExtraFields{ + Latency: time.Since(startTime).Milliseconds(), + } + if providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse) { + result.ExtraFields.RawResponse = rawResponse + } + return &result, nil + } + + return nil, lastErr } -// BatchDelete is not supported by Vertex AI provider. -func (provider *VertexProvider) BatchDelete(_ *schemas.BifrostContext, _ []schemas.Key, _ *schemas.BifrostBatchDeleteRequest) (*schemas.BifrostBatchDeleteResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchDeleteRequest, provider.GetProviderKey()) +// vertexGetBatchJob fetches a BatchPredictionJob resource. The returned rawResponse is +// the raw response payload when raw-response capture is enabled (nil otherwise); it is a +// GET, so there is no raw request to capture. +func (provider *VertexProvider) vertexGetBatchJob(ctx *schemas.BifrostContext, key schemas.Key, batchID string) (*VertexBatchPredictionJob, interface{}, *schemas.BifrostError) { + jobURL, cfgErr := vertexBatchJobURL(key, batchID) + if cfgErr != nil { + return nil, nil, cfgErr + } + + authHeader, authErr := gcsGetAuthHeader(key) + if authErr != nil { + return nil, nil, providerUtils.NewBifrostOperationError(authErr.Error(), nil) + } + + req := fasthttp.AcquireRequest() + resp := fasthttp.AcquireResponse() + defer fasthttp.ReleaseRequest(req) + defer fasthttp.ReleaseResponse(resp) + + req.SetRequestURI(jobURL) + req.Header.SetMethod(http.MethodGet) + providerUtils.SetExtraHeaders(ctx, req, provider.networkConfig.ExtraHeaders, nil) + req.Header.Set("Authorization", authHeader) + + _, bifrostErr, wait := providerUtils.MakeRequestWithContext(ctx, provider.client, req, resp) + defer wait() + if bifrostErr != nil { + return nil, nil, providerUtils.EnrichError(ctx, bifrostErr, nil, nil, provider.sendBackRawRequest, provider.sendBackRawResponse) + } + + if resp.StatusCode() != fasthttp.StatusOK { + if resp.StatusCode() == fasthttp.StatusUnauthorized || resp.StatusCode() == fasthttp.StatusForbidden { + removeVertexClient(key.VertexKeyConfig.AuthCredentials.GetValue()) + } + return nil, nil, providerUtils.EnrichError(ctx, parseVertexJobAPIError(resp.Body(), resp.StatusCode(), "batch retrieve"), nil, resp.Body(), provider.sendBackRawRequest, provider.sendBackRawResponse) + } + + var job VertexBatchPredictionJob + _, rawResponse, parseErr := providerUtils.HandleProviderResponse(resp.Body(), &job, nil, false, providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse)) + if parseErr != nil { + return nil, nil, providerUtils.EnrichError(ctx, parseErr, nil, resp.Body(), provider.sendBackRawRequest, provider.sendBackRawResponse) + } + return &job, rawResponse, nil } -// BatchResults is not supported by Vertex AI provider. -func (provider *VertexProvider) BatchResults(_ *schemas.BifrostContext, _ []schemas.Key, _ *schemas.BifrostBatchResultsRequest) (*schemas.BifrostBatchResultsResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchResultsRequest, provider.GetProviderKey()) +// BatchCancel cancels a running Vertex AI batch prediction job. +func (provider *VertexProvider) BatchCancel(ctx *schemas.BifrostContext, keys []schemas.Key, request *schemas.BifrostBatchCancelRequest) (*schemas.BifrostBatchCancelResponse, *schemas.BifrostError) { + if len(keys) == 0 { + return nil, providerUtils.NewBifrostOperationError("no keys provided for Vertex BatchCancel", nil) + } + + // A job ID is scoped to the project/region of the key that created it, so try each key + // until the cancel succeeds; return the last error only if all keys fail. + var lastErr *schemas.BifrostError + for _, key := range keys { + resp, bifrostErr := provider.batchCancelByKey(ctx, key, request) + if bifrostErr == nil { + return resp, nil + } + lastErr = bifrostErr + } + return nil, lastErr +} + +// batchCancelByKey cancels a batch prediction job using a single Vertex key. +func (provider *VertexProvider) batchCancelByKey(ctx *schemas.BifrostContext, key schemas.Key, request *schemas.BifrostBatchCancelRequest) (*schemas.BifrostBatchCancelResponse, *schemas.BifrostError) { + jobURL, cfgErr := vertexBatchJobURL(key, request.BatchID) + if cfgErr != nil { + return nil, cfgErr + } + + authHeader, authErr := gcsGetAuthHeader(key) + if authErr != nil { + return nil, providerUtils.NewBifrostOperationError(authErr.Error(), nil) + } + + req := fasthttp.AcquireRequest() + resp := fasthttp.AcquireResponse() + defer fasthttp.ReleaseRequest(req) + defer fasthttp.ReleaseResponse(resp) + + req.SetRequestURI(jobURL + ":cancel") + req.Header.SetMethod(http.MethodPost) + req.Header.SetContentType("application/json") + providerUtils.SetExtraHeaders(ctx, req, provider.networkConfig.ExtraHeaders, nil) + req.Header.Set("Authorization", authHeader) + + startTime := time.Now() + _, bifrostErr, wait := providerUtils.MakeRequestWithContext(ctx, provider.client, req, resp) + defer wait() + if bifrostErr != nil { + return nil, providerUtils.EnrichError(ctx, bifrostErr, nil, nil, provider.sendBackRawRequest, provider.sendBackRawResponse) + } + + if resp.StatusCode() != fasthttp.StatusOK { + if resp.StatusCode() == fasthttp.StatusUnauthorized || resp.StatusCode() == fasthttp.StatusForbidden { + removeVertexClient(key.VertexKeyConfig.AuthCredentials.GetValue()) + } + return nil, providerUtils.EnrichError(ctx, parseVertexJobAPIError(resp.Body(), resp.StatusCode(), "batch cancel"), nil, resp.Body(), provider.sendBackRawRequest, provider.sendBackRawResponse) + } + + return &schemas.BifrostBatchCancelResponse{ + ID: vertexBatchJobIDFromName(request.BatchID), + Object: "batch", + Status: schemas.BatchStatusCancelling, + CancellingAt: schemas.Ptr(startTime.Unix()), + ExtraFields: schemas.BifrostResponseExtraFields{ + Latency: time.Since(startTime).Milliseconds(), + }, + }, nil +} + +// BatchDelete deletes a finished Vertex AI batch prediction job. +func (provider *VertexProvider) BatchDelete(ctx *schemas.BifrostContext, keys []schemas.Key, request *schemas.BifrostBatchDeleteRequest) (*schemas.BifrostBatchDeleteResponse, *schemas.BifrostError) { + if len(keys) == 0 { + return nil, providerUtils.NewBifrostOperationError("no keys provided for Vertex BatchDelete", nil) + } + + // A job ID is scoped to the project/region of the key that created it, so try each key + // until the delete succeeds; return the last error only if all keys fail. + var lastErr *schemas.BifrostError + for _, key := range keys { + resp, bifrostErr := provider.batchDeleteByKey(ctx, key, request) + if bifrostErr == nil { + return resp, nil + } + lastErr = bifrostErr + } + return nil, lastErr +} + +// batchDeleteByKey deletes a batch prediction job using a single Vertex key. +func (provider *VertexProvider) batchDeleteByKey(ctx *schemas.BifrostContext, key schemas.Key, request *schemas.BifrostBatchDeleteRequest) (*schemas.BifrostBatchDeleteResponse, *schemas.BifrostError) { + jobURL, cfgErr := vertexBatchJobURL(key, request.BatchID) + if cfgErr != nil { + return nil, cfgErr + } + + authHeader, authErr := gcsGetAuthHeader(key) + if authErr != nil { + return nil, providerUtils.NewBifrostOperationError(authErr.Error(), nil) + } + + req := fasthttp.AcquireRequest() + resp := fasthttp.AcquireResponse() + defer fasthttp.ReleaseRequest(req) + defer fasthttp.ReleaseResponse(resp) + + req.SetRequestURI(jobURL) + req.Header.SetMethod(http.MethodDelete) + providerUtils.SetExtraHeaders(ctx, req, provider.networkConfig.ExtraHeaders, nil) + req.Header.Set("Authorization", authHeader) + + startTime := time.Now() + _, bifrostErr, wait := providerUtils.MakeRequestWithContext(ctx, provider.client, req, resp) + defer wait() + if bifrostErr != nil { + return nil, providerUtils.EnrichError(ctx, bifrostErr, nil, nil, provider.sendBackRawRequest, provider.sendBackRawResponse) + } + + if resp.StatusCode() != fasthttp.StatusOK { + if resp.StatusCode() == fasthttp.StatusUnauthorized || resp.StatusCode() == fasthttp.StatusForbidden { + removeVertexClient(key.VertexKeyConfig.AuthCredentials.GetValue()) + } + return nil, providerUtils.EnrichError(ctx, parseVertexJobAPIError(resp.Body(), resp.StatusCode(), "batch delete"), nil, resp.Body(), provider.sendBackRawRequest, provider.sendBackRawResponse) + } + + return &schemas.BifrostBatchDeleteResponse{ + ID: vertexBatchJobIDFromName(request.BatchID), + Object: "batch", + Status: schemas.BatchStatusDeleted, + ExtraFields: schemas.BifrostResponseExtraFields{ + Latency: time.Since(startTime).Milliseconds(), + }, + }, nil +} + +// BatchResults reads the predictions-*.jsonl files a finished job wrote to its GCS +// output directory and maps each line to a Bifrost batch result item. The custom_id +// is recovered from the echoed request labels. +func (provider *VertexProvider) BatchResults(ctx *schemas.BifrostContext, keys []schemas.Key, request *schemas.BifrostBatchResultsRequest) (*schemas.BifrostBatchResultsResponse, *schemas.BifrostError) { + if len(keys) == 0 { + return nil, providerUtils.NewBifrostOperationError("no keys provided for Vertex BatchResults", nil) + } + + // A job ID is scoped to the project/region of the key that created it, so try each key + // until one resolves the job and reads its results; return the last error if all fail. + var lastErr *schemas.BifrostError + for _, key := range keys { + resp, bifrostErr := provider.batchResultsByKey(ctx, key, request) + if bifrostErr == nil { + return resp, nil + } + lastErr = bifrostErr + } + return nil, lastErr +} + +// batchResultsByKey reads a finished job's GCS output using a single Vertex key. +func (provider *VertexProvider) batchResultsByKey(ctx *schemas.BifrostContext, key schemas.Key, request *schemas.BifrostBatchResultsRequest) (*schemas.BifrostBatchResultsResponse, *schemas.BifrostError) { + startTime := time.Now() + job, _, bifrostErr := provider.vertexGetBatchJob(ctx, key, request.BatchID) + if bifrostErr != nil { + return nil, bifrostErr + } + if job.OutputInfo == nil || job.OutputInfo.GcsOutputDirectory == "" { + return nil, providerUtils.NewBifrostOperationError(fmt.Sprintf("batch output is not available yet (job state: %s)", job.State), nil) + } + + bucket, dirKey, parseErr := parseGCSURI(job.OutputInfo.GcsOutputDirectory) + if parseErr != nil { + return nil, providerUtils.NewBifrostOperationError(parseErr.Error(), nil) + } + + authHeader, authErr := gcsGetAuthHeader(key) + if authErr != nil { + return nil, providerUtils.NewBifrostOperationError(authErr.Error(), nil) + } + + objects, listErr := provider.gcsListAllObjects(ctx, authHeader, bucket, strings.Trim(dirKey, "/")+"/") + if listErr != nil { + return nil, listErr + } + + results := []schemas.BatchResultItem{} + for _, obj := range objects { + name := obj.Name + if idx := strings.LastIndexByte(name, '/'); idx >= 0 { + name = name[idx+1:] + } + if !strings.HasPrefix(name, "predictions") { + continue + } + + content, downloadErr := provider.gcsDownloadObject(ctx, authHeader, bucket, obj.Name) + if downloadErr != nil { + return nil, downloadErr + } + + for _, rawLine := range bytes.Split(content, []byte("\n")) { + if len(bytes.TrimSpace(rawLine)) == 0 { + continue + } + var line VertexBatchOutputLine + if err := sonic.Unmarshal(rawLine, &line); err != nil { + continue // skip malformed lines rather than failing the whole result set + } + item := schemas.BatchResultItem{ + CustomID: line.Request.Labels[vertexBatchCustomIDLabel], + } + if line.Response != nil { + item.Response = &schemas.BatchResultResponse{ + StatusCode: 200, + Body: line.Response, + } + } else { + item.Error = &schemas.BatchResultError{Message: line.Status} + } + results = append(results, item) + } + } + + return &schemas.BifrostBatchResultsResponse{ + BatchID: request.BatchID, + Results: results, + ExtraFields: schemas.BifrostResponseExtraFields{ + Latency: time.Since(startTime).Milliseconds(), + }, + }, nil +} + +// gcsListAllObjects lists every object under a prefix, following pagination. +func (provider *VertexProvider) gcsListAllObjects(ctx *schemas.BifrostContext, authHeader, bucket, prefix string) ([]gcsObjectMetadata, *schemas.BifrostError) { + var objects []gcsObjectMetadata + pageToken := "" + for { + params := url.Values{} + params.Set("prefix", prefix) + params.Set("maxResults", "1000") + if pageToken != "" { + params.Set("pageToken", pageToken) + } + + req := fasthttp.AcquireRequest() + resp := fasthttp.AcquireResponse() + + req.SetRequestURI(fmt.Sprintf("%s/b/%s/o?%s", gcsStorageBase, url.PathEscape(bucket), params.Encode())) + req.Header.SetMethod(http.MethodGet) + providerUtils.SetExtraHeaders(ctx, req, provider.networkConfig.ExtraHeaders, nil) + req.Header.Set("Authorization", authHeader) + + _, bifrostErr, wait := providerUtils.MakeRequestWithContext(ctx, provider.client, req, resp) + if bifrostErr != nil { + wait() + fasthttp.ReleaseRequest(req) + fasthttp.ReleaseResponse(resp) + return nil, bifrostErr + } + + statusCode := resp.StatusCode() + var listResp gcsObjectListResponse + var unmarshalErr error + if statusCode == fasthttp.StatusOK { + unmarshalErr = sonic.Unmarshal(resp.Body(), &listResp) + } + var apiErr *schemas.BifrostError + if statusCode != fasthttp.StatusOK { + apiErr = parseGCSAPIError(resp.Body(), statusCode, "list") + } + wait() + fasthttp.ReleaseRequest(req) + fasthttp.ReleaseResponse(resp) + + if apiErr != nil { + return nil, apiErr + } + if unmarshalErr != nil { + return nil, providerUtils.NewBifrostOperationError("failed to parse GCS list response", unmarshalErr) + } + + objects = append(objects, listResp.Items...) + if listResp.NextPageToken == "" { + return objects, nil + } + pageToken = listResp.NextPageToken + } +} + +// gcsDownloadObject downloads the raw bytes of a GCS object. +func (provider *VertexProvider) gcsDownloadObject(ctx *schemas.BifrostContext, authHeader, bucket, objectKey string) ([]byte, *schemas.BifrostError) { + req := fasthttp.AcquireRequest() + resp := fasthttp.AcquireResponse() + defer fasthttp.ReleaseRequest(req) + defer fasthttp.ReleaseResponse(resp) + + req.SetRequestURI(fmt.Sprintf("%s/b/%s/o/%s?alt=media", gcsStorageBase, url.PathEscape(bucket), gcsEncodeObjectName(objectKey))) + req.Header.SetMethod(http.MethodGet) + providerUtils.SetExtraHeaders(ctx, req, provider.networkConfig.ExtraHeaders, nil) + req.Header.Set("Authorization", authHeader) + + _, bifrostErr, wait := providerUtils.MakeRequestWithContext(ctx, provider.client, req, resp) + defer wait() + if bifrostErr != nil { + return nil, bifrostErr + } + + if resp.StatusCode() != fasthttp.StatusOK { + return nil, parseGCSAPIError(resp.Body(), resp.StatusCode(), "content download") + } + + content := make([]byte, len(resp.Body())) + copy(content, resp.Body()) + return content, nil } const ( @@ -2879,7 +3525,8 @@ func (provider *VertexProvider) gcsFileUploadDirect( StorageBackend: schemas.FileStorageGCS, StorageURI: gcsURI, ExtraFields: schemas.BifrostResponseExtraFields{ - Latency: time.Since(startTime).Milliseconds(), + Latency: time.Since(startTime).Milliseconds(), + ProviderResponseHeaders: providerUtils.ExtractProviderResponseHeaders(resp), }, }, nil } @@ -2943,7 +3590,8 @@ func (provider *VertexProvider) gcsFileUploadResumable( StorageURI: gcsURI, UploadURL: &sessionURL, ExtraFields: schemas.BifrostResponseExtraFields{ - Latency: time.Since(startTime).Milliseconds(), + Latency: time.Since(startTime).Milliseconds(), + ProviderResponseHeaders: providerUtils.ExtractProviderResponseHeaders(resp), }, }, nil } diff --git a/core/providers/vertex/vertex_test.go b/core/providers/vertex/vertex_test.go index d754f33d22..d29499f232 100644 --- a/core/providers/vertex/vertex_test.go +++ b/core/providers/vertex/vertex_test.go @@ -25,6 +25,27 @@ func TestVertex(t *testing.T) { rerankModel := strings.TrimSpace(os.Getenv("VERTEX_RERANK_MODEL")) + // Vertex file operations are GCS-backed: the bucket/prefix are passed via the typed + // StorageConfig (VERTEX_GCS_BUCKET, optional VERTEX_GCS_PREFIX), not extra_params. + var fileStorageConfig *schemas.FileStorageConfig + var batchOutputFolder *schemas.BatchOutputFolder + if gcsBucket := strings.TrimSpace(os.Getenv("VERTEX_GCS_BUCKET")); gcsBucket != "" { + gcsPrefix := strings.TrimSpace(os.Getenv("VERTEX_GCS_PREFIX")) + fileStorageConfig = &schemas.FileStorageConfig{ + GCS: &schemas.GCSStorageConfig{ + Bucket: gcsBucket, + Prefix: gcsPrefix, + }, + } + // Batch output is a gs:// prefix; Vertex writes results into its own subdirectory under it. + outputURI := "gs://" + gcsBucket + if gcsPrefix != "" { + outputURI += "/" + strings.Trim(gcsPrefix, "/") + } + outputURI += "/batch-output" + batchOutputFolder = &schemas.BatchOutputFolder{URL: outputURI} + } + testConfig := llmtests.ComprehensiveTestConfig{ Provider: schemas.Vertex, ChatModel: "gemini-2.5-pro", @@ -37,41 +58,54 @@ func TestVertex(t *testing.T) { ImageGenerationModel: "gemini-2.5-flash-image", ImageEditModel: "imagen-3.0-capability-001", VideoGenerationModel: "veo-3.1-generate-preview", + FileStorageConfig: fileStorageConfig, + BatchOutputFolder: batchOutputFolder, Scenarios: llmtests.TestScenarios{ - TextCompletion: false, // Not supported - SimpleChat: true, - CompletionStream: true, - MultiTurnConversation: true, - ToolCalls: true, - ToolCallsStreaming: true, - MultipleToolCalls: true, - MultipleToolCallsStreaming: true, - End2EndToolCalling: true, - AutomaticFunctionCall: true, - ImageURL: false, - ImageBase64: true, - ImageGeneration: true, - ImageGenerationStream: false, - ImageEdit: true, - VideoGeneration: false, // disabled for now because of long running operations - VideoRetrieve: false, - VideoRemix: false, - VideoDownload: false, - VideoList: false, - VideoDelete: false, - MultipleImages: true, - CompleteEnd2End: true, - FileBase64: true, - Embedding: true, - Rerank: rerankModel != "", - Reasoning: true, - PromptCaching: true, - ListModels: false, - CountTokens: true, - StructuredOutputs: true, // Structured outputs with nullable enum support + TextCompletion: false, // Not supported + SimpleChat: true, + CompletionStream: true, + MultiTurnConversation: true, + ToolCalls: true, + ToolCallsStreaming: true, + MultipleToolCalls: true, + MultipleToolCallsStreaming: true, + End2EndToolCalling: true, + AutomaticFunctionCall: true, + ImageURL: false, + ImageBase64: true, + ImageGeneration: true, + ImageGenerationStream: false, + ImageEdit: true, + VideoGeneration: false, // disabled for now because of long running operations + VideoRetrieve: false, + VideoRemix: false, + VideoDownload: false, + VideoList: false, + VideoDelete: false, + MultipleImages: true, + CompleteEnd2End: true, + FileBase64: true, + Embedding: true, + Rerank: rerankModel != "", + Reasoning: true, + PromptCaching: true, + ListModels: false, + CountTokens: true, + StructuredOutputs: true, // Structured outputs with nullable enum support InterleavedThinking: true, EagerInputStreaming: true, // fine-grained-tool-streaming-2025-05-14 (GA on Vertex) ServerToolsViaOpenAIEndpoint: true, // web_search only on Vertex per Table 20 (web_fetch/code_execution skip) + FileUpload: true, + FileList: true, + FileRetrieve: true, + FileDelete: true, + FileContent: true, + FileBatchInput: true, + BatchCreate: true, + BatchList: true, + BatchRetrieve: true, + BatchCancel: true, + BatchResults: true, }, } diff --git a/core/schemas/batch.go b/core/schemas/batch.go index 0f7bdd2f6c..92c4f32adf 100644 --- a/core/schemas/batch.go +++ b/core/schemas/batch.go @@ -79,6 +79,7 @@ type BifrostBatchCreateRequest struct { OutputFolder *BatchOutputFolder `json:"output_folder,omitempty"` // Common fields + DisplayName *string `json:"display_name,omitempty"` // Human-readable job name (e.g. Vertex displayName) Endpoint BatchEndpoint `json:"endpoint,omitempty"` // Target endpoint for batch requests CompletionWindow string `json:"completion_window,omitempty"` // Time window (e.g., "24h") Metadata map[string]string `json:"metadata,omitempty"` // User-provided metadata @@ -106,7 +107,8 @@ func (request *BifrostBatchCreateRequest) GetRawRequestBody() []byte { // BifrostBatchCreateResponse represents the response from creating a batch job. type BifrostBatchCreateResponse struct { ID string `json:"id"` - Object string `json:"object,omitempty"` // "batch" for OpenAI + Object string `json:"object,omitempty"` // "batch" for OpenAI + DisplayName *string `json:"display_name,omitempty"` // Human-readable job name (e.g. Vertex displayName) Endpoint string `json:"endpoint,omitempty"` InputFileID string `json:"input_file_id,omitempty"` CompletionWindow string `json:"completion_window,omitempty"` @@ -188,6 +190,7 @@ func (request *BifrostBatchRetrieveRequest) GetRawRequestBody() []byte { type BifrostBatchRetrieveResponse struct { ID string `json:"id"` Object string `json:"object,omitempty"` + DisplayName *string `json:"display_name,omitempty"` // Human-readable job name (e.g. Vertex displayName) Endpoint string `json:"endpoint,omitempty"` InputFileID string `json:"input_file_id,omitempty"` CompletionWindow string `json:"completion_window,omitempty"` diff --git a/tests/integrations/python/config.yml b/tests/integrations/python/config.yml index 1d1be6e3aa..0a5d51615d 100644 --- a/tests/integrations/python/config.yml +++ b/tests/integrations/python/config.yml @@ -165,6 +165,12 @@ providers: streaming: "gemini-2.5-flash" count_tokens: "claude-sonnet-4-5" video: "veo-3.1-generate-preview" + batch_create: "gemini-2.5-flash" + batch_inline: "gemini-2.5-flash" + batch_file_upload: "gemini-2.5-flash" + batch_list: "gemini-2.5-flash" + batch_retrieve: "gemini-2.5-flash" + batch_cancel: "gemini-2.5-flash" bedrock: chat: "global.anthropic.claude-sonnet-4-20250514-v1:0" vision: "global.anthropic.claude-sonnet-4-20250514-v1:0" diff --git a/tests/integrations/python/tests/test_google.py b/tests/integrations/python/tests/test_google.py index 2ff99ecce4..d7fc219bda 100644 --- a/tests/integrations/python/tests/test_google.py +++ b/tests/integrations/python/tests/test_google.py @@ -109,17 +109,34 @@ get_api_key, get_provider_voice, get_provider_voices, + # Vertex batch GCS utilities + get_vertex_batch_dest_uri, + get_vertex_project, + get_vertex_location, + get_bifrost_base_url, + is_vertex_gcs_configured, + skip_if_no_vertex_gcs, + skip_if_no_vertex_native_batch, + stage_vertex_batch_input, skip_if_no_api_key, ) -from .utils.config_loader import get_model +from .utils.config_loader import get_config, get_model from .utils.parametrize import ( format_provider_model, get_cross_provider_params_for_scenario, ) -def get_provider_google_client(provider: str = "gemini", passthrough: bool = False): - """Create Google GenAI client with x-model-provider header for given provider""" +def get_provider_google_client( + provider: str = "gemini", + passthrough: bool = False, + extra_headers: Dict[str, str] | None = None, +): + """Create Google GenAI client with x-model-provider header for given provider. + + extra_headers: optional additional HTTP headers forwarded on every request + (e.g. to carry provider-specific routing/config the SDK doesn't model natively). + """ from .utils.config_loader import get_config, get_integration_url api_key = get_api_key(provider) @@ -135,8 +152,11 @@ def get_provider_google_client(provider: str = "gemini", passthrough: bool = Fal } # Add base URL support, timeout, and x-model-provider header through HttpOptions + headers = {"x-model-provider": provider} + if extra_headers: + headers.update(extra_headers) http_options_kwargs = { - "headers": {"x-model-provider": provider}, + "headers": headers, } if base_url: http_options_kwargs["base_url"] = base_url @@ -148,6 +168,58 @@ def get_provider_google_client(provider: str = "gemini", passthrough: bool = Fal return genai.Client(**client_kwargs) +def get_vertex_job_service_client(): + """Build a native Vertex AI JobServiceClient pointed at the Bifrost gateway. + + Vertex batch prediction is a Vertex-native (aiplatform) API — not the Gemini + Developer batches surface — so these tests use the aiplatform gapic + JobServiceClient with the regional batchPredictionJobs methods, routed through + Bifrost via the gateway base URL. Auth is anonymous because Bifrost injects the + real Vertex credentials from its key config; Bifrost detects Vertex routing from + the /projects/{p}/locations/{l}/... request path. + """ + from google.cloud import aiplatform + from google.api_core.client_options import ClientOptions + from google.auth.credentials import AnonymousCredentials + + # Route through Bifrost's genai integration (the Vertex batch routes are mounted + # under the /genai prefix alongside the other GenAI endpoints). + api_endpoint = get_bifrost_base_url().rstrip("/") + "/genai" + return aiplatform.gapic.JobServiceClient( + client_options=ClientOptions(api_endpoint=api_endpoint), + transport="rest", + credentials=AnonymousCredentials(), + ) + + +def build_vertex_batch_prediction_job( + display_name: str, + model: str, + gcs_source_uri: str, + gcs_destination_output_uri_prefix: str, +) -> Dict[str, Any]: + """Build a native Vertex BatchPredictionJob request body (jsonl GCS in/out). + + Mirrors the official aiplatform create_batch_prediction_job sample. Gemini + publisher models do not require dedicated_resources/machine_spec, so those are + omitted (they apply to custom-trained models). + """ + if "/" not in model: + model = "publishers/google/models/" + model + return { + "display_name": display_name, + "model": model, + "input_config": { + "instances_format": "jsonl", + "gcs_source": {"uris": [gcs_source_uri]}, + }, + "output_config": { + "predictions_format": "jsonl", + "gcs_destination": {"output_uri_prefix": gcs_destination_output_uri_prefix}, + }, + } + + @pytest.fixture def google_client(): """Configure Google GenAI client for testing with default gemini provider""" @@ -3046,6 +3118,187 @@ def test_39_batch_e2e_file_api(self, test_config, provider, model): except Exception as e: print(f"Cleanup warning: Failed to delete file: {e}") + # ========================================================================= + # VERTEX AI BATCH API TEST CASES (native aiplatform JobServiceClient) + # + # Vertex batch prediction is a Vertex-native (aiplatform) API, distinct from + # the Gemini Developer batches surface. These tests use the aiplatform gapic + # JobServiceClient with the regional batchPredictionJobs methods, routed + # through Bifrost. Inputs/outputs live in GCS (instances_format / + # predictions_format = jsonl). Requires VERTEX_PROJECT_ID + VERTEX_GCS_BUCKET + # (and ADC for staging the input object); otherwise the tests skip. + # ========================================================================= + + @staticmethod + def _vertex_parent(): + return f"projects/{get_vertex_project()}/locations/{get_vertex_location()}" + + @staticmethod + def _cleanup_vertex_job(client, job_name): + """Best-effort cancel + delete of a native Vertex batch prediction job.""" + if not job_name: + return + try: + client.cancel_batch_prediction_job(name=job_name) + except Exception as e: + print(f"Cleanup info: Could not cancel job: {e}") + try: + client.delete_batch_prediction_job(name=job_name) + except Exception as e: + print(f"Cleanup info: Could not delete job: {e}") + + def test_vertex_batch_create(self, test_config): + """Vertex Batch: create a batch prediction job (jsonl GCS in/out).""" + skip_if_no_vertex_native_batch() + client = get_vertex_job_service_client() + model = get_config().get_provider_model("vertex", "batch_create") + + gcs_source_uri = stage_vertex_batch_input( + create_google_batch_json_content(model=model, num_requests=2) + ) + body = build_vertex_batch_prediction_job( + display_name="bifrost-vertex-batch-create", + model=model, + gcs_source_uri=gcs_source_uri, + gcs_destination_output_uri_prefix=get_vertex_batch_dest_uri(), + ) + + job = None + try: + job = client.create_batch_prediction_job( + parent=self._vertex_parent(), batch_prediction_job=body + ) + assert job.name, "Created job should have a resource name" + print(f"Success: Created Vertex batch job {job.name}, state: {job.state.name}") + finally: + self._cleanup_vertex_job(client, job.name if job else None) + + @skip_if_no_api_key("vertex") + def test_vertex_batch_get(self, test_config): + """Vertex Batch: retrieve a batch prediction job by name.""" + skip_if_no_vertex_native_batch() + client = get_vertex_job_service_client() + model = get_config().get_provider_model("vertex", "batch_retrieve") + + gcs_source_uri = stage_vertex_batch_input( + create_google_batch_json_content(model=model, num_requests=1) + ) + body = build_vertex_batch_prediction_job( + display_name="bifrost-vertex-batch-get", + model=model, + gcs_source_uri=gcs_source_uri, + gcs_destination_output_uri_prefix=get_vertex_batch_dest_uri(), + ) + + job = None + try: + job = client.create_batch_prediction_job( + parent=self._vertex_parent(), batch_prediction_job=body + ) + retrieved = client.get_batch_prediction_job(name=job.name) + assert retrieved.name == job.name, ( + f"Retrieved job name should match: expected {job.name}, got {retrieved.name}" + ) + print(f"Success: Retrieved Vertex batch job {retrieved.name}, state: {retrieved.state.name}") + finally: + self._cleanup_vertex_job(client, job.name if job else None) + + @skip_if_no_api_key("vertex") + def test_vertex_batch_list(self, test_config): + """Vertex Batch: list batch prediction jobs for the project/location.""" + skip_if_no_vertex_native_batch() + client = get_vertex_job_service_client() + model = get_config().get_provider_model("vertex", "batch_create") + + gcs_source_uri = stage_vertex_batch_input( + create_google_batch_json_content(model=model, num_requests=1) + ) + body = build_vertex_batch_prediction_job( + display_name="bifrost-vertex-batch-list", + model=model, + gcs_source_uri=gcs_source_uri, + gcs_destination_output_uri_prefix=get_vertex_batch_dest_uri(), + ) + + job = None + try: + job = client.create_batch_prediction_job( + parent=self._vertex_parent(), batch_prediction_job=body + ) + found = any( + listed.name == job.name + for listed in client.list_batch_prediction_jobs(parent=self._vertex_parent()) + ) + assert found, f"Created job {job.name} should appear in the listing" + print(f"Success: Found created Vertex batch job {job.name} in listing") + finally: + self._cleanup_vertex_job(client, job.name if job else None) + + @skip_if_no_api_key("vertex") + def test_vertex_batch_cancel(self, test_config): + """Vertex Batch: cancel a running batch prediction job.""" + skip_if_no_vertex_native_batch() + client = get_vertex_job_service_client() + model = get_config().get_provider_model("vertex", "batch_cancel") + + gcs_source_uri = stage_vertex_batch_input( + create_google_batch_json_content(model=model, num_requests=2) + ) + body = build_vertex_batch_prediction_job( + display_name="bifrost-vertex-batch-cancel", + model=model, + gcs_source_uri=gcs_source_uri, + gcs_destination_output_uri_prefix=get_vertex_batch_dest_uri(), + ) + + job = None + try: + job = client.create_batch_prediction_job( + parent=self._vertex_parent(), batch_prediction_job=body + ) + client.cancel_batch_prediction_job(name=job.name) + + retrieved = client.get_batch_prediction_job(name=job.name) + assert retrieved.state.name in ("JOB_STATE_CANCELLING", "JOB_STATE_CANCELLED"), ( + f"Job state should be cancelling/cancelled, got {retrieved.state.name}" + ) + print(f"Success: Cancelled Vertex batch job {job.name}, state: {retrieved.state.name}") + finally: + # Already cancelled above; just delete. + if job: + try: + client.delete_batch_prediction_job(name=job.name) + except Exception as e: + print(f"Cleanup info: Could not delete job: {e}") + + @skip_if_no_api_key("vertex") + def test_vertex_batch_delete(self, test_config): + """Vertex Batch: delete a batch prediction job.""" + skip_if_no_vertex_native_batch() + client = get_vertex_job_service_client() + model = get_config().get_provider_model("vertex", "batch_create") + + gcs_source_uri = stage_vertex_batch_input( + create_google_batch_json_content(model=model, num_requests=1) + ) + body = build_vertex_batch_prediction_job( + display_name="bifrost-vertex-batch-delete", + model=model, + gcs_source_uri=gcs_source_uri, + gcs_destination_output_uri_prefix=get_vertex_batch_dest_uri(), + ) + + job = client.create_batch_prediction_job( + parent=self._vertex_parent(), batch_prediction_job=body + ) + # Cancel first so the job is deletable, then delete (returns an LRO). + try: + client.cancel_batch_prediction_job(name=job.name) + except Exception as e: + print(f"Info: Could not cancel before delete: {e}") + client.delete_batch_prediction_job(name=job.name) + print(f"Success: Deleted Vertex batch job {job.name}") + # ========================================================================= # INPUT TOKENS / TOKEN COUNTING TEST CASES # ========================================================================= diff --git a/tests/integrations/python/tests/utils/common.py b/tests/integrations/python/tests/utils/common.py index 30ca7dc234..dbcfe48c89 100644 --- a/tests/integrations/python/tests/utils/common.py +++ b/tests/integrations/python/tests/utils/common.py @@ -2535,6 +2535,168 @@ def skip_if_no_bedrock_s3(): pytest.skip("Bedrock S3 tests require AWS_S3_BUCKET environment variable") +def get_vertex_gcs_config() -> Dict[str, Optional[str]]: + """ + Get Vertex AI batch GCS configuration from environment variables. + + Vertex batch prediction reads inputs from / writes outputs to Google Cloud Storage, + so a bucket must be provided to exercise the batch API end-to-end. + + Returns: + Dictionary with GCS configuration: + - bucket: GCS bucket name (from VERTEX_GCS_BUCKET) + - prefix: Output object prefix (from VERTEX_GCS_PREFIX or a default) + """ + return { + "bucket": os.environ.get("VERTEX_GCS_BUCKET"), + "prefix": os.environ.get("VERTEX_GCS_PREFIX", "bifrost-batch-tests/"), + } + + +def is_vertex_gcs_configured() -> bool: + """ + Check if Vertex AI batch GCS configuration is available. + + Returns: + True if VERTEX_GCS_BUCKET is set, False otherwise + """ + config = get_vertex_gcs_config() + return config["bucket"] is not None and len(config["bucket"]) > 0 + + +def get_vertex_batch_dest_uri() -> str: + """ + Build the GCS output destination URI (gs:// prefix) for a Vertex batch job. + + Returns: + GCS URI string (e.g., gs://bucket/bifrost-batch-tests/output) + + Raises: + ValueError if VERTEX_GCS_BUCKET is not configured + """ + config = get_vertex_gcs_config() + if not config["bucket"]: + raise ValueError( + "VERTEX_GCS_BUCKET environment variable is required for Vertex batch API" + ) + prefix = (config["prefix"] or "").strip("/") + base = f"gs://{config['bucket']}" + if prefix: + base = f"{base}/{prefix}" + return f"{base}/output" + + +def skip_if_no_vertex_gcs(): + """ + Pytest skip helper for tests requiring Vertex GCS configuration. + Call skip_if_no_vertex_gcs() at the start of a test. + """ + import pytest + + if not is_vertex_gcs_configured(): + pytest.skip("Vertex batch tests require VERTEX_GCS_BUCKET environment variable") + + +def get_vertex_project() -> Optional[str]: + """Vertex project id for native batch prediction (from VERTEX_PROJECT_ID).""" + return os.environ.get("VERTEX_PROJECT_ID") + + +def get_vertex_location() -> str: + """Vertex regional location for native batch prediction (from GOOGLE_LOCATION).""" + return os.environ.get("GOOGLE_LOCATION", "us-central1") + + +def get_bifrost_base_url() -> str: + """Base URL of the Bifrost gateway (from BIFROST_BASE_URL).""" + return os.environ.get("BIFROST_BASE_URL", "http://localhost:8080") + + +def skip_if_no_vertex_native_batch(): + """ + Pytest skip helper for native Vertex batch tests (aiplatform JobServiceClient). + Requires both a project id and a GCS bucket. + """ + import pytest + + if not get_vertex_project(): + pytest.skip("Vertex native batch tests require VERTEX_PROJECT_ID environment variable") + if not is_vertex_gcs_configured(): + pytest.skip("Vertex native batch tests require VERTEX_GCS_BUCKET environment variable") + + +def get_vertex_google_credentials(scopes: Optional[List[str]] = None): + """ + Build google-auth credentials for direct GCS/Vertex calls in tests. + + The Vertex service-account key is provided to Bifrost via VERTEX_CREDENTIALS, which + may hold either the service-account JSON *content* or a path to a JSON file. ADC + (GOOGLE_APPLICATION_CREDENTIALS) only accepts a file path, so when the content is + inlined we must construct credentials explicitly instead of relying on ADC. + + Returns None if no usable credentials are found (caller falls back to ADC). + """ + import json + + from google.oauth2 import service_account + + raw = os.environ.get("VERTEX_CREDENTIALS") or os.environ.get("GOOGLE_APPLICATION_CREDENTIALS") + if not raw: + return None + raw = raw.strip() + + if os.path.isfile(raw): + creds = service_account.Credentials.from_service_account_file(raw) + else: + try: + info = json.loads(raw) + except (ValueError, TypeError): + return None + creds = service_account.Credentials.from_service_account_info(info) + + if scopes: + creds = creds.with_scopes(scopes) + return creds + + +def stage_vertex_batch_input(content: str, filename: str | None = None) -> str: + """ + Upload JSONL batch input to GCS and return its gs:// URI. + + Vertex batch prediction reads inputs from Cloud Storage, so the input file must + exist in GCS before creating the job (the native API has no inline mode). + + Args: + content: Newline-delimited JSON batch input + filename: Optional object filename (auto-generated if not provided) + + Returns: + gs:// URI of the uploaded input object + """ + import time + + from google.cloud import storage + + cfg = get_vertex_gcs_config() + if not cfg["bucket"]: + raise ValueError("VERTEX_GCS_BUCKET environment variable is required for Vertex batch API") + + if filename is None: + filename = f"batch-input-{int(time.time())}.jsonl" + prefix = (cfg["prefix"] or "").strip("/") + blob_name = f"{prefix}/input/{filename}" if prefix else f"input/{filename}" + + # Build credentials from VERTEX_CREDENTIALS (JSON content or path); fall back to ADC. + creds = get_vertex_google_credentials( + scopes=["https://www.googleapis.com/auth/cloud-platform"] + ) + client = storage.Client(project=get_vertex_project(), credentials=creds) + bucket = client.bucket(cfg["bucket"]) + blob = bucket.blob(blob_name) + blob.upload_from_string(content, content_type="application/jsonl") + return f"gs://{cfg['bucket']}/{blob_name}" + + def get_content_string_with_summary(response: Any) -> tuple[str, bool]: """ Extract content from response, handling both OpenAI API responses and LangChain AIMessage objects. diff --git a/tests/integrations/python/tests/utils/config_loader.py b/tests/integrations/python/tests/utils/config_loader.py index eed542fb13..394b6eab33 100644 --- a/tests/integrations/python/tests/utils/config_loader.py +++ b/tests/integrations/python/tests/utils/config_loader.py @@ -23,6 +23,7 @@ "pydanticai": "openai", # Pydantic AI defaults to OpenAI "bedrock": "bedrock", # Bedrock defaults to Amazon provider "azure": "azure", + "vertex": "vertex", } @dataclass diff --git a/transports/bifrost-http/handlers/inference.go b/transports/bifrost-http/handlers/inference.go index 9bba6d3cac..60b52fe6e7 100644 --- a/transports/bifrost-http/handlers/inference.go +++ b/transports/bifrost-http/handlers/inference.go @@ -341,6 +341,7 @@ var batchCreateParamsKnownFields = map[string]bool{ "input_file_id": true, "input_blob": true, "output_folder": true, + "display_name": true, "requests": true, "endpoint": true, "completion_window": true, @@ -575,6 +576,7 @@ type BatchCreateRequest struct { Requests []schemas.BatchRequestItem `json:"requests,omitempty"` // Anthropic-style inline requests InputBlob *string `json:"input_blob,omitempty"` // Azure-style blob storage input OutputFolder *schemas.BatchOutputFolder `json:"output_folder,omitempty"` // Azure-style output destination + DisplayName *string `json:"display_name,omitempty"` // Human-readable job name (e.g. Vertex displayName) Endpoint string `json:"endpoint,omitempty"` // e.g., "/v1/chat/completions" CompletionWindow string `json:"completion_window,omitempty"` // e.g., "24h" Metadata map[string]string `json:"metadata,omitempty"` @@ -2773,6 +2775,7 @@ func (h *CompletionHandler) batchCreate(ctx *fasthttp.RequestCtx) { InputFileID: req.InputFileID, InputBlob: req.InputBlob, OutputFolder: req.OutputFolder, + DisplayName: req.DisplayName, Requests: req.Requests, Endpoint: schemas.BatchEndpoint(req.Endpoint), CompletionWindow: req.CompletionWindow, diff --git a/transports/bifrost-http/integrations/genai.go b/transports/bifrost-http/integrations/genai.go index 5ac45fd198..58967c281e 100644 --- a/transports/bifrost-http/integrations/genai.go +++ b/transports/bifrost-http/integrations/genai.go @@ -676,6 +676,220 @@ func CreateGenAIBatchRouteConfigs(pathPrefix string, handlerStore lib.HandlerSto return routes } +// CreateVertexBatchRouteConfigs creates route configurations for the native Vertex AI +// batchPredictionJobs API (as used by the aiplatform JobServiceClient). Unlike the Gemini +// Developer batches surface, Vertex batch prediction is GCS-backed and addressed by the +// regional resource path projects/{project}/locations/{location}/batchPredictionJobs. +// Key/project selection happens in Bifrost from the vertex key config, so the project and +// location in the path are placeholders used only for routing the request shape. +func CreateVertexBatchRouteConfigs(pathPrefix string) []RouteConfig { + var routes []RouteConfig + + collectionPath := pathPrefix + "/v1/projects/{project}/locations/{location}/batchPredictionJobs" + itemPath := collectionPath + "/{batch_id}" + + // Create batch prediction job - POST .../batchPredictionJobs + routes = append(routes, RouteConfig{ + Type: RouteConfigTypeGenAI, + Path: collectionPath, + Method: "POST", + GetHTTPRequestType: func(ctx *fasthttp.RequestCtx) schemas.RequestType { + return schemas.BatchCreateRequest + }, + GetRequestTypeInstance: func(ctx context.Context) interface{} { + return &vertex.VertexBatchPredictionJob{} + }, + BatchRequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*BatchRequest, error) { + if job, ok := req.(*vertex.VertexBatchPredictionJob); ok { + createReq := vertex.ToBifrostBatchCreateRequest(job) + // The native body is already a Vertex BatchPredictionJob; carry it verbatim + // so BigQuery IO, non-JSONL formats and multi-URI inputs round-trip losslessly. + createReq.RawRequestBody = getGenAIRawRequestBody(ctx) + return &BatchRequest{ + Type: schemas.BatchCreateRequest, + CreateRequest: createReq, + }, nil + } + return nil, errors.New("invalid vertex batch create request type") + }, + BatchCreateResponseConverter: func(ctx *schemas.BifrostContext, resp *schemas.BifrostBatchCreateResponse) (interface{}, error) { + if resp.ExtraFields.Provider == schemas.Vertex && resp.ExtraFields.RawResponse != nil { + return resp.ExtraFields.RawResponse, nil + } + return vertex.ToVertexBatchCreateResponse(resp), nil + }, + ErrorConverter: func(ctx *schemas.BifrostContext, err *schemas.BifrostError) interface{} { + return gemini.ToGeminiError(err) + }, + // Native Vertex batch bodies pass through verbatim (see RawRequestBody above). + PreCallback: func(ctx *fasthttp.RequestCtx, bifrostCtx *schemas.BifrostContext, req interface{}) error { + setGenAIRawRequestBodyFromRequest(ctx, bifrostCtx) + bifrostCtx.SetValue(schemas.BifrostContextKeyUseRawRequestBody, true) + return nil + }, + }) + + // List batch prediction jobs - GET .../batchPredictionJobs + routes = append(routes, RouteConfig{ + Type: RouteConfigTypeGenAI, + Path: collectionPath, + Method: "GET", + GetHTTPRequestType: func(ctx *fasthttp.RequestCtx) schemas.RequestType { + return schemas.BatchListRequest + }, + GetRequestTypeInstance: func(ctx context.Context) interface{} { + return &schemas.BifrostBatchListRequest{} + }, + BatchRequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*BatchRequest, error) { + if listReq, ok := req.(*schemas.BifrostBatchListRequest); ok { + return &BatchRequest{Type: schemas.BatchListRequest, ListRequest: listReq}, nil + } + return nil, errors.New("invalid vertex batch list request type") + }, + BatchListResponseConverter: func(ctx *schemas.BifrostContext, resp *schemas.BifrostBatchListResponse) (interface{}, error) { + if resp.ExtraFields.Provider == schemas.Vertex && resp.ExtraFields.RawResponse != nil { + return resp.ExtraFields.RawResponse, nil + } + return vertex.ToVertexBatchListResponse(resp), nil + }, + ErrorConverter: func(ctx *schemas.BifrostContext, err *schemas.BifrostError) interface{} { + return gemini.ToGeminiError(err) + }, + PreCallback: extractVertexBatchPathParams, + }) + + // Retrieve batch prediction job - GET .../batchPredictionJobs/{batch_id} + routes = append(routes, RouteConfig{ + Type: RouteConfigTypeGenAI, + Path: itemPath, + Method: "GET", + GetHTTPRequestType: func(ctx *fasthttp.RequestCtx) schemas.RequestType { + return schemas.BatchRetrieveRequest + }, + GetRequestTypeInstance: func(ctx context.Context) interface{} { + return &schemas.BifrostBatchRetrieveRequest{} + }, + BatchRequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*BatchRequest, error) { + if retrieveReq, ok := req.(*schemas.BifrostBatchRetrieveRequest); ok { + return &BatchRequest{Type: schemas.BatchRetrieveRequest, RetrieveRequest: retrieveReq}, nil + } + return nil, errors.New("invalid vertex batch retrieve request type") + }, + BatchRetrieveResponseConverter: func(ctx *schemas.BifrostContext, resp *schemas.BifrostBatchRetrieveResponse) (interface{}, error) { + if resp.ExtraFields.Provider == schemas.Vertex && resp.ExtraFields.RawResponse != nil { + return resp.ExtraFields.RawResponse, nil + } + return vertex.ToVertexBatchRetrieveResponse(resp), nil + }, + ErrorConverter: func(ctx *schemas.BifrostContext, err *schemas.BifrostError) interface{} { + return gemini.ToGeminiError(err) + }, + PreCallback: extractVertexBatchPathParams, + }) + + // Cancel batch prediction job - POST .../batchPredictionJobs/{batch_id}:cancel + routes = append(routes, RouteConfig{ + Type: RouteConfigTypeGenAI, + Path: itemPath, + Method: "POST", + GetHTTPRequestType: func(ctx *fasthttp.RequestCtx) schemas.RequestType { + return schemas.BatchCancelRequest + }, + GetRequestTypeInstance: func(ctx context.Context) interface{} { + return &schemas.BifrostBatchCancelRequest{} + }, + BatchRequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*BatchRequest, error) { + if cancelReq, ok := req.(*schemas.BifrostBatchCancelRequest); ok { + return &BatchRequest{Type: schemas.BatchCancelRequest, CancelRequest: cancelReq}, nil + } + return nil, errors.New("invalid vertex batch cancel request type") + }, + BatchCancelResponseConverter: func(ctx *schemas.BifrostContext, resp *schemas.BifrostBatchCancelResponse) (interface{}, error) { + if resp.ExtraFields.Provider == schemas.Vertex && resp.ExtraFields.RawResponse != nil { + return resp.ExtraFields.RawResponse, nil + } + // Vertex batchPredictionJobs.cancel returns google.protobuf.Empty. + return map[string]interface{}{}, nil + }, + ErrorConverter: func(ctx *schemas.BifrostContext, err *schemas.BifrostError) interface{} { + return gemini.ToGeminiError(err) + }, + PreCallback: extractVertexBatchPathParams, + }) + + // Delete batch prediction job - DELETE .../batchPredictionJobs/{batch_id} + routes = append(routes, RouteConfig{ + Type: RouteConfigTypeGenAI, + Path: itemPath, + Method: "DELETE", + GetHTTPRequestType: func(ctx *fasthttp.RequestCtx) schemas.RequestType { + return schemas.BatchDeleteRequest + }, + GetRequestTypeInstance: func(ctx context.Context) interface{} { + return &schemas.BifrostBatchDeleteRequest{} + }, + BatchRequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*BatchRequest, error) { + if deleteReq, ok := req.(*schemas.BifrostBatchDeleteRequest); ok { + return &BatchRequest{Type: schemas.BatchDeleteRequest, DeleteRequest: deleteReq}, nil + } + return nil, errors.New("invalid vertex batch delete request type") + }, + BatchDeleteResponseConverter: func(ctx *schemas.BifrostContext, resp *schemas.BifrostBatchDeleteResponse) (interface{}, error) { + if resp.ExtraFields.Provider == schemas.Vertex && resp.ExtraFields.RawResponse != nil { + return resp.ExtraFields.RawResponse, nil + } + // Vertex batchPredictionJobs.delete returns a long-running Operation. + return map[string]interface{}{"done": true}, nil + }, + ErrorConverter: func(ctx *schemas.BifrostContext, err *schemas.BifrostError) interface{} { + return gemini.ToGeminiError(err) + }, + PreCallback: extractVertexBatchPathParams, + }) + + return routes +} + +// extractVertexBatchPathParams pins the provider to Vertex and extracts the bare batch_id +// (stripping any :cancel action suffix) for the native Vertex batch routes. The job ID is +// passed bare so the provider resolves project/region from its key config. +func extractVertexBatchPathParams(ctx *fasthttp.RequestCtx, bifrostCtx *schemas.BifrostContext, req interface{}) error { + batchID, _ := ctx.UserValue("batch_id").(string) + batchID = strings.TrimSuffix(batchID, ":cancel") + + switch r := req.(type) { + case *schemas.BifrostBatchListRequest: + r.Provider = schemas.Vertex + if pageSizeStr := string(ctx.QueryArgs().Peek("pageSize")); pageSizeStr != "" { + if pageSize, err := strconv.Atoi(pageSizeStr); err == nil { + r.Limit = pageSize + } + } + if pageToken := string(ctx.QueryArgs().Peek("pageToken")); pageToken != "" { + r.After = &pageToken + } + case *schemas.BifrostBatchRetrieveRequest: + if batchID == "" { + return errors.New("batch_id is required") + } + r.Provider = schemas.Vertex + r.BatchID = batchID + case *schemas.BifrostBatchCancelRequest: + if batchID == "" { + return errors.New("batch_id is required") + } + r.Provider = schemas.Vertex + r.BatchID = batchID + case *schemas.BifrostBatchDeleteRequest: + if batchID == "" { + return errors.New("batch_id is required") + } + r.Provider = schemas.Vertex + r.BatchID = batchID + } + return nil +} + // extractGeminiBatchIDFromPath extracts batch_id from path parameters for Gemini func extractGeminiBatchIDFromPath(ctx *fasthttp.RequestCtx, bifrostCtx *schemas.BifrostContext, req interface{}) error { provider := getProviderFromHeader(ctx, schemas.Gemini) @@ -1110,6 +1324,7 @@ func NewGenAIRouter(client *bifrost.Bifrost, handlerStore lib.HandlerStore, logg routes := CreateGenAIRouteConfigs("/genai") routes = append(routes, CreateGenAIFileRouteConfigs("/genai", handlerStore)...) routes = append(routes, CreateGenAIBatchRouteConfigs("/genai", handlerStore)...) + routes = append(routes, CreateVertexBatchRouteConfigs("/genai")...) routes = append(routes, CreateGenAICachedContentRouteConfigs("/genai", handlerStore)...) return &GenAIRouter{ diff --git a/transports/bifrost-http/integrations/openai.go b/transports/bifrost-http/integrations/openai.go index 0657a40372..bd4c42d1c1 100644 --- a/transports/bifrost-http/integrations/openai.go +++ b/transports/bifrost-http/integrations/openai.go @@ -1503,13 +1503,15 @@ func CreateOpenAIBatchRouteConfigs(pathPrefix string, handlerStore lib.HandlerSt } } - // For Azure, extract inline requests from raw body - if createReq.Provider == schemas.Azure { + // Azure (input_blob + output_folder) and Vertex (output_folder, a gs:// prefix) + // carry their storage location in the request body rather than a managed file. + if createReq.Provider == schemas.Azure || createReq.Provider == schemas.Vertex { var extraFields map[string]interface{} if err := json.Unmarshal(ctx.Request.Body(), &extraFields); err == nil { - // Extract requests array for inline batching - if inputBlob, ok := extraFields["input_blob"].(string); ok { - createReq.InputBlob = &inputBlob + if createReq.Provider == schemas.Azure { + if inputBlob, ok := extraFields["input_blob"].(string); ok { + createReq.InputBlob = &inputBlob + } } if outputFolder, ok := extraFields["output_folder"].(map[string]interface{}); ok { outputURL, ok := outputFolder["url"].(string)