diff --git a/core/bifrost.go b/core/bifrost.go index 35b364a1a..2ec2d3061 100644 --- a/core/bifrost.go +++ b/core/bifrost.go @@ -34,6 +34,7 @@ import ( providerUtils "github.com/maximhq/bifrost/core/providers/utils" "github.com/maximhq/bifrost/core/providers/vertex" schemas "github.com/maximhq/bifrost/core/schemas" + "github.com/valyala/fasthttp" ) // ChannelMessage represents a message passed through the request channel. @@ -908,42 +909,15 @@ func (bifrost *Bifrost) BatchCreateRequest(ctx context.Context, req *schemas.Bif } } - config, err := bifrost.account.GetConfigForProvider(req.Provider) - if err != nil { - return nil, newBifrostErrorFromMsg(fmt.Sprintf("failed to get config for provider %s: %v", req.Provider, err.Error())) - } - if config == nil { - return nil, newBifrostErrorFromMsg(fmt.Sprintf("config is nil for provider %s", req.Provider)) - } - - // Determine the base provider type for key requirement checks - baseProvider := req.Provider - if config.CustomProviderConfig != nil && config.CustomProviderConfig.BaseProviderType != "" { - baseProvider = config.CustomProviderConfig.BaseProviderType - } - - var key schemas.Key - if providerRequiresKey(baseProvider, config.CustomProviderConfig) { - keys, keyErr := bifrost.getAllSupportedKeys(&ctx, req.Provider, baseProvider) - if keyErr != nil { - return nil, newBifrostError(keyErr) - } - if len(keys) > 0 { - key = keys[0] - } - } + bifrostReq := bifrost.getBifrostRequest() + bifrostReq.RequestType = schemas.BatchCreateRequest + bifrostReq.BatchCreateRequest = req - response, bifrostErr := executeRequestWithRetries(&ctx, config, func() (*schemas.BifrostBatchCreateResponse, *schemas.BifrostError) { - return provider.BatchCreate(ctx, key, req) - }, schemas.BatchCreateRequest, req.Provider, req.Model) - if bifrostErr != nil { - bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{ - RequestType: schemas.BatchCreateRequest, - Provider: req.Provider, - } - return nil, bifrostErr + response, err := bifrost.handleRequest(ctx, bifrostReq) + if err != nil { + return nil, err } - return response, nil + return response.BatchCreateResponse, nil } // BatchListRequest lists batch jobs for the specified provider. @@ -968,49 +942,15 @@ func (bifrost *Bifrost) BatchListRequest(ctx context.Context, req *schemas.Bifro ctx = bifrost.ctx } - provider := bifrost.getProviderByKey(req.Provider) - if provider == nil { - return nil, &schemas.BifrostError{ - IsBifrostError: false, - Error: &schemas.ErrorField{ - Message: "provider not found for batch list request", - }, - } - } + bifrostReq := bifrost.getBifrostRequest() + bifrostReq.RequestType = schemas.BatchListRequest + bifrostReq.BatchListRequest = req - config, err := bifrost.account.GetConfigForProvider(req.Provider) + response, err := bifrost.handleRequest(ctx, bifrostReq) if err != nil { - return nil, newBifrostErrorFromMsg(fmt.Sprintf("failed to get config for provider %s: %v", req.Provider, err.Error())) - } - if config == nil { - return nil, newBifrostErrorFromMsg(fmt.Sprintf("config is nil for provider %s", req.Provider)) - } - - // Determine the base provider type for key requirement checks - baseProvider := req.Provider - if config.CustomProviderConfig != nil && config.CustomProviderConfig.BaseProviderType != "" { - baseProvider = config.CustomProviderConfig.BaseProviderType - } - - var keys []schemas.Key - if providerRequiresKey(baseProvider, config.CustomProviderConfig) { - keys, err = bifrost.getAllSupportedKeys(&ctx, req.Provider, baseProvider) - if err != nil { - return nil, newBifrostError(err) - } - } - - response, bifrostErr := executeRequestWithRetries(&ctx, config, func() (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { - return provider.BatchList(ctx, keys, req) - }, schemas.BatchListRequest, req.Provider, "") - if bifrostErr != nil { - bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{ - RequestType: schemas.BatchListRequest, - Provider: req.Provider, - } - return nil, bifrostErr + return nil, err } - return response, nil + return response.BatchListResponse, nil } // BatchRetrieveRequest retrieves a specific batch job. @@ -1043,52 +983,15 @@ func (bifrost *Bifrost) BatchRetrieveRequest(ctx context.Context, req *schemas.B ctx = bifrost.ctx } - provider := bifrost.getProviderByKey(req.Provider) - if provider == nil { - return nil, &schemas.BifrostError{ - IsBifrostError: false, - Error: &schemas.ErrorField{ - Message: "provider not found for batch retrieve request", - }, - } - } + bifrostReq := bifrost.getBifrostRequest() + bifrostReq.RequestType = schemas.BatchRetrieveRequest + bifrostReq.BatchRetrieveRequest = req - config, err := bifrost.account.GetConfigForProvider(req.Provider) + response, err := bifrost.handleRequest(ctx, bifrostReq) if err != nil { - return nil, newBifrostErrorFromMsg(fmt.Sprintf("failed to get config for provider %s: %v", req.Provider, err.Error())) - } - if config == nil { - return nil, newBifrostErrorFromMsg(fmt.Sprintf("config is nil for provider %s", req.Provider)) - } - - // Determine the base provider type for key requirement checks - baseProvider := req.Provider - if config.CustomProviderConfig != nil && config.CustomProviderConfig.BaseProviderType != "" { - baseProvider = config.CustomProviderConfig.BaseProviderType - } - - var key schemas.Key - if providerRequiresKey(baseProvider, config.CustomProviderConfig) { - keys, keyErr := bifrost.getAllSupportedKeys(&ctx, req.Provider, baseProvider) - if keyErr != nil { - return nil, newBifrostError(keyErr) - } - if len(keys) > 0 { - key = keys[0] - } - } - - response, bifrostErr := executeRequestWithRetries(&ctx, config, func() (*schemas.BifrostBatchRetrieveResponse, *schemas.BifrostError) { - return provider.BatchRetrieve(ctx, key, req) - }, schemas.BatchRetrieveRequest, req.Provider, "") - if bifrostErr != nil { - bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{ - RequestType: schemas.BatchRetrieveRequest, - Provider: req.Provider, - } - return nil, bifrostErr + return nil, err } - return response, nil + return response.BatchRetrieveResponse, nil } // BatchCancelRequest cancels a batch job. @@ -1121,52 +1024,15 @@ func (bifrost *Bifrost) BatchCancelRequest(ctx context.Context, req *schemas.Bif ctx = bifrost.ctx } - provider := bifrost.getProviderByKey(req.Provider) - if provider == nil { - return nil, &schemas.BifrostError{ - IsBifrostError: false, - Error: &schemas.ErrorField{ - Message: "provider not found for batch cancel request", - }, - } - } + bifrostReq := bifrost.getBifrostRequest() + bifrostReq.RequestType = schemas.BatchCancelRequest + bifrostReq.BatchCancelRequest = req - config, err := bifrost.account.GetConfigForProvider(req.Provider) + response, err := bifrost.handleRequest(ctx, bifrostReq) if err != nil { - return nil, newBifrostErrorFromMsg(fmt.Sprintf("failed to get config for provider %s: %v", req.Provider, err.Error())) - } - if config == nil { - return nil, newBifrostErrorFromMsg(fmt.Sprintf("config is nil for provider %s", req.Provider)) - } - - // Determine the base provider type for key requirement checks - baseProvider := req.Provider - if config.CustomProviderConfig != nil && config.CustomProviderConfig.BaseProviderType != "" { - baseProvider = config.CustomProviderConfig.BaseProviderType - } - - var key schemas.Key - if providerRequiresKey(baseProvider, config.CustomProviderConfig) { - keys, keyErr := bifrost.getAllSupportedKeys(&ctx, req.Provider, baseProvider) - if keyErr != nil { - return nil, newBifrostError(keyErr) - } - if len(keys) > 0 { - key = keys[0] - } - } - - response, bifrostErr := executeRequestWithRetries(&ctx, config, func() (*schemas.BifrostBatchCancelResponse, *schemas.BifrostError) { - return provider.BatchCancel(ctx, key, req) - }, schemas.BatchCancelRequest, req.Provider, "") - if bifrostErr != nil { - bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{ - RequestType: schemas.BatchCancelRequest, - Provider: req.Provider, - } - return nil, bifrostErr + return nil, err } - return response, nil + return response.BatchCancelResponse, nil } // BatchResultsRequest retrieves results from a completed batch job. @@ -1209,56 +1075,15 @@ func (bifrost *Bifrost) BatchResultsRequest(ctx context.Context, req *schemas.Bi ctx = bifrost.ctx } - provider := bifrost.getProviderByKey(req.Provider) - if provider == nil { - return nil, &schemas.BifrostError{ - IsBifrostError: false, - Error: &schemas.ErrorField{ - Message: "provider not found for batch results request", - }, - ExtraFields: schemas.BifrostErrorExtraFields{ - RequestType: schemas.BatchResultsRequest, - Provider: req.Provider, - }, - } - } + bifrostReq := bifrost.getBifrostRequest() + bifrostReq.RequestType = schemas.BatchResultsRequest + bifrostReq.BatchResultsRequest = req - config, err := bifrost.account.GetConfigForProvider(req.Provider) + response, err := bifrost.handleRequest(ctx, bifrostReq) if err != nil { - return nil, newBifrostErrorFromMsg(fmt.Sprintf("failed to get config for provider %s: %v", req.Provider, err.Error())) - } - if config == nil { - return nil, newBifrostErrorFromMsg(fmt.Sprintf("config is nil for provider %s", req.Provider)) - } - - // Determine the base provider type for key requirement checks - baseProvider := req.Provider - if config.CustomProviderConfig != nil && config.CustomProviderConfig.BaseProviderType != "" { - baseProvider = config.CustomProviderConfig.BaseProviderType - } - - var key schemas.Key - if providerRequiresKey(baseProvider, config.CustomProviderConfig) { - keys, keyErr := bifrost.getAllSupportedKeys(&ctx, req.Provider, baseProvider) - if keyErr != nil { - return nil, newBifrostError(keyErr) - } - if len(keys) > 0 { - key = keys[0] - } - } - - response, bifrostErr := executeRequestWithRetries(&ctx, config, func() (*schemas.BifrostBatchResultsResponse, *schemas.BifrostError) { - return provider.BatchResults(ctx, key, req) - }, schemas.BatchResultsRequest, req.Provider, "") - if bifrostErr != nil { - bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{ - RequestType: schemas.BatchResultsRequest, - Provider: req.Provider, - } - return nil, bifrostErr + return nil, err } - return response, nil + return response.BatchResultsResponse, nil } // FileUploadRequest uploads a file to the specified provider. @@ -1301,52 +1126,15 @@ func (bifrost *Bifrost) FileUploadRequest(ctx context.Context, req *schemas.Bifr ctx = bifrost.ctx } - provider := bifrost.getProviderByKey(req.Provider) - if provider == nil { - return nil, &schemas.BifrostError{ - IsBifrostError: false, - Error: &schemas.ErrorField{ - Message: "provider not found for file upload request", - }, - } - } + bifrostReq := bifrost.getBifrostRequest() + bifrostReq.RequestType = schemas.FileUploadRequest + bifrostReq.FileUploadRequest = req - config, err := bifrost.account.GetConfigForProvider(req.Provider) + response, err := bifrost.handleRequest(ctx, bifrostReq) if err != nil { - return nil, newBifrostErrorFromMsg(fmt.Sprintf("failed to get config for provider %s: %v", req.Provider, err.Error())) - } - if config == nil { - return nil, newBifrostErrorFromMsg(fmt.Sprintf("config is nil for provider %s", req.Provider)) - } - - // Determine the base provider type for key requirement checks - baseProvider := req.Provider - if config.CustomProviderConfig != nil && config.CustomProviderConfig.BaseProviderType != "" { - baseProvider = config.CustomProviderConfig.BaseProviderType - } - - var key schemas.Key - if providerRequiresKey(baseProvider, config.CustomProviderConfig) { - keys, keyErr := bifrost.getAllSupportedKeys(&ctx, req.Provider, baseProvider) - if keyErr != nil { - return nil, newBifrostError(keyErr) - } - if len(keys) > 0 { - key = keys[0] - } - } - - response, bifrostErr := executeRequestWithRetries(&ctx, config, func() (*schemas.BifrostFileUploadResponse, *schemas.BifrostError) { - return provider.FileUpload(ctx, key, req) - }, schemas.FileUploadRequest, req.Provider, "") - if bifrostErr != nil { - bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{ - RequestType: schemas.FileUploadRequest, - Provider: req.Provider, - } - return nil, bifrostErr + return nil, err } - return response, nil + return response.FileUploadResponse, nil } // FileListRequest lists files from the specified provider. @@ -1377,54 +1165,15 @@ func (bifrost *Bifrost) FileListRequest(ctx context.Context, req *schemas.Bifros ctx = bifrost.ctx } - provider := bifrost.getProviderByKey(req.Provider) - if provider == nil { - return nil, &schemas.BifrostError{ - IsBifrostError: false, - Error: &schemas.ErrorField{ - Message: "provider not found for file list request", - }, - ExtraFields: schemas.BifrostErrorExtraFields{ - RequestType: schemas.FileListRequest, - Provider: req.Provider, - }, - } - } + bifrostReq := bifrost.getBifrostRequest() + bifrostReq.RequestType = schemas.FileListRequest + bifrostReq.FileListRequest = req - config, err := bifrost.account.GetConfigForProvider(req.Provider) + response, err := bifrost.handleRequest(ctx, bifrostReq) if err != nil { - return nil, newBifrostErrorFromMsg(fmt.Sprintf("failed to get config for provider %s: %v", req.Provider, err.Error())) - } - if config == nil { - return nil, newBifrostErrorFromMsg(fmt.Sprintf("config is nil for provider %s", req.Provider)) - } - - // Determine the base provider type for key requirement checks - baseProvider := req.Provider - if config.CustomProviderConfig != nil && config.CustomProviderConfig.BaseProviderType != "" { - baseProvider = config.CustomProviderConfig.BaseProviderType - } - - var keys []schemas.Key - var keyErr error - if providerRequiresKey(baseProvider, config.CustomProviderConfig) { - keys, keyErr = bifrost.getAllSupportedKeys(&ctx, req.Provider, baseProvider) - if keyErr != nil { - return nil, newBifrostError(keyErr) - } - } - - response, bifrostErr := executeRequestWithRetries(&ctx, config, func() (*schemas.BifrostFileListResponse, *schemas.BifrostError) { - return provider.FileList(ctx, keys, req) - }, schemas.FileListRequest, req.Provider, "") - if bifrostErr != nil { - bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{ - RequestType: schemas.FileListRequest, - Provider: req.Provider, - } - return nil, bifrostErr + return nil, err } - return response, nil + return response.FileListResponse, nil } // FileRetrieveRequest retrieves file metadata from the specified provider. @@ -1457,52 +1206,15 @@ func (bifrost *Bifrost) FileRetrieveRequest(ctx context.Context, req *schemas.Bi ctx = bifrost.ctx } - provider := bifrost.getProviderByKey(req.Provider) - if provider == nil { - return nil, &schemas.BifrostError{ - IsBifrostError: false, - Error: &schemas.ErrorField{ - Message: "provider not found for file retrieve request", - }, - } - } + bifrostReq := bifrost.getBifrostRequest() + bifrostReq.RequestType = schemas.FileRetrieveRequest + bifrostReq.FileRetrieveRequest = req - config, err := bifrost.account.GetConfigForProvider(req.Provider) + response, err := bifrost.handleRequest(ctx, bifrostReq) if err != nil { - return nil, newBifrostErrorFromMsg(fmt.Sprintf("failed to get config for provider %s: %v", req.Provider, err.Error())) - } - if config == nil { - return nil, newBifrostErrorFromMsg(fmt.Sprintf("config is nil for provider %s", req.Provider)) - } - - // Determine the base provider type for key requirement checks - baseProvider := req.Provider - if config.CustomProviderConfig != nil && config.CustomProviderConfig.BaseProviderType != "" { - baseProvider = config.CustomProviderConfig.BaseProviderType - } - - var key schemas.Key - if providerRequiresKey(baseProvider, config.CustomProviderConfig) { - keys, keyErr := bifrost.getAllSupportedKeys(&ctx, req.Provider, baseProvider) - if keyErr != nil { - return nil, newBifrostError(keyErr) - } - if len(keys) > 0 { - key = keys[0] - } - } - - response, bifrostErr := executeRequestWithRetries(&ctx, config, func() (*schemas.BifrostFileRetrieveResponse, *schemas.BifrostError) { - return provider.FileRetrieve(ctx, key, req) - }, schemas.FileRetrieveRequest, req.Provider, "") - if bifrostErr != nil { - bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{ - RequestType: schemas.FileRetrieveRequest, - Provider: req.Provider, - } - return nil, bifrostErr + return nil, err } - return response, nil + return response.FileRetrieveResponse, nil } // FileDeleteRequest deletes a file from the specified provider. @@ -1535,52 +1247,15 @@ func (bifrost *Bifrost) FileDeleteRequest(ctx context.Context, req *schemas.Bifr ctx = bifrost.ctx } - provider := bifrost.getProviderByKey(req.Provider) - if provider == nil { - return nil, &schemas.BifrostError{ - IsBifrostError: false, - Error: &schemas.ErrorField{ - Message: "provider not found for file delete request", - }, - } - } + bifrostReq := bifrost.getBifrostRequest() + bifrostReq.RequestType = schemas.FileDeleteRequest + bifrostReq.FileDeleteRequest = req - config, err := bifrost.account.GetConfigForProvider(req.Provider) + response, err := bifrost.handleRequest(ctx, bifrostReq) if err != nil { - return nil, newBifrostErrorFromMsg(fmt.Sprintf("failed to get config for provider %s: %v", req.Provider, err.Error())) - } - if config == nil { - return nil, newBifrostErrorFromMsg(fmt.Sprintf("config is nil for provider %s", req.Provider)) - } - - // Determine the base provider type for key requirement checks - baseProvider := req.Provider - if config.CustomProviderConfig != nil && config.CustomProviderConfig.BaseProviderType != "" { - baseProvider = config.CustomProviderConfig.BaseProviderType - } - - var key schemas.Key - if providerRequiresKey(baseProvider, config.CustomProviderConfig) { - keys, keyErr := bifrost.getAllSupportedKeys(&ctx, req.Provider, baseProvider) - if keyErr != nil { - return nil, newBifrostError(keyErr) - } - if len(keys) > 0 { - key = keys[0] - } - } - - response, bifrostErr := executeRequestWithRetries(&ctx, config, func() (*schemas.BifrostFileDeleteResponse, *schemas.BifrostError) { - return provider.FileDelete(ctx, key, req) - }, schemas.FileDeleteRequest, req.Provider, "") - if bifrostErr != nil { - bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{ - RequestType: schemas.FileDeleteRequest, - Provider: req.Provider, - } - return nil, bifrostErr + return nil, err } - return response, nil + return response.FileDeleteResponse, nil } // FileContentRequest downloads file content from the specified provider. @@ -1613,52 +1288,15 @@ func (bifrost *Bifrost) FileContentRequest(ctx context.Context, req *schemas.Bif ctx = bifrost.ctx } - provider := bifrost.getProviderByKey(req.Provider) - if provider == nil { - return nil, &schemas.BifrostError{ - IsBifrostError: false, - Error: &schemas.ErrorField{ - Message: "provider not found for file content request", - }, - } - } + bifrostReq := bifrost.getBifrostRequest() + bifrostReq.RequestType = schemas.FileContentRequest + bifrostReq.FileContentRequest = req - config, err := bifrost.account.GetConfigForProvider(req.Provider) + response, err := bifrost.handleRequest(ctx, bifrostReq) if err != nil { - return nil, newBifrostErrorFromMsg(fmt.Sprintf("failed to get config for provider %s: %v", req.Provider, err.Error())) - } - if config == nil { - return nil, newBifrostErrorFromMsg(fmt.Sprintf("config is nil for provider %s", req.Provider)) - } - - // Determine the base provider type for key requirement checks - baseProvider := req.Provider - if config.CustomProviderConfig != nil && config.CustomProviderConfig.BaseProviderType != "" { - baseProvider = config.CustomProviderConfig.BaseProviderType - } - - var key schemas.Key - if providerRequiresKey(baseProvider, config.CustomProviderConfig) { - keys, keyErr := bifrost.getAllSupportedKeys(&ctx, req.Provider, baseProvider) - if keyErr != nil { - return nil, newBifrostError(keyErr) - } - if len(keys) > 0 { - key = keys[0] - } - } - - response, bifrostErr := executeRequestWithRetries(&ctx, config, func() (*schemas.BifrostFileContentResponse, *schemas.BifrostError) { - return provider.FileContent(ctx, key, req) - }, schemas.FileContentRequest, req.Provider, "") - if bifrostErr != nil { - bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{ - RequestType: schemas.FileContentRequest, - Provider: req.Provider, - } - return nil, bifrostErr + return nil, err } - return response, nil + return response.FileContentResponse, nil } // RemovePlugin removes a plugin from the server. @@ -2613,6 +2251,7 @@ func (bifrost *Bifrost) handleStreamRequest(ctx context.Context, req *schemas.Bi Provider: provider, ModelRequested: model, } + err.StatusCode = schemas.Ptr(fasthttp.StatusBadRequest) return nil, err } @@ -3220,6 +2859,66 @@ func (bifrost *Bifrost) handleProviderRequest(provider schemas.Provider, req *Ch return nil, bifrostError } response.TranscriptionResponse = transcriptionResponse + case schemas.FileUploadRequest: + fileUploadResponse, bifrostError := provider.FileUpload(req.Context, key, req.BifrostRequest.FileUploadRequest) + if bifrostError != nil { + return nil, bifrostError + } + response.FileUploadResponse = fileUploadResponse + case schemas.FileListRequest: + fileListResponse, bifrostError := provider.FileList(req.Context, key, req.BifrostRequest.FileListRequest) + if bifrostError != nil { + return nil, bifrostError + } + response.FileListResponse = fileListResponse + case schemas.FileRetrieveRequest: + fileRetrieveResponse, bifrostError := provider.FileRetrieve(req.Context, key, req.BifrostRequest.FileRetrieveRequest) + if bifrostError != nil { + return nil, bifrostError + } + response.FileRetrieveResponse = fileRetrieveResponse + case schemas.FileDeleteRequest: + fileDeleteResponse, bifrostError := provider.FileDelete(req.Context, key, req.BifrostRequest.FileDeleteRequest) + if bifrostError != nil { + return nil, bifrostError + } + response.FileDeleteResponse = fileDeleteResponse + case schemas.FileContentRequest: + fileContentResponse, bifrostError := provider.FileContent(req.Context, key, req.BifrostRequest.FileContentRequest) + if bifrostError != nil { + return nil, bifrostError + } + response.FileContentResponse = fileContentResponse + case schemas.BatchCreateRequest: + batchCreateResponse, bifrostError := provider.BatchCreate(req.Context, key, req.BifrostRequest.BatchCreateRequest) + if bifrostError != nil { + return nil, bifrostError + } + response.BatchCreateResponse = batchCreateResponse + case schemas.BatchListRequest: + batchListResponse, bifrostError := provider.BatchList(req.Context, key, req.BifrostRequest.BatchListRequest) + if bifrostError != nil { + return nil, bifrostError + } + response.BatchListResponse = batchListResponse + case schemas.BatchRetrieveRequest: + batchRetrieveResponse, bifrostError := provider.BatchRetrieve(req.Context, key, req.BifrostRequest.BatchRetrieveRequest) + if bifrostError != nil { + return nil, bifrostError + } + response.BatchRetrieveResponse = batchRetrieveResponse + case schemas.BatchCancelRequest: + batchCancelResponse, bifrostError := provider.BatchCancel(req.Context, key, req.BifrostRequest.BatchCancelRequest) + if bifrostError != nil { + return nil, bifrostError + } + response.BatchCancelResponse = batchCancelResponse + case schemas.BatchResultsRequest: + batchResultsResponse, bifrostError := provider.BatchResults(req.Context, key, req.BifrostRequest.BatchResultsRequest) + if bifrostError != nil { + return nil, bifrostError + } + response.BatchResultsResponse = batchResultsResponse default: _, model, _ := req.BifrostRequest.GetRequestFields() return nil, &schemas.BifrostError{ @@ -3420,12 +3119,23 @@ func (bifrost *Bifrost) releaseChannelMessage(msg *ChannelMessage) { // resetBifrostRequest resets a BifrostRequest instance for reuse func resetBifrostRequest(req *schemas.BifrostRequest) { req.RequestType = "" + req.ListModelsRequest = nil req.TextCompletionRequest = nil req.ChatRequest = nil req.ResponsesRequest = nil req.EmbeddingRequest = nil req.SpeechRequest = nil req.TranscriptionRequest = nil + req.FileUploadRequest = nil + req.FileListRequest = nil + req.FileRetrieveRequest = nil + req.FileDeleteRequest = nil + req.FileContentRequest = nil + req.BatchCreateRequest = nil + req.BatchListRequest = nil + req.BatchRetrieveRequest = nil + req.BatchCancelRequest = nil + req.BatchResultsRequest = nil } // getBifrostRequest gets a BifrostRequest from the pool diff --git a/core/changelog.md b/core/changelog.md index cbd2ac094..ce943b98c 100644 --- a/core/changelog.md +++ b/core/changelog.md @@ -1,2 +1,5 @@ -fix: vertex and bedrock usage aggregation improvements for streaming -fix: choice index fixed to 0 for anthropic and bedrock streaming \ No newline at end of file +- feat: adds batch and files API support for bedrock, openai, anthropic and gemini +- feat: new provider support - nebius +- feat: structured output support +- fix: vertex and bedrock usage aggregation improvements for streaming +- fix: choice index fixed to 0 for anthropic and bedrock streaming \ No newline at end of file diff --git a/core/providers/anthropic/anthropic.go b/core/providers/anthropic/anthropic.go index 56f83abd8..62217a58d 100644 --- a/core/providers/anthropic/anthropic.go +++ b/core/providers/anthropic/anthropic.go @@ -1110,7 +1110,7 @@ func (provider *AnthropicProvider) BatchCreate(ctx context.Context, key schemas. } // BatchList lists batch jobs. -func (provider *AnthropicProvider) BatchList(ctx context.Context, keys []schemas.Key, request *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { +func (provider *AnthropicProvider) BatchList(ctx context.Context, key schemas.Key, request *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { if err := providerUtils.CheckOperationAllowed(schemas.Anthropic, provider.customProviderConfig, schemas.BatchListRequest); err != nil { return nil, err } @@ -1147,8 +1147,8 @@ func (provider *AnthropicProvider) BatchList(ctx context.Context, keys []schemas req.Header.SetContentType("application/json") // Use first key if available - if len(keys) > 0 && keys[0].Value != "" { - req.Header.Set("x-api-key", keys[0].Value) + if key.Value != "" { + req.Header.Set("x-api-key", key.Value) } req.Header.Set("anthropic-version", provider.apiVersion) @@ -1579,19 +1579,13 @@ func (provider *AnthropicProvider) FileUpload(ctx context.Context, key schemas.K } // FileList lists files from Anthropic's Files API. -func (provider *AnthropicProvider) FileList(ctx context.Context, keys []schemas.Key, request *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { +func (provider *AnthropicProvider) FileList(ctx context.Context, key schemas.Key, request *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { if err := providerUtils.CheckOperationAllowed(schemas.Anthropic, provider.customProviderConfig, schemas.FileListRequest); err != nil { return nil, err } providerName := provider.GetProviderKey() - - if len(keys) == 0 { - return nil, providerUtils.NewConfigurationError("no keys provided", providerName) - } - - key := keys[0] - + // Create request req := fasthttp.AcquireRequest() resp := fasthttp.AcquireResponse() @@ -1599,15 +1593,14 @@ func (provider *AnthropicProvider) FileList(ctx context.Context, keys []schemas. defer fasthttp.ReleaseResponse(resp) // Build URL with query params - baseURL := provider.buildRequestURL(ctx, "/v1/files", schemas.FileListRequest) + requestURL := provider.buildRequestURL(ctx, "/v1/files", schemas.FileListRequest) values := url.Values{} if request.Limit > 0 { values.Set("limit", fmt.Sprintf("%d", request.Limit)) } if request.After != nil && *request.After != "" { values.Set("after_id", *request.After) - } - requestURL := baseURL + } if encodedValues := values.Encode(); encodedValues != "" { requestURL += "?" + encodedValues } diff --git a/core/providers/anthropic/batch.go b/core/providers/anthropic/batch.go index 8125882c7..405738330 100644 --- a/core/providers/anthropic/batch.go +++ b/core/providers/anthropic/batch.go @@ -255,9 +255,15 @@ func ToAnthropicBatchCreateResponse(resp *schemas.BifrostBatchCreateResponse) *A Type: "message_batch", ProcessingStatus: toAnthropicProcessingStatus(resp.Status), CreatedAt: formatAnthropicTimestamp(resp.CreatedAt), - ExpiresAt: formatAnthropicTimestamp(*resp.ExpiresAt), ResultsURL: resp.ResultsURL, } + if resp.ExpiresAt != nil { + result.ExpiresAt = formatAnthropicTimestamp(*resp.ExpiresAt) + } else { + // This is a fallback for worst case scenario where expires_at is not available + // Which is never expected to happen, but just in case. + result.ExpiresAt = formatAnthropicTimestamp(time.Now().Add(24 * time.Hour).Unix()) + } if resp.RequestCounts.Total > 0 { result.RequestCounts = &AnthropicBatchRequestCounts{ Processing: resp.RequestCounts.Pending, diff --git a/core/providers/azure/azure.go b/core/providers/azure/azure.go index 070675b8b..26157235f 100644 --- a/core/providers/azure/azure.go +++ b/core/providers/azure/azure.go @@ -941,12 +941,8 @@ func (provider *AzureProvider) FileUpload(ctx context.Context, key schemas.Key, } // FileList lists files from Azure OpenAI. -func (provider *AzureProvider) FileList(ctx context.Context, keys []schemas.Key, request *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { - if len(keys) == 0 { - return nil, providerUtils.NewConfigurationError("no keys provided", provider.GetProviderKey()) - } - - key := keys[0] +func (provider *AzureProvider) FileList(ctx context.Context, key schemas.Key, request *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { + if err := provider.validateKeyConfigForFiles(key); err != nil { return nil, err } @@ -966,13 +962,15 @@ func (provider *AzureProvider) FileList(ctx context.Context, keys []schemas.Key, defer fasthttp.ReleaseResponse(resp) // Build URL with query params - baseURL := fmt.Sprintf("%s/openai/files", key.AzureKeyConfig.Endpoint) + requestURL := fmt.Sprintf("%s/openai/files", key.AzureKeyConfig.Endpoint) values := url.Values{} values.Set("api-version", *apiVersion) if request.Purpose != "" { values.Set("purpose", string(request.Purpose)) } - requestURL := baseURL + "?" + values.Encode() + if encodedValues := values.Encode(); encodedValues != "" { + requestURL += "?" + encodedValues + } // Set headers providerUtils.SetExtraHeaders(ctx, req, provider.networkConfig.ExtraHeaders, nil) @@ -1396,12 +1394,8 @@ func (provider *AzureProvider) BatchCreate(ctx context.Context, key schemas.Key, } // BatchList lists batch jobs from Azure OpenAI. -func (provider *AzureProvider) BatchList(ctx context.Context, keys []schemas.Key, request *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { - if len(keys) == 0 { - return nil, providerUtils.NewConfigurationError("no keys provided", provider.GetProviderKey()) - } - - key := keys[0] +func (provider *AzureProvider) BatchList(ctx context.Context, key schemas.Key, request *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { + if err := provider.validateKeyConfigForFiles(key); err != nil { return nil, err } diff --git a/core/providers/bedrock/batch.go b/core/providers/bedrock/batch.go index 4f5809ddc..4303ef5ab 100644 --- a/core/providers/bedrock/batch.go +++ b/core/providers/bedrock/batch.go @@ -332,7 +332,10 @@ func toBedrockBatchStatus(status schemas.BatchStatus) string { func ToBifrostBatchListRequest(req *BedrockBatchListRequest, provider schemas.ModelProvider) *schemas.BifrostBatchListRequest { result := &schemas.BifrostBatchListRequest{ Provider: provider, - Limit: req.MaxResults, + // We add a dummy model to avoid validation errors + // This model is never used in any of the provider flows + Model: "dummy-model", + Limit: req.MaxResults, } if req.NextToken != nil { diff --git a/core/providers/bedrock/bedrock.go b/core/providers/bedrock/bedrock.go index adb3cdc40..4f056553f 100644 --- a/core/providers/bedrock/bedrock.go +++ b/core/providers/bedrock/bedrock.go @@ -1291,8 +1291,11 @@ func (provider *BedrockProvider) TranscriptionStream(ctx context.Context, postHo // FileUpload uploads a file to S3 for Bedrock batch processing. func (provider *BedrockProvider) FileUpload(ctx context.Context, key schemas.Key, request *schemas.BifrostFileUploadRequest) (*schemas.BifrostFileUploadResponse, *schemas.BifrostError) { + if err := providerUtils.CheckOperationAllowed(schemas.Bedrock, provider.customProviderConfig, schemas.FileUploadRequest); err != nil { - provider.logger.Error("file upload operation not allowed: %s", err.Error.Message) + if err.Error != nil { + provider.logger.Error("file upload operation not allowed: %s", err.Error.Message) + } return nil, err } @@ -1417,22 +1420,13 @@ func (provider *BedrockProvider) FileUpload(ctx context.Context, key schemas.Key } // FileList lists files in the S3 bucket used for Bedrock batch processing. -func (provider *BedrockProvider) FileList(ctx context.Context, keys []schemas.Key, request *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { +func (provider *BedrockProvider) FileList(ctx context.Context, key schemas.Key, request *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { if err := providerUtils.CheckOperationAllowed(schemas.Bedrock, provider.customProviderConfig, schemas.FileListRequest); err != nil { return nil, err } providerName := provider.GetProviderKey() - if len(keys) == 0 { - return nil, providerUtils.NewConfigurationError("no keys provided", providerName) - } - - key := keys[0] - if key.BedrockKeyConfig == nil { - return nil, providerUtils.NewConfigurationError("bedrock key config is not provided", providerName) - } - // Get S3 bucket from storage config or extra params s3Bucket := "" s3Prefix := "" @@ -1478,9 +1472,9 @@ func (provider *BedrockProvider) FileList(ctx context.Context, keys []schemas.Ke params.Set("continuation-token", *request.After) } - reqURL := fmt.Sprintf("https://%s.s3.%s.amazonaws.com/?%s", bucketName, region, params.Encode()) + requestURL := fmt.Sprintf("https://%s.s3.%s.amazonaws.com/?%s", bucketName, region, params.Encode()) - httpReq, err := http.NewRequestWithContext(ctx, http.MethodGet, reqURL, nil) + httpReq, err := http.NewRequestWithContext(ctx, http.MethodGet, requestURL, nil) if err != nil { return nil, providerUtils.NewBifrostOperationError("error creating request", err, providerName) } @@ -2052,18 +2046,13 @@ func (provider *BedrockProvider) BatchCreate(ctx context.Context, key schemas.Ke } // BatchList lists batch inference jobs from AWS Bedrock. -func (provider *BedrockProvider) BatchList(ctx context.Context, keys []schemas.Key, request *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { +func (provider *BedrockProvider) BatchList(ctx context.Context, key schemas.Key, request *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { if err := providerUtils.CheckOperationAllowed(schemas.Bedrock, provider.customProviderConfig, schemas.BatchListRequest); err != nil { return nil, err } providerName := provider.GetProviderKey() - if len(keys) == 0 { - return nil, providerUtils.NewConfigurationError("no keys provided", providerName) - } - - key := keys[0] if key.BedrockKeyConfig == nil { return nil, providerUtils.NewConfigurationError("bedrock key config is not provided", providerName) } @@ -2508,7 +2497,7 @@ func (provider *BedrockProvider) BatchResults(ctx context.Context, key schemas.K allFiles []schemas.FileObject ) for { - listResp, bifrostErr = provider.FileList(ctx, []schemas.Key{key}, &schemas.BifrostFileListRequest{ + listResp, bifrostErr = provider.FileList(ctx, key, &schemas.BifrostFileListRequest{ Provider: request.Provider, StorageConfig: &schemas.FileStorageConfig{ S3: &schemas.S3StorageConfig{ diff --git a/core/providers/cerebras/cerebras.go b/core/providers/cerebras/cerebras.go index 5a25ff6f9..aeed94fee 100644 --- a/core/providers/cerebras/cerebras.go +++ b/core/providers/cerebras/cerebras.go @@ -217,7 +217,7 @@ func (provider *CerebrasProvider) FileUpload(_ context.Context, _ schemas.Key, _ } // FileList is not supported by Cerebras provider. -func (provider *CerebrasProvider) FileList(_ context.Context, _ []schemas.Key, _ *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { +func (provider *CerebrasProvider) FileList(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { return nil, providerUtils.NewUnsupportedOperationError(schemas.FileListRequest, provider.GetProviderKey()) } @@ -242,7 +242,7 @@ func (provider *CerebrasProvider) BatchCreate(_ context.Context, _ schemas.Key, } // BatchList is not supported by Cerebras provider. -func (provider *CerebrasProvider) BatchList(_ context.Context, _ []schemas.Key, _ *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { +func (provider *CerebrasProvider) BatchList(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchListRequest, provider.GetProviderKey()) } diff --git a/core/providers/cohere/cohere.go b/core/providers/cohere/cohere.go index 973cf2bae..c87794f46 100644 --- a/core/providers/cohere/cohere.go +++ b/core/providers/cohere/cohere.go @@ -847,7 +847,7 @@ func (provider *CohereProvider) BatchCreate(_ context.Context, _ schemas.Key, _ } // BatchList is not supported by Cohere provider. -func (provider *CohereProvider) BatchList(_ context.Context, _ []schemas.Key, _ *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { +func (provider *CohereProvider) BatchList(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchListRequest, provider.GetProviderKey()) } @@ -872,7 +872,7 @@ func (provider *CohereProvider) FileUpload(_ context.Context, _ schemas.Key, _ * } // FileList is not supported by Cohere provider. -func (provider *CohereProvider) FileList(_ context.Context, _ []schemas.Key, _ *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { +func (provider *CohereProvider) FileList(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { return nil, providerUtils.NewUnsupportedOperationError(schemas.FileListRequest, provider.GetProviderKey()) } diff --git a/core/providers/elevenlabs/elevenlabs.go b/core/providers/elevenlabs/elevenlabs.go index 095d61b24..da937b30d 100644 --- a/core/providers/elevenlabs/elevenlabs.go +++ b/core/providers/elevenlabs/elevenlabs.go @@ -718,7 +718,7 @@ func (provider *ElevenlabsProvider) BatchCreate(_ context.Context, _ schemas.Key } // BatchList is not supported by Elevenlabs provider. -func (provider *ElevenlabsProvider) BatchList(_ context.Context, _ []schemas.Key, _ *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { +func (provider *ElevenlabsProvider) BatchList(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchListRequest, provider.GetProviderKey()) } @@ -743,7 +743,7 @@ func (provider *ElevenlabsProvider) FileUpload(_ context.Context, _ schemas.Key, } // FileList is not supported by Elevenlabs provider. -func (provider *ElevenlabsProvider) FileList(_ context.Context, _ []schemas.Key, _ *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { +func (provider *ElevenlabsProvider) FileList(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { return nil, providerUtils.NewUnsupportedOperationError(schemas.FileListRequest, provider.GetProviderKey()) } diff --git a/core/providers/gemini/gemini.go b/core/providers/gemini/gemini.go index b32b6bf09..a6f847101 100644 --- a/core/providers/gemini/gemini.go +++ b/core/providers/gemini/gemini.go @@ -1624,19 +1624,13 @@ func (provider *GeminiProvider) BatchCreate(ctx context.Context, key schemas.Key // BatchList lists batch jobs for Gemini. // Note: The consumer API may have limited list functionality. -func (provider *GeminiProvider) BatchList(ctx context.Context, keys []schemas.Key, request *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { +func (provider *GeminiProvider) BatchList(ctx context.Context, key schemas.Key, request *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { if err := providerUtils.CheckOperationAllowed(schemas.Gemini, provider.customProviderConfig, schemas.BatchListRequest); err != nil { return nil, err } providerName := provider.GetProviderKey() - // Select a key for the request - if len(keys) == 0 { - return nil, providerUtils.NewBifrostOperationError("at least one API key is required", nil, providerName) - } - key := keys[0] - // Create HTTP request req := fasthttp.AcquireRequest() resp := fasthttp.AcquireResponse() @@ -2238,19 +2232,13 @@ func (provider *GeminiProvider) FileUpload(ctx context.Context, key schemas.Key, } // FileList lists files from Gemini. -func (provider *GeminiProvider) FileList(ctx context.Context, keys []schemas.Key, request *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { +func (provider *GeminiProvider) FileList(ctx context.Context, key schemas.Key, request *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { if err := providerUtils.CheckOperationAllowed(schemas.Gemini, provider.customProviderConfig, schemas.FileListRequest); err != nil { return nil, err } providerName := provider.GetProviderKey() - if len(keys) == 0 { - return nil, providerUtils.NewConfigurationError("no keys provided", providerName) - } - - key := keys[0] - // Create request req := fasthttp.AcquireRequest() resp := fasthttp.AcquireResponse() @@ -2258,7 +2246,7 @@ func (provider *GeminiProvider) FileList(ctx context.Context, keys []schemas.Key defer fasthttp.ReleaseResponse(resp) // Build URL with pagination - baseURL := fmt.Sprintf("%s/files", provider.networkConfig.BaseURL) + requestURL := fmt.Sprintf("%s/files", provider.networkConfig.BaseURL) values := url.Values{} if request.Limit > 0 { values.Set("pageSize", fmt.Sprintf("%d", request.Limit)) @@ -2266,7 +2254,6 @@ func (provider *GeminiProvider) FileList(ctx context.Context, keys []schemas.Key if request.After != nil && *request.After != "" { values.Set("pageToken", *request.After) } - requestURL := baseURL if encodedValues := values.Encode(); encodedValues != "" { requestURL += "?" + encodedValues } @@ -2315,6 +2302,10 @@ func (provider *GeminiProvider) FileList(ctx context.Context, keys []schemas.Key }, } + if geminiResp.NextPageToken != "" { + bifrostResp.After = &geminiResp.NextPageToken + } + for i, file := range geminiResp.Files { var sizeBytes int64 fmt.Sscanf(file.SizeBytes, "%d", &sizeBytes) diff --git a/core/providers/groq/groq.go b/core/providers/groq/groq.go index 240792534..9796b4556 100644 --- a/core/providers/groq/groq.go +++ b/core/providers/groq/groq.go @@ -255,7 +255,7 @@ func (provider *GroqProvider) BatchCreate(_ context.Context, _ schemas.Key, _ *s } // BatchList is not supported by Groq provider. -func (provider *GroqProvider) BatchList(_ context.Context, _ []schemas.Key, _ *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { +func (provider *GroqProvider) BatchList(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchListRequest, provider.GetProviderKey()) } @@ -280,7 +280,7 @@ func (provider *GroqProvider) FileUpload(_ context.Context, _ schemas.Key, _ *sc } // FileList is not supported by Groq provider. -func (provider *GroqProvider) FileList(_ context.Context, _ []schemas.Key, _ *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { +func (provider *GroqProvider) FileList(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { return nil, providerUtils.NewUnsupportedOperationError(schemas.FileListRequest, provider.GetProviderKey()) } diff --git a/core/providers/mistral/batch.go b/core/providers/mistral/batch.go deleted file mode 100644 index e368910db..000000000 --- a/core/providers/mistral/batch.go +++ /dev/null @@ -1,34 +0,0 @@ -package mistral - -import ( - "context" - - providerUtils "github.com/maximhq/bifrost/core/providers/utils" - "github.com/maximhq/bifrost/core/schemas" -) - -// BatchCreate is not supported by Mistral provider. -func (provider *MistralProvider) BatchCreate(ctx context.Context, key schemas.Key, request *schemas.BifrostBatchCreateRequest) (*schemas.BifrostBatchCreateResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchCreateRequest, provider.GetProviderKey()) -} - -// BatchList is not supported by Mistral provider. -func (provider *MistralProvider) BatchList(ctx context.Context, keys []schemas.Key, request *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchListRequest, provider.GetProviderKey()) -} - -// BatchRetrieve is not supported by Mistral provider. -func (provider *MistralProvider) BatchRetrieve(ctx context.Context, key schemas.Key, request *schemas.BifrostBatchRetrieveRequest) (*schemas.BifrostBatchRetrieveResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchRetrieveRequest, provider.GetProviderKey()) -} - -// BatchCancel is not supported by Mistral provider. -func (provider *MistralProvider) BatchCancel(ctx context.Context, key schemas.Key, request *schemas.BifrostBatchCancelRequest) (*schemas.BifrostBatchCancelResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchCancelRequest, provider.GetProviderKey()) -} - -// BatchResults is not supported by Mistral provider. -func (provider *MistralProvider) BatchResults(ctx context.Context, key schemas.Key, request *schemas.BifrostBatchResultsRequest) (*schemas.BifrostBatchResultsResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchResultsRequest, provider.GetProviderKey()) -} - diff --git a/core/providers/mistral/files.go b/core/providers/mistral/files.go deleted file mode 100644 index d34906fa4..000000000 --- a/core/providers/mistral/files.go +++ /dev/null @@ -1,34 +0,0 @@ -package mistral - -import ( - "context" - - providerUtils "github.com/maximhq/bifrost/core/providers/utils" - "github.com/maximhq/bifrost/core/schemas" -) - -// FileUpload is not supported by Mistral provider. -func (provider *MistralProvider) FileUpload(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileUploadRequest) (*schemas.BifrostFileUploadResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.FileUploadRequest, provider.GetProviderKey()) -} - -// FileList is not supported by Mistral provider. -func (provider *MistralProvider) FileList(_ context.Context, _ []schemas.Key, _ *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.FileListRequest, provider.GetProviderKey()) -} - -// FileRetrieve is not supported by Mistral provider. -func (provider *MistralProvider) FileRetrieve(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileRetrieveRequest) (*schemas.BifrostFileRetrieveResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.FileRetrieveRequest, provider.GetProviderKey()) -} - -// FileDelete is not supported by Mistral provider. -func (provider *MistralProvider) FileDelete(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileDeleteRequest) (*schemas.BifrostFileDeleteResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.FileDeleteRequest, provider.GetProviderKey()) -} - -// FileContent is not supported by Mistral provider. -func (provider *MistralProvider) FileContent(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileContentRequest) (*schemas.BifrostFileContentResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.FileContentRequest, provider.GetProviderKey()) -} - diff --git a/core/providers/mistral/mistral.go b/core/providers/mistral/mistral.go index 6e093f710..d33e9541e 100644 --- a/core/providers/mistral/mistral.go +++ b/core/providers/mistral/mistral.go @@ -256,3 +256,53 @@ func (provider *MistralProvider) Transcription(ctx context.Context, key schemas. func (provider *MistralProvider) TranscriptionStream(ctx context.Context, postHookRunner schemas.PostHookRunner, key schemas.Key, request *schemas.BifrostTranscriptionRequest) (chan *schemas.BifrostStream, *schemas.BifrostError) { return nil, providerUtils.NewUnsupportedOperationError(schemas.TranscriptionStreamRequest, provider.GetProviderKey()) } + +// BatchCreate is not supported by Mistral provider. +func (provider *MistralProvider) BatchCreate(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchCreateRequest) (*schemas.BifrostBatchCreateResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchCreateRequest, provider.GetProviderKey()) +} + +// BatchList is not supported by Mistral provider. +func (provider *MistralProvider) BatchList(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchListRequest, provider.GetProviderKey()) +} + +// BatchRetrieve is not supported by Mistral provider. +func (provider *MistralProvider) BatchRetrieve(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchRetrieveRequest) (*schemas.BifrostBatchRetrieveResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchRetrieveRequest, provider.GetProviderKey()) +} + +// BatchCancel is not supported by Mistral provider. +func (provider *MistralProvider) BatchCancel(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchCancelRequest) (*schemas.BifrostBatchCancelResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchCancelRequest, provider.GetProviderKey()) +} + +// BatchResults is not supported by Mistral provider. +func (provider *MistralProvider) BatchResults(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchResultsRequest) (*schemas.BifrostBatchResultsResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchResultsRequest, provider.GetProviderKey()) +} + +// FileUpload is not supported by Mistral provider. +func (provider *MistralProvider) FileUpload(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileUploadRequest) (*schemas.BifrostFileUploadResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.FileUploadRequest, provider.GetProviderKey()) +} + +// FileList is not supported by Mistral provider. +func (provider *MistralProvider) FileList(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.FileListRequest, provider.GetProviderKey()) +} + +// FileRetrieve is not supported by Mistral provider. +func (provider *MistralProvider) FileRetrieve(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileRetrieveRequest) (*schemas.BifrostFileRetrieveResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.FileRetrieveRequest, provider.GetProviderKey()) +} + +// FileDelete is not supported by Mistral provider. +func (provider *MistralProvider) FileDelete(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileDeleteRequest) (*schemas.BifrostFileDeleteResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.FileDeleteRequest, provider.GetProviderKey()) +} + +// FileContent is not supported by Mistral provider. +func (provider *MistralProvider) FileContent(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileContentRequest) (*schemas.BifrostFileContentResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.FileContentRequest, provider.GetProviderKey()) +} diff --git a/core/providers/nebius/nebius.go b/core/providers/nebius/nebius.go index b9519b4aa..0574616d4 100644 --- a/core/providers/nebius/nebius.go +++ b/core/providers/nebius/nebius.go @@ -18,6 +18,7 @@ type NebiusProvider struct { logger schemas.Logger // Logger for provider operations client *fasthttp.Client // HTTP client for API requests networkConfig schemas.NetworkConfig // Network configuration including extra headers + sendBackRawRequest bool // Whether to include raw request in BifrostResponse sendBackRawResponse bool // Whether to include raw response in BifrostResponse } @@ -67,6 +68,7 @@ func (provider *NebiusProvider) ListModels(ctx context.Context, keys []schemas.K keys, provider.networkConfig.ExtraHeaders, provider.GetProviderKey(), + providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest), providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse), provider.logger, ) @@ -84,6 +86,7 @@ func (provider *NebiusProvider) TextCompletion(ctx context.Context, key schemas. key, provider.networkConfig.ExtraHeaders, provider.GetProviderKey(), + providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest), providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse), provider.logger, ) @@ -105,6 +108,7 @@ func (provider *NebiusProvider) TextCompletionStream(ctx context.Context, postHo request, authHeader, provider.networkConfig.ExtraHeaders, + providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest), providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse), provider.GetProviderKey(), postHookRunner, @@ -133,6 +137,7 @@ func (provider *NebiusProvider) ChatCompletion(ctx context.Context, key schemas. request, key, provider.networkConfig.ExtraHeaders, + providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest), providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse), provider.GetProviderKey(), provider.logger, @@ -157,6 +162,7 @@ func (provider *NebiusProvider) ChatCompletionStream(ctx context.Context, postHo request, authHeader, provider.networkConfig.ExtraHeaders, + providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest), providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse), provider.GetProviderKey(), postHookRunner, @@ -204,6 +210,7 @@ func (provider *NebiusProvider) Embedding(ctx context.Context, key schemas.Key, key, provider.networkConfig.ExtraHeaders, provider.GetProviderKey(), + providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest), providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse), provider.logger) } @@ -227,3 +234,53 @@ func (provider *NebiusProvider) Transcription(ctx context.Context, key schemas.K func (provider *NebiusProvider) TranscriptionStream(ctx context.Context, postHookRunner schemas.PostHookRunner, key schemas.Key, request *schemas.BifrostTranscriptionRequest) (chan *schemas.BifrostStream, *schemas.BifrostError) { return nil, providerUtils.NewUnsupportedOperationError(schemas.TranscriptionStreamRequest, provider.GetProviderKey()) } + +// BatchCreate is not supported by Nebius provider. +func (provider *NebiusProvider) BatchCreate(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchCreateRequest) (*schemas.BifrostBatchCreateResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchCreateRequest, provider.GetProviderKey()) +} + +// BatchList is not supported by Nebius provider. +func (provider *NebiusProvider) BatchList(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchListRequest, provider.GetProviderKey()) +} + +// BatchRetrieve is not supported by Nebius provider. +func (provider *NebiusProvider) BatchRetrieve(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchRetrieveRequest) (*schemas.BifrostBatchRetrieveResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchRetrieveRequest, provider.GetProviderKey()) +} + +// BatchCancel is not supported by Nebius provider. +func (provider *NebiusProvider) BatchCancel(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchCancelRequest) (*schemas.BifrostBatchCancelResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchCancelRequest, provider.GetProviderKey()) +} + +// BatchResults is not supported by Nebius provider. +func (provider *NebiusProvider) BatchResults(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchResultsRequest) (*schemas.BifrostBatchResultsResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchResultsRequest, provider.GetProviderKey()) +} + +// FileUpload is not supported by Nebius provider. +func (provider *NebiusProvider) FileUpload(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileUploadRequest) (*schemas.BifrostFileUploadResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.FileUploadRequest, provider.GetProviderKey()) +} + +// FileList is not supported by Nebius provider. +func (provider *NebiusProvider) FileList(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.FileListRequest, provider.GetProviderKey()) +} + +// FileRetrieve is not supported by Nebius provider. +func (provider *NebiusProvider) FileRetrieve(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileRetrieveRequest) (*schemas.BifrostFileRetrieveResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.FileRetrieveRequest, provider.GetProviderKey()) +} + +// FileDelete is not supported by Nebius provider. +func (provider *NebiusProvider) FileDelete(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileDeleteRequest) (*schemas.BifrostFileDeleteResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.FileDeleteRequest, provider.GetProviderKey()) +} + +// FileContent is not supported by Nebius provider. +func (provider *NebiusProvider) FileContent(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileContentRequest) (*schemas.BifrostFileContentResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.FileContentRequest, provider.GetProviderKey()) +} diff --git a/core/providers/ollama/ollama.go b/core/providers/ollama/ollama.go index de13e9a66..cb02bb3ab 100644 --- a/core/providers/ollama/ollama.go +++ b/core/providers/ollama/ollama.go @@ -229,7 +229,7 @@ func (provider *OllamaProvider) BatchCreate(_ context.Context, _ schemas.Key, _ } // BatchList is not supported by Ollama provider. -func (provider *OllamaProvider) BatchList(_ context.Context, _ []schemas.Key, _ *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { +func (provider *OllamaProvider) BatchList(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchListRequest, provider.GetProviderKey()) } @@ -254,7 +254,7 @@ func (provider *OllamaProvider) FileUpload(_ context.Context, _ schemas.Key, _ * } // FileList is not supported by Ollama provider. -func (provider *OllamaProvider) FileList(_ context.Context, _ []schemas.Key, _ *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { +func (provider *OllamaProvider) FileList(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { return nil, providerUtils.NewUnsupportedOperationError(schemas.FileListRequest, provider.GetProviderKey()) } diff --git a/core/providers/openai/openai.go b/core/providers/openai/openai.go index a95432cb2..0d636ead6 100644 --- a/core/providers/openai/openai.go +++ b/core/providers/openai/openai.go @@ -2242,7 +2242,7 @@ func (provider *OpenAIProvider) FileUpload(ctx context.Context, key schemas.Key, } // FileList lists files from OpenAI. -func (provider *OpenAIProvider) FileList(ctx context.Context, keys []schemas.Key, request *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { +func (provider *OpenAIProvider) FileList(ctx context.Context, key schemas.Key, request *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { if err := providerUtils.CheckOperationAllowed(schemas.OpenAI, provider.customProviderConfig, schemas.FileListRequest); err != nil { return nil, err } @@ -2256,7 +2256,7 @@ func (provider *OpenAIProvider) FileList(ctx context.Context, keys []schemas.Key defer fasthttp.ReleaseResponse(resp) // Build URL with query params - baseUrl := provider.buildRequestURL(ctx, "/v1/files", schemas.FileListRequest) + requestURL := provider.buildRequestURL(ctx, "/v1/files", schemas.FileListRequest) values := url.Values{} if request.Purpose != "" { values.Set("purpose", string(request.Purpose)) @@ -2271,18 +2271,18 @@ func (provider *OpenAIProvider) FileList(ctx context.Context, keys []schemas.Key values.Set("order", *request.Order) } if encoded := values.Encode(); encoded != "" { - baseUrl += "?" + encoded + requestURL += "?" + encoded } // Set headers providerUtils.SetExtraHeaders(ctx, req, provider.networkConfig.ExtraHeaders, nil) - req.SetRequestURI(baseUrl) + req.SetRequestURI(requestURL) req.Header.SetMethod(http.MethodGet) req.Header.SetContentType("application/json") // Use first key if available - if len(keys) > 0 && keys[0].Value != "" { - req.Header.Set("Authorization", "Bearer "+keys[0].Value) + if key.Value != "" { + req.Header.Set("Authorization", "Bearer "+key.Value) } // Make request @@ -2644,7 +2644,7 @@ func (provider *OpenAIProvider) BatchCreate(ctx context.Context, key schemas.Key } // BatchList lists batch jobs. -func (provider *OpenAIProvider) BatchList(ctx context.Context, keys []schemas.Key, request *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { +func (provider *OpenAIProvider) BatchList(ctx context.Context, key schemas.Key, request *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { if err := providerUtils.CheckOperationAllowed(schemas.OpenAI, provider.customProviderConfig, schemas.BatchListRequest); err != nil { return nil, err } @@ -2678,8 +2678,8 @@ func (provider *OpenAIProvider) BatchList(ctx context.Context, keys []schemas.Ke req.Header.SetContentType("application/json") // Use first key if available - if len(keys) > 0 && keys[0].Value != "" { - req.Header.Set("Authorization", "Bearer "+keys[0].Value) + if key.Value != "" { + req.Header.Set("Authorization", "Bearer "+key.Value) } // Make request diff --git a/core/providers/openrouter/openrouter.go b/core/providers/openrouter/openrouter.go index 5885ef451..e7342cf1c 100644 --- a/core/providers/openrouter/openrouter.go +++ b/core/providers/openrouter/openrouter.go @@ -288,7 +288,7 @@ func (provider *OpenRouterProvider) BatchCreate(_ context.Context, _ schemas.Key } // BatchList is not supported by OpenRouter provider. -func (provider *OpenRouterProvider) BatchList(_ context.Context, _ []schemas.Key, _ *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { +func (provider *OpenRouterProvider) BatchList(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchListRequest, provider.GetProviderKey()) } @@ -308,26 +308,26 @@ func (provider *OpenRouterProvider) BatchResults(_ context.Context, _ schemas.Ke } // FileUpload is not supported by OpenRouter provider. -func (provider *OpenRouterProvider) FileUpload(ctx context.Context, key schemas.Key, request *schemas.BifrostFileUploadRequest) (*schemas.BifrostFileUploadResponse, *schemas.BifrostError) { +func (provider *OpenRouterProvider) FileUpload(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileUploadRequest) (*schemas.BifrostFileUploadResponse, *schemas.BifrostError) { return nil, providerUtils.NewUnsupportedOperationError(schemas.FileUploadRequest, provider.GetProviderKey()) } // FileList is not supported by OpenRouter provider. -func (provider *OpenRouterProvider) FileList(ctx context.Context, keys []schemas.Key, request *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { +func (provider *OpenRouterProvider) FileList(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { return nil, providerUtils.NewUnsupportedOperationError(schemas.FileListRequest, provider.GetProviderKey()) } // FileRetrieve is not supported by OpenRouter provider. -func (provider *OpenRouterProvider) FileRetrieve(ctx context.Context, key schemas.Key, request *schemas.BifrostFileRetrieveRequest) (*schemas.BifrostFileRetrieveResponse, *schemas.BifrostError) { +func (provider *OpenRouterProvider) FileRetrieve(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileRetrieveRequest) (*schemas.BifrostFileRetrieveResponse, *schemas.BifrostError) { return nil, providerUtils.NewUnsupportedOperationError(schemas.FileRetrieveRequest, provider.GetProviderKey()) } // FileDelete is not supported by OpenRouter provider. -func (provider *OpenRouterProvider) FileDelete(ctx context.Context, key schemas.Key, request *schemas.BifrostFileDeleteRequest) (*schemas.BifrostFileDeleteResponse, *schemas.BifrostError) { +func (provider *OpenRouterProvider) FileDelete(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileDeleteRequest) (*schemas.BifrostFileDeleteResponse, *schemas.BifrostError) { return nil, providerUtils.NewUnsupportedOperationError(schemas.FileDeleteRequest, provider.GetProviderKey()) } // FileContent is not supported by OpenRouter provider. -func (provider *OpenRouterProvider) FileContent(ctx context.Context, key schemas.Key, request *schemas.BifrostFileContentRequest) (*schemas.BifrostFileContentResponse, *schemas.BifrostError) { +func (provider *OpenRouterProvider) FileContent(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileContentRequest) (*schemas.BifrostFileContentResponse, *schemas.BifrostError) { return nil, providerUtils.NewUnsupportedOperationError(schemas.FileContentRequest, provider.GetProviderKey()) } diff --git a/core/providers/parasail/batch.go b/core/providers/parasail/batch.go deleted file mode 100644 index 67c9cacfc..000000000 --- a/core/providers/parasail/batch.go +++ /dev/null @@ -1,34 +0,0 @@ -package parasail - -import ( - "context" - - providerUtils "github.com/maximhq/bifrost/core/providers/utils" - "github.com/maximhq/bifrost/core/schemas" -) - -// BatchCreate is not supported by Parasail provider. -func (provider *ParasailProvider) BatchCreate(ctx context.Context, key schemas.Key, request *schemas.BifrostBatchCreateRequest) (*schemas.BifrostBatchCreateResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchCreateRequest, provider.GetProviderKey()) -} - -// BatchList is not supported by Parasail provider. -func (provider *ParasailProvider) BatchList(ctx context.Context, keys []schemas.Key, request *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchListRequest, provider.GetProviderKey()) -} - -// BatchRetrieve is not supported by Parasail provider. -func (provider *ParasailProvider) BatchRetrieve(ctx context.Context, key schemas.Key, request *schemas.BifrostBatchRetrieveRequest) (*schemas.BifrostBatchRetrieveResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchRetrieveRequest, provider.GetProviderKey()) -} - -// BatchCancel is not supported by Parasail provider. -func (provider *ParasailProvider) BatchCancel(ctx context.Context, key schemas.Key, request *schemas.BifrostBatchCancelRequest) (*schemas.BifrostBatchCancelResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchCancelRequest, provider.GetProviderKey()) -} - -// BatchResults is not supported by Parasail provider. -func (provider *ParasailProvider) BatchResults(ctx context.Context, key schemas.Key, request *schemas.BifrostBatchResultsRequest) (*schemas.BifrostBatchResultsResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchResultsRequest, provider.GetProviderKey()) -} - diff --git a/core/providers/parasail/files.go b/core/providers/parasail/files.go deleted file mode 100644 index 5f9972128..000000000 --- a/core/providers/parasail/files.go +++ /dev/null @@ -1,34 +0,0 @@ -package parasail - -import ( - "context" - - providerUtils "github.com/maximhq/bifrost/core/providers/utils" - "github.com/maximhq/bifrost/core/schemas" -) - -// FileUpload is not supported by Parasail provider. -func (provider *ParasailProvider) FileUpload(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileUploadRequest) (*schemas.BifrostFileUploadResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.FileUploadRequest, provider.GetProviderKey()) -} - -// FileList is not supported by Parasail provider. -func (provider *ParasailProvider) FileList(_ context.Context, _ []schemas.Key, _ *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.FileListRequest, provider.GetProviderKey()) -} - -// FileRetrieve is not supported by Parasail provider. -func (provider *ParasailProvider) FileRetrieve(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileRetrieveRequest) (*schemas.BifrostFileRetrieveResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.FileRetrieveRequest, provider.GetProviderKey()) -} - -// FileDelete is not supported by Parasail provider. -func (provider *ParasailProvider) FileDelete(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileDeleteRequest) (*schemas.BifrostFileDeleteResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.FileDeleteRequest, provider.GetProviderKey()) -} - -// FileContent is not supported by Parasail provider. -func (provider *ParasailProvider) FileContent(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileContentRequest) (*schemas.BifrostFileContentResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.FileContentRequest, provider.GetProviderKey()) -} - diff --git a/core/providers/parasail/parasail.go b/core/providers/parasail/parasail.go index df7764ec2..82f00f65f 100644 --- a/core/providers/parasail/parasail.go +++ b/core/providers/parasail/parasail.go @@ -181,3 +181,53 @@ func (provider *ParasailProvider) Transcription(ctx context.Context, key schemas func (provider *ParasailProvider) TranscriptionStream(ctx context.Context, postHookRunner schemas.PostHookRunner, key schemas.Key, request *schemas.BifrostTranscriptionRequest) (chan *schemas.BifrostStream, *schemas.BifrostError) { return nil, providerUtils.NewUnsupportedOperationError(schemas.TranscriptionStreamRequest, provider.GetProviderKey()) } + +// FileUpload is not supported by Parasail provider. +func (provider *ParasailProvider) FileUpload(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileUploadRequest) (*schemas.BifrostFileUploadResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.FileUploadRequest, provider.GetProviderKey()) +} + +// FileList is not supported by Parasail provider. +func (provider *ParasailProvider) FileList(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.FileListRequest, provider.GetProviderKey()) +} + +// FileRetrieve is not supported by Parasail provider. +func (provider *ParasailProvider) FileRetrieve(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileRetrieveRequest) (*schemas.BifrostFileRetrieveResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.FileRetrieveRequest, provider.GetProviderKey()) +} + +// FileDelete is not supported by Parasail provider. +func (provider *ParasailProvider) FileDelete(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileDeleteRequest) (*schemas.BifrostFileDeleteResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.FileDeleteRequest, provider.GetProviderKey()) +} + +// FileContent is not supported by Parasail provider. +func (provider *ParasailProvider) FileContent(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileContentRequest) (*schemas.BifrostFileContentResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.FileContentRequest, provider.GetProviderKey()) +} + +// BatchCreate is not supported by Parasail provider. +func (provider *ParasailProvider) BatchCreate(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchCreateRequest) (*schemas.BifrostBatchCreateResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchCreateRequest, provider.GetProviderKey()) +} + +// BatchList is not supported by Parasail provider. +func (provider *ParasailProvider) BatchList(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchListRequest, provider.GetProviderKey()) +} + +// BatchRetrieve is not supported by Parasail provider. +func (provider *ParasailProvider) BatchRetrieve(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchRetrieveRequest) (*schemas.BifrostBatchRetrieveResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchRetrieveRequest, provider.GetProviderKey()) +} + +// BatchCancel is not supported by Parasail provider. +func (provider *ParasailProvider) BatchCancel(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchCancelRequest) (*schemas.BifrostBatchCancelResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchCancelRequest, provider.GetProviderKey()) +} + +// BatchResults is not supported by Parasail provider. +func (provider *ParasailProvider) BatchResults(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchResultsRequest) (*schemas.BifrostBatchResultsResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchResultsRequest, provider.GetProviderKey()) +} diff --git a/core/providers/perplexity/perplexity.go b/core/providers/perplexity/perplexity.go index f7719b79f..060159e81 100644 --- a/core/providers/perplexity/perplexity.go +++ b/core/providers/perplexity/perplexity.go @@ -258,7 +258,7 @@ func (provider *PerplexityProvider) BatchCreate(_ context.Context, _ schemas.Key } // BatchList is not supported by Perplexity provider. -func (provider *PerplexityProvider) BatchList(_ context.Context, _ []schemas.Key, _ *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { +func (provider *PerplexityProvider) BatchList(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchListRequest, provider.GetProviderKey()) } @@ -283,7 +283,7 @@ func (provider *PerplexityProvider) FileUpload(_ context.Context, _ schemas.Key, } // FileList is not supported by Perplexity provider. -func (provider *PerplexityProvider) FileList(_ context.Context, _ []schemas.Key, _ *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { +func (provider *PerplexityProvider) FileList(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { return nil, providerUtils.NewUnsupportedOperationError(schemas.FileListRequest, provider.GetProviderKey()) } diff --git a/core/providers/sgl/batch.go b/core/providers/sgl/batch.go deleted file mode 100644 index c142fa973..000000000 --- a/core/providers/sgl/batch.go +++ /dev/null @@ -1,34 +0,0 @@ -package sgl - -import ( - "context" - - providerUtils "github.com/maximhq/bifrost/core/providers/utils" - "github.com/maximhq/bifrost/core/schemas" -) - -// BatchCreate is not supported by SGL provider. -func (provider *SGLProvider) BatchCreate(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchCreateRequest) (*schemas.BifrostBatchCreateResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchCreateRequest, provider.GetProviderKey()) -} - -// BatchList is not supported by SGL provider. -func (provider *SGLProvider) BatchList(_ context.Context, _ []schemas.Key, _ *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchListRequest, provider.GetProviderKey()) -} - -// BatchRetrieve is not supported by SGL provider. -func (provider *SGLProvider) BatchRetrieve(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchRetrieveRequest) (*schemas.BifrostBatchRetrieveResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchRetrieveRequest, provider.GetProviderKey()) -} - -// BatchCancel is not supported by SGL provider. -func (provider *SGLProvider) BatchCancel(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchCancelRequest) (*schemas.BifrostBatchCancelResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchCancelRequest, provider.GetProviderKey()) -} - -// BatchResults is not supported by SGL provider. -func (provider *SGLProvider) BatchResults(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchResultsRequest) (*schemas.BifrostBatchResultsResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchResultsRequest, provider.GetProviderKey()) -} - diff --git a/core/providers/sgl/files.go b/core/providers/sgl/files.go deleted file mode 100644 index b90e1943d..000000000 --- a/core/providers/sgl/files.go +++ /dev/null @@ -1,34 +0,0 @@ -package sgl - -import ( - "context" - - providerUtils "github.com/maximhq/bifrost/core/providers/utils" - "github.com/maximhq/bifrost/core/schemas" -) - -// FileUpload is not supported by SGL provider. -func (provider *SGLProvider) FileUpload(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileUploadRequest) (*schemas.BifrostFileUploadResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.FileUploadRequest, provider.GetProviderKey()) -} - -// FileList is not supported by SGL provider. -func (provider *SGLProvider) FileList(_ context.Context, _ []schemas.Key, _ *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.FileListRequest, provider.GetProviderKey()) -} - -// FileRetrieve is not supported by SGL provider. -func (provider *SGLProvider) FileRetrieve(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileRetrieveRequest) (*schemas.BifrostFileRetrieveResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.FileRetrieveRequest, provider.GetProviderKey()) -} - -// FileDelete is not supported by SGL provider. -func (provider *SGLProvider) FileDelete(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileDeleteRequest) (*schemas.BifrostFileDeleteResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.FileDeleteRequest, provider.GetProviderKey()) -} - -// FileContent is not supported by SGL provider. -func (provider *SGLProvider) FileContent(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileContentRequest) (*schemas.BifrostFileContentResponse, *schemas.BifrostError) { - return nil, providerUtils.NewUnsupportedOperationError(schemas.FileContentRequest, provider.GetProviderKey()) -} - diff --git a/core/providers/sgl/sgl.go b/core/providers/sgl/sgl.go index af33fc720..3bce9d423 100644 --- a/core/providers/sgl/sgl.go +++ b/core/providers/sgl/sgl.go @@ -219,3 +219,53 @@ func (provider *SGLProvider) Transcription(ctx context.Context, key schemas.Key, func (provider *SGLProvider) TranscriptionStream(ctx context.Context, postHookRunner schemas.PostHookRunner, key schemas.Key, request *schemas.BifrostTranscriptionRequest) (chan *schemas.BifrostStream, *schemas.BifrostError) { return nil, providerUtils.NewUnsupportedOperationError(schemas.TranscriptionStreamRequest, provider.GetProviderKey()) } + +// FileUpload is not supported by SGL provider. +func (provider *SGLProvider) FileUpload(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileUploadRequest) (*schemas.BifrostFileUploadResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.FileUploadRequest, provider.GetProviderKey()) +} + +// FileList is not supported by SGL provider. +func (provider *SGLProvider) FileList(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.FileListRequest, provider.GetProviderKey()) +} + +// FileRetrieve is not supported by SGL provider. +func (provider *SGLProvider) FileRetrieve(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileRetrieveRequest) (*schemas.BifrostFileRetrieveResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.FileRetrieveRequest, provider.GetProviderKey()) +} + +// FileDelete is not supported by SGL provider. +func (provider *SGLProvider) FileDelete(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileDeleteRequest) (*schemas.BifrostFileDeleteResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.FileDeleteRequest, provider.GetProviderKey()) +} + +// FileContent is not supported by SGL provider. +func (provider *SGLProvider) FileContent(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileContentRequest) (*schemas.BifrostFileContentResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.FileContentRequest, provider.GetProviderKey()) +} + +// BatchCreate is not supported by SGL provider. +func (provider *SGLProvider) BatchCreate(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchCreateRequest) (*schemas.BifrostBatchCreateResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchCreateRequest, provider.GetProviderKey()) +} + +// BatchList is not supported by SGL provider. +func (provider *SGLProvider) BatchList(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchListRequest, provider.GetProviderKey()) +} + +// BatchRetrieve is not supported by SGL provider. +func (provider *SGLProvider) BatchRetrieve(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchRetrieveRequest) (*schemas.BifrostBatchRetrieveResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchRetrieveRequest, provider.GetProviderKey()) +} + +// BatchCancel is not supported by SGL provider. +func (provider *SGLProvider) BatchCancel(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchCancelRequest) (*schemas.BifrostBatchCancelResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchCancelRequest, provider.GetProviderKey()) +} + +// BatchResults is not supported by SGL provider. +func (provider *SGLProvider) BatchResults(_ context.Context, _ schemas.Key, _ *schemas.BifrostBatchResultsRequest) (*schemas.BifrostBatchResultsResponse, *schemas.BifrostError) { + return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchResultsRequest, provider.GetProviderKey()) +} diff --git a/core/providers/utils/utils.go b/core/providers/utils/utils.go index 7dff53129..c1a67647e 100644 --- a/core/providers/utils/utils.go +++ b/core/providers/utils/utils.go @@ -457,6 +457,7 @@ func NewUnsupportedOperationError(requestType schemas.RequestType, providerName IsBifrostError: false, Error: &schemas.ErrorField{ Message: fmt.Sprintf("%s is not supported by %s provider", requestType, providerName), + Code: schemas.Ptr("unsupported_operation"), }, ExtraFields: schemas.BifrostErrorExtraFields{ Provider: providerName, diff --git a/core/providers/vertex/vertex.go b/core/providers/vertex/vertex.go index c95018c63..7fcd1b10d 100644 --- a/core/providers/vertex/vertex.go +++ b/core/providers/vertex/vertex.go @@ -1465,7 +1465,7 @@ func (provider *VertexProvider) BatchCreate(ctx context.Context, key schemas.Key } // BatchList is not supported by Vertex AI provider. -func (provider *VertexProvider) BatchList(ctx context.Context, keys []schemas.Key, request *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { +func (provider *VertexProvider) BatchList(ctx context.Context, keys schemas.Key, request *schemas.BifrostBatchListRequest) (*schemas.BifrostBatchListResponse, *schemas.BifrostError) { return nil, providerUtils.NewUnsupportedOperationError(schemas.BatchListRequest, provider.GetProviderKey()) } @@ -1491,7 +1491,7 @@ func (provider *VertexProvider) FileUpload(_ context.Context, _ schemas.Key, _ * } // FileList is not yet implemented for Vertex AI provider. -func (provider *VertexProvider) FileList(_ context.Context, _ []schemas.Key, _ *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { +func (provider *VertexProvider) FileList(_ context.Context, _ schemas.Key, _ *schemas.BifrostFileListRequest) (*schemas.BifrostFileListResponse, *schemas.BifrostError) { return nil, providerUtils.NewUnsupportedOperationError(schemas.FileListRequest, provider.GetProviderKey()) } diff --git a/core/schemas/batch.go b/core/schemas/batch.go index 23ceac505..f8c308fde 100644 --- a/core/schemas/batch.go +++ b/core/schemas/batch.go @@ -65,6 +65,7 @@ type BatchError struct { type BifrostBatchCreateRequest struct { Provider ModelProvider `json:"provider"` Model string `json:"model,omitempty"` // Model hint for routing (optional for file-based) + RawRequestBody []byte `json:"-"` // Raw request body (not serialized) // OpenAI-style: file-based batching InputFileID string `json:"input_file_id,omitempty"` // ID of uploaded JSONL file @@ -81,6 +82,11 @@ type BifrostBatchCreateRequest struct { ExtraParams map[string]interface{} `json:"-"` } +// GetRawRequestBody returns the raw request body. +func (request *BifrostBatchCreateRequest) GetRawRequestBody() []byte { + return request.RawRequestBody +} + // BifrostBatchCreateResponse represents the response from creating a batch job. type BifrostBatchCreateResponse struct { ID string `json:"id"` @@ -111,6 +117,7 @@ type BifrostBatchCreateResponse struct { // BifrostBatchListRequest represents a request to list batch jobs. type BifrostBatchListRequest struct { Provider ModelProvider `json:"provider"` + Model string `json:"model"` // Pagination Limit int `json:"limit,omitempty"` // Max results to return @@ -142,12 +149,20 @@ type BifrostBatchListResponse struct { // BifrostBatchRetrieveRequest represents a request to retrieve a batch job. type BifrostBatchRetrieveRequest struct { Provider ModelProvider `json:"provider"` + Model string `json:"model"` BatchID string `json:"batch_id"` // ID of the batch to retrieve + RawRequestBody []byte `json:"-"` // Raw request body (not serialized) + // Extra parameters for provider-specific features ExtraParams map[string]interface{} `json:"-"` } +// GetRawRequestBody returns the raw request body. +func (request *BifrostBatchRetrieveRequest) GetRawRequestBody() []byte { + return request.RawRequestBody +} + // BifrostBatchRetrieveResponse represents the response from retrieving a batch job. type BifrostBatchRetrieveResponse struct { ID string `json:"id"` @@ -189,12 +204,20 @@ type BifrostBatchRetrieveResponse struct { // BifrostBatchCancelRequest represents a request to cancel a batch job. type BifrostBatchCancelRequest struct { Provider ModelProvider `json:"provider"` + Model string `json:"model"` BatchID string `json:"batch_id"` // ID of the batch to cancel + RawRequestBody []byte `json:"-"` // Raw request body (not serialized) + // Extra parameters for provider-specific features ExtraParams map[string]interface{} `json:"-"` } +// GetRawRequestBody returns the raw request body. +func (request *BifrostBatchCancelRequest) GetRawRequestBody() []byte { + return request.RawRequestBody +} + // BifrostBatchCancelResponse represents the response from cancelling a batch job. type BifrostBatchCancelResponse struct { ID string `json:"id"` @@ -210,8 +233,11 @@ type BifrostBatchCancelResponse struct { // BifrostBatchResultsRequest represents a request to retrieve batch results. type BifrostBatchResultsRequest struct { Provider ModelProvider `json:"provider"` + Model string `json:"model"` BatchID string `json:"batch_id"` // ID of the batch to get results for + RawRequestBody []byte `json:"-"` // Raw request body (not serialized) + // For OpenAI, results are retrieved via output_file_id (file download) // For Anthropic, results are streamed from a dedicated endpoint @@ -219,6 +245,11 @@ type BifrostBatchResultsRequest struct { ExtraParams map[string]interface{} `json:"-"` } +// GetRawRequestBody returns the raw request body. +func (request *BifrostBatchResultsRequest) GetRawRequestBody() []byte { + return request.RawRequestBody +} + // BatchResultItem represents a single result from a batch request. type BatchResultItem struct { CustomID string `json:"custom_id"` diff --git a/core/schemas/bifrost.go b/core/schemas/bifrost.go index 39e7cdb00..712da41e9 100644 --- a/core/schemas/bifrost.go +++ b/core/schemas/bifrost.go @@ -167,6 +167,16 @@ type BifrostRequest struct { EmbeddingRequest *BifrostEmbeddingRequest SpeechRequest *BifrostSpeechRequest TranscriptionRequest *BifrostTranscriptionRequest + FileUploadRequest *BifrostFileUploadRequest + FileListRequest *BifrostFileListRequest + FileRetrieveRequest *BifrostFileRetrieveRequest + FileDeleteRequest *BifrostFileDeleteRequest + FileContentRequest *BifrostFileContentRequest + BatchCreateRequest *BifrostBatchCreateRequest + BatchListRequest *BifrostBatchListRequest + BatchRetrieveRequest *BifrostBatchRetrieveRequest + BatchCancelRequest *BifrostBatchCancelRequest + BatchResultsRequest *BifrostBatchResultsRequest } // GetRequestFields returns the provider, model, and fallbacks from the request. @@ -183,7 +193,27 @@ func (br *BifrostRequest) GetRequestFields() (provider ModelProvider, model stri case br.SpeechRequest != nil: return br.SpeechRequest.Provider, br.SpeechRequest.Model, br.SpeechRequest.Fallbacks case br.TranscriptionRequest != nil: - return br.TranscriptionRequest.Provider, br.TranscriptionRequest.Model, br.TranscriptionRequest.Fallbacks + return br.TranscriptionRequest.Provider, br.TranscriptionRequest.Model, br.TranscriptionRequest.Fallbacks + case br.FileUploadRequest != nil: + return br.FileUploadRequest.Provider, br.FileUploadRequest.Model, nil + case br.FileListRequest != nil: + return br.FileListRequest.Provider, br.FileListRequest.Model, nil + case br.FileRetrieveRequest != nil: + return br.FileRetrieveRequest.Provider, br.FileRetrieveRequest.Model, nil + case br.FileDeleteRequest != nil: + return br.FileDeleteRequest.Provider, br.FileDeleteRequest.Model, nil + case br.FileContentRequest != nil: + return br.FileContentRequest.Provider, br.FileContentRequest.Model, nil + case br.BatchCreateRequest != nil: + return br.BatchCreateRequest.Provider, br.BatchCreateRequest.Model, nil + case br.BatchListRequest != nil: + return br.BatchListRequest.Provider, br.BatchListRequest.Model, nil + case br.BatchRetrieveRequest != nil: + return br.BatchRetrieveRequest.Provider, br.BatchRetrieveRequest.Model, nil + case br.BatchCancelRequest != nil: + return br.BatchCancelRequest.Provider, br.BatchCancelRequest.Model, nil + case br.BatchResultsRequest != nil: + return br.BatchResultsRequest.Provider, br.BatchResultsRequest.Model, nil } return "", "", nil @@ -270,6 +300,16 @@ type BifrostResponse struct { SpeechStreamResponse *BifrostSpeechStreamResponse TranscriptionResponse *BifrostTranscriptionResponse TranscriptionStreamResponse *BifrostTranscriptionStreamResponse + FileUploadResponse *BifrostFileUploadResponse + FileListResponse *BifrostFileListResponse + FileRetrieveResponse *BifrostFileRetrieveResponse + FileDeleteResponse *BifrostFileDeleteResponse + FileContentResponse *BifrostFileContentResponse + BatchCreateResponse *BifrostBatchCreateResponse + BatchListResponse *BifrostBatchListResponse + BatchRetrieveResponse *BifrostBatchRetrieveResponse + BatchCancelResponse *BifrostBatchCancelResponse + BatchResultsResponse *BifrostBatchResultsResponse } func (r *BifrostResponse) GetExtraFields() *BifrostResponseExtraFields { @@ -292,6 +332,26 @@ func (r *BifrostResponse) GetExtraFields() *BifrostResponseExtraFields { return &r.TranscriptionResponse.ExtraFields case r.TranscriptionStreamResponse != nil: return &r.TranscriptionStreamResponse.ExtraFields + case r.FileUploadResponse != nil: + return &r.FileUploadResponse.ExtraFields + case r.FileListResponse != nil: + return &r.FileListResponse.ExtraFields + case r.FileRetrieveResponse != nil: + return &r.FileRetrieveResponse.ExtraFields + case r.FileDeleteResponse != nil: + return &r.FileDeleteResponse.ExtraFields + case r.FileContentResponse != nil: + return &r.FileContentResponse.ExtraFields + case r.BatchCreateResponse != nil: + return &r.BatchCreateResponse.ExtraFields + case r.BatchListResponse != nil: + return &r.BatchListResponse.ExtraFields + case r.BatchRetrieveResponse != nil: + return &r.BatchRetrieveResponse.ExtraFields + case r.BatchCancelResponse != nil: + return &r.BatchCancelResponse.ExtraFields + case r.BatchResultsResponse != nil: + return &r.BatchResultsResponse.ExtraFields } return &BifrostResponseExtraFields{} diff --git a/core/schemas/chatcompletions.go b/core/schemas/chatcompletions.go index 15a774896..bfb70d5f7 100644 --- a/core/schemas/chatcompletions.go +++ b/core/schemas/chatcompletions.go @@ -18,6 +18,7 @@ type BifrostChatRequest struct { RawRequestBody []byte `json:"-"` // set bifrost-use-raw-request-body to true in ctx to use the raw request body. Bifrost will directly send this to the downstream provider. } +// GetRawRequestBody returns the raw request body func (r *BifrostChatRequest) GetRawRequestBody() []byte { return r.RawRequestBody } diff --git a/core/schemas/files.go b/core/schemas/files.go index f776db457..d09f11940 100644 --- a/core/schemas/files.go +++ b/core/schemas/files.go @@ -52,6 +52,7 @@ type FileObject struct { // BifrostFileUploadRequest represents a request to upload a file. type BifrostFileUploadRequest struct { Provider ModelProvider `json:"provider"` + Model string `json:"model"` // File content File []byte `json:"-"` // Raw file content (not serialized) @@ -107,6 +108,9 @@ type BifrostFileUploadResponse struct { // BifrostFileListRequest represents a request to list files. type BifrostFileListRequest struct { Provider ModelProvider `json:"provider"` + Model string `json:"model"` + + RawRequestBody []byte `json:"-"` // Raw request body (not serialized) // Filters Purpose FilePurpose `json:"purpose,omitempty"` // Filter by purpose @@ -123,6 +127,11 @@ type BifrostFileListRequest struct { ExtraParams map[string]interface{} `json:"-"` } +// GetRawRequestBody returns the raw request body. +func (request *BifrostFileListRequest) GetRawRequestBody() []byte { + return request.RawRequestBody +} + // BifrostFileListResponse represents the response from listing files. type BifrostFileListResponse struct { Object string `json:"object,omitempty"` // "list" @@ -136,7 +145,11 @@ type BifrostFileListResponse struct { // BifrostFileRetrieveRequest represents a request to retrieve file metadata. type BifrostFileRetrieveRequest struct { Provider ModelProvider `json:"provider"` - FileID string `json:"file_id"` // ID of the file to retrieve + Model string `json:"model"` + + RawRequestBody []byte `json:"-"` // Raw request body (not serialized) + + FileID string `json:"file_id"` // ID of the file to retrieve // Storage configuration (for S3/GCS backends) StorageConfig *FileStorageConfig `json:"storage_config,omitempty"` @@ -145,6 +158,11 @@ type BifrostFileRetrieveRequest struct { ExtraParams map[string]interface{} `json:"-"` } +// GetRawRequestBody returns the raw request body. +func (request *BifrostFileRetrieveRequest) GetRawRequestBody() []byte { + return request.RawRequestBody +} + // BifrostFileRetrieveResponse represents the response from retrieving file metadata. type BifrostFileRetrieveResponse struct { ID string `json:"id"` @@ -167,8 +185,11 @@ type BifrostFileRetrieveResponse struct { // BifrostFileDeleteRequest represents a request to delete a file. type BifrostFileDeleteRequest struct { Provider ModelProvider `json:"provider"` + Model string `json:"model"` FileID string `json:"file_id"` // ID of the file to delete + RawRequestBody []byte `json:"-"` // Raw request body (not serialized) + // Storage configuration (for S3/GCS backends) StorageConfig *FileStorageConfig `json:"storage_config,omitempty"` @@ -176,6 +197,11 @@ type BifrostFileDeleteRequest struct { ExtraParams map[string]interface{} `json:"-"` } +// GetRawRequestBody returns the raw request body. +func (request *BifrostFileDeleteRequest) GetRawRequestBody() []byte { + return request.RawRequestBody +} + // BifrostFileDeleteResponse represents the response from deleting a file. type BifrostFileDeleteResponse struct { ID string `json:"id"` @@ -188,8 +214,11 @@ type BifrostFileDeleteResponse struct { // BifrostFileContentRequest represents a request to download file content. type BifrostFileContentRequest struct { Provider ModelProvider `json:"provider"` + Model string `json:"model"` FileID string `json:"file_id"` // ID of the file to download + RawRequestBody []byte `json:"-"` // Raw request body (not serialized) + // Storage configuration (for S3/GCS backends) StorageConfig *FileStorageConfig `json:"storage_config,omitempty"` @@ -197,6 +226,11 @@ type BifrostFileContentRequest struct { ExtraParams map[string]interface{} `json:"-"` } +// GetRawRequestBody returns the raw request body. +func (request *BifrostFileContentRequest) GetRawRequestBody() []byte { + return request.RawRequestBody +} + // BifrostFileContentResponse represents the response from downloading file content. type BifrostFileContentResponse struct { FileID string `json:"file_id"` diff --git a/core/schemas/provider.go b/core/schemas/provider.go index f4a6a01b1..86074848b 100644 --- a/core/schemas/provider.go +++ b/core/schemas/provider.go @@ -339,7 +339,7 @@ type Provider interface { // BatchCreate creates a new batch job for asynchronous processing BatchCreate(ctx context.Context, key Key, request *BifrostBatchCreateRequest) (*BifrostBatchCreateResponse, *BifrostError) // BatchList lists batch jobs - BatchList(ctx context.Context, keys []Key, request *BifrostBatchListRequest) (*BifrostBatchListResponse, *BifrostError) + BatchList(ctx context.Context, key Key, request *BifrostBatchListRequest) (*BifrostBatchListResponse, *BifrostError) // BatchRetrieve retrieves a specific batch job BatchRetrieve(ctx context.Context, key Key, request *BifrostBatchRetrieveRequest) (*BifrostBatchRetrieveResponse, *BifrostError) // BatchCancel cancels a batch job @@ -349,7 +349,7 @@ type Provider interface { // FileUpload uploads a file to the provider FileUpload(ctx context.Context, key Key, request *BifrostFileUploadRequest) (*BifrostFileUploadResponse, *BifrostError) // FileList lists files from the provider - FileList(ctx context.Context, keys []Key, request *BifrostFileListRequest) (*BifrostFileListResponse, *BifrostError) + FileList(ctx context.Context, key Key, request *BifrostFileListRequest) (*BifrostFileListResponse, *BifrostError) // FileRetrieve retrieves file metadata from the provider FileRetrieve(ctx context.Context, key Key, request *BifrostFileRetrieveRequest) (*BifrostFileRetrieveResponse, *BifrostError) // FileDelete deletes a file from the provider diff --git a/core/utils.go b/core/utils.go index 340ca1b9b..27f511b1a 100644 --- a/core/utils.go +++ b/core/utils.go @@ -47,6 +47,11 @@ var rateLimitPatterns = []string{ "concurrent requests limit", } +// IsModelRequired returns true if the request type requires a model +func IsModelRequired(reqType schemas.RequestType) bool { + return reqType == schemas.TextCompletionRequest || reqType == schemas.TextCompletionStreamRequest || reqType == schemas.ChatCompletionRequest || reqType == schemas.ChatCompletionStreamRequest || reqType == schemas.ResponsesRequest || reqType == schemas.ResponsesStreamRequest || reqType == schemas.SpeechRequest || reqType == schemas.SpeechStreamRequest || reqType == schemas.TranscriptionRequest || reqType == schemas.TranscriptionStreamRequest || reqType == schemas.EmbeddingRequest +} + // Ptr returns a pointer to the given value. func Ptr[T any](v T) *T { return &v @@ -92,7 +97,7 @@ func validateRequest(req *schemas.BifrostRequest) *schemas.BifrostError { if provider == "" { return newBifrostErrorFromMsg("provider is required") } - if model == "" { + if IsModelRequired(req.RequestType) && model == "" { return newBifrostErrorFromMsg("model is required") } diff --git a/core/version b/core/version index f480e7993..adf1ebc44 100644 --- a/core/version +++ b/core/version @@ -1 +1 @@ -1.2.37 \ No newline at end of file +1.2.38 \ No newline at end of file diff --git a/docs/docs.json b/docs/docs.json index cdb604586..9d3cc64e9 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -75,10 +75,37 @@ "icon": "plug", "pages": [ "integrations/what-is-an-integration", - "integrations/openai-sdk", - "integrations/anthropic-sdk", - "integrations/bedrock-sdk", - "integrations/genai-sdk", + { + "group": "OpenAI SDK", + "icon": "o", + "pages": [ + "integrations/openai-sdk/overview", + "integrations/openai-sdk/files-and-batch" + ] + }, + { + "group": "Anthropic SDK", + "icon": "a", + "pages": [ + "integrations/anthropic-sdk/overview", + "integrations/anthropic-sdk/files-and-batch" + ] + }, + { + "group": "Bedrock SDK", + "icon": "b", + "pages": [ + "integrations/bedrock-sdk/overview", + "integrations/bedrock-sdk/files-and-batch" + ] + }, + { + "group": "GenAI SDK", + "icon": "g", + "pages": [ + "integrations/genai-sdk/overview" + ] + }, "integrations/litellm-sdk", "integrations/langchain-sdk", "integrations/pydanticai-sdk" @@ -344,4 +371,4 @@ "linkedin": "https://linkedin.com/company/maxim-ai" } } -} +} \ No newline at end of file diff --git a/docs/integrations/anthropic-sdk/files-and-batch.mdx b/docs/integrations/anthropic-sdk/files-and-batch.mdx new file mode 100644 index 000000000..032e85e14 --- /dev/null +++ b/docs/integrations/anthropic-sdk/files-and-batch.mdx @@ -0,0 +1,618 @@ +--- +title: "Files and Batch API" +tag: "Beta" +description: "Upload files and create batch jobs for asynchronous processing using the Anthropic SDK through Bifrost across multiple providers." +icon: "folder-open" +--- + +## Overview + +Bifrost supports the Anthropic Files API and Batch API (via the `beta` namespace) with **cross-provider routing**. This means you can use the Anthropic SDK to manage files and batch jobs across multiple providers including Anthropic, OpenAI, and Gemini. + +The provider is specified using the `x-model-provider` header in `default_headers`. + + +**Bedrock Limitation:** Bedrock batch operations require file-based input with S3 storage, which is not supported via the Anthropic SDK's inline batch API. For Bedrock batch operations, use the [Bedrock SDK](../bedrock-sdk/files-and-batch) directly. + + +--- + +## Client Setup + + +In API Key section, you can either send virtual key or a dummy key to escape client side validation. + + +### Anthropic Provider (Default) + +```python +import anthropic + +client = anthropic.Anthropic( + base_url="http://localhost:8080/anthropic", + api_key="virtual-key-or-dummy-key" +) +``` + +### Cross-Provider Client + +To route requests to a different provider, set the `x-model-provider` header: + + + + +```python +import anthropic + +client = anthropic.Anthropic( + base_url="http://localhost:8080/anthropic", + api_key="virtual-key-or-dummy-key", + default_headers={"x-model-provider": "openai"} +) +``` + + + + +```python +import anthropic + +client = anthropic.Anthropic( + base_url="http://localhost:8080/anthropic", + api_key="virtual-key-or-dummy-key", + default_headers={"x-model-provider": "bedrock"} +) +``` + + +Bedrock can be used for chat completions via the Anthropic SDK, but **batch operations are not supported**. Bedrock requires file-based batch input with S3 storage. Use the [Bedrock SDK](../bedrock-sdk/files-and-batch) for batch operations. + + + + + +```python +import anthropic + +client = anthropic.Anthropic( + base_url="http://localhost:8080/anthropic", + api_key="virtual-key-or-dummy-key", + default_headers={"x-model-provider": "gemini"} +) +``` + + + + +--- + +## Files API + +The Files API is accessed through the `beta.files` namespace. Note that file support varies by provider. + +### Upload a File + + + + +Upload a text file for use with Anthropic: + +```python +import anthropic + +client = anthropic.Anthropic( + base_url="http://localhost:8080/anthropic", + api_key="virtual-key-or-dummy-key" +) + +# Upload a text file +text_content = b"This is a test file for Files API integration." + +response = client.beta.files.upload( + file=("test_upload.txt", text_content, "text/plain"), +) + +print(f"File ID: {response.id}") +print(f"Filename: {response.filename}") +``` + + + + +Upload a JSONL file for OpenAI batch processing: + +```python +import anthropic + +# Client configured for OpenAI provider +client = anthropic.Anthropic( + base_url="http://localhost:8080/anthropic", + api_key="virtual-key-or-dummy-key", + default_headers={"x-model-provider": "openai"} +) + +# Create JSONL content in OpenAI batch format +jsonl_content = b'''{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "Hello!"}], "max_tokens": 100}} +{"custom_id": "request-2", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "How are you?"}], "max_tokens": 100}}''' + +response = client.beta.files.upload( + file=("batch_input.jsonl", jsonl_content, "application/jsonl"), +) + +print(f"File ID: {response.id}") +``` + + + + +### List Files + + + + +```python +import anthropic + +client = anthropic.Anthropic( + base_url="http://localhost:8080/anthropic", + api_key="virtual-key-or-dummy-key" +) + +# List all files +response = client.beta.files.list() + +for file in response.data: + print(f"File ID: {file.id}") + print(f"Filename: {file.filename}") + print(f"Size: {file.size} bytes") + print("---") +``` + + + + +```python +import anthropic + +# Client configured for OpenAI provider +client = anthropic.Anthropic( + base_url="http://localhost:8080/anthropic", + api_key="virtual-key-or-dummy-key", + default_headers={"x-model-provider": "openai"} +) + +# List all files from OpenAI +response = client.beta.files.list() + +for file in response.data: + print(f"File ID: {file.id}, Name: {file.filename}") +``` + + + + +### Delete a File + +```python +import anthropic + +client = anthropic.Anthropic( + base_url="http://localhost:8080/anthropic", + api_key="virtual-key-or-dummy-key", + default_headers={"x-model-provider": "openai"} # or omit for anthropic +) + +# Delete a file +file_id = "file-abc123" +response = client.beta.files.delete(file_id) + +print(f"Deleted file: {file_id}") +``` + +### Download File Content + +Note: Anthropic only allows downloading files created by certain tools (like code execution). OpenAI allows downloading batch output files. + +```python +import anthropic + +client = anthropic.Anthropic( + base_url="http://localhost:8080/anthropic", + api_key="virtual-key-or-dummy-key", + default_headers={"x-model-provider": "openai"} +) + +# Download file content +file_id = "file-abc123" +response = client.beta.files.download(file_id) + +content = response.text() +print(f"File content:\n{content}") +``` + +--- + +## Batch API + +The Anthropic Batch API is accessed through `beta.messages.batches`. Anthropic's batch API uses **inline requests** rather than file uploads. + +### Create a Batch with Inline Requests + + + + +```python +import anthropic + +client = anthropic.Anthropic( + base_url="http://localhost:8080/anthropic", + api_key="virtual-key-or-dummy-key" +) + +# Create batch with inline requests +batch_requests = [ + { + "custom_id": "request-1", + "params": { + "model": "claude-3-sonnet-20240229", + "max_tokens": 100, + "messages": [ + {"role": "user", "content": "What is 2+2?"} + ] + } + }, + { + "custom_id": "request-2", + "params": { + "model": "claude-3-sonnet-20240229", + "max_tokens": 100, + "messages": [ + {"role": "user", "content": "What is the capital of France?"} + ] + } + } +] + +batch = client.beta.messages.batches.create(requests=batch_requests) + +print(f"Batch ID: {batch.id}") +print(f"Status: {batch.processing_status}") +``` + + + + +When routing to OpenAI, use OpenAI-compatible models: + +```python +import anthropic + +# Client configured for OpenAI provider +client = anthropic.Anthropic( + base_url="http://localhost:8080/anthropic", + api_key="virtual-key-or-dummy-key", + default_headers={"x-model-provider": "openai"} +) + +# Create batch with inline requests (using OpenAI models) +batch_requests = [ + { + "custom_id": "request-1", + "params": { + "model": "gpt-4o-mini", + "max_tokens": 100, + "messages": [ + {"role": "user", "content": "What is 2+2?"} + ] + } + }, + { + "custom_id": "request-2", + "params": { + "model": "gpt-4o-mini", + "max_tokens": 100, + "messages": [ + {"role": "user", "content": "What is the capital of France?"} + ] + } + } +] + +batch = client.beta.messages.batches.create(requests=batch_requests) + +print(f"Batch ID: {batch.id}") +print(f"Status: {batch.processing_status}") +``` + + + + +When routing to Gemini: + +```python +import anthropic + +# Client configured for Gemini provider +client = anthropic.Anthropic( + base_url="http://localhost:8080/anthropic", + api_key="virtual-key-or-dummy-key", + default_headers={"x-model-provider": "gemini"} +) + +# Create batch with inline requests (using Gemini models) +batch_requests = [ + { + "custom_id": "request-1", + "params": { + "model": "gemini-1.5-flash", + "max_tokens": 100, + "messages": [ + {"role": "user", "content": "What is 2+2?"} + ] + } + }, + { + "custom_id": "request-2", + "params": { + "model": "gemini-1.5-flash", + "max_tokens": 100, + "messages": [ + {"role": "user", "content": "What is the capital of France?"} + ] + } + } +] + +batch = client.beta.messages.batches.create(requests=batch_requests) + +print(f"Batch ID: {batch.id}") +print(f"Status: {batch.processing_status}") +``` + + + + + +**Bedrock Note:** Bedrock requires file-based batch creation with S3 storage. When routing to Bedrock from the Anthropic SDK, you'll need to use the Bedrock SDK directly for batch operations. See the [Bedrock SDK documentation](../bedrock-sdk/files-and-batch) for details. + + +### List Batches + +```python +import anthropic + +client = anthropic.Anthropic( + base_url="http://localhost:8080/anthropic", + api_key="virtual-key-or-dummy-key", + default_headers={"x-model-provider": "anthropic"} # or "openai", "gemini" +) + +# List batches +response = client.beta.messages.batches.list(limit=10) + +for batch in response.data: + print(f"Batch ID: {batch.id}") + print(f"Status: {batch.processing_status}") + if batch.request_counts: + print(f"Processing: {batch.request_counts.processing}") + print(f"Succeeded: {batch.request_counts.succeeded}") + print(f"Errored: {batch.request_counts.errored}") + print("---") +``` + +### Retrieve Batch Status + +```python +import anthropic + +client = anthropic.Anthropic( + base_url="http://localhost:8080/anthropic", + api_key="virtual-key-or-dummy-key", + default_headers={"x-model-provider": "anthropic"} # or "openai", "gemini" +) + +# Retrieve batch status +batch_id = "batch-abc123" +batch = client.beta.messages.batches.retrieve(batch_id) + +print(f"Batch ID: {batch.id}") +print(f"Status: {batch.processing_status}") + +if batch.request_counts: + print(f"Processing: {batch.request_counts.processing}") + print(f"Succeeded: {batch.request_counts.succeeded}") + print(f"Errored: {batch.request_counts.errored}") +``` + +### Cancel a Batch + +```python +import anthropic + +client = anthropic.Anthropic( + base_url="http://localhost:8080/anthropic", + api_key="virtual-key-or-dummy-key", + default_headers={"x-model-provider": "anthropic"} # or "openai", "gemini" +) + +# Cancel batch +batch_id = "batch-abc123" +batch = client.beta.messages.batches.cancel(batch_id) + +print(f"Batch ID: {batch.id}") +print(f"Status: {batch.processing_status}") # "canceling" or "ended" +``` + +### Get Batch Results + +```python +import anthropic + +client = anthropic.Anthropic( + base_url="http://localhost:8080/anthropic", + api_key="virtual-key-or-dummy-key" +) + +# Get batch results (only available after batch is completed) +batch_id = "batch-abc123" +results = client.beta.messages.batches.results(batch_id) + +# Iterate over results +for result in results: + print(f"Custom ID: {result.custom_id}") + if result.result.type == "succeeded": + message = result.result.message + print(f"Response: {message.content[0].text}") + elif result.result.type == "errored": + print(f"Error: {result.result.error}") + print("---") +``` + +--- + +## End-to-End Workflows + +### Anthropic Batch Workflow + +```python +import time +import anthropic + +client = anthropic.Anthropic( + base_url="http://localhost:8080/anthropic", + api_key="virtual-key-or-dummy-key" +) + +# Step 1: Create batch with inline requests +print("Step 1: Creating batch...") +batch_requests = [ + { + "custom_id": "math-question", + "params": { + "model": "claude-3-sonnet-20240229", + "max_tokens": 100, + "messages": [{"role": "user", "content": "What is 15 * 7?"}] + } + }, + { + "custom_id": "geography-question", + "params": { + "model": "claude-3-sonnet-20240229", + "max_tokens": 100, + "messages": [{"role": "user", "content": "What is the largest ocean?"}] + } + } +] + +batch = client.beta.messages.batches.create(requests=batch_requests) +print(f" Created batch: {batch.id}, status: {batch.processing_status}") + +# Step 2: Poll for completion +print("Step 2: Polling batch status...") +for i in range(20): + batch = client.beta.messages.batches.retrieve(batch.id) + print(f" Poll {i+1}: status = {batch.processing_status}") + + if batch.processing_status == "ended": + print(" Batch completed!") + break + + if batch.request_counts: + print(f" Processing: {batch.request_counts.processing}") + print(f" Succeeded: {batch.request_counts.succeeded}") + + time.sleep(5) + +# Step 3: Verify batch is in list +print("Step 3: Verifying batch in list...") +batch_list = client.beta.messages.batches.list(limit=20) +batch_ids = [b.id for b in batch_list.data] +assert batch.id in batch_ids, f"Batch {batch.id} should be in list" +print(f" Verified batch {batch.id} is in list") + +# Step 4: Get results (if completed) +if batch.processing_status == "ended": + print("Step 4: Getting results...") + try: + results = client.beta.messages.batches.results(batch.id) + for result in results: + print(f" {result.custom_id}: ", end="") + if result.result.type == "succeeded": + print(result.result.message.content[0].text[:50] + "...") + else: + print(f"Error: {result.result.error}") + except Exception as e: + print(f" Results not yet available: {e}") + +print(f"\nSuccess! Batch {batch.id} workflow completed.") +``` + +### Cross-Provider Batch Workflow (OpenAI via Anthropic SDK) + +```python +import time +import anthropic + +# Create client with OpenAI provider header +client = anthropic.Anthropic( + base_url="http://localhost:8080/anthropic", + api_key="virtual-key-or-dummy-key", + default_headers={"x-model-provider": "openai"} +) + +# Step 1: Create batch with OpenAI models +print("Step 1: Creating batch for OpenAI provider...") +batch_requests = [ + { + "custom_id": "openai-request-1", + "params": { + "model": "gpt-4o-mini", + "max_tokens": 100, + "messages": [{"role": "user", "content": "Explain AI in one sentence."}] + } + }, + { + "custom_id": "openai-request-2", + "params": { + "model": "gpt-4o-mini", + "max_tokens": 100, + "messages": [{"role": "user", "content": "What is machine learning?"}] + } + } +] + +batch = client.beta.messages.batches.create(requests=batch_requests) +print(f" Created batch: {batch.id}, status: {batch.processing_status}") + +# Step 2: Poll for completion +print("Step 2: Polling batch status...") +for i in range(10): + batch = client.beta.messages.batches.retrieve(batch.id) + print(f" Poll {i+1}: status = {batch.processing_status}") + + if batch.processing_status in ["ended", "completed"]: + break + + time.sleep(5) + +print(f"\nSuccess! Cross-provider batch {batch.id} completed via Anthropic SDK.") +``` + +--- + +## Provider-Specific Notes + +| Provider | Header Value | File Upload | Batch Type | Models | +|----------|--------------|-------------|------------|--------| +| **Anthropic** | `anthropic` or omit | ✅ Beta API | Inline requests | `claude-3-*` | +| **OpenAI** | `openai` | ✅ Beta API | Inline requests | `gpt-4o-*`, `gpt-4-*` | +| **Gemini** | `gemini` | ✅ Beta API | Inline requests | `gemini-1.5-*` | +| **Bedrock** | `bedrock` | ❌ Use Bedrock SDK | File-based (S3) | `anthropic.claude-*` | + +--- + +## Next Steps + +- **[Overview](./overview)** - Anthropic SDK integration basics +- **[Configuration](../../quickstart/gateway/configuration)** - Bifrost setup and configuration +- **[Core Features](../../features/)** - Governance, semantic caching, and more diff --git a/docs/integrations/anthropic-sdk.mdx b/docs/integrations/anthropic-sdk/overview.mdx similarity index 93% rename from docs/integrations/anthropic-sdk.mdx rename to docs/integrations/anthropic-sdk/overview.mdx index fbcd1a2fc..ae7c6d256 100644 --- a/docs/integrations/anthropic-sdk.mdx +++ b/docs/integrations/anthropic-sdk/overview.mdx @@ -1,7 +1,7 @@ --- -title: "Anthropic SDK" +title: "Overview" description: "Use Bifrost as a drop-in replacement for Anthropic API with full compatibility and enhanced features." -icon: "a" +icon: "book" --- ## Overview @@ -223,7 +223,7 @@ const response = await anthropic.messages.create({ Pass API keys directly in requests to bypass Bifrost's load balancing. You can pass any provider's API key (OpenAI, Anthropic, Mistral, etc.) since Bifrost only looks for `Authorization` or `x-api-key` headers. This requires the **Allow Direct API keys** option to be enabled in Bifrost configuration. -> **Learn more:** See [Quickstart Configuration](../quickstart/README) for enabling direct API key usage. +> **Learn more:** See [Quickstart Configuration](../../quickstart/README) for enabling direct API key usage. @@ -324,13 +324,15 @@ const openaiResponseWithHeader = await anthropic.messages.create({ ## Supported Features -The Anthropic integration supports all features that are available in both the Anthropic SDK and Bifrost core functionality. If the Anthropic SDK supports a feature and Bifrost supports it, the integration will work seamlessly. 😄 +The Anthropic integration supports all features that are available in both the Anthropic SDK and Bifrost core functionality. If the Anthropic SDK supports a feature and Bifrost supports it, the integration will work seamlessly. --- ## Next Steps -- **[OpenAI SDK](./openai-sdk)** - GPT integration patterns -- **[Google GenAI SDK](./genai-sdk)** - Gemini integration patterns -- **[Configuration](../quickstart/README)** - Bifrost setup and configuration -- **[Core Features](../features/)** - Advanced Bifrost capabilities +- **[Files and Batch API](./files-and-batch)** - File uploads and batch processing +- **[OpenAI SDK](../openai-sdk/overview)** - GPT integration patterns +- **[Google GenAI SDK](../genai-sdk)** - Gemini integration patterns +- **[Configuration](../../quickstart/README)** - Bifrost setup and configuration +- **[Core Features](../../features/)** - Advanced Bifrost capabilities + diff --git a/docs/integrations/bedrock-sdk/files-and-batch.mdx b/docs/integrations/bedrock-sdk/files-and-batch.mdx new file mode 100644 index 000000000..e3e68f1a6 --- /dev/null +++ b/docs/integrations/bedrock-sdk/files-and-batch.mdx @@ -0,0 +1,1162 @@ +--- +title: "Files and Batch API" +description: "Manage S3-based files and batch inference jobs using the AWS Bedrock SDK (boto3) through Bifrost across multiple providers." +tag: "Beta" +icon: "folder-open" +--- + +## Overview + +Bifrost supports AWS Bedrock's Files and Batch APIs with **cross-provider routing**. This means you can use boto3 to manage files and batch jobs across multiple providers including Bedrock, OpenAI, and Gemini. + +For Bedrock SDK: +- **Files** are managed through an S3-compatible API +- **Batches** are managed through the Bedrock service API +- **Provider routing** is done via the `x-model-provider` header + + +**Anthropic Limitation:** Anthropic does not support file upload via the S3-compatible API. For Anthropic batch operations, use the [Anthropic SDK](../anthropic-sdk/files-and-batch) with inline requests instead. + + +--- + +## Client Setup + +### Default Bedrock Client + +```python +import boto3 + +# S3 client for file operations +s3_client = boto3.client( + "s3", + region_name="us-west-2", + endpoint_url="http://localhost:8080/bedrock/files", +) + +# Bedrock client for batch operations +bedrock_client = boto3.client( + "bedrock", + region_name="us-west-2", + endpoint_url="http://localhost:8080/bedrock", +) +``` + +### Cross-Provider Client Setup + +To route requests to different providers, add the `x-model-provider` header using boto3 events: + + + + +```python +import boto3 + +def add_bedrock_header(request, **kwargs): + request.headers["x-model-provider"] = "bedrock" + +# S3 client for Bedrock files +s3_client = boto3.client( + "s3", + region_name="us-west-2", + endpoint_url="http://localhost:8080/bedrock/files", +) +s3_client.meta.events.register("before-send", add_bedrock_header) + +# Bedrock client for batches +bedrock_client = boto3.client( + "bedrock", + region_name="us-west-2", + endpoint_url="http://localhost:8080/bedrock", +) +bedrock_client.meta.events.register("before-send", add_bedrock_header) +``` + + + + +```python +import boto3 + +def add_openai_header(request, **kwargs): + request.headers["x-model-provider"] = "openai" + +# S3 client for OpenAI files +s3_client = boto3.client( + "s3", + region_name="us-west-2", + endpoint_url="http://localhost:8080/bedrock/files", +) +s3_client.meta.events.register("before-send", add_openai_header) + +# Bedrock client for OpenAI batches +bedrock_client = boto3.client( + "bedrock", + region_name="us-west-2", + endpoint_url="http://localhost:8080/bedrock", +) +bedrock_client.meta.events.register("before-send", add_openai_header) +``` + + + + + +Anthropic does not support S3-based file uploads. For Anthropic operations, use the [Anthropic SDK](../anthropic-sdk/files-and-batch) with inline batch requests instead. + + +```python +import boto3 + +def add_anthropic_header(request, **kwargs): + request.headers["x-model-provider"] = "anthropic" + +# Note: File operations are NOT supported for Anthropic +# Use Anthropic SDK with inline requests instead + +# Bedrock client for Anthropic (limited to non-batch operations) +bedrock_client = boto3.client( + "bedrock", + region_name="us-west-2", + endpoint_url="http://localhost:8080/bedrock", +) +bedrock_client.meta.events.register("before-send", add_anthropic_header) +``` + + + + +```python +import boto3 + +def add_gemini_header(request, **kwargs): + request.headers["x-model-provider"] = "gemini" + +# S3 client for Gemini files +s3_client = boto3.client( + "s3", + region_name="us-west-2", + endpoint_url="http://localhost:8080/bedrock/files", +) +s3_client.meta.events.register("before-send", add_gemini_header) + +# Bedrock client for Gemini batches +bedrock_client = boto3.client( + "bedrock", + region_name="us-west-2", + endpoint_url="http://localhost:8080/bedrock", +) +bedrock_client.meta.events.register("before-send", add_gemini_header) +``` + + + + +### Helper Function for Provider-Specific Clients + +```python +import boto3 + +def create_provider_header_handler(provider: str): + """Create a header handler function for a specific provider""" + def add_provider_header(request, **kwargs): + request.headers["x-model-provider"] = provider + return add_provider_header + +def get_provider_s3_client(provider: str): + """Create S3 client with x-model-provider header""" + client = boto3.client( + "s3", + region_name="us-west-2", + endpoint_url="http://localhost:8080/bedrock/files", + ) + client.meta.events.register("before-send", create_provider_header_handler(provider)) + return client + +def get_provider_bedrock_client(provider: str): + """Create Bedrock batch client with x-model-provider header""" + client = boto3.client( + "bedrock", + region_name="us-west-2", + endpoint_url="http://localhost:8080/bedrock", + ) + client.meta.events.register("before-send", create_provider_header_handler(provider)) + return client +``` + +--- + +## Files API (S3-Compatible) + +Files are managed through Bifrost's S3-compatible endpoint. + +### Upload a File + + + + +```python +import boto3 +import json +import time + +def add_bedrock_header(request, **kwargs): + request.headers["x-model-provider"] = "bedrock" + +s3_client = boto3.client( + "s3", + region_name="us-west-2", + endpoint_url="http://localhost:8080/bedrock/files", +) +s3_client.meta.events.register("before-send", add_bedrock_header) + +# Create JSONL content for Bedrock batch format +def create_bedrock_batch_jsonl(model_id: str, num_requests: int = 2) -> str: + lines = [] + for i in range(num_requests): + record = { + "recordId": f"request-{i+1}", + "modelInput": { + "messages": [ + { + "role": "user", + "content": [ + {"text": f"Hello, this is test message {i+1}. Say hi back briefly."} + ], + } + ], + "inferenceConfig": {"maxTokens": 100}, + }, + } + lines.append(json.dumps(record)) + return "\n".join(lines) + +# Create content +jsonl_content = create_bedrock_batch_jsonl("anthropic.claude-3-sonnet-20240229-v1:0") + +# Upload to S3 +s3_bucket = "your-s3-bucket" +s3_key = f"bifrost-batch-input/batch_input_{int(time.time())}.jsonl" + +response = s3_client.put_object( + Bucket=s3_bucket, + Key=s3_key, + Body=jsonl_content.encode(), + ContentType="application/jsonl", +) + +# Extract file ID from ETag header +file_id = response.get("ETag", "").strip('"') +print(f"Uploaded file ID: {file_id}") +print(f"S3 URI: s3://{s3_bucket}/{s3_key}") +``` + + + + + +**No S3 configuration required.** Files are stored in OpenAI's native storage. The bucket/key values are identifiers used by Bifrost for routing. + + +```python +import boto3 +import json +import time + +def add_openai_header(request, **kwargs): + request.headers["x-model-provider"] = "openai" + +s3_client = boto3.client( + "s3", + region_name="us-west-2", + endpoint_url="http://localhost:8080/bedrock/files", +) +s3_client.meta.events.register("before-send", add_openai_header) + +# Create JSONL content for OpenAI batch format +def create_openai_batch_jsonl(model_id: str, num_requests: int = 2) -> str: + lines = [] + for i in range(num_requests): + record = { + "custom_id": f"request-{i+1}", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": model_id, + "messages": [ + {"role": "user", "content": f"Hello, this is test message {i+1}. Say hi back briefly."} + ], + "max_tokens": 100, + }, + } + lines.append(json.dumps(record)) + return "\n".join(lines) + +# Create content +jsonl_content = create_openai_batch_jsonl("gpt-4o-mini") + +# Upload file (bucket/key are routing identifiers, not actual S3 paths) +response = s3_client.put_object( + Bucket="openai-files", + Key=f"batch_input_{int(time.time())}.jsonl", + Body=jsonl_content.encode(), + ContentType="application/jsonl", +) + +# Extract file ID from ETag header +file_id = response.get("ETag", "").strip('"') +print(f"Uploaded file ID: {file_id}") +``` + + + + + +**No S3 configuration required.** Files are stored in Google Cloud Storage. The bucket/key values are identifiers used by Bifrost for routing. + + +```python +import boto3 +import json +import time + +def add_gemini_header(request, **kwargs): + request.headers["x-model-provider"] = "gemini" + +s3_client = boto3.client( + "s3", + region_name="us-west-2", + endpoint_url="http://localhost:8080/bedrock/files", +) +s3_client.meta.events.register("before-send", add_gemini_header) + +# Create JSONL content for Gemini batch format +def create_gemini_batch_jsonl(model_id: str, num_requests: int = 2) -> str: + lines = [] + for i in range(num_requests): + record = { + "request": { + "contents": [ + { + "role": "user", + "parts": [ + {"text": f"Hello, this is test message {i+1}. Say hi back briefly."} + ], + } + ], + "generationConfig": {"maxOutputTokens": 100}, + }, + "metadata": {"key": f"request-{i+1}"}, + } + lines.append(json.dumps(record)) + return "\n".join(lines) + +# Create content +jsonl_content = create_gemini_batch_jsonl("gemini-1.5-flash") + +# Upload file (bucket/key are routing identifiers, not actual S3 paths) +response = s3_client.put_object( + Bucket="gemini-files", + Key=f"batch_input_{int(time.time())}.jsonl", + Body=jsonl_content.encode(), + ContentType="application/jsonl", +) + +file_id = response.get("ETag", "").strip('"') +print(f"Uploaded file ID: {file_id}") +``` + + + + +### List Files + + +For **OpenAI** and **Gemini**, use any bucket name as an identifier—files are stored in the provider's native storage and listed by file ID. + + +```python +import boto3 + +def add_provider_header(request, **kwargs): + request.headers["x-model-provider"] = "bedrock" + +s3_client = boto3.client( + "s3", + region_name="us-west-2", + endpoint_url="http://localhost:8080/bedrock/files", +) +s3_client.meta.events.register("before-send", add_provider_header) + +# List files (S3 bucket required for Bedrock only) +s3_bucket = "your-s3-bucket" +response = s3_client.list_objects_v2( + Bucket=s3_bucket, + Prefix="bifrost-batch-input/" +) + +if "Contents" in response: + for obj in response["Contents"]: + print(f"Key: {obj['Key']}") + print(f"Size: {obj['Size']} bytes") + print(f"Last Modified: {obj['LastModified']}") + print("---") +``` + +### Retrieve File Metadata + +```python +import boto3 + +def add_provider_header(request, **kwargs): + request.headers["x-model-provider"] = "bedrock" + +s3_client = boto3.client( + "s3", + region_name="us-west-2", + endpoint_url="http://localhost:8080/bedrock/files", +) +s3_client.meta.events.register("before-send", add_provider_header) + +# Retrieve file metadata (HEAD request) +# For OpenAI/Gemini: use any bucket name, file_id from upload +s3_bucket = "your-s3-bucket" +s3_key = "bifrost-batch-input/batch_input.jsonl" +file_id = "abc123" # ETag from upload + +response = s3_client.head_object( + Bucket=s3_bucket, + Key=s3_key, + IfMatch=file_id +) + +print(f"Content Length: {response['ContentLength']} bytes") +print(f"Content Type: {response['ContentType']}") +print(f"ETag (File ID): {response['ETag']}") +``` + +### Delete a File + +```python +import boto3 + +def add_provider_header(request, **kwargs): + request.headers["x-model-provider"] = "bedrock" + +s3_client = boto3.client( + "s3", + region_name="us-west-2", + endpoint_url="http://localhost:8080/bedrock/files", +) +s3_client.meta.events.register("before-send", add_provider_header) + +# Delete file +# For OpenAI/Gemini: use any bucket name, file_id from upload +s3_bucket = "your-s3-bucket" +s3_key = "bifrost-batch-input/batch_input.jsonl" +file_id = "abc123" + +s3_client.delete_object( + Bucket=s3_bucket, + Key=s3_key, + IfMatch=file_id +) + +print(f"Deleted file: {s3_key}") +``` + +### Download File Content + + +File content download is **only supported for Bedrock**. For OpenAI and Gemini, use their native SDKs to download file content. + + +```python +import boto3 + +def add_provider_header(request, **kwargs): + request.headers["x-model-provider"] = "bedrock" + +s3_client = boto3.client( + "s3", + region_name="us-west-2", + endpoint_url="http://localhost:8080/bedrock/files", +) +s3_client.meta.events.register("before-send", add_provider_header) + +# Download file content (Bedrock only) +s3_bucket = "your-s3-bucket" +s3_key = "bifrost-batch-input/batch_input.jsonl" +file_id = "abc123" + +response = s3_client.get_object( + Bucket=s3_bucket, + Key=s3_key, + IfMatch=file_id +) + +content = response["Body"].read().decode("utf-8") +print(f"File content:\n{content}") +``` + +--- + +## Batch API + +The Bedrock Batch API uses `create_model_invocation_job` and related methods. + +### Create a Batch Job + + + + +```python +import boto3 +import time + +def add_bedrock_header(request, **kwargs): + request.headers["x-model-provider"] = "bedrock" + +bedrock_client = boto3.client( + "bedrock", + region_name="us-west-2", + endpoint_url="http://localhost:8080/bedrock", +) +bedrock_client.meta.events.register("before-send", add_bedrock_header) + +# Configuration +s3_bucket = "your-s3-bucket" +role_arn = "arn:aws:iam::123456789:role/BedrockBatchRole" +model_id = "anthropic.claude-3-sonnet-20240229-v1:0" + +# Input/output URIs (file should already be uploaded) +input_uri = f"s3://{s3_bucket}/bifrost-batch-input/batch_input.jsonl" +output_uri = f"s3://{s3_bucket}/bifrost-batch-output/" + +# Create batch job +response = bedrock_client.create_model_invocation_job( + jobName=f"bifrost-batch-{int(time.time())}", + modelId=model_id, + roleArn=role_arn, + inputDataConfig={ + "s3InputDataConfig": { + "s3Uri": input_uri, + "s3InputFormat": "JSONL" + } + }, + outputDataConfig={ + "s3OutputDataConfig": { + "s3Uri": output_uri + } + }, + tags=[ + {"key": "endpoint", "value": "/v1/chat/completions"}, + {"key": "source", "value": "bifrost-docs"}, + ], +) + +job_arn = response["jobArn"] +print(f"Created batch job: {job_arn}") +``` + + + + + +**No S3 or IAM configuration required.** Files are stored in OpenAI's native storage. The S3 URIs are routing identifiers used by Bifrost. + + +```python +import boto3 +import time + +def add_openai_header(request, **kwargs): + request.headers["x-model-provider"] = "openai" + +bedrock_client = boto3.client( + "bedrock", + region_name="us-west-2", + endpoint_url="http://localhost:8080/bedrock", +) +bedrock_client.meta.events.register("before-send", add_openai_header) + +# Configuration (no S3 bucket or IAM role needed) +model_id = "gpt-4o-mini" +file_id = "file-abc123" # File ID from upload step + +# Create batch job +response = bedrock_client.create_model_invocation_job( + jobName=f"openai-batch-{int(time.time())}", + modelId=model_id, + roleArn="not-required-for-openai", + inputDataConfig={ + "s3InputDataConfig": { + "s3Uri": f"s3://openai-files/{file_id}", # Routing identifier + "s3InputFormat": "JSONL" + } + }, + outputDataConfig={ + "s3OutputDataConfig": { + "s3Uri": "s3://openai-output/" + } + }, + tags=[ + {"key": "endpoint", "value": "/v1/chat/completions"}, + {"key": "file_id", "value": file_id}, + ], +) + +job_arn = response["jobArn"] +print(f"Created OpenAI batch job: {job_arn}") +``` + + + + + +**No S3 or IAM configuration required.** Files are stored in Google Cloud Storage. The S3 URIs are routing identifiers used by Bifrost. + + +```python +import boto3 +import time + +def add_gemini_header(request, **kwargs): + request.headers["x-model-provider"] = "gemini" + +bedrock_client = boto3.client( + "bedrock", + region_name="us-west-2", + endpoint_url="http://localhost:8080/bedrock", +) +bedrock_client.meta.events.register("before-send", add_gemini_header) + +# Configuration (no S3 bucket or IAM role needed) +model_id = "gemini-1.5-flash" +file_id = "file-xyz789" # File ID from upload step + +# Create batch job +response = bedrock_client.create_model_invocation_job( + jobName=f"gemini-batch-{int(time.time())}", + modelId=model_id, + roleArn="not-required-for-gemini", + inputDataConfig={ + "s3InputDataConfig": { + "s3Uri": f"s3://gemini-files/{file_id}", # Routing identifier + "s3InputFormat": "JSONL" + } + }, + outputDataConfig={ + "s3OutputDataConfig": { + "s3Uri": "s3://gemini-output/" + } + }, +) + +job_arn = response["jobArn"] +print(f"Created Gemini batch job: {job_arn}") +``` + + + + + +**Anthropic Note:** Anthropic prefers inline batch requests rather than file-based batching. When targeting Anthropic from the Bedrock SDK, consider using the [Anthropic SDK](../anthropic-sdk/files-and-batch) directly for better compatibility. + + +### List Batch Jobs + +```python +import boto3 + +def add_provider_header(request, **kwargs): + request.headers["x-model-provider"] = "bedrock" # or "gemini" + +bedrock_client = boto3.client( + "bedrock", + region_name="us-west-2", + endpoint_url="http://localhost:8080/bedrock", +) +bedrock_client.meta.events.register("before-send", add_provider_header) + +# List batch jobs +response = bedrock_client.list_model_invocation_jobs(maxResults=10) + +if "invocationJobSummaries" in response: + for job in response["invocationJobSummaries"]: + print(f"Job ARN: {job['jobArn']}") + print(f"Job Name: {job['jobName']}") + print(f"Status: {job['status']}") + print(f"Model ID: {job.get('modelId', 'N/A')}") + print("---") +``` + +### Retrieve Batch Job Status + +```python +import boto3 + +def add_provider_header(request, **kwargs): + request.headers["x-model-provider"] = "bedrock" + +bedrock_client = boto3.client( + "bedrock", + region_name="us-west-2", + endpoint_url="http://localhost:8080/bedrock", +) +bedrock_client.meta.events.register("before-send", add_provider_header) + +# Get batch job status +job_arn = "arn:aws:bedrock:us-west-2:123456789:model-invocation-job/abc123" + +response = bedrock_client.get_model_invocation_job(jobIdentifier=job_arn) + +print(f"Job ARN: {response['jobArn']}") +print(f"Job Name: {response['jobName']}") +print(f"Status: {response['status']}") +print(f"Model ID: {response['modelId']}") + +if response["status"] == "Completed" and "statistics" in response: + stats = response["statistics"] + print(f"Total Records: {stats.get('totalRecordCount', 'N/A')}") + print(f"Successful: {stats.get('successfulRecordCount', 'N/A')}") + print(f"Failed: {stats.get('failedRecordCount', 'N/A')}") +``` + +### Stop a Batch Job + +```python +import boto3 + +def add_provider_header(request, **kwargs): + request.headers["x-model-provider"] = "bedrock" + +bedrock_client = boto3.client( + "bedrock", + region_name="us-west-2", + endpoint_url="http://localhost:8080/bedrock", +) +bedrock_client.meta.events.register("before-send", add_provider_header) + +# Stop batch job +job_arn = "arn:aws:bedrock:us-west-2:123456789:model-invocation-job/abc123" + +bedrock_client.stop_model_invocation_job(jobIdentifier=job_arn) + +print(f"Stopped job: {job_arn}") +``` + +--- + +## End-to-End Batch Workflow + +### Bedrock Provider + +```python +import boto3 +import json +import time + +# Configuration +region = "us-west-2" +s3_bucket = "your-s3-bucket" +role_arn = "arn:aws:iam::123456789:role/BedrockBatchRole" +model_id = "anthropic.claude-3-sonnet-20240229-v1:0" +provider = "bedrock" + +# Provider header handler +def add_provider_header(request, **kwargs): + request.headers["x-model-provider"] = provider + +# Setup clients +s3_client = boto3.client( + "s3", + region_name=region, + endpoint_url="http://localhost:8080/bedrock/files", +) +s3_client.meta.events.register("before-send", add_provider_header) + +bedrock_client = boto3.client( + "bedrock", + region_name=region, + endpoint_url="http://localhost:8080/bedrock", +) +bedrock_client.meta.events.register("before-send", add_provider_header) + +# Step 1: Create JSONL content +print("Step 1: Creating batch input file...") + +def create_batch_jsonl(num_requests: int) -> str: + lines = [] + for i in range(num_requests): + record = { + "recordId": f"request-{i+1}", + "modelInput": { + "messages": [ + { + "role": "user", + "content": [{"text": f"What is {i+1} + {i+1}? Answer briefly."}], + } + ], + "inferenceConfig": {"maxTokens": 100}, + }, + } + lines.append(json.dumps(record)) + return "\n".join(lines) + +jsonl_content = create_batch_jsonl(num_requests=3) + +# Step 2: Upload input file to S3 +print("Step 2: Uploading input file to S3...") +timestamp = int(time.time()) +s3_key = f"bifrost-batch-input/batch_{timestamp}.jsonl" + +upload_response = s3_client.put_object( + Bucket=s3_bucket, + Key=s3_key, + Body=jsonl_content.encode(), + ContentType="application/jsonl", +) +file_id = upload_response.get("ETag", "").strip('"') +input_uri = f"s3://{s3_bucket}/{s3_key}" +print(f" Uploaded: {input_uri}") + +# Step 3: Create batch job +print("Step 3: Creating batch job...") +output_uri = f"s3://{s3_bucket}/bifrost-batch-output/" + +job_response = bedrock_client.create_model_invocation_job( + jobName=f"bifrost-e2e-{timestamp}", + modelId=model_id, + roleArn=role_arn, + inputDataConfig={ + "s3InputDataConfig": {"s3Uri": input_uri, "s3InputFormat": "JSONL"} + }, + outputDataConfig={ + "s3OutputDataConfig": {"s3Uri": output_uri} + }, + tags=[ + {"key": "endpoint", "value": "/v1/chat/completions"}, + {"key": "file_id", "value": file_id}, + ], +) +job_arn = job_response["jobArn"] +print(f" Created job: {job_arn}") + +# Step 4: Poll for completion +print("Step 4: Polling job status...") +for i in range(20): + status_response = bedrock_client.get_model_invocation_job(jobIdentifier=job_arn) + status = status_response["status"] + print(f" Poll {i+1}: status = {status}") + + if status in ["Completed", "Failed", "Stopped"]: + print(f" Job reached terminal state: {status}") + if status == "Completed" and "statistics" in status_response: + stats = status_response["statistics"] + print(f" Total: {stats.get('totalRecordCount')}") + print(f" Successful: {stats.get('successfulRecordCount')}") + print(f" Failed: {stats.get('failedRecordCount')}") + break + + time.sleep(10) + +# Step 5: Verify job is in list +print("Step 5: Verifying job in list...") +list_response = bedrock_client.list_model_invocation_jobs(maxResults=20) +job_arns = [job["jobArn"] for job in list_response.get("invocationJobSummaries", [])] +assert job_arn in job_arns, f"Job {job_arn} should be in list" +print(f" Verified job is in list") + +print(f"\nSuccess! Batch workflow completed for job {job_arn}") +``` + +### OpenAI Provider + + +**No S3 configuration required.** Files are stored in OpenAI's native storage. The bucket/key values are routing identifiers used by Bifrost. + + +```python +import boto3 +import json +import time + +# Configuration (no S3 bucket needed for OpenAI) +region = "us-west-2" +model_id = "gpt-4o-mini" +provider = "openai" + +# Provider header handler +def add_provider_header(request, **kwargs): + request.headers["x-model-provider"] = provider + +# Setup clients +s3_client = boto3.client( + "s3", + region_name=region, + endpoint_url="http://localhost:8080/bedrock/files", +) +s3_client.meta.events.register("before-send", add_provider_header) + +bedrock_client = boto3.client( + "bedrock", + region_name=region, + endpoint_url="http://localhost:8080/bedrock", +) +bedrock_client.meta.events.register("before-send", add_provider_header) + +# Step 1: Create OpenAI JSONL content +print("Step 1: Creating OpenAI batch input file...") + +def create_openai_jsonl(num_requests: int) -> str: + lines = [] + for i in range(num_requests): + record = { + "custom_id": f"request-{i+1}", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": model_id, + "messages": [ + {"role": "user", "content": f"What is {i+1} + {i+1}? Answer briefly."} + ], + "max_tokens": 100, + }, + } + lines.append(json.dumps(record)) + return "\n".join(lines) + +jsonl_content = create_openai_jsonl(num_requests=3) + +# Step 2: Upload input file (bucket/key are routing identifiers) +print("Step 2: Uploading input file...") +timestamp = int(time.time()) + +upload_response = s3_client.put_object( + Bucket="openai-files", # Routing identifier, not actual S3 + Key=f"batch_{timestamp}.jsonl", + Body=jsonl_content.encode(), + ContentType="application/jsonl", +) +file_id = upload_response.get("ETag", "").strip('"') +print(f" Uploaded file ID: {file_id}") + +# Step 3: Create batch job +print("Step 3: Creating OpenAI batch job...") + +job_response = bedrock_client.create_model_invocation_job( + jobName=f"openai-e2e-{timestamp}", + modelId=model_id, + roleArn="not-required-for-openai", # Not used for OpenAI + inputDataConfig={ + "s3InputDataConfig": {"s3Uri": f"s3://openai-files/{file_id}", "s3InputFormat": "JSONL"} + }, + outputDataConfig={ + "s3OutputDataConfig": {"s3Uri": "s3://openai-output/"} + }, + tags=[ + {"key": "endpoint", "value": "/v1/chat/completions"}, + {"key": "file_id", "value": file_id}, + ], +) +job_arn = job_response["jobArn"] +print(f" Created job: {job_arn}") + +# Step 4: Poll for completion +print("Step 4: Polling job status...") +for i in range(20): + status_response = bedrock_client.get_model_invocation_job(jobIdentifier=job_arn) + status = status_response["status"] + print(f" Poll {i+1}: status = {status}") + + if status in ["Completed", "Failed", "Stopped"]: + print(f" Job reached terminal state: {status}") + break + + time.sleep(10) + +print(f"\nSuccess! OpenAI batch workflow completed for job {job_arn}") +``` + +### Gemini Provider + + +**No S3 configuration required.** Files are stored in Google Cloud Storage. The bucket/key values are routing identifiers used by Bifrost. + + +```python +import boto3 +import json +import time + +# Configuration (no S3 bucket needed for Gemini) +region = "us-west-2" +model_id = "gemini-1.5-flash" +provider = "gemini" + +# Provider header handler +def add_provider_header(request, **kwargs): + request.headers["x-model-provider"] = provider + +# Setup clients +s3_client = boto3.client( + "s3", + region_name=region, + endpoint_url="http://localhost:8080/bedrock/files", +) +s3_client.meta.events.register("before-send", add_provider_header) + +bedrock_client = boto3.client( + "bedrock", + region_name=region, + endpoint_url="http://localhost:8080/bedrock", +) +bedrock_client.meta.events.register("before-send", add_provider_header) + +# Step 1: Create Gemini JSONL content +print("Step 1: Creating Gemini batch input file...") + +def create_gemini_jsonl(num_requests: int) -> str: + lines = [] + for i in range(num_requests): + record = { + "request": { + "contents": [ + { + "role": "user", + "parts": [{"text": f"What is {i+1} + {i+1}? Answer briefly."}], + } + ], + "generationConfig": {"maxOutputTokens": 100}, + }, + "metadata": {"key": f"request-{i+1}"}, + } + lines.append(json.dumps(record)) + return "\n".join(lines) + +jsonl_content = create_gemini_jsonl(num_requests=3) + +# Step 2: Upload input file (bucket/key are routing identifiers) +print("Step 2: Uploading input file...") +timestamp = int(time.time()) + +upload_response = s3_client.put_object( + Bucket="gemini-files", # Routing identifier, not actual S3 + Key=f"batch_{timestamp}.jsonl", + Body=jsonl_content.encode(), + ContentType="application/jsonl", +) +file_id = upload_response.get("ETag", "").strip('"') +print(f" Uploaded file ID: {file_id}") + +# Step 3: Create batch job +print("Step 3: Creating Gemini batch job...") + +job_response = bedrock_client.create_model_invocation_job( + jobName=f"gemini-e2e-{timestamp}", + modelId=model_id, + roleArn="not-required-for-gemini", # Not used for Gemini + inputDataConfig={ + "s3InputDataConfig": {"s3Uri": f"s3://gemini-files/{file_id}", "s3InputFormat": "JSONL"} + }, + outputDataConfig={ + "s3OutputDataConfig": {"s3Uri": "s3://gemini-output/"} + }, +) +job_arn = job_response["jobArn"] +print(f" Created job: {job_arn}") + +# Step 4: Poll for completion (same as Bedrock) +# ... (same polling logic as above) + +print(f"\nSuccess! Gemini batch workflow completed.") +``` + +--- + +## JSONL Format Reference + +### Bedrock Format + +```json +{"recordId": "request-1", "modelInput": {"messages": [{"role": "user", "content": [{"text": "Hello!"}]}], "inferenceConfig": {"maxTokens": 100}}} +``` + +### OpenAI Format + +```json +{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "Hello!"}], "max_tokens": 100}} +``` + +### Gemini Format + +```json +{"request": {"contents": [{"role": "user", "parts": [{"text": "Hello!"}]}], "generationConfig": {"maxOutputTokens": 100}}, "metadata": {"key": "request-1"}} +``` + +--- + +## Provider-Specific Notes + +| Provider | Header Value | File Storage | S3 Config Required | IAM Role Required | +|----------|--------------|--------------|-------------------|-------------------| +| **Bedrock** | `bedrock` | AWS S3 | ✅ Yes | ✅ Yes | +| **OpenAI** | `openai` | OpenAI storage | ❌ No | ❌ No | +| **Gemini** | `gemini` | Google Cloud Storage | ❌ No | ❌ No | +| **Anthropic** | `anthropic` | N/A | N/A | N/A | + + +**Bedrock Provider:** Requires S3 bucket and IAM role configuration. You can use Anthropic models deployed on Bedrock for batch and files APIs. + +**OpenAI & Gemini Providers:** No AWS infrastructure needed. Files are stored in the provider's native storage. The S3 bucket/key values in the examples are routing identifiers used by Bifrost. + +**Anthropic Provider:** Does not support S3-based file uploads. Use the [Anthropic SDK](../anthropic-sdk/files-and-batch) with inline batch requests instead. + + +--- + +## IAM Role Requirements (Bedrock Only) + +For **Bedrock provider** batch jobs, you need an IAM role with these permissions: + +```json +{ + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Action": ["s3:GetObject", "s3:PutObject", "s3:ListBucket"], + "Resource": ["arn:aws:s3:::your-bucket", "arn:aws:s3:::your-bucket/*"] + }, + { + "Effect": "Allow", + "Action": ["bedrock:InvokeModel"], + "Resource": "*" + } + ] +} +``` + +Trust relationship: + +```json +{ + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Principal": {"Service": "bedrock.amazonaws.com"}, + "Action": "sts:AssumeRole" + } + ] +} +``` + +--- + +## Next Steps + +- **[Overview](./overview)** - Bedrock SDK integration basics +- **[Configuration](../../quickstart/gateway/configuration)** - Bifrost setup and configuration +- **[Core Features](../../features/)** - Governance, semantic caching, and more diff --git a/docs/integrations/bedrock-sdk.mdx b/docs/integrations/bedrock-sdk/overview.mdx similarity index 93% rename from docs/integrations/bedrock-sdk.mdx rename to docs/integrations/bedrock-sdk/overview.mdx index 7d9f02e9c..32c085a2f 100644 --- a/docs/integrations/bedrock-sdk.mdx +++ b/docs/integrations/bedrock-sdk/overview.mdx @@ -1,7 +1,7 @@ --- -title: "AWS Bedrock Integration" +title: "Overview" description: "Use Bifrost as a Bedrock-compatible gateway for the Converse and Invoke APIs, with Bifrost features on top." -icon: "b" +icon: "book" --- ## Overview @@ -91,8 +91,6 @@ mistral_invoke_response = client.invoke_model( --- ---- - ## Streaming Examples ### Converse Stream @@ -174,7 +172,7 @@ for event in response.get("body"): Pass AWS credentials or Bedrock API keys directly in requests to bypass Bifrost's load balancing. This requires the **Allow Direct API keys** option to be enabled in Bifrost configuration. -> **Learn more:** See [Quickstart Configuration](../quickstart/README) for enabling direct API key usage. +> **Learn more:** See [Quickstart Configuration](../../../quickstart) for enabling direct API key usage. When direct keys are enabled, you can pass your AWS credentials directly to the boto3 client instead of using dummy credentials. @@ -221,8 +219,8 @@ The Bedrock integration currently supports: ## Next Steps -- **[What is an integration?](./what-is-an-integration)** – Core integration concepts -- **[Configuration](../quickstart/README)** – Bedrock provider setup and API key management -- **[Core Features](../features/)** – Governance, semantic caching, and more - +- **[Files and Batch API](./files-and-batch)** - S3-based file operations and batch processing +- **[What is an integration?](../what-is-an-integration)** - Core integration concepts +- **[Configuration](../../../quickstart)** - Bedrock provider setup and API key management +- **[Core Features](../../features/)** - Governance, semantic caching, and more diff --git a/docs/integrations/genai-sdk/files-and-batch.mdx b/docs/integrations/genai-sdk/files-and-batch.mdx new file mode 100644 index 000000000..91a932638 --- /dev/null +++ b/docs/integrations/genai-sdk/files-and-batch.mdx @@ -0,0 +1,1022 @@ +--- +title: "Files and Batch API" +description: "Upload files and create batch jobs for asynchronous processing using the Google GenAI SDK through Bifrost across multiple providers." +icon: "folder-open" +--- + +## Overview + +Bifrost supports the Google GenAI Files API and Batch API with **cross-provider routing**. This means you can use the Google GenAI SDK to manage files and batch jobs across multiple providers including Gemini, OpenAI, Anthropic, and Bedrock. + +The provider is specified using the `x-model-provider` header in `HttpOptions`. + + +**Bedrock Limitation:** Bedrock batch operations require file-based input with S3 storage, which is not fully supported via the GenAI SDK's batch API. For Bedrock batch operations, use the [Bedrock SDK](../bedrock-sdk/files-and-batch) directly. + + +--- + +## Client Setup + +### Gemini Provider (Default) + +```python +from google import genai +from google.genai.types import HttpOptions + +client = genai.Client( + api_key="your-gemini-api-key", + http_options=HttpOptions(base_url="http://localhost:8080/genai") +) +``` + +### Cross-Provider Client Setup + +To route requests to different providers, add the `x-model-provider` header: + + + + +```python +from google import genai +from google.genai.types import HttpOptions + +client = genai.Client( + api_key="your-gemini-api-key", + http_options=HttpOptions( + base_url="http://localhost:8080/genai", + headers={"x-model-provider": "gemini"} + ) +) +``` + + + + +```python +from google import genai +from google.genai.types import HttpOptions + +client = genai.Client( + api_key="your-openai-api-key", + http_options=HttpOptions( + base_url="http://localhost:8080/genai", + headers={"x-model-provider": "openai"} + ) +) +``` + + + + +```python +from google import genai +from google.genai.types import HttpOptions + +client = genai.Client( + api_key="your-anthropic-api-key", + http_options=HttpOptions( + base_url="http://localhost:8080/genai", + headers={"x-model-provider": "anthropic"} + ) +) +``` + + +Anthropic batch operations use inline requests. File uploads for batch processing are not supported for Anthropic. + + + + + +```python +from google import genai +from google.genai.types import HttpOptions + +client = genai.Client( + api_key="your-api-key", + http_options=HttpOptions( + base_url="http://localhost:8080/genai", + headers={"x-model-provider": "bedrock"} + ) +) +``` + + +Bedrock requires S3-based file storage for batch operations. Use the [Bedrock SDK](../bedrock-sdk/files-and-batch) for full batch support. + + + + + +### Helper Function for Provider-Specific Clients + +```python +from google import genai +from google.genai.types import HttpOptions + +def get_provider_client(provider: str, api_key: str): + """Create GenAI client with x-model-provider header""" + return genai.Client( + api_key=api_key, + http_options=HttpOptions( + base_url="http://localhost:8080/genai", + headers={"x-model-provider": provider} + ) + ) + +# Usage +gemini_client = get_provider_client("gemini", "your-gemini-key") +openai_client = get_provider_client("openai", "your-openai-key") +``` + +--- + +## Files API + +Files are managed through the `client.files` namespace. + +### Upload a File + + + + +```python +from google import genai +from google.genai.types import HttpOptions, UploadFileConfig +import json +import tempfile + +def add_gemini_header(): + return {"x-model-provider": "gemini"} + +client = genai.Client( + api_key="your-gemini-api-key", + http_options=HttpOptions( + base_url="http://localhost:8080/genai", + headers=add_gemini_header() + ) +) + +# Create JSON content for Gemini batch format +def create_gemini_batch_json(num_requests: int = 2) -> str: + requests_list = [] + for i in range(num_requests): + request = { + "key": f"request_{i+1}", + "request": { + "contents": [ + { + "parts": [{"text": f"Hello, this is test message {i+1}. Say hi back briefly."}], + "role": "user" + } + ] + } + } + requests_list.append(json.dumps(request)) + return "\n".join(requests_list) + +# Write content to a temporary file +json_content = create_gemini_batch_json(num_requests=2) +with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f: + f.write(json_content) + temp_file_path = f.name + +# Upload the file +response = client.files.upload( + file=temp_file_path, + config=UploadFileConfig(display_name='batch_input_gemini') +) + +print(f"File name: {response.name}") +print(f"Display name: {response.display_name}") +``` + + + + +```python +from google import genai +from google.genai.types import HttpOptions, UploadFileConfig +import json +import tempfile + +client = genai.Client( + api_key="your-openai-api-key", + http_options=HttpOptions( + base_url="http://localhost:8080/genai", + headers={"x-model-provider": "openai"} + ) +) + +# Create JSONL content for OpenAI batch format +def create_openai_batch_jsonl(model_id: str, num_requests: int = 2) -> str: + lines = [] + for i in range(num_requests): + record = { + "custom_id": f"request-{i+1}", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": model_id, + "messages": [ + {"role": "user", "content": f"Hello, this is test message {i+1}. Say hi back briefly."} + ], + "max_tokens": 100, + }, + } + lines.append(json.dumps(record)) + return "\n".join(lines) + +# Write content to a temporary file +jsonl_content = create_openai_batch_jsonl("gpt-4o-mini") +with tempfile.NamedTemporaryFile(mode='w', suffix='.jsonl', delete=False) as f: + f.write(jsonl_content) + temp_file_path = f.name + +# Upload the file +response = client.files.upload( + file=temp_file_path, + config=UploadFileConfig(display_name='batch_input_openai') +) + +print(f"File name: {response.name}") +``` + + + + + +Anthropic does not support file uploads for batch processing via this API. Use inline batch requests instead (see Batch API section). + + + + + + +Bedrock requires S3-based file storage. Use the [Bedrock SDK](../bedrock-sdk/files-and-batch) for file operations. + + + + + +### List Files + + + + +```python +from google import genai +from google.genai.types import HttpOptions + +client = genai.Client( + api_key="your-gemini-api-key", + http_options=HttpOptions( + base_url="http://localhost:8080/genai", + headers={"x-model-provider": "gemini"} + ) +) + +# List files +for file in client.files.list(config={'page_size': 50}): + print(f"File name: {file.name}") + print(f"Display name: {file.display_name}") + if hasattr(file, 'size_bytes'): + print(f"Size: {file.size_bytes} bytes") + print("---") +``` + + + + +```python +from google import genai +from google.genai.types import HttpOptions + +client = genai.Client( + api_key="your-openai-api-key", + http_options=HttpOptions( + base_url="http://localhost:8080/genai", + headers={"x-model-provider": "openai"} + ) +) + +# List files from OpenAI +for file in client.files.list(config={'page_size': 50}): + print(f"File name: {file.name}") + print(f"Display name: {file.display_name}") + print("---") +``` + + + + + +File listing is not supported for Anthropic via this API. + + + + + + +Use the [Bedrock SDK](../bedrock-sdk/files-and-batch) for file listing with S3 storage. + + + + + +### Retrieve File Metadata + +```python +from google import genai +from google.genai.types import HttpOptions + +client = genai.Client( + api_key="your-api-key", + http_options=HttpOptions( + base_url="http://localhost:8080/genai", + headers={"x-model-provider": "gemini"} # or "openai" + ) +) + +# Retrieve file metadata by name +file_name = "files/abc123" +response = client.files.get(name=file_name) + +print(f"File name: {response.name}") +print(f"Display name: {response.display_name}") +if hasattr(response, 'size_bytes'): + print(f"Size: {response.size_bytes} bytes") +``` + +### Delete a File + +```python +from google import genai +from google.genai.types import HttpOptions + +client = genai.Client( + api_key="your-api-key", + http_options=HttpOptions( + base_url="http://localhost:8080/genai", + headers={"x-model-provider": "gemini"} # or "openai" + ) +) + +# Delete a file +file_name = "files/abc123" +client.files.delete(name=file_name) + +print(f"Deleted file: {file_name}") +``` + +--- + +## Batch API + +Batches are managed through the `client.batches` namespace. The GenAI SDK supports both file-based and inline batch creation. + +### Create a Batch with File + + + + +```python +from google import genai +from google.genai.types import HttpOptions, UploadFileConfig +import json +import tempfile + +client = genai.Client( + api_key="your-gemini-api-key", + http_options=HttpOptions( + base_url="http://localhost:8080/genai", + headers={"x-model-provider": "gemini"} + ) +) + +# Create batch input JSON content (Gemini format) +batch_request = json.dumps({ + "key": "request_1", + "request": { + "contents": [ + {"parts": [{"text": "Hello! Say hi back briefly."}], "role": "user"} + ] + } +}) + +# Write to temporary file and upload +with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f: + f.write(batch_request) + temp_file_path = f.name + +uploaded_file = client.files.upload( + file=temp_file_path, + config=UploadFileConfig(display_name='batch_input_gemini') +) + +# Create batch job using file reference +batch_job = client.batches.create( + model="gemini-1.5-flash", + src=uploaded_file.name, +) + +print(f"Batch name: {batch_job.name}") +print(f"State: {batch_job.state}") +``` + + + + +```python +from google import genai +from google.genai.types import HttpOptions, UploadFileConfig +import json +import tempfile + +client = genai.Client( + api_key="your-openai-api-key", + http_options=HttpOptions( + base_url="http://localhost:8080/genai", + headers={"x-model-provider": "openai"} + ) +) + +# Create batch input JSONL content (OpenAI format) +batch_request = json.dumps({ + "custom_id": "request-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "Hello! Say hi back briefly."}], + "max_tokens": 100 + } +}) + +# Write to temporary file and upload +with tempfile.NamedTemporaryFile(mode='w', suffix='.jsonl', delete=False) as f: + f.write(batch_request) + temp_file_path = f.name + +uploaded_file = client.files.upload( + file=temp_file_path, + config=UploadFileConfig(display_name='batch_input_openai') +) + +# Create batch job using file reference +batch_job = client.batches.create( + model="gpt-4o-mini", + src=uploaded_file.name, +) + +print(f"Batch name: {batch_job.name}") +print(f"State: {batch_job.state}") +``` + + + + + +Anthropic does not support file-based batch creation. Use inline requests instead. + + + + + + +Use the [Bedrock SDK](../bedrock-sdk/files-and-batch) for Bedrock batch operations with S3 storage. + + + + + +### Create a Batch with Inline Requests + + + + +```python +from google import genai +from google.genai.types import HttpOptions + +client = genai.Client( + api_key="your-gemini-api-key", + http_options=HttpOptions( + base_url="http://localhost:8080/genai", + headers={"x-model-provider": "gemini"} + ) +) + +# Create inline requests +inline_requests = [ + { + "contents": [ + { + "parts": [{"text": "What is 2+2?"}], + "role": "user" + } + ], + "config": {"response_modalities": ["TEXT"]} + }, + { + "contents": [ + { + "parts": [{"text": "What is the capital of France?"}], + "role": "user" + } + ], + "config": {"response_modalities": ["TEXT"]} + } +] + +# Create batch job with inline requests +batch_job = client.batches.create( + model="gemini-1.5-flash", + src=inline_requests, +) + +print(f"Batch name: {batch_job.name}") +print(f"State: {batch_job.state}") +``` + + + + +```python +from google import genai +from google.genai.types import HttpOptions + +client = genai.Client( + api_key="your-openai-api-key", + http_options=HttpOptions( + base_url="http://localhost:8080/genai", + headers={"x-model-provider": "openai"} + ) +) + +# Create inline requests (OpenAI format via Bifrost translation) +inline_requests = [ + { + "contents": [ + { + "parts": [{"text": "What is 2+2?"}], + "role": "user" + } + ], + "config": {"response_modalities": ["TEXT"]} + }, + { + "contents": [ + { + "parts": [{"text": "What is the capital of France?"}], + "role": "user" + } + ], + "config": {"response_modalities": ["TEXT"]} + } +] + +batch_job = client.batches.create( + model="gpt-4o-mini", + src=inline_requests, +) + +print(f"Batch name: {batch_job.name}") +print(f"State: {batch_job.state}") +``` + + + + +```python +from google import genai +from google.genai.types import HttpOptions + +client = genai.Client( + api_key="your-anthropic-api-key", + http_options=HttpOptions( + base_url="http://localhost:8080/genai", + headers={"x-model-provider": "anthropic"} + ) +) + +# Create inline requests for Anthropic +inline_requests = [ + { + "contents": [ + { + "parts": [{"text": "What is 2+2?"}], + "role": "user" + } + ], + "config": {"response_modalities": ["TEXT"]} + }, + { + "contents": [ + { + "parts": [{"text": "What is the capital of France?"}], + "role": "user" + } + ], + "config": {"response_modalities": ["TEXT"]} + } +] + +batch_job = client.batches.create( + model="claude-3-sonnet-20240229", + src=inline_requests, +) + +print(f"Batch name: {batch_job.name}") +print(f"State: {batch_job.state}") +``` + + + + + +Use the [Bedrock SDK](../bedrock-sdk/files-and-batch) for Bedrock batch operations. + + + + + +### List Batches + +```python +from google import genai +from google.genai.types import HttpOptions, ListBatchJobsConfig + +client = genai.Client( + api_key="your-api-key", + http_options=HttpOptions( + base_url="http://localhost:8080/genai", + headers={"x-model-provider": "gemini"} # or "openai", "anthropic" + ) +) + +# List batch jobs +for job in client.batches.list(config=ListBatchJobsConfig(page_size=10)): + print(f"Batch name: {job.name}") + print(f"State: {job.state}") + print("---") +``` + +### Retrieve Batch Status + +```python +from google import genai +from google.genai.types import HttpOptions + +client = genai.Client( + api_key="your-api-key", + http_options=HttpOptions( + base_url="http://localhost:8080/genai", + headers={"x-model-provider": "gemini"} # or "openai", "anthropic" + ) +) + +# Get batch job status +batch_name = "batches/abc123" +batch_job = client.batches.get(name=batch_name) + +print(f"Batch name: {batch_job.name}") +print(f"State: {batch_job.state}") +``` + +### Cancel a Batch + +```python +from google import genai +from google.genai.types import HttpOptions + +client = genai.Client( + api_key="your-api-key", + http_options=HttpOptions( + base_url="http://localhost:8080/genai", + headers={"x-model-provider": "gemini"} # or "openai", "anthropic" + ) +) + +# Cancel batch job +batch_name = "batches/abc123" +cancelled_job = client.batches.cancel(name=batch_name) + +print(f"Batch name: {cancelled_job.name}") +print(f"State: {cancelled_job.state}") # JOB_STATE_CANCELLING or JOB_STATE_CANCELLED +``` + +### Delete a Batch + +```python +from google import genai +from google.genai.types import HttpOptions + +client = genai.Client( + api_key="your-api-key", + http_options=HttpOptions( + base_url="http://localhost:8080/genai", + headers={"x-model-provider": "gemini"} # or "openai", "anthropic" + ) +) + +# Delete batch job +batch_name = "batches/abc123" +client.batches.delete(name=batch_name) + +print(f"Deleted batch: {batch_name}") +``` + +--- + +## End-to-End Workflows + +### Gemini Batch Workflow + +```python +import time +from google import genai +from google.genai.types import HttpOptions, UploadFileConfig, ListBatchJobsConfig +import json +import tempfile +import os + +# Configuration +provider = "gemini" +model = "gemini-1.5-flash" + +client = genai.Client( + api_key="your-gemini-api-key", + http_options=HttpOptions( + base_url="http://localhost:8080/genai", + headers={"x-model-provider": provider} + ) +) + +# Step 1: Create batch input file +print("Step 1: Creating batch input file...") + +def create_gemini_batch_json(num_requests: int) -> str: + requests_list = [] + for i in range(num_requests): + request = { + "key": f"request_{i+1}", + "request": { + "contents": [ + { + "parts": [{"text": f"What is {i+1} + {i+1}? Answer briefly."}], + "role": "user" + } + ] + } + } + requests_list.append(json.dumps(request)) + return "\n".join(requests_list) + +json_content = create_gemini_batch_json(num_requests=3) + +with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f: + f.write(json_content) + temp_file_path = f.name + +# Step 2: Upload batch input file +print("Step 2: Uploading batch input file...") +uploaded_file = client.files.upload( + file=temp_file_path, + config=UploadFileConfig(display_name='batch_e2e_gemini') +) +print(f" Uploaded file: {uploaded_file.name}") + +# Step 3: Create batch job +print("Step 3: Creating batch job...") +batch_job = client.batches.create( + model=model, + src=uploaded_file.name, +) +print(f" Created batch: {batch_job.name}, state: {batch_job.state}") + +# Step 4: Poll for completion +print("Step 4: Polling batch status...") +terminal_states = ["JOB_STATE_SUCCEEDED", "JOB_STATE_FAILED", "JOB_STATE_CANCELLED"] + +for i in range(20): + batch_job = client.batches.get(name=batch_job.name) + print(f" Poll {i+1}: state = {batch_job.state}") + + if batch_job.state in terminal_states: + print(f" Batch reached terminal state: {batch_job.state}") + break + + time.sleep(5) + +# Step 5: Verify batch is in list +print("Step 5: Verifying batch in list...") +found = False +for job in client.batches.list(config=ListBatchJobsConfig(page_size=20)): + if job.name == batch_job.name: + found = True + break + +assert found, f"Batch {batch_job.name} should be in list" +print(f" Verified batch {batch_job.name} is in list") + +# Cleanup +os.remove(temp_file_path) +try: + client.files.delete(name=uploaded_file.name) + client.batches.delete(name=batch_job.name) +except Exception as e: + print(f"Cleanup note: {e}") + +print(f"\nSuccess! Batch workflow completed for {batch_job.name}") +``` + +### OpenAI via GenAI SDK Workflow + +```python +import time +from google import genai +from google.genai.types import HttpOptions, ListBatchJobsConfig + +# Configuration +provider = "openai" +model = "gpt-4o-mini" + +client = genai.Client( + api_key="your-openai-api-key", + http_options=HttpOptions( + base_url="http://localhost:8080/genai", + headers={"x-model-provider": provider} + ) +) + +# Step 1: Create inline requests +print("Step 1: Creating inline requests...") +inline_requests = [ + { + "contents": [ + {"parts": [{"text": "What is 2+2?"}], "role": "user"} + ], + "config": {"response_modalities": ["TEXT"]} + }, + { + "contents": [ + {"parts": [{"text": "What is the capital of France?"}], "role": "user"} + ], + "config": {"response_modalities": ["TEXT"]} + } +] +print(f" Created {len(inline_requests)} inline requests") + +# Step 2: Create batch job +print("Step 2: Creating batch job...") +batch_job = client.batches.create( + model=model, + src=inline_requests, +) +print(f" Created batch: {batch_job.name}, state: {batch_job.state}") + +# Step 3: Poll for completion +print("Step 3: Polling batch status...") +terminal_states = ["JOB_STATE_SUCCEEDED", "JOB_STATE_FAILED", "JOB_STATE_CANCELLED"] + +for i in range(10): + batch_job = client.batches.get(name=batch_job.name) + print(f" Poll {i+1}: state = {batch_job.state}") + + if batch_job.state in terminal_states: + break + + time.sleep(5) + +# Cleanup +try: + client.batches.delete(name=batch_job.name) +except Exception as e: + print(f"Cleanup note: {e}") + +print(f"\nSuccess! Cross-provider batch {batch_job.name} completed via GenAI SDK.") +``` + +### Anthropic via GenAI SDK Workflow + +```python +import time +from google import genai +from google.genai.types import HttpOptions + +# Configuration +provider = "anthropic" +model = "claude-3-sonnet-20240229" + +client = genai.Client( + api_key="your-anthropic-api-key", + http_options=HttpOptions( + base_url="http://localhost:8080/genai", + headers={"x-model-provider": provider} + ) +) + +# Step 1: Create inline requests +print("Step 1: Creating inline requests...") +inline_requests = [ + { + "contents": [ + {"parts": [{"text": "What is 15 * 7?"}], "role": "user"} + ], + "config": {"response_modalities": ["TEXT"]} + }, + { + "contents": [ + {"parts": [{"text": "What is the largest ocean?"}], "role": "user"} + ], + "config": {"response_modalities": ["TEXT"]} + } +] +print(f" Created {len(inline_requests)} inline requests") + +# Step 2: Create batch job +print("Step 2: Creating batch job...") +batch_job = client.batches.create( + model=model, + src=inline_requests, +) +print(f" Created batch: {batch_job.name}, state: {batch_job.state}") + +# Step 3: Poll for completion +print("Step 3: Polling batch status...") +terminal_states = ["JOB_STATE_SUCCEEDED", "JOB_STATE_FAILED", "JOB_STATE_CANCELLED", "ended"] + +for i in range(10): + batch_job = client.batches.get(name=batch_job.name) + print(f" Poll {i+1}: state = {batch_job.state}") + + if batch_job.state in terminal_states: + break + + time.sleep(5) + +print(f"\nSuccess! Anthropic batch {batch_job.name} completed via GenAI SDK.") +``` + +--- + +## Batch Job States + +| State | Description | +|-------|-------------| +| `JOB_STATE_QUEUED` | Job is queued and waiting to start | +| `JOB_STATE_PENDING` | Job is pending processing | +| `JOB_STATE_RUNNING` | Job is currently running | +| `JOB_STATE_SUCCEEDED` | Job completed successfully | +| `JOB_STATE_FAILED` | Job failed | +| `JOB_STATE_CANCELLING` | Job is being cancelled | +| `JOB_STATE_CANCELLED` | Job was cancelled | + +--- + +## JSON Format Reference + +### Gemini Batch Format + +```json +{"key": "request-1", "request": {"contents": [{"parts": [{"text": "Hello!"}], "role": "user"}]}} +``` + +### Inline Request Format + +```python +{ + "contents": [ + {"parts": [{"text": "Hello!"}], "role": "user"} + ], + "config": {"response_modalities": ["TEXT"]} +} +``` + +--- + +## Provider-Specific Notes + +| Provider | Header Value | File Upload | Batch Type | Models | +|----------|--------------|-------------|------------|--------| +| **Gemini** | `gemini` or omit | Native storage | File or Inline | `gemini-1.5-*` | +| **OpenAI** | `openai` | Native storage | File or Inline | `gpt-4o-*`, `gpt-4-*` | +| **Anthropic** | `anthropic` | Not supported | Inline only | `claude-3-*` | +| **Bedrock** | `bedrock` | Use Bedrock SDK | Use Bedrock SDK | `anthropic.claude-*` | + + +- **Gemini** and **OpenAI** support both file-based and inline batch creation +- **Anthropic** only supports inline batch requests via this SDK +- **Bedrock** requires the [Bedrock SDK](../bedrock-sdk/files-and-batch) for full batch support + + +--- + +## Next Steps + +- **[Overview](./overview)** - GenAI SDK integration basics +- **[Configuration](../../quickstart/gateway/configuration)** - Bifrost setup and configuration +- **[Core Features](../../features/)** - Governance, semantic caching, and more + diff --git a/docs/integrations/genai-sdk.mdx b/docs/integrations/genai-sdk/overview.mdx similarity index 94% rename from docs/integrations/genai-sdk.mdx rename to docs/integrations/genai-sdk/overview.mdx index 7bd95617c..ab125a6bc 100644 --- a/docs/integrations/genai-sdk.mdx +++ b/docs/integrations/genai-sdk/overview.mdx @@ -1,7 +1,7 @@ --- -title: "Google GenAI SDK" +title: "Overview" description: "Use Bifrost as a drop-in replacement for Google GenAI API with full compatibility and enhanced features." -icon: "g" +icon: "book" --- ## Overview @@ -199,7 +199,7 @@ const response = await model.generateContent("Hello with custom headers!"); Pass API keys directly in requests to bypass Bifrost's load balancing. You can pass any provider's API key (OpenAI, Anthropic, Mistral, etc.) since Bifrost only looks for `Authorization`, `x-api-key` and `x-goog-api-key` headers. This requires the **Allow Direct API keys** option to be enabled in Bifrost configuration. -> **Learn more:** See [Quickstart Configuration](../quickstart/README) for enabling direct API key usage. +> **Learn more:** See [Quickstart Configuration](../../quickstart/README) for enabling direct API key usage. @@ -282,13 +282,13 @@ const gptResponse = await gptModel.generateContent("Hello GPT!"); ## Supported Features -The Google GenAI integration supports all features that are available in both the Google GenAI SDK and Bifrost core functionality. If the Google GenAI SDK supports a feature and Bifrost supports it, the integration will work seamlessly. 😄 +The Google GenAI integration supports all features that are available in both the Google GenAI SDK and Bifrost core functionality. If the Google GenAI SDK supports a feature and Bifrost supports it, the integration will work seamlessly. --- ## Next Steps -- **[OpenAI SDK](./openai-sdk)** - GPT integration patterns -- **[Anthropic SDK](./anthropic-sdk)** - Claude integration patterns -- **[Configuration](../quickstart/README)** - Bifrost setup and configuration -- **[Core Features](../features/)** - Advanced Bifrost capabilities +- **[OpenAI SDK](../openai-sdk/overview)** - GPT integration patterns +- **[Configuration](../../../quickstart)** - Bifrost setup and configuration +- **[Core Features](../../features/)** - Advanced Bifrost capabilities + diff --git a/docs/integrations/openai-sdk/files-and-batch.mdx b/docs/integrations/openai-sdk/files-and-batch.mdx new file mode 100644 index 000000000..cb61a8caa --- /dev/null +++ b/docs/integrations/openai-sdk/files-and-batch.mdx @@ -0,0 +1,673 @@ +--- +title: "Files and Batch API" +description: "Upload files and create batch jobs for asynchronous processing using the OpenAI SDK through Bifrost across multiple providers." +tag: "Beta" +icon: "folder-open" +--- + +## Overview + +Bifrost supports the OpenAI Files API and Batch API with **cross-provider routing**. This means you can use the familiar OpenAI SDK to manage files and batch jobs across multiple providers including OpenAI, Anthropic, Bedrock, and Gemini. + +The provider is specified using `extra_body` (for POST requests) or `extra_query` (for GET requests) parameters. + +--- + +## Client Setup + +The base client setup is the same for all providers. The provider is specified per-request: + +```python +from openai import OpenAI + +client = OpenAI( + base_url="http://localhost:8080/openai", + api_key="your-api-key" # Your actual API key +) +``` + +--- + +## Files API + +### Upload a File + + +**Bedrock** requires S3 storage configuration. OpenAI and Gemini use their native file storage. Anthropic uses inline requests (no file upload). + + + + + +```python +from openai import OpenAI + +client = OpenAI( + base_url="http://localhost:8080/openai", + api_key="your-openai-api-key" +) + +# Create JSONL content for OpenAI batch format +jsonl_content = '''{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "Hello!"}], "max_tokens": 100}} +{"custom_id": "request-2", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "How are you?"}], "max_tokens": 100}}''' + +# Upload file (uses OpenAI's native file storage) +response = client.files.create( + file=("batch_input.jsonl", jsonl_content.encode(), "application/jsonl"), + purpose="batch", + extra_body={"provider": "openai"}, +) + +print(f"Uploaded file ID: {response.id}") +``` + + + + +For Bedrock, you need to provide S3 storage configuration: + +```python +from openai import OpenAI + +client = OpenAI( + base_url="http://localhost:8080/openai", + api_key="your-api-key" +) + +# Create JSONL content using OpenAI-style format (Bifrost converts to Bedrock format internally) +jsonl_content = '''{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "anthropic.claude-3-sonnet-20240229-v1:0", "messages": [{"role": "user", "content": "Hello!"}], "max_tokens": 100}} +{"custom_id": "request-2", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "anthropic.claude-3-sonnet-20240229-v1:0", "messages": [{"role": "user", "content": "How are you?"}], "max_tokens": 100}}''' + +# Upload file with S3 storage configuration +response = client.files.create( + file=("batch_input.jsonl", jsonl_content.encode(), "application/jsonl"), + purpose="batch", + extra_body={ + "provider": "bedrock", + "storage_config": { + "s3": { + "bucket": "your-s3-bucket", + "region": "us-west-2", + "prefix": "bifrost-batch-output", + }, + }, + }, +) + +print(f"Uploaded file ID: {response.id}") +``` + + + + +Anthropic uses inline requests for batching (no file upload needed). See the Batch API section below. + + + + +```python +from openai import OpenAI + +client = OpenAI( + base_url="http://localhost:8080/openai", + api_key="your-api-key" +) + +# Create JSONL content using OpenAI-style format (Bifrost converts to Gemini format internally) +jsonl_content = '''{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gemini-1.5-flash", "messages": [{"role": "user", "content": "Hello!"}], "max_tokens": 100}} +{"custom_id": "request-2", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gemini-1.5-flash", "messages": [{"role": "user", "content": "How are you?"}], "max_tokens": 100}}''' + +# Upload file (uses Gemini's native file storage) +response = client.files.create( + file=("batch_input.jsonl", jsonl_content.encode(), "application/jsonl"), + purpose="batch", + extra_body={"provider": "gemini"}, +) + +print(f"Uploaded file ID: {response.id}") +``` + + + + +### List Files + +```python +# List files for OpenAI or Gemini (no S3 config needed) +response = client.files.list( + extra_query={"provider": "openai"} # or "gemini" +) + +for file in response.data: + print(f"File ID: {file.id}, Name: {file.filename}") + +# For Bedrock (requires S3 config) +response = client.files.list( + extra_query={ + "provider": "bedrock", + "storage_config": { + "s3": { + "bucket": "your-s3-bucket", + "region": "us-west-2", + "prefix": "bifrost-batch-output", + }, + }, + } +) +``` + +### Retrieve File Metadata + +```python +# Retrieve file metadata (specify provider) +file_id = "file-abc123" +response = client.files.retrieve( + file_id, + extra_query={"provider": "bedrock"} # or "openai", "gemini" +) + +print(f"File ID: {response.id}") +print(f"Filename: {response.filename}") +print(f"Purpose: {response.purpose}") +print(f"Bytes: {response.bytes}") +``` + +### Delete a File + +```python +# Delete file (specify provider) +file_id = "file-abc123" +response = client.files.delete( + file_id, + extra_query={"provider": "bedrock"} # or "openai", "gemini" +) + +print(f"Deleted: {response.deleted}") +``` + +### Download File Content + +```python +# Download file content (specify provider) +file_id = "file-abc123" +response = client.files.content( + file_id, + extra_query={"provider": "bedrock"} # or "openai", "gemini" +) + +# Handle different response types +if hasattr(response, "read"): + content = response.read() +elif hasattr(response, "content"): + content = response.content +else: + content = response + +# Decode bytes to string if needed +if isinstance(content, bytes): + content = content.decode("utf-8") + +print(f"File content:\n{content}") +``` + +--- + +## Batch API + +### Create a Batch + + + + +For native OpenAI batching: + +```python +from openai import OpenAI + +client = OpenAI( + base_url="http://localhost:8080/openai", + api_key="your-openai-api-key" +) + +# First upload a file (see Files API section) +# Then create batch using the file ID + +batch = client.batches.create( + input_file_id="file-abc123", + endpoint="/v1/chat/completions", + completion_window="24h", + extra_body={"provider": "openai"}, +) + +print(f"Batch ID: {batch.id}") +print(f"Status: {batch.status}") +``` + + + + +For Bedrock, you need to provide role ARN and output S3 URI: + +```python +from openai import OpenAI + +client = OpenAI( + base_url="http://localhost:8080/openai", + api_key="your-api-key" +) + +# First upload a file with S3 config (see Files API section) +# Then create batch using the file ID + +batch = client.batches.create( + input_file_id="file-abc123", + endpoint="/v1/chat/completions", + completion_window="24h", + extra_body={ + "provider": "bedrock", + "model": "anthropic.claude-3-sonnet-20240229-v1:0", + "role_arn": "arn:aws:iam::123456789:role/BedrockBatchRole", + "output_s3_uri": "s3://your-bucket/batch-output", + }, +) + +print(f"Batch ID: {batch.id}") +print(f"Status: {batch.status}") +``` + + + + +Anthropic supports inline requests (no file upload required): + +```python +from openai import OpenAI + +client = OpenAI( + base_url="http://localhost:8080/openai", + api_key="your-anthropic-api-key" +) + +# Create inline requests for Anthropic +requests = [ + { + "custom_id": "request-1", + "params": { + "model": "claude-3-sonnet-20240229", + "max_tokens": 100, + "messages": [{"role": "user", "content": "Hello!"}] + } + }, + { + "custom_id": "request-2", + "params": { + "model": "claude-3-sonnet-20240229", + "max_tokens": 100, + "messages": [{"role": "user", "content": "How are you?"}] + } + } +] + +# Create batch with inline requests (no file ID needed) +batch = client.batches.create( + input_file_id="", # Empty for inline requests + endpoint="/v1/chat/completions", + completion_window="24h", + extra_body={ + "provider": "anthropic", + "requests": requests, + }, +) + +print(f"Batch ID: {batch.id}") +print(f"Status: {batch.status}") +``` + + + + +```python +from openai import OpenAI + +client = OpenAI( + base_url="http://localhost:8080/openai", + api_key="your-api-key" +) + +# First upload a file with Gemini format (see Files API section) +# Then create batch using the file ID + +batch = client.batches.create( + input_file_id="file-abc123", + endpoint="/v1/chat/completions", + completion_window="24h", + extra_body={ + "provider": "gemini", + "model": "gemini-1.5-flash", + }, +) + +print(f"Batch ID: {batch.id}") +print(f"Status: {batch.status}") +``` + + + + +### List Batches + +```python +# List batches (specify provider) +response = client.batches.list( + limit=10, + extra_query={ + "provider": "bedrock", # or "openai", "anthropic", "gemini" + "model": "anthropic.claude-3-sonnet-20240229-v1:0", # Required for bedrock + "role_arn": "arn:aws:iam::123456789:role/BedrockBatchRole", # Required for bedrock + } +) + +for batch in response.data: + print(f"Batch ID: {batch.id}, Status: {batch.status}") +``` + +### Retrieve Batch Status + +```python +# Retrieve batch status (specify provider) +batch_id = "batch-abc123" +batch = client.batches.retrieve( + batch_id, + extra_query={"provider": "bedrock"} # or "openai", "anthropic", "gemini" +) + +print(f"Batch ID: {batch.id}") +print(f"Status: {batch.status}") + +if batch.request_counts: + print(f"Total: {batch.request_counts.total}") + print(f"Completed: {batch.request_counts.completed}") + print(f"Failed: {batch.request_counts.failed}") +``` + +### Cancel a Batch + +```python +# Cancel batch (specify provider) +batch_id = "batch-abc123" +batch = client.batches.cancel( + batch_id, + extra_body={"provider": "bedrock"} # or "openai", "anthropic", "gemini" +) + +print(f"Batch ID: {batch.id}") +print(f"Status: {batch.status}") # "cancelling" or "cancelled" +``` + +--- + +## End-to-End Workflows + +### OpenAI Batch Workflow + +```python +import time +from openai import OpenAI + +client = OpenAI( + base_url="http://localhost:8080/openai", + api_key="your-openai-api-key" +) + +# Configuration +provider = "openai" + +# Step 1: Create OpenAI JSONL content +jsonl_content = '''{"custom_id": "req-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "What is 2+2?"}], "max_tokens": 100}} +{"custom_id": "req-2", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "What is the capital of France?"}], "max_tokens": 100}}''' + +# Step 2: Upload file (uses OpenAI's native file storage) +print("Step 1: Uploading batch input file...") +uploaded_file = client.files.create( + file=("batch_e2e.jsonl", jsonl_content.encode(), "application/jsonl"), + purpose="batch", + extra_body={"provider": provider}, +) +print(f" Uploaded file: {uploaded_file.id}") + +# Step 3: Create batch +print("Step 2: Creating batch job...") +batch = client.batches.create( + input_file_id=uploaded_file.id, + endpoint="/v1/chat/completions", + completion_window="24h", + extra_body={"provider": provider}, +) +print(f" Created batch: {batch.id}, status: {batch.status}") + +# Step 4: Poll for completion +print("Step 3: Polling batch status...") +for i in range(10): + batch = client.batches.retrieve(batch.id, extra_query={"provider": provider}) + print(f" Poll {i+1}: status = {batch.status}") + + if batch.status in ["completed", "failed", "expired", "cancelled"]: + break + + if batch.request_counts: + print(f" Completed: {batch.request_counts.completed}/{batch.request_counts.total}") + + time.sleep(5) + +print(f"\nSuccess! Batch {batch.id} workflow completed.") +``` + +### Bedrock Batch Workflow + +```python +import time +from openai import OpenAI + +client = OpenAI( + base_url="http://localhost:8080/openai", + api_key="your-api-key" +) + +# Configuration +provider = "bedrock" +s3_bucket = "your-s3-bucket" +s3_region = "us-west-2" +role_arn = "arn:aws:iam::123456789:role/BedrockBatchRole" +model = "anthropic.claude-3-sonnet-20240229-v1:0" + +# Step 1: Create JSONL content using OpenAI-style format (Bifrost converts to Bedrock format internally) +jsonl_content = '''{"custom_id": "req-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "anthropic.claude-3-sonnet-20240229-v1:0", "messages": [{"role": "user", "content": "What is 2+2?"}], "max_tokens": 100}} +{"custom_id": "req-2", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "anthropic.claude-3-sonnet-20240229-v1:0", "messages": [{"role": "user", "content": "What is the capital of France?"}], "max_tokens": 100}}''' + +# Step 2: Upload file +print("Step 1: Uploading batch input file...") +uploaded_file = client.files.create( + file=("batch_e2e.jsonl", jsonl_content.encode(), "application/jsonl"), + purpose="batch", + extra_body={ + "provider": provider, + "storage_config": { + "s3": {"bucket": s3_bucket, "region": s3_region, "prefix": "batch-input"}, + }, + }, +) +print(f" Uploaded file: {uploaded_file.id}") + +# Step 3: Create batch +print("Step 2: Creating batch job...") +batch = client.batches.create( + input_file_id=uploaded_file.id, + endpoint="/v1/chat/completions", + completion_window="24h", + extra_body={ + "provider": provider, + "model": model, + "role_arn": role_arn, + "output_s3_uri": f"s3://{s3_bucket}/batch-output", + }, +) +print(f" Created batch: {batch.id}, status: {batch.status}") + +# Step 4: Poll for completion +print("Step 3: Polling batch status...") +for i in range(10): + batch = client.batches.retrieve(batch.id, extra_query={"provider": provider}) + print(f" Poll {i+1}: status = {batch.status}") + + if batch.status in ["completed", "failed", "expired", "cancelled"]: + break + + if batch.request_counts: + print(f" Completed: {batch.request_counts.completed}/{batch.request_counts.total}") + + time.sleep(5) + +print(f"\nSuccess! Batch {batch.id} workflow completed.") +``` + +### Anthropic Inline Batch Workflow + +```python +import time +from openai import OpenAI + +client = OpenAI( + base_url="http://localhost:8080/openai", + api_key="your-anthropic-api-key" +) + +provider = "anthropic" + +# Step 1: Create inline requests +print("Step 1: Creating inline requests...") +requests = [ + { + "custom_id": "math-question", + "params": { + "model": "claude-3-sonnet-20240229", + "max_tokens": 100, + "messages": [{"role": "user", "content": "What is 15 * 7?"}] + } + }, + { + "custom_id": "geography-question", + "params": { + "model": "claude-3-sonnet-20240229", + "max_tokens": 100, + "messages": [{"role": "user", "content": "What is the largest ocean?"}] + } + } +] +print(f" Created {len(requests)} inline requests") + +# Step 2: Create batch +print("Step 2: Creating batch job...") +batch = client.batches.create( + input_file_id="", + endpoint="/v1/chat/completions", + completion_window="24h", + extra_body={"provider": provider, "requests": requests}, +) +print(f" Created batch: {batch.id}, status: {batch.status}") + +# Step 3: Poll for completion +print("Step 3: Polling batch status...") +for i in range(10): + batch = client.batches.retrieve(batch.id, extra_query={"provider": provider}) + print(f" Poll {i+1}: status = {batch.status}") + + if batch.status in ["completed", "failed", "expired", "cancelled", "ended"]: + break + + time.sleep(5) + +print(f"\nSuccess! Batch {batch.id} workflow completed.") +``` + +### Gemini Batch Workflow + +```python +import time +from openai import OpenAI + +client = OpenAI( + base_url="http://localhost:8080/openai", + api_key="your-api-key" +) + +# Configuration +provider = "gemini" +model = "gemini-1.5-flash" + +# Step 1: Create JSONL content using OpenAI-style format (Bifrost converts to Gemini format internally) +jsonl_content = '''{"custom_id": "req-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gemini-1.5-flash", "messages": [{"role": "user", "content": "What is 2+2?"}], "max_tokens": 100}} +{"custom_id": "req-2", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gemini-1.5-flash", "messages": [{"role": "user", "content": "What is the capital of France?"}], "max_tokens": 100}}''' + +# Step 2: Upload file (uses Gemini's native file storage) +print("Step 1: Uploading batch input file...") +uploaded_file = client.files.create( + file=("batch_e2e.jsonl", jsonl_content.encode(), "application/jsonl"), + purpose="batch", + extra_body={"provider": provider}, +) +print(f" Uploaded file: {uploaded_file.id}") + +# Step 3: Create batch +print("Step 2: Creating batch job...") +batch = client.batches.create( + input_file_id=uploaded_file.id, + endpoint="/v1/chat/completions", + completion_window="24h", + extra_body={ + "provider": provider, + "model": model, + }, +) +print(f" Created batch: {batch.id}, status: {batch.status}") + +# Step 4: Poll for completion +print("Step 3: Polling batch status...") +for i in range(10): + batch = client.batches.retrieve(batch.id, extra_query={"provider": provider}) + print(f" Poll {i+1}: status = {batch.status}") + + if batch.status in ["completed", "failed", "expired", "cancelled"]: + break + + if batch.request_counts: + print(f" Completed: {batch.request_counts.completed}/{batch.request_counts.total}") + + time.sleep(5) + +print(f"\nSuccess! Batch {batch.id} workflow completed.") +``` + +--- + +## Provider-Specific Notes + +| Provider | File Upload | Batch Creation | Extra Configuration | +|----------|-------------|----------------|---------------------| +| **OpenAI** | ✅ Native storage | ✅ File-based | None | +| **Bedrock** | ✅ S3-based | ✅ File-based | `storage_config`, `role_arn`, `output_s3_uri` | +| **Anthropic** | ❌ Not supported | ✅ Inline requests | `requests` array in `extra_body` | +| **Gemini** | ✅ Native storage | ✅ File-based | `model` in `extra_body` | + + +- **OpenAI** and **Gemini** use their native file storage - no S3 configuration needed +- **Bedrock** requires S3 storage configuration (`storage_config`, `role_arn`, `output_s3_uri`) +- **Anthropic** does not support file-based batch operations - use inline requests instead + + +--- + +## Next Steps + +- **[Overview](./overview)** - OpenAI SDK integration basics +- **[Configuration](../../quickstart/gateway/configuration)** - Bifrost setup and configuration +- **[Core Features](../../features/)** - Governance, semantic caching, and more diff --git a/docs/integrations/openai-sdk.mdx b/docs/integrations/openai-sdk/overview.mdx similarity index 94% rename from docs/integrations/openai-sdk.mdx rename to docs/integrations/openai-sdk/overview.mdx index dfa201553..b3addf624 100644 --- a/docs/integrations/openai-sdk.mdx +++ b/docs/integrations/openai-sdk/overview.mdx @@ -1,7 +1,7 @@ --- -title: "OpenAI SDK" +title: "Overview" description: "Use Bifrost as a drop-in replacement for OpenAI API with full compatibility and enhanced features." -icon: "o" +icon: "book" --- ## Overview @@ -209,7 +209,7 @@ const response = await openai.chat.completions.create({ Pass API keys directly in requests to bypass Bifrost's load balancing. You can pass any provider's API key (OpenAI, Anthropic, Mistral, etc.) since Bifrost only looks for `Authorization` or `x-api-key` headers. This requires the **Allow Direct API keys** option to be enabled in Bifrost configuration. -> **Learn more:** See [Quickstart Configuration](../quickstart/README) for enabling direct API key usage. +> **Learn more:** See [Quickstart Configuration](../../quickstart/README) for enabling direct API key usage. @@ -371,13 +371,15 @@ console.log(azureResponse.choices[0].message.content); ## Supported Features -The OpenAI integration supports all features that are available in both the OpenAI SDK and Bifrost core functionality. If the OpenAI SDK supports a feature and Bifrost supports it, the integration will work seamlessly. 😄 +The OpenAI integration supports all features that are available in both the OpenAI SDK and Bifrost core functionality. If the OpenAI SDK supports a feature and Bifrost supports it, the integration will work seamlessly. --- ## Next Steps -- **[Anthropic SDK](./anthropic-sdk)** - Claude integration patterns -- **[Google GenAI SDK](./genai-sdk)** - Gemini integration patterns -- **[Configuration](../quickstart/README)** - Bifrost setup and configuration -- **[Core Features](../features/)** - Advanced Bifrost capabilities \ No newline at end of file +- **[Files and Batch API](./files-and-batch)** - File uploads and batch processing +- **[Anthropic SDK](../anthropic-sdk/overview)** - Claude integration patterns +- **[Google GenAI SDK](../genai-sdk)** - Gemini integration patterns +- **[Configuration](../../quickstart/README)** - Bifrost setup and configuration +- **[Core Features](../../features/)** - Advanced Bifrost capabilities + diff --git a/framework/changelog.md b/framework/changelog.md index e69de29bb..bd9b266bb 100644 --- a/framework/changelog.md +++ b/framework/changelog.md @@ -0,0 +1,2 @@ +- feat: adds logging support for batch and file requests +- chore: upgrades core to 1.2.38 and framework to 1.1.48 \ No newline at end of file diff --git a/framework/configstore/migrations.go b/framework/configstore/migrations.go index 0ca229f67..19e84e924 100644 --- a/framework/configstore/migrations.go +++ b/framework/configstore/migrations.go @@ -86,9 +86,6 @@ func triggerMigrations(ctx context.Context, db *gorm.DB) error { if err := migrationAddBatchAndCachePricingColumns(ctx, db); err != nil { return err } - if err := migrationAdd200kTokenPricingColumns(ctx, db); err != nil { - return err - } if err := migrationMoveKeysToProviderConfig(ctx, db); err != nil { return err } @@ -107,6 +104,9 @@ func triggerMigrations(ctx context.Context, db *gorm.DB) error { if err := migrationAddAdditionalConfigHashColumns(ctx, db); err != nil { return err } + if err := migrationAdd200kTokenPricingColumns(ctx, db); err != nil { + return err + } return nil } @@ -1191,54 +1191,6 @@ func migrationAddBatchAndCachePricingColumns(ctx context.Context, db *gorm.DB) e return m.Migrate() } -// migrationAdd200kTokenPricingColumns adds pricing columns for 200k token tier models -func migrationAdd200kTokenPricingColumns(ctx context.Context, db *gorm.DB) error { - m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{ - ID: "add_200k_token_pricing_columns", - Migrate: func(tx *gorm.DB) error { - tx = tx.WithContext(ctx) - migrator := tx.Migrator() - - columns := []string{ - "input_cost_per_token_above_200k_tokens", - "output_cost_per_token_above_200k_tokens", - "cache_creation_input_token_cost_above_200k_tokens", - "cache_read_input_token_cost_above_200k_tokens", - } - - for _, field := range columns { - if !migrator.HasColumn(&tables.TableModelPricing{}, field) { - if err := migrator.AddColumn(&tables.TableModelPricing{}, field); err != nil { - return fmt.Errorf("failed to add column %s: %w", field, err) - } - } - } - return nil - }, - Rollback: func(tx *gorm.DB) error { - tx = tx.WithContext(ctx) - migrator := tx.Migrator() - - columns := []string{ - "input_cost_per_token_above_200k_tokens", - "output_cost_per_token_above_200k_tokens", - "cache_creation_input_token_cost_above_200k_tokens", - "cache_read_input_token_cost_above_200k_tokens", - } - - for _, field := range columns { - if migrator.HasColumn(&tables.TableModelPricing{}, field) { - if err := migrator.DropColumn(&tables.TableModelPricing{}, field); err != nil { - return fmt.Errorf("failed to drop column %s: %w", field, err) - } - } - } - return nil - }, - }}) - return m.Migrate() -} - // migrationMoveKeysToProviderConfig migrates keys from virtual key level to provider config level func migrationMoveKeysToProviderConfig(ctx context.Context, db *gorm.DB) error { m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{ @@ -1774,3 +1726,51 @@ func migrationAddAdditionalConfigHashColumns(ctx context.Context, db *gorm.DB) e } return nil } + +// migrationAdd200kTokenPricingColumns adds pricing columns for 200k token tier models +func migrationAdd200kTokenPricingColumns(ctx context.Context, db *gorm.DB) error { + m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{ + ID: "add_200k_token_pricing_columns", + Migrate: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + migrator := tx.Migrator() + + columns := []string{ + "input_cost_per_token_above_200k_tokens", + "output_cost_per_token_above_200k_tokens", + "cache_creation_input_token_cost_above_200k_tokens", + "cache_read_input_token_cost_above_200k_tokens", + } + + for _, field := range columns { + if !migrator.HasColumn(&tables.TableModelPricing{}, field) { + if err := migrator.AddColumn(&tables.TableModelPricing{}, field); err != nil { + return fmt.Errorf("failed to add column %s: %w", field, err) + } + } + } + return nil + }, + Rollback: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + migrator := tx.Migrator() + + columns := []string{ + "input_cost_per_token_above_200k_tokens", + "output_cost_per_token_above_200k_tokens", + "cache_creation_input_token_cost_above_200k_tokens", + "cache_read_input_token_cost_above_200k_tokens", + } + + for _, field := range columns { + if migrator.HasColumn(&tables.TableModelPricing{}, field) { + if err := migrator.DropColumn(&tables.TableModelPricing{}, field); err != nil { + return fmt.Errorf("failed to drop column %s: %w", field, err) + } + } + } + return nil + }, + }}) + return m.Migrate() +} diff --git a/framework/configstore/tables/modelpricing.go b/framework/configstore/tables/modelpricing.go index 06cf9f8ad..f818f02e1 100644 --- a/framework/configstore/tables/modelpricing.go +++ b/framework/configstore/tables/modelpricing.go @@ -28,16 +28,16 @@ type TableModelPricing struct { OutputCostPerCharacterAbove128kTokens *float64 `gorm:"default:null" json:"output_cost_per_character_above_128k_tokens,omitempty"` //Pricing above 200k tokens (for gemini and claude models) - InputCostPerTokenAbove200kTokens *float64 `gorm:"default:null" json:"input_cost_per_token_above_200k_tokens,omitempty"` - OutputCostPerTokenAbove200kTokens *float64 `gorm:"default:null" json:"output_cost_per_token_above_200k_tokens,omitempty"` - CacheCreationInputTokenCostAbove200kTokens *float64 `gorm:"default:null" json:"cache_creation_input_token_cost_above_200k_tokens,omitempty"` - CacheReadInputTokenCostAbove200kTokens *float64 `gorm:"default:null" json:"cache_read_input_token_cost_above_200k_tokens,omitempty"` + InputCostPerTokenAbove200kTokens *float64 `gorm:"default:null;column:input_cost_per_token_above_200k_tokens" json:"input_cost_per_token_above_200k_tokens,omitempty"` + OutputCostPerTokenAbove200kTokens *float64 `gorm:"default:null;column:output_cost_per_token_above_200k_tokens" json:"output_cost_per_token_above_200k_tokens,omitempty"` + CacheCreationInputTokenCostAbove200kTokens *float64 `gorm:"default:null;column:cache_creation_input_token_cost_above_200k_tokens" json:"cache_creation_input_token_cost_above_200k_tokens,omitempty"` + CacheReadInputTokenCostAbove200kTokens *float64 `gorm:"default:null;column:cache_read_input_token_cost_above_200k_tokens" json:"cache_read_input_token_cost_above_200k_tokens,omitempty"` // Cache and batch pricing - CacheReadInputTokenCost *float64 `gorm:"default:null" json:"cache_read_input_token_cost,omitempty"` - CacheCreationInputTokenCost *float64 `gorm:"default:null" json:"cache_creation_input_token_cost,omitempty"` - InputCostPerTokenBatches *float64 `gorm:"default:null" json:"input_cost_per_token_batches,omitempty"` - OutputCostPerTokenBatches *float64 `gorm:"default:null" json:"output_cost_per_token_batches,omitempty"` + CacheReadInputTokenCost *float64 `gorm:"default:null;column:cache_read_input_token_cost" json:"cache_read_input_token_cost,omitempty"` + CacheCreationInputTokenCost *float64 `gorm:"default:null;column:cache_creation_input_token_cost" json:"cache_creation_input_token_cost,omitempty"` + InputCostPerTokenBatches *float64 `gorm:"default:null;column:input_cost_per_token_batches" json:"input_cost_per_token_batches,omitempty"` + OutputCostPerTokenBatches *float64 `gorm:"default:null;column:output_cost_per_token_batches" json:"output_cost_per_token_batches,omitempty"` } // TableName sets the table name for each model diff --git a/framework/version b/framework/version index e20060401..3361394de 100644 --- a/framework/version +++ b/framework/version @@ -1 +1 @@ -1.1.47 \ No newline at end of file +1.1.48 \ No newline at end of file diff --git a/plugins/governance/changelog.md b/plugins/governance/changelog.md index e69de29bb..c9b9868c3 100644 --- a/plugins/governance/changelog.md +++ b/plugins/governance/changelog.md @@ -0,0 +1 @@ +- chore: upgrades core to 1.2.38 and framework to 1.1.48 \ No newline at end of file diff --git a/plugins/governance/go.mod b/plugins/governance/go.mod index ae2b18f3b..6190fc384 100644 --- a/plugins/governance/go.mod +++ b/plugins/governance/go.mod @@ -84,7 +84,7 @@ require ( github.com/twitchyliquid64/golang-asm v0.15.1 // indirect github.com/valyala/bytebufferpool v1.0.0 // indirect github.com/valyala/fasthttp v1.67.0 // indirect - github.com/weaviate/weaviate v1.33.1 // indirect + github.com/weaviate/weaviate v1.33.4 // indirect github.com/weaviate/weaviate-go-client/v5 v5.5.0 // indirect github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect github.com/yosida95/uritemplate/v3 v3.0.2 // indirect diff --git a/plugins/governance/go.sum b/plugins/governance/go.sum index 5dc9c29ab..9c25418f9 100644 --- a/plugins/governance/go.sum +++ b/plugins/governance/go.sum @@ -185,8 +185,7 @@ github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6Kllzaw github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc= github.com/valyala/fasthttp v1.67.0 h1:tqKlJMUP6iuNG8hGjK/s9J4kadH7HLV4ijEcPGsezac= github.com/valyala/fasthttp v1.67.0/go.mod h1:qYSIpqt/0XNmShgo/8Aq8E3UYWVVwNS2QYmzd8WIEPM= -github.com/weaviate/weaviate v1.33.1 h1:fV69ffJSH0aO3LvLiKYlVZ8wFa94oQ1g3uMyZGTb838= -github.com/weaviate/weaviate v1.33.1/go.mod h1:SnxXSIoiusZttZ/gI9knXhFAu0UYqn9N/ekgsNnXbNw= +github.com/weaviate/weaviate v1.33.4 h1:eA37l538+3pEBJAZ3/mFBHG0IEWcj5/aAJSzFhrYuUA= github.com/weaviate/weaviate-go-client/v5 v5.5.0 h1:+5qkHodrL3/Qc7kXvMXnDaIxSBN5+djivLqzmCx7VS4= github.com/weaviate/weaviate-go-client/v5 v5.5.0/go.mod h1:Zdm2MEXG27I0Nf6fM0FZ3P2vLR4JM0iJZrOxwc+Zj34= github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc= diff --git a/plugins/governance/version b/plugins/governance/version index bba60c093..8229aaebc 100644 --- a/plugins/governance/version +++ b/plugins/governance/version @@ -1 +1 @@ -1.3.48 \ No newline at end of file +1.3.49 \ No newline at end of file diff --git a/plugins/jsonparser/changelog.md b/plugins/jsonparser/changelog.md index e69de29bb..c9b9868c3 100644 --- a/plugins/jsonparser/changelog.md +++ b/plugins/jsonparser/changelog.md @@ -0,0 +1 @@ +- chore: upgrades core to 1.2.38 and framework to 1.1.48 \ No newline at end of file diff --git a/plugins/jsonparser/version b/plugins/jsonparser/version index bba60c093..8229aaebc 100644 --- a/plugins/jsonparser/version +++ b/plugins/jsonparser/version @@ -1 +1 @@ -1.3.48 \ No newline at end of file +1.3.49 \ No newline at end of file diff --git a/plugins/logging/changelog.md b/plugins/logging/changelog.md index e69de29bb..c9b9868c3 100644 --- a/plugins/logging/changelog.md +++ b/plugins/logging/changelog.md @@ -0,0 +1 @@ +- chore: upgrades core to 1.2.38 and framework to 1.1.48 \ No newline at end of file diff --git a/plugins/logging/go.mod b/plugins/logging/go.mod index a7a7b9464..283e827a1 100644 --- a/plugins/logging/go.mod +++ b/plugins/logging/go.mod @@ -82,7 +82,7 @@ require ( github.com/twitchyliquid64/golang-asm v0.15.1 // indirect github.com/valyala/bytebufferpool v1.0.0 // indirect github.com/valyala/fasthttp v1.67.0 // indirect - github.com/weaviate/weaviate v1.33.1 // indirect + github.com/weaviate/weaviate v1.33.4 // indirect github.com/weaviate/weaviate-go-client/v5 v5.5.0 // indirect github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect github.com/yosida95/uritemplate/v3 v3.0.2 // indirect diff --git a/plugins/logging/go.sum b/plugins/logging/go.sum index 5dc9c29ab..9c25418f9 100644 --- a/plugins/logging/go.sum +++ b/plugins/logging/go.sum @@ -185,8 +185,7 @@ github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6Kllzaw github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc= github.com/valyala/fasthttp v1.67.0 h1:tqKlJMUP6iuNG8hGjK/s9J4kadH7HLV4ijEcPGsezac= github.com/valyala/fasthttp v1.67.0/go.mod h1:qYSIpqt/0XNmShgo/8Aq8E3UYWVVwNS2QYmzd8WIEPM= -github.com/weaviate/weaviate v1.33.1 h1:fV69ffJSH0aO3LvLiKYlVZ8wFa94oQ1g3uMyZGTb838= -github.com/weaviate/weaviate v1.33.1/go.mod h1:SnxXSIoiusZttZ/gI9knXhFAu0UYqn9N/ekgsNnXbNw= +github.com/weaviate/weaviate v1.33.4 h1:eA37l538+3pEBJAZ3/mFBHG0IEWcj5/aAJSzFhrYuUA= github.com/weaviate/weaviate-go-client/v5 v5.5.0 h1:+5qkHodrL3/Qc7kXvMXnDaIxSBN5+djivLqzmCx7VS4= github.com/weaviate/weaviate-go-client/v5 v5.5.0/go.mod h1:Zdm2MEXG27I0Nf6fM0FZ3P2vLR4JM0iJZrOxwc+Zj34= github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc= diff --git a/plugins/logging/main.go b/plugins/logging/main.go index 0ede4aca3..297290ae6 100644 --- a/plugins/logging/main.go +++ b/plugins/logging/main.go @@ -264,7 +264,7 @@ func (p *LoggerPlugin) PreHook(ctx *schemas.BifrostContext, req *schemas.Bifrost initialData.SpeechInput = req.SpeechRequest.Input case schemas.TranscriptionRequest, schemas.TranscriptionStreamRequest: initialData.Params = req.TranscriptionRequest.Params - initialData.TranscriptionInput = req.TranscriptionRequest.Input + initialData.TranscriptionInput = req.TranscriptionRequest.Input } } diff --git a/plugins/logging/version b/plugins/logging/version index bba60c093..8229aaebc 100644 --- a/plugins/logging/version +++ b/plugins/logging/version @@ -1 +1 @@ -1.3.48 \ No newline at end of file +1.3.49 \ No newline at end of file diff --git a/plugins/maxim/changelog.md b/plugins/maxim/changelog.md index e69de29bb..c9b9868c3 100644 --- a/plugins/maxim/changelog.md +++ b/plugins/maxim/changelog.md @@ -0,0 +1 @@ +- chore: upgrades core to 1.2.38 and framework to 1.1.48 \ No newline at end of file diff --git a/plugins/maxim/go.mod b/plugins/maxim/go.mod index c9c8f11f7..4d1e87973 100644 --- a/plugins/maxim/go.mod +++ b/plugins/maxim/go.mod @@ -84,7 +84,7 @@ require ( github.com/twitchyliquid64/golang-asm v0.15.1 // indirect github.com/valyala/bytebufferpool v1.0.0 // indirect github.com/valyala/fasthttp v1.67.0 // indirect - github.com/weaviate/weaviate v1.33.1 // indirect + github.com/weaviate/weaviate v1.33.4 // indirect github.com/weaviate/weaviate-go-client/v5 v5.5.0 // indirect github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect github.com/yosida95/uritemplate/v3 v3.0.2 // indirect diff --git a/plugins/maxim/go.sum b/plugins/maxim/go.sum index ac6ce47b9..353445df4 100644 --- a/plugins/maxim/go.sum +++ b/plugins/maxim/go.sum @@ -187,8 +187,7 @@ github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6Kllzaw github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc= github.com/valyala/fasthttp v1.67.0 h1:tqKlJMUP6iuNG8hGjK/s9J4kadH7HLV4ijEcPGsezac= github.com/valyala/fasthttp v1.67.0/go.mod h1:qYSIpqt/0XNmShgo/8Aq8E3UYWVVwNS2QYmzd8WIEPM= -github.com/weaviate/weaviate v1.33.1 h1:fV69ffJSH0aO3LvLiKYlVZ8wFa94oQ1g3uMyZGTb838= -github.com/weaviate/weaviate v1.33.1/go.mod h1:SnxXSIoiusZttZ/gI9knXhFAu0UYqn9N/ekgsNnXbNw= +github.com/weaviate/weaviate v1.33.4 h1:eA37l538+3pEBJAZ3/mFBHG0IEWcj5/aAJSzFhrYuUA= github.com/weaviate/weaviate-go-client/v5 v5.5.0 h1:+5qkHodrL3/Qc7kXvMXnDaIxSBN5+djivLqzmCx7VS4= github.com/weaviate/weaviate-go-client/v5 v5.5.0/go.mod h1:Zdm2MEXG27I0Nf6fM0FZ3P2vLR4JM0iJZrOxwc+Zj34= github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc= diff --git a/plugins/maxim/version b/plugins/maxim/version index 15539864f..0750769ee 100644 --- a/plugins/maxim/version +++ b/plugins/maxim/version @@ -1 +1 @@ -1.4.48 \ No newline at end of file +1.4.49 \ No newline at end of file diff --git a/plugins/mocker/changelog.md b/plugins/mocker/changelog.md index e69de29bb..c9b9868c3 100644 --- a/plugins/mocker/changelog.md +++ b/plugins/mocker/changelog.md @@ -0,0 +1 @@ +- chore: upgrades core to 1.2.38 and framework to 1.1.48 \ No newline at end of file diff --git a/plugins/mocker/version b/plugins/mocker/version index fea650a6e..bba60c093 100644 --- a/plugins/mocker/version +++ b/plugins/mocker/version @@ -1 +1 @@ -1.3.47 \ No newline at end of file +1.3.48 \ No newline at end of file diff --git a/plugins/otel/changelog.md b/plugins/otel/changelog.md index e69de29bb..1dfed1bf4 100644 --- a/plugins/otel/changelog.md +++ b/plugins/otel/changelog.md @@ -0,0 +1,2 @@ +- feat: add batch and file request logging support; refactor centralized request handling +- chore: upgrades core to 1.2.38 and framework to 1.1.48 \ No newline at end of file diff --git a/plugins/otel/converter.go b/plugins/otel/converter.go index 1c27b44a5..d113b2bf5 100644 --- a/plugins/otel/converter.go +++ b/plugins/otel/converter.go @@ -367,6 +367,188 @@ func getResponsesRequestParams(req *schemas.BifrostResponsesRequest) []*KeyValue return params } +// getFileUploadRequestParams handles the file upload request +func getFileUploadRequestParams(req *schemas.BifrostFileUploadRequest) []*KeyValue { + params := []*KeyValue{} + if req.Filename != "" { + params = append(params, kvStr("gen_ai.file.filename", req.Filename)) + } + if req.Purpose != "" { + params = append(params, kvStr("gen_ai.file.purpose", string(req.Purpose))) + } + if len(req.File) > 0 { + params = append(params, kvInt("gen_ai.file.bytes", int64(len(req.File)))) + } + if req.ExtraParams != nil { + for k, v := range req.ExtraParams { + params = append(params, kvStr(k, fmt.Sprintf("%v", v))) + } + } + return params +} + +// getFileListRequestParams handles the file list request +func getFileListRequestParams(req *schemas.BifrostFileListRequest) []*KeyValue { + params := []*KeyValue{} + if req.Purpose != "" { + params = append(params, kvStr("gen_ai.file.purpose", string(req.Purpose))) + } + if req.Limit > 0 { + params = append(params, kvInt("gen_ai.file.limit", int64(req.Limit))) + } + if req.After != nil { + params = append(params, kvStr("gen_ai.file.after", *req.After)) + } + if req.Order != nil { + params = append(params, kvStr("gen_ai.file.order", *req.Order)) + } + if req.ExtraParams != nil { + for k, v := range req.ExtraParams { + params = append(params, kvStr(k, fmt.Sprintf("%v", v))) + } + } + return params +} + +// getFileRetrieveRequestParams handles the file retrieve request +func getFileRetrieveRequestParams(req *schemas.BifrostFileRetrieveRequest) []*KeyValue { + params := []*KeyValue{} + if req.FileID != "" { + params = append(params, kvStr("gen_ai.file.file_id", req.FileID)) + } + if req.ExtraParams != nil { + for k, v := range req.ExtraParams { + params = append(params, kvStr(k, fmt.Sprintf("%v", v))) + } + } + return params +} + +// getFileDeleteRequestParams handles the file delete request +func getFileDeleteRequestParams(req *schemas.BifrostFileDeleteRequest) []*KeyValue { + params := []*KeyValue{} + if req.FileID != "" { + params = append(params, kvStr("gen_ai.file.file_id", req.FileID)) + } + if req.ExtraParams != nil { + for k, v := range req.ExtraParams { + params = append(params, kvStr(k, fmt.Sprintf("%v", v))) + } + } + return params +} + +// getFileContentRequestParams handles the file content request +func getFileContentRequestParams(req *schemas.BifrostFileContentRequest) []*KeyValue { + params := []*KeyValue{} + if req.FileID != "" { + params = append(params, kvStr("gen_ai.file.file_id", req.FileID)) + } + if req.ExtraParams != nil { + for k, v := range req.ExtraParams { + params = append(params, kvStr(k, fmt.Sprintf("%v", v))) + } + } + return params +} + +// getBatchCreateRequestParams handles the batch create request +func getBatchCreateRequestParams(req *schemas.BifrostBatchCreateRequest) []*KeyValue { + params := []*KeyValue{} + if req.InputFileID != "" { + params = append(params, kvStr("gen_ai.batch.input_file_id", req.InputFileID)) + } + if req.Endpoint != "" { + params = append(params, kvStr("gen_ai.batch.endpoint", string(req.Endpoint))) + } + if req.CompletionWindow != "" { + params = append(params, kvStr("gen_ai.batch.completion_window", req.CompletionWindow)) + } + if len(req.Requests) > 0 { + params = append(params, kvInt("gen_ai.batch.requests_count", int64(len(req.Requests)))) + } + if len(req.Metadata) > 0 { + params = append(params, kvStr("gen_ai.batch.metadata", fmt.Sprintf("%v", req.Metadata))) + } + if req.ExtraParams != nil { + for k, v := range req.ExtraParams { + params = append(params, kvStr(k, fmt.Sprintf("%v", v))) + } + } + return params +} + +// getBatchListRequestParams handles the batch list request +func getBatchListRequestParams(req *schemas.BifrostBatchListRequest) []*KeyValue { + params := []*KeyValue{} + if req.Limit > 0 { + params = append(params, kvInt("gen_ai.batch.limit", int64(req.Limit))) + } + if req.After != nil { + params = append(params, kvStr("gen_ai.batch.after", *req.After)) + } + if req.BeforeID != nil { + params = append(params, kvStr("gen_ai.batch.before_id", *req.BeforeID)) + } + if req.AfterID != nil { + params = append(params, kvStr("gen_ai.batch.after_id", *req.AfterID)) + } + if req.PageToken != nil { + params = append(params, kvStr("gen_ai.batch.page_token", *req.PageToken)) + } + if req.PageSize > 0 { + params = append(params, kvInt("gen_ai.batch.page_size", int64(req.PageSize))) + } + if req.ExtraParams != nil { + for k, v := range req.ExtraParams { + params = append(params, kvStr(k, fmt.Sprintf("%v", v))) + } + } + return params +} + +// getBatchRetrieveRequestParams handles the batch retrieve request +func getBatchRetrieveRequestParams(req *schemas.BifrostBatchRetrieveRequest) []*KeyValue { + params := []*KeyValue{} + if req.BatchID != "" { + params = append(params, kvStr("gen_ai.batch.batch_id", req.BatchID)) + } + if req.ExtraParams != nil { + for k, v := range req.ExtraParams { + params = append(params, kvStr(k, fmt.Sprintf("%v", v))) + } + } + return params +} + +// getBatchCancelRequestParams handles the batch cancel request +func getBatchCancelRequestParams(req *schemas.BifrostBatchCancelRequest) []*KeyValue { + params := []*KeyValue{} + if req.BatchID != "" { + params = append(params, kvStr("gen_ai.batch.batch_id", req.BatchID)) + } + if req.ExtraParams != nil { + for k, v := range req.ExtraParams { + params = append(params, kvStr(k, fmt.Sprintf("%v", v))) + } + } + return params +} + +// getBatchResultsRequestParams handles the batch results request +func getBatchResultsRequestParams(req *schemas.BifrostBatchResultsRequest) []*KeyValue { + params := []*KeyValue{} + if req.BatchID != "" { + params = append(params, kvStr("gen_ai.batch.batch_id", req.BatchID)) + } + if req.ExtraParams != nil { + for k, v := range req.ExtraParams { + params = append(params, kvStr(k, fmt.Sprintf("%v", v))) + } + } + return params +} + // createResourceSpan creates a new resource span for a Bifrost request func (p *OtelPlugin) createResourceSpan(traceID, spanID string, timestamp time.Time, req *schemas.BifrostRequest) *ResourceSpan { provider, model, _ := req.GetRequestFields() @@ -396,6 +578,36 @@ func (p *OtelPlugin) createResourceSpan(traceID, spanID string, timestamp time.T case schemas.ResponsesRequest, schemas.ResponsesStreamRequest: spanName = "gen_ai.responses" params = append(params, getResponsesRequestParams(req.ResponsesRequest)...) + case schemas.BatchCreateRequest: + spanName = "gen_ai.batch.create" + params = append(params, getBatchCreateRequestParams(req.BatchCreateRequest)...) + case schemas.BatchListRequest: + spanName = "gen_ai.batch.list" + params = append(params, getBatchListRequestParams(req.BatchListRequest)...) + case schemas.BatchRetrieveRequest: + spanName = "gen_ai.batch.retrieve" + params = append(params, getBatchRetrieveRequestParams(req.BatchRetrieveRequest)...) + case schemas.BatchCancelRequest: + spanName = "gen_ai.batch.cancel" + params = append(params, getBatchCancelRequestParams(req.BatchCancelRequest)...) + case schemas.BatchResultsRequest: + spanName = "gen_ai.batch.results" + params = append(params, getBatchResultsRequestParams(req.BatchResultsRequest)...) + case schemas.FileUploadRequest: + spanName = "gen_ai.file.upload" + params = append(params, getFileUploadRequestParams(req.FileUploadRequest)...) + case schemas.FileListRequest: + spanName = "gen_ai.file.list" + params = append(params, getFileListRequestParams(req.FileListRequest)...) + case schemas.FileRetrieveRequest: + spanName = "gen_ai.file.retrieve" + params = append(params, getFileRetrieveRequestParams(req.FileRetrieveRequest)...) + case schemas.FileDeleteRequest: + spanName = "gen_ai.file.delete" + params = append(params, getFileDeleteRequestParams(req.FileDeleteRequest)...) + case schemas.FileContentRequest: + spanName = "gen_ai.file.content" + params = append(params, getFileContentRequestParams(req.FileContentRequest)...) } attributes := append(p.attributesFromEnvironment, kvStr("service.name", p.serviceName), kvStr("service.version", p.bifrostVersion)) // Preparing final resource span @@ -660,6 +872,171 @@ func completeResourceSpan( params = append(params, kvInt("gen_ai.usage.input_token_details.audio_tokens", int64(resp.TranscriptionResponse.Usage.InputTokenDetails.AudioTokens))) } } + case resp.BatchCreateResponse != nil: + params = append(params, kvStr("gen_ai.batch.id", resp.BatchCreateResponse.ID)) + params = append(params, kvStr("gen_ai.batch.status", string(resp.BatchCreateResponse.Status))) + if resp.BatchCreateResponse.Object != "" { + params = append(params, kvStr("gen_ai.batch.object", resp.BatchCreateResponse.Object)) + } + if resp.BatchCreateResponse.Endpoint != "" { + params = append(params, kvStr("gen_ai.batch.endpoint", resp.BatchCreateResponse.Endpoint)) + } + if resp.BatchCreateResponse.InputFileID != "" { + params = append(params, kvStr("gen_ai.batch.input_file_id", resp.BatchCreateResponse.InputFileID)) + } + if resp.BatchCreateResponse.CompletionWindow != "" { + params = append(params, kvStr("gen_ai.batch.completion_window", resp.BatchCreateResponse.CompletionWindow)) + } + if resp.BatchCreateResponse.CreatedAt != 0 { + params = append(params, kvInt("gen_ai.batch.created_at", resp.BatchCreateResponse.CreatedAt)) + } + if resp.BatchCreateResponse.ExpiresAt != nil { + params = append(params, kvInt("gen_ai.batch.expires_at", *resp.BatchCreateResponse.ExpiresAt)) + } + if resp.BatchCreateResponse.OutputFileID != nil { + params = append(params, kvStr("gen_ai.batch.output_file_id", *resp.BatchCreateResponse.OutputFileID)) + } + if resp.BatchCreateResponse.ErrorFileID != nil { + params = append(params, kvStr("gen_ai.batch.error_file_id", *resp.BatchCreateResponse.ErrorFileID)) + } + params = append(params, kvInt("gen_ai.batch.request_counts.total", int64(resp.BatchCreateResponse.RequestCounts.Total))) + params = append(params, kvInt("gen_ai.batch.request_counts.completed", int64(resp.BatchCreateResponse.RequestCounts.Completed))) + params = append(params, kvInt("gen_ai.batch.request_counts.failed", int64(resp.BatchCreateResponse.RequestCounts.Failed))) + case resp.BatchListResponse != nil: + if resp.BatchListResponse.Object != "" { + params = append(params, kvStr("gen_ai.batch.object", resp.BatchListResponse.Object)) + } + params = append(params, kvInt("gen_ai.batch.data_count", int64(len(resp.BatchListResponse.Data)))) + params = append(params, kvBool("gen_ai.batch.has_more", resp.BatchListResponse.HasMore)) + if resp.BatchListResponse.FirstID != nil { + params = append(params, kvStr("gen_ai.batch.first_id", *resp.BatchListResponse.FirstID)) + } + if resp.BatchListResponse.LastID != nil { + params = append(params, kvStr("gen_ai.batch.last_id", *resp.BatchListResponse.LastID)) + } + case resp.BatchRetrieveResponse != nil: + params = append(params, kvStr("gen_ai.batch.id", resp.BatchRetrieveResponse.ID)) + params = append(params, kvStr("gen_ai.batch.status", string(resp.BatchRetrieveResponse.Status))) + if resp.BatchRetrieveResponse.Object != "" { + params = append(params, kvStr("gen_ai.batch.object", resp.BatchRetrieveResponse.Object)) + } + if resp.BatchRetrieveResponse.Endpoint != "" { + params = append(params, kvStr("gen_ai.batch.endpoint", resp.BatchRetrieveResponse.Endpoint)) + } + if resp.BatchRetrieveResponse.InputFileID != "" { + params = append(params, kvStr("gen_ai.batch.input_file_id", resp.BatchRetrieveResponse.InputFileID)) + } + if resp.BatchRetrieveResponse.CompletionWindow != "" { + params = append(params, kvStr("gen_ai.batch.completion_window", resp.BatchRetrieveResponse.CompletionWindow)) + } + if resp.BatchRetrieveResponse.CreatedAt != 0 { + params = append(params, kvInt("gen_ai.batch.created_at", resp.BatchRetrieveResponse.CreatedAt)) + } + if resp.BatchRetrieveResponse.ExpiresAt != nil { + params = append(params, kvInt("gen_ai.batch.expires_at", *resp.BatchRetrieveResponse.ExpiresAt)) + } + if resp.BatchRetrieveResponse.InProgressAt != nil { + params = append(params, kvInt("gen_ai.batch.in_progress_at", *resp.BatchRetrieveResponse.InProgressAt)) + } + if resp.BatchRetrieveResponse.FinalizingAt != nil { + params = append(params, kvInt("gen_ai.batch.finalizing_at", *resp.BatchRetrieveResponse.FinalizingAt)) + } + if resp.BatchRetrieveResponse.CompletedAt != nil { + params = append(params, kvInt("gen_ai.batch.completed_at", *resp.BatchRetrieveResponse.CompletedAt)) + } + if resp.BatchRetrieveResponse.FailedAt != nil { + params = append(params, kvInt("gen_ai.batch.failed_at", *resp.BatchRetrieveResponse.FailedAt)) + } + if resp.BatchRetrieveResponse.ExpiredAt != nil { + params = append(params, kvInt("gen_ai.batch.expired_at", *resp.BatchRetrieveResponse.ExpiredAt)) + } + if resp.BatchRetrieveResponse.CancellingAt != nil { + params = append(params, kvInt("gen_ai.batch.cancelling_at", *resp.BatchRetrieveResponse.CancellingAt)) + } + if resp.BatchRetrieveResponse.CancelledAt != nil { + params = append(params, kvInt("gen_ai.batch.cancelled_at", *resp.BatchRetrieveResponse.CancelledAt)) + } + if resp.BatchRetrieveResponse.OutputFileID != nil { + params = append(params, kvStr("gen_ai.batch.output_file_id", *resp.BatchRetrieveResponse.OutputFileID)) + } + if resp.BatchRetrieveResponse.ErrorFileID != nil { + params = append(params, kvStr("gen_ai.batch.error_file_id", *resp.BatchRetrieveResponse.ErrorFileID)) + } + params = append(params, kvInt("gen_ai.batch.request_counts.total", int64(resp.BatchRetrieveResponse.RequestCounts.Total))) + params = append(params, kvInt("gen_ai.batch.request_counts.completed", int64(resp.BatchRetrieveResponse.RequestCounts.Completed))) + params = append(params, kvInt("gen_ai.batch.request_counts.failed", int64(resp.BatchRetrieveResponse.RequestCounts.Failed))) + case resp.BatchCancelResponse != nil: + params = append(params, kvStr("gen_ai.batch.id", resp.BatchCancelResponse.ID)) + params = append(params, kvStr("gen_ai.batch.status", string(resp.BatchCancelResponse.Status))) + if resp.BatchCancelResponse.Object != "" { + params = append(params, kvStr("gen_ai.batch.object", resp.BatchCancelResponse.Object)) + } + if resp.BatchCancelResponse.CancellingAt != nil { + params = append(params, kvInt("gen_ai.batch.cancelling_at", *resp.BatchCancelResponse.CancellingAt)) + } + if resp.BatchCancelResponse.CancelledAt != nil { + params = append(params, kvInt("gen_ai.batch.cancelled_at", *resp.BatchCancelResponse.CancelledAt)) + } + params = append(params, kvInt("gen_ai.batch.request_counts.total", int64(resp.BatchCancelResponse.RequestCounts.Total))) + params = append(params, kvInt("gen_ai.batch.request_counts.completed", int64(resp.BatchCancelResponse.RequestCounts.Completed))) + params = append(params, kvInt("gen_ai.batch.request_counts.failed", int64(resp.BatchCancelResponse.RequestCounts.Failed))) + case resp.BatchResultsResponse != nil: + params = append(params, kvStr("gen_ai.batch.batch_id", resp.BatchResultsResponse.BatchID)) + params = append(params, kvInt("gen_ai.batch.results_count", int64(len(resp.BatchResultsResponse.Results)))) + params = append(params, kvBool("gen_ai.batch.has_more", resp.BatchResultsResponse.HasMore)) + if resp.BatchResultsResponse.NextCursor != nil { + params = append(params, kvStr("gen_ai.batch.next_cursor", *resp.BatchResultsResponse.NextCursor)) + } + case resp.FileUploadResponse != nil: + params = append(params, kvStr("gen_ai.file.id", resp.FileUploadResponse.ID)) + if resp.FileUploadResponse.Object != "" { + params = append(params, kvStr("gen_ai.file.object", resp.FileUploadResponse.Object)) + } + params = append(params, kvInt("gen_ai.file.bytes", resp.FileUploadResponse.Bytes)) + params = append(params, kvInt("gen_ai.file.created_at", resp.FileUploadResponse.CreatedAt)) + params = append(params, kvStr("gen_ai.file.filename", resp.FileUploadResponse.Filename)) + params = append(params, kvStr("gen_ai.file.purpose", string(resp.FileUploadResponse.Purpose))) + if resp.FileUploadResponse.Status != "" { + params = append(params, kvStr("gen_ai.file.status", string(resp.FileUploadResponse.Status))) + } + if resp.FileUploadResponse.StorageBackend != "" { + params = append(params, kvStr("gen_ai.file.storage_backend", string(resp.FileUploadResponse.StorageBackend))) + } + case resp.FileListResponse != nil: + if resp.FileListResponse.Object != "" { + params = append(params, kvStr("gen_ai.file.object", resp.FileListResponse.Object)) + } + params = append(params, kvInt("gen_ai.file.data_count", int64(len(resp.FileListResponse.Data)))) + params = append(params, kvBool("gen_ai.file.has_more", resp.FileListResponse.HasMore)) + case resp.FileRetrieveResponse != nil: + params = append(params, kvStr("gen_ai.file.id", resp.FileRetrieveResponse.ID)) + if resp.FileRetrieveResponse.Object != "" { + params = append(params, kvStr("gen_ai.file.object", resp.FileRetrieveResponse.Object)) + } + params = append(params, kvInt("gen_ai.file.bytes", resp.FileRetrieveResponse.Bytes)) + params = append(params, kvInt("gen_ai.file.created_at", resp.FileRetrieveResponse.CreatedAt)) + params = append(params, kvStr("gen_ai.file.filename", resp.FileRetrieveResponse.Filename)) + params = append(params, kvStr("gen_ai.file.purpose", string(resp.FileRetrieveResponse.Purpose))) + if resp.FileRetrieveResponse.Status != "" { + params = append(params, kvStr("gen_ai.file.status", string(resp.FileRetrieveResponse.Status))) + } + if resp.FileRetrieveResponse.StorageBackend != "" { + params = append(params, kvStr("gen_ai.file.storage_backend", string(resp.FileRetrieveResponse.StorageBackend))) + } + case resp.FileDeleteResponse != nil: + params = append(params, kvStr("gen_ai.file.id", resp.FileDeleteResponse.ID)) + if resp.FileDeleteResponse.Object != "" { + params = append(params, kvStr("gen_ai.file.object", resp.FileDeleteResponse.Object)) + } + params = append(params, kvBool("gen_ai.file.deleted", resp.FileDeleteResponse.Deleted)) + case resp.FileContentResponse != nil: + params = append(params, kvStr("gen_ai.file.file_id", resp.FileContentResponse.FileID)) + if resp.FileContentResponse.ContentType != "" { + params = append(params, kvStr("gen_ai.file.content_type", resp.FileContentResponse.ContentType)) + } + if len(resp.FileContentResponse.Content) > 0 { + params = append(params, kvInt("gen_ai.file.content_bytes", int64(len(resp.FileContentResponse.Content)))) + } } } diff --git a/plugins/otel/go.mod b/plugins/otel/go.mod index afb87b626..89e48776b 100644 --- a/plugins/otel/go.mod +++ b/plugins/otel/go.mod @@ -80,7 +80,7 @@ require ( github.com/spf13/cast v1.10.0 // indirect github.com/valyala/bytebufferpool v1.0.0 // indirect github.com/valyala/fasthttp v1.67.0 // indirect - github.com/weaviate/weaviate v1.33.1 // indirect + github.com/weaviate/weaviate v1.33.4 // indirect github.com/weaviate/weaviate-go-client/v5 v5.5.0 // indirect github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect github.com/yosida95/uritemplate/v3 v3.0.2 // indirect diff --git a/plugins/otel/go.sum b/plugins/otel/go.sum index c1bdc80dd..7906b2d79 100644 --- a/plugins/otel/go.sum +++ b/plugins/otel/go.sum @@ -187,8 +187,7 @@ github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6Kllzaw github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc= github.com/valyala/fasthttp v1.67.0 h1:tqKlJMUP6iuNG8hGjK/s9J4kadH7HLV4ijEcPGsezac= github.com/valyala/fasthttp v1.67.0/go.mod h1:qYSIpqt/0XNmShgo/8Aq8E3UYWVVwNS2QYmzd8WIEPM= -github.com/weaviate/weaviate v1.33.1 h1:fV69ffJSH0aO3LvLiKYlVZ8wFa94oQ1g3uMyZGTb838= -github.com/weaviate/weaviate v1.33.1/go.mod h1:SnxXSIoiusZttZ/gI9knXhFAu0UYqn9N/ekgsNnXbNw= +github.com/weaviate/weaviate v1.33.4 h1:eA37l538+3pEBJAZ3/mFBHG0IEWcj5/aAJSzFhrYuUA= github.com/weaviate/weaviate-go-client/v5 v5.5.0 h1:+5qkHodrL3/Qc7kXvMXnDaIxSBN5+djivLqzmCx7VS4= github.com/weaviate/weaviate-go-client/v5 v5.5.0/go.mod h1:Zdm2MEXG27I0Nf6fM0FZ3P2vLR4JM0iJZrOxwc+Zj34= github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc= diff --git a/plugins/otel/version b/plugins/otel/version index de0434d0b..c30197954 100644 --- a/plugins/otel/version +++ b/plugins/otel/version @@ -1 +1 @@ -1.0.47 \ No newline at end of file +1.0.48 \ No newline at end of file diff --git a/plugins/semanticcache/changelog.md b/plugins/semanticcache/changelog.md index e69de29bb..c9b9868c3 100644 --- a/plugins/semanticcache/changelog.md +++ b/plugins/semanticcache/changelog.md @@ -0,0 +1 @@ +- chore: upgrades core to 1.2.38 and framework to 1.1.48 \ No newline at end of file diff --git a/plugins/semanticcache/go.mod b/plugins/semanticcache/go.mod index c2957643d..078e2837b 100644 --- a/plugins/semanticcache/go.mod +++ b/plugins/semanticcache/go.mod @@ -84,7 +84,7 @@ require ( github.com/twitchyliquid64/golang-asm v0.15.1 // indirect github.com/valyala/bytebufferpool v1.0.0 // indirect github.com/valyala/fasthttp v1.67.0 // indirect - github.com/weaviate/weaviate v1.33.1 // indirect + github.com/weaviate/weaviate v1.33.4 // indirect github.com/weaviate/weaviate-go-client/v5 v5.5.0 // indirect github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect github.com/yosida95/uritemplate/v3 v3.0.2 // indirect diff --git a/plugins/semanticcache/go.sum b/plugins/semanticcache/go.sum index 6063cd879..8839f86f6 100644 --- a/plugins/semanticcache/go.sum +++ b/plugins/semanticcache/go.sum @@ -189,8 +189,7 @@ github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6Kllzaw github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc= github.com/valyala/fasthttp v1.67.0 h1:tqKlJMUP6iuNG8hGjK/s9J4kadH7HLV4ijEcPGsezac= github.com/valyala/fasthttp v1.67.0/go.mod h1:qYSIpqt/0XNmShgo/8Aq8E3UYWVVwNS2QYmzd8WIEPM= -github.com/weaviate/weaviate v1.33.1 h1:fV69ffJSH0aO3LvLiKYlVZ8wFa94oQ1g3uMyZGTb838= -github.com/weaviate/weaviate v1.33.1/go.mod h1:SnxXSIoiusZttZ/gI9knXhFAu0UYqn9N/ekgsNnXbNw= +github.com/weaviate/weaviate v1.33.4 h1:eA37l538+3pEBJAZ3/mFBHG0IEWcj5/aAJSzFhrYuUA= github.com/weaviate/weaviate-go-client/v5 v5.5.0 h1:+5qkHodrL3/Qc7kXvMXnDaIxSBN5+djivLqzmCx7VS4= github.com/weaviate/weaviate-go-client/v5 v5.5.0/go.mod h1:Zdm2MEXG27I0Nf6fM0FZ3P2vLR4JM0iJZrOxwc+Zj34= github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc= diff --git a/plugins/semanticcache/version b/plugins/semanticcache/version index fea650a6e..bba60c093 100644 --- a/plugins/semanticcache/version +++ b/plugins/semanticcache/version @@ -1 +1 @@ -1.3.47 \ No newline at end of file +1.3.48 \ No newline at end of file diff --git a/plugins/telemetry/changelog.md b/plugins/telemetry/changelog.md index e69de29bb..bd9b266bb 100644 --- a/plugins/telemetry/changelog.md +++ b/plugins/telemetry/changelog.md @@ -0,0 +1,2 @@ +- feat: adds logging support for batch and file requests +- chore: upgrades core to 1.2.38 and framework to 1.1.48 \ No newline at end of file diff --git a/plugins/telemetry/go.mod b/plugins/telemetry/go.mod index e805c2793..406bdafbc 100644 --- a/plugins/telemetry/go.mod +++ b/plugins/telemetry/go.mod @@ -88,7 +88,7 @@ require ( github.com/spf13/cast v1.10.0 // indirect github.com/twitchyliquid64/golang-asm v0.15.1 // indirect github.com/valyala/bytebufferpool v1.0.0 // indirect - github.com/weaviate/weaviate v1.33.1 // indirect + github.com/weaviate/weaviate v1.33.4 // indirect github.com/weaviate/weaviate-go-client/v5 v5.5.0 // indirect github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect github.com/yosida95/uritemplate/v3 v3.0.2 // indirect diff --git a/plugins/telemetry/go.sum b/plugins/telemetry/go.sum index 39404a855..42e248b34 100644 --- a/plugins/telemetry/go.sum +++ b/plugins/telemetry/go.sum @@ -197,8 +197,7 @@ github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6Kllzaw github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc= github.com/valyala/fasthttp v1.67.0 h1:tqKlJMUP6iuNG8hGjK/s9J4kadH7HLV4ijEcPGsezac= github.com/valyala/fasthttp v1.67.0/go.mod h1:qYSIpqt/0XNmShgo/8Aq8E3UYWVVwNS2QYmzd8WIEPM= -github.com/weaviate/weaviate v1.33.1 h1:fV69ffJSH0aO3LvLiKYlVZ8wFa94oQ1g3uMyZGTb838= -github.com/weaviate/weaviate v1.33.1/go.mod h1:SnxXSIoiusZttZ/gI9knXhFAu0UYqn9N/ekgsNnXbNw= +github.com/weaviate/weaviate v1.33.4 h1:eA37l538+3pEBJAZ3/mFBHG0IEWcj5/aAJSzFhrYuUA= github.com/weaviate/weaviate-go-client/v5 v5.5.0 h1:+5qkHodrL3/Qc7kXvMXnDaIxSBN5+djivLqzmCx7VS4= github.com/weaviate/weaviate-go-client/v5 v5.5.0/go.mod h1:Zdm2MEXG27I0Nf6fM0FZ3P2vLR4JM0iJZrOxwc+Zj34= github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc= diff --git a/plugins/telemetry/version b/plugins/telemetry/version index fea650a6e..bba60c093 100644 --- a/plugins/telemetry/version +++ b/plugins/telemetry/version @@ -1 +1 @@ -1.3.47 \ No newline at end of file +1.3.48 \ No newline at end of file diff --git a/tests/integrations/config.json b/tests/integrations/config.json index 50fd0d0b0..a12220579 100644 --- a/tests/integrations/config.json +++ b/tests/integrations/config.json @@ -37,9 +37,9 @@ "vertex": { "keys": [ { - "vertex_key_config": { - "projectId": "env.GOOGLE_PROJECT_ID", - "location": "env.GOOGLE_LOCATION" + "vertex_key_config": { + "project_id": "env.GOOGLE_PROJECT_ID", + "region": "env.GOOGLE_LOCATION" }, "weight": 1 } @@ -135,7 +135,7 @@ "bedrock_key_config": { "access_key": "env.AWS_ACCESS_KEY_ID", "secret_key": "env.AWS_SECRET_ACCESS_KEY", - "region": "env.AWS_REGION" + "region": "env.AWS_REGION" }, "weight": 1 } @@ -148,10 +148,14 @@ "config_store": { "enabled": false }, - "log_store": { - "enabled": false + "logs_store": { + "enabled": true, + "type": "sqlite", + "config": { + "path": "./logs.db" + } }, - "client_config": { + "client": { "drop_excess_requests": false, "initial_pool_size": 300, "allowed_origins": [ @@ -160,7 +164,7 @@ "enable_logging": true, "enable_governance": false, "enforce_governance_header": false, - "allow_direct_keys": true, + "allow_direct_keys": false, "max_request_body_size_mb": 100, "enable_litellm_fallbacks": false } diff --git a/tests/integrations/tests/test_anthropic.py b/tests/integrations/tests/test_anthropic.py index ec8521f7d..c9e817c47 100644 --- a/tests/integrations/tests/test_anthropic.py +++ b/tests/integrations/tests/test_anthropic.py @@ -1139,7 +1139,7 @@ def test_22_batch_create_inline(self, anthropic_client, test_config, provider, m """ if provider == "_no_providers_" or model == "_no_model_": pytest.skip("No providers configured for batch_inline scenario") - + # Get provider-specific client client = get_provider_anthropic_client(provider) diff --git a/tests/integrations/tests/test_bedrock.py b/tests/integrations/tests/test_bedrock.py index aff0babaa..4df2d3167 100644 --- a/tests/integrations/tests/test_bedrock.py +++ b/tests/integrations/tests/test_bedrock.py @@ -1493,8 +1493,12 @@ def test_19_batch_cancel(self, test_config, provider, model): assert create_response["jobArn"], "jobArn should not be empty" assert create_response["jobArn"].startswith("arn:"), "jobArn should be a valid ARN" + print(f"create_response: {create_response}") + job_arn = create_response["jobArn"] + print(f"stopping the job") + # Cancel the job bedrock_client.stop_model_invocation_job(jobIdentifier=job_arn) @@ -1547,6 +1551,9 @@ def test_20_batch_e2e(self, test_config, provider, model): # Get provider-specific clients with x-model-provider header s3_client = get_provider_s3_client(provider) + + print(f"getting the bedrock client for provider {provider}") + bedrock_client = get_provider_bedrock_batch_client(provider) # Step 1: Upload input file to S3 diff --git a/transports/bifrost-http/handlers/inference.go b/transports/bifrost-http/handlers/inference.go index 78ac172b3..d0806c3d3 100644 --- a/transports/bifrost-http/handlers/inference.go +++ b/transports/bifrost-http/handlers/inference.go @@ -156,26 +156,6 @@ var batchCreateParamsKnownFields = map[string]bool{ "metadata": true, } -var batchListParamsKnownFields = map[string]bool{ - "provider": true, - "limit": true, - "after": true, - "before": true, -} - -var fileUploadParamsKnownFields = map[string]bool{ - "purpose": true, - "file": true, -} - -var fileListParamsKnownFields = map[string]bool{ - "provider": true, - "purpose": true, - "limit": true, - "after": true, - "order": true, -} - type BifrostParams struct { Model string `json:"model"` // Model to use in "provider/model" format Fallbacks []string `json:"fallbacks"` // Fallback providers and models in "provider/model" format diff --git a/transports/bifrost-http/handlers/middlewares.go b/transports/bifrost-http/handlers/middlewares.go index 28aa9e436..e2c4e7737 100644 --- a/transports/bifrost-http/handlers/middlewares.go +++ b/transports/bifrost-http/handlers/middlewares.go @@ -79,14 +79,12 @@ func TransportInterceptorMiddleware(config *lib.Config) lib.BifrostHTTPMiddlewar return true }) - - // Unmarshal request body requestBody := make(map[string]any) bodyBytes := ctx.Request.Body() - if len(bodyBytes) > 0 { + if len(bodyBytes) > 0 && strings.HasPrefix(string(ctx.Request.Header.Peek("Content-Type")), "application/json") { if err := json.Unmarshal(bodyBytes, &requestBody); err != nil { // If body is not valid JSON, log warning and continue without interception - logger.Warn(fmt.Sprintf("TransportInterceptor: Failed to unmarshal request body: %v, skipping interceptor", err)) + logger.Warn(fmt.Sprintf("[transportInterceptor]: Failed to unmarshal request body: %v, skipping interceptor", err)) next(ctx) return } diff --git a/transports/bifrost-http/integrations/anthropic.go b/transports/bifrost-http/integrations/anthropic.go index 91f6d3b57..523c201b3 100644 --- a/transports/bifrost-http/integrations/anthropic.go +++ b/transports/bifrost-http/integrations/anthropic.go @@ -287,13 +287,13 @@ func CreateAnthropicBatchRouteConfigs(pathPrefix string, handlerStore lib.Handle GetRequestTypeInstance: func() any { return &anthropic.AnthropicBatchCreateRequest{} }, - BatchCreateRequestConverter: func(ctx *context.Context, req any) (*BatchRequest, error) { + BatchRequestConverter: func(ctx *context.Context, req any) (*BatchRequest, error) { if anthropicReq, ok := req.(*anthropic.AnthropicBatchCreateRequest); ok { // Convert Anthropic batch request items to Bifrost format isNonAnthropicProvider := false var provider schemas.ModelProvider var ok bool - if provider, ok = (*ctx).Value(schemas.BifrostContextKey("batch_provider")).(schemas.ModelProvider); ok && provider != schemas.Anthropic { + if provider, ok = (*ctx).Value(bifrostContextKeyProvider).(schemas.ModelProvider); ok && provider != schemas.Anthropic { isNonAnthropicProvider = true } var model string @@ -362,7 +362,7 @@ func CreateAnthropicBatchRouteConfigs(pathPrefix string, handlerStore lib.Handle GetRequestTypeInstance: func() interface{} { return &anthropic.AnthropicBatchListRequest{} }, - BatchCreateRequestConverter: func(ctx *context.Context, req interface{}) (*BatchRequest, error) { + BatchRequestConverter: func(ctx *context.Context, req interface{}) (*BatchRequest, error) { if listReq, ok := req.(*anthropic.AnthropicBatchListRequest); ok { provider, ok := (*ctx).Value(bifrostContextKeyProvider).(schemas.ModelProvider) if !ok { @@ -404,7 +404,7 @@ func CreateAnthropicBatchRouteConfigs(pathPrefix string, handlerStore lib.Handle GetRequestTypeInstance: func() interface{} { return &anthropic.AnthropicBatchRetrieveRequest{} }, - BatchCreateRequestConverter: func(ctx *context.Context, req interface{}) (*BatchRequest, error) { + BatchRequestConverter: func(ctx *context.Context, req interface{}) (*BatchRequest, error) { if retrieveReq, ok := req.(*anthropic.AnthropicBatchRetrieveRequest); ok { provider := (*ctx).Value(bifrostContextKeyProvider).(schemas.ModelProvider) if provider == schemas.Gemini { @@ -440,7 +440,7 @@ func CreateAnthropicBatchRouteConfigs(pathPrefix string, handlerStore lib.Handle GetRequestTypeInstance: func() any { return &anthropic.AnthropicBatchCancelRequest{} }, - BatchCreateRequestConverter: func(ctx *context.Context, req interface{}) (*BatchRequest, error) { + BatchRequestConverter: func(ctx *context.Context, req interface{}) (*BatchRequest, error) { if cancelReq, ok := req.(*anthropic.AnthropicBatchCancelRequest); ok { provider := (*ctx).Value(bifrostContextKeyProvider).(schemas.ModelProvider) if provider == schemas.Gemini { @@ -476,7 +476,7 @@ func CreateAnthropicBatchRouteConfigs(pathPrefix string, handlerStore lib.Handle GetRequestTypeInstance: func() interface{} { return &anthropic.AnthropicBatchResultsRequest{} }, - BatchCreateRequestConverter: func(ctx *context.Context, req interface{}) (*BatchRequest, error) { + BatchRequestConverter: func(ctx *context.Context, req interface{}) (*BatchRequest, error) { if resultsReq, ok := req.(*anthropic.AnthropicBatchResultsRequest); ok { provider := (*ctx).Value(bifrostContextKeyProvider).(schemas.ModelProvider) if provider == schemas.Gemini { @@ -515,7 +515,7 @@ func extractAnthropicBatchCreateParams(ctx *fasthttp.RequestCtx, bifrostCtx *con provider = string(schemas.Anthropic) } // Store provider in context for batch create converter to use - *bifrostCtx = context.WithValue(*bifrostCtx, schemas.BifrostContextKey("batch_provider"), schemas.ModelProvider(provider)) + *bifrostCtx = context.WithValue(*bifrostCtx, bifrostContextKeyProvider, schemas.ModelProvider(provider)) return nil } diff --git a/transports/bifrost-http/integrations/bedrock.go b/transports/bifrost-http/integrations/bedrock.go index 74f2fef07..696ba3b67 100644 --- a/transports/bifrost-http/integrations/bedrock.go +++ b/transports/bifrost-http/integrations/bedrock.go @@ -201,7 +201,7 @@ func createBedrockBatchRouteConfigs(pathPrefix string, handlerStore lib.HandlerS GetRequestTypeInstance: func() interface{} { return &bedrock.BedrockBatchJobRequest{} }, - BatchCreateRequestConverter: func(ctx *context.Context, req interface{}) (*BatchRequest, error) { + BatchRequestConverter: func(ctx *context.Context, req interface{}) (*BatchRequest, error) { if bedrockReq, ok := req.(*bedrock.BedrockBatchJobRequest); ok { provider := (*ctx).Value(bifrostContextKeyProvider).(schemas.ModelProvider) @@ -300,7 +300,7 @@ func createBedrockBatchRouteConfigs(pathPrefix string, handlerStore lib.HandlerS GetRequestTypeInstance: func() interface{} { return &bedrock.BedrockBatchListRequest{} }, - BatchCreateRequestConverter: func(ctx *context.Context, req interface{}) (*BatchRequest, error) { + BatchRequestConverter: func(ctx *context.Context, req interface{}) (*BatchRequest, error) { if bedrockReq, ok := req.(*bedrock.BedrockBatchListRequest); ok { provider := (*ctx).Value(bifrostContextKeyProvider).(schemas.ModelProvider) bifrostReq := bedrock.ToBifrostBatchListRequest(bedrockReq, provider) @@ -342,7 +342,7 @@ func createBedrockBatchRouteConfigs(pathPrefix string, handlerStore lib.HandlerS GetRequestTypeInstance: func() interface{} { return &bedrock.BedrockBatchRetrieveRequest{} }, - BatchCreateRequestConverter: func(ctx *context.Context, req interface{}) (*BatchRequest, error) { + BatchRequestConverter: func(ctx *context.Context, req interface{}) (*BatchRequest, error) { if bedrockReq, ok := req.(*bedrock.BedrockBatchRetrieveRequest); ok { provider := (*ctx).Value(bifrostContextKeyProvider).(schemas.ModelProvider) bifrostReq := bedrock.ToBifrostBatchRetrieveRequest(bedrockReq, provider) @@ -384,7 +384,7 @@ func createBedrockBatchRouteConfigs(pathPrefix string, handlerStore lib.HandlerS GetRequestTypeInstance: func() interface{} { return &bedrock.BedrockBatchCancelRequest{} }, - BatchCreateRequestConverter: func(ctx *context.Context, req interface{}) (*BatchRequest, error) { + BatchRequestConverter: func(ctx *context.Context, req interface{}) (*BatchRequest, error) { if bedrockReq, ok := req.(*bedrock.BedrockBatchCancelRequest); ok { provider := (*ctx).Value(bifrostContextKeyProvider).(schemas.ModelProvider) bifrostReq := bedrock.ToBifrostBatchCancelRequest(bedrockReq, provider) diff --git a/transports/bifrost-http/integrations/bedrock_test.go b/transports/bifrost-http/integrations/bedrock_test.go index 0aecdfcfe..2556676b0 100644 --- a/transports/bifrost-http/integrations/bedrock_test.go +++ b/transports/bifrost-http/integrations/bedrock_test.go @@ -188,7 +188,7 @@ func Test_createBedrockBatchRouteConfigs(t *testing.T) { assert.Equal(t, expected.method, routes[i].Method, "batch route %d method mismatch", i) assert.Equal(t, RouteConfigTypeBedrock, routes[i].Type, "batch route %d type mismatch", i) assert.NotNil(t, routes[i].GetRequestTypeInstance, "batch route %d GetRequestTypeInstance should not be nil", i) - assert.NotNil(t, routes[i].BatchCreateRequestConverter, "batch route %d BatchCreateRequestConverter should not be nil", i) + assert.NotNil(t, routes[i].BatchRequestConverter, "batch route %d BatchCreateRequestConverter should not be nil", i) assert.NotNil(t, routes[i].ErrorConverter, "batch route %d ErrorConverter should not be nil", i) assert.NotNil(t, routes[i].PreCallback, "batch route %d PreCallback should not be nil", i) } diff --git a/transports/bifrost-http/integrations/openai.go b/transports/bifrost-http/integrations/openai.go index c30f8b70c..ebdfaf069 100644 --- a/transports/bifrost-http/integrations/openai.go +++ b/transports/bifrost-http/integrations/openai.go @@ -459,7 +459,7 @@ func CreateOpenAIBatchRouteConfigs(pathPrefix string, handlerStore lib.HandlerSt GetRequestTypeInstance: func() interface{} { return &schemas.BifrostBatchCreateRequest{} }, - BatchCreateRequestConverter: func(ctx *context.Context, req interface{}) (*BatchRequest, error) { + BatchRequestConverter: func(ctx *context.Context, req interface{}) (*BatchRequest, error) { if openaiReq, ok := req.(*schemas.BifrostBatchCreateRequest); ok { switch openaiReq.Provider { case schemas.Gemini: @@ -568,7 +568,7 @@ func CreateOpenAIBatchRouteConfigs(pathPrefix string, handlerStore lib.HandlerSt GetRequestTypeInstance: func() interface{} { return &schemas.BifrostBatchListRequest{} }, - BatchCreateRequestConverter: func(ctx *context.Context, req interface{}) (*BatchRequest, error) { + BatchRequestConverter: func(ctx *context.Context, req interface{}) (*BatchRequest, error) { if listReq, ok := req.(*schemas.BifrostBatchListRequest); ok { if listReq.Provider == "" { listReq.Provider = schemas.OpenAI @@ -614,7 +614,7 @@ func CreateOpenAIBatchRouteConfigs(pathPrefix string, handlerStore lib.HandlerSt GetRequestTypeInstance: func() interface{} { return &schemas.BifrostBatchRetrieveRequest{} }, - BatchCreateRequestConverter: func(ctx *context.Context, req interface{}) (*BatchRequest, error) { + BatchRequestConverter: func(ctx *context.Context, req interface{}) (*BatchRequest, error) { if retrieveReq, ok := req.(*schemas.BifrostBatchRetrieveRequest); ok { if retrieveReq.Provider == "" { retrieveReq.Provider = schemas.OpenAI @@ -665,7 +665,7 @@ func CreateOpenAIBatchRouteConfigs(pathPrefix string, handlerStore lib.HandlerSt GetRequestTypeInstance: func() interface{} { return &schemas.BifrostBatchCancelRequest{} }, - BatchCreateRequestConverter: func(ctx *context.Context, req interface{}) (*BatchRequest, error) { + BatchRequestConverter: func(ctx *context.Context, req interface{}) (*BatchRequest, error) { if cancelReq, ok := req.(*schemas.BifrostBatchCancelRequest); ok { if cancelReq.Provider == "" { cancelReq.Provider = schemas.OpenAI diff --git a/transports/bifrost-http/integrations/router.go b/transports/bifrost-http/integrations/router.go index 8eff01453..d784334c3 100644 --- a/transports/bifrost-http/integrations/router.go +++ b/transports/bifrost-http/integrations/router.go @@ -97,8 +97,8 @@ type FileRequest struct { ContentRequest *schemas.BifrostFileContentRequest } -// BatchCreateRequestConverter is a function that converts integration-specific batch requests to Bifrost format. -type BatchCreateRequestConverter func(ctx *context.Context, req interface{}) (*BatchRequest, error) +// BatchRequestConverter is a function that converts integration-specific batch requests to Bifrost format. +type BatchRequestConverter func(ctx *context.Context, req interface{}) (*BatchRequest, error) // FileRequestConverter is a function that converts integration-specific file requests to Bifrost format. type FileRequestConverter func(ctx *context.Context, req interface{}) (*FileRequest, error) @@ -269,7 +269,7 @@ type RouteConfig struct { GetRequestTypeInstance func() interface{} // Factory function to create request instance (SHOULD NOT BE NIL) RequestParser RequestParser // Optional: custom request parsing (e.g., multipart/form-data) RequestConverter RequestConverter // Function to convert request to BifrostRequest (for inference requests) - BatchCreateRequestConverter BatchCreateRequestConverter // Function to convert request to BatchRequest (for batch operations) + BatchRequestConverter BatchRequestConverter // Function to convert request to BatchRequest (for batch operations) FileRequestConverter FileRequestConverter // Function to convert request to FileRequest (for file operations) ListModelsResponseConverter ListModelsResponseConverter // Function to convert BifrostListModelsResponse to integration format (SHOULD NOT BE NIL) TextResponseConverter TextResponseConverter // Function to convert BifrostTextCompletionResponse to integration format (SHOULD NOT BE NIL) @@ -334,7 +334,7 @@ func (g *GenericRouter) RegisterRoutes(r *router.Router, middlewares ...lib.Bifr } // Determine route type: inference, batch, or file - isBatchRoute := route.BatchCreateRequestConverter != nil + isBatchRoute := route.BatchRequestConverter != nil isFileRoute := route.FileRequestConverter != nil isInferenceRoute := !isBatchRoute && !isFileRoute @@ -432,9 +432,9 @@ func (g *GenericRouter) createHandler(config RouteConfig) fasthttp.RequestHandle } // Handle batch requests if BatchRequestConverter is set - if config.BatchCreateRequestConverter != nil { + if config.BatchRequestConverter != nil { defer cancel() - batchReq, err := config.BatchCreateRequestConverter(bifrostCtx, req) + batchReq, err := config.BatchRequestConverter(bifrostCtx, req) if err != nil { g.sendError(ctx, bifrostCtx, config.ErrorConverter, newBifrostError(err, "failed to convert batch request")) return diff --git a/transports/changelog.md b/transports/changelog.md index 8f2746e97..4a486d36e 100644 --- a/transports/changelog.md +++ b/transports/changelog.md @@ -1 +1,7 @@ -feat: add support for enabling/disabling provider keys without deletion \ No newline at end of file +- feat: add support for enabling/disabling provider keys without deletion. +- feat: add batch api support for OpenAI, Anthropic, Google Gemini and Bedrock Beta. +- feat: new provider support - nebius. +- feat: force refresh datasheet support. +- fix: fixed minor issues with structured output support for Gemini and Bedrock. +- fix: proper cost compute for gemini models (>200k and <200k token costs are now considered). +- chore: CORS policy now allows `x-stainless-timeout` \ No newline at end of file diff --git a/transports/go.mod b/transports/go.mod index 3bee70db3..8bbbca8b2 100644 --- a/transports/go.mod +++ b/transports/go.mod @@ -106,7 +106,7 @@ require ( github.com/spf13/cast v1.10.0 // indirect github.com/twitchyliquid64/golang-asm v0.15.1 // indirect github.com/valyala/bytebufferpool v1.0.0 // indirect - github.com/weaviate/weaviate v1.33.1 // indirect + github.com/weaviate/weaviate v1.33.4 // indirect github.com/weaviate/weaviate-go-client/v5 v5.5.0 // indirect github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect github.com/yosida95/uritemplate/v3 v3.0.2 // indirect diff --git a/transports/go.sum b/transports/go.sum index 8b6a63b13..6780af1b9 100644 --- a/transports/go.sum +++ b/transports/go.sum @@ -225,8 +225,7 @@ github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6Kllzaw github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc= github.com/valyala/fasthttp v1.67.0 h1:tqKlJMUP6iuNG8hGjK/s9J4kadH7HLV4ijEcPGsezac= github.com/valyala/fasthttp v1.67.0/go.mod h1:qYSIpqt/0XNmShgo/8Aq8E3UYWVVwNS2QYmzd8WIEPM= -github.com/weaviate/weaviate v1.33.1 h1:fV69ffJSH0aO3LvLiKYlVZ8wFa94oQ1g3uMyZGTb838= -github.com/weaviate/weaviate v1.33.1/go.mod h1:SnxXSIoiusZttZ/gI9knXhFAu0UYqn9N/ekgsNnXbNw= +github.com/weaviate/weaviate v1.33.4 h1:eA37l538+3pEBJAZ3/mFBHG0IEWcj5/aAJSzFhrYuUA= github.com/weaviate/weaviate-go-client/v5 v5.5.0 h1:+5qkHodrL3/Qc7kXvMXnDaIxSBN5+djivLqzmCx7VS4= github.com/weaviate/weaviate-go-client/v5 v5.5.0/go.mod h1:Zdm2MEXG27I0Nf6fM0FZ3P2vLR4JM0iJZrOxwc+Zj34= github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc= diff --git a/transports/version b/transports/version index bba60c093..8229aaebc 100644 --- a/transports/version +++ b/transports/version @@ -1 +1 @@ -1.3.48 \ No newline at end of file +1.3.49 \ No newline at end of file diff --git a/ui/lib/constants/logs.ts b/ui/lib/constants/logs.ts index 2d47afcfe..bc6db7c7b 100644 --- a/ui/lib/constants/logs.ts +++ b/ui/lib/constants/logs.ts @@ -57,7 +57,7 @@ export const ProviderLabels: Record = { cerebras: "Cerebras", gemini: "Gemini", openrouter: "OpenRouter", - nebius: "Nebius Token Factory" + nebius: "Nebius Token Factory", } as const; // Helper function to get provider label, supporting custom providers @@ -103,6 +103,20 @@ export const RequestTypeLabels = { speech_stream: "Speech Stream", transcription: "Transcription", transcription_stream: "Transcription Stream", + + batch_create: "Batch Create", + batch_list: "Batch List", + batch_retrieve: "Batch Retrieve", + batch_cancel: "Batch Cancel", + batch_results: "Batch Results", + + file_upload: "File Upload", + file_list: "File List", + file_retrieve: "File Retrieve", + file_delete: "File Delete", + file_content: "File Content", + + } as const; export const RequestTypeColors = { @@ -130,6 +144,18 @@ export const RequestTypeColors = { speech_stream: "bg-pink-100 text-pink-800", transcription: "bg-orange-100 text-orange-800", transcription_stream: "bg-lime-100 text-lime-800", + + batch_create: "bg-green-100 text-green-800", + batch_list: "bg-blue-100 text-blue-800", + batch_retrieve: "bg-red-100 text-red-800", + batch_cancel: "bg-yellow-100 text-yellow-800", + batch_results: "bg-purple-100 text-purple-800", + + file_upload: "bg-pink-100 text-pink-800", + file_list: "bg-lime-100 text-lime-800", + file_retrieve: "bg-orange-100 text-orange-800", + file_delete: "bg-red-100 text-red-800", + file_content: "bg-blue-100 text-blue-800", } as const; export type Status = (typeof Statuses)[number];