bifrost/core/providers/openai/batch.go

package openai

import (
	"time"

	"github.com/maximhq/bifrost/core/schemas"
)

// OpenAI Batch API Types

// OpenAIBatchRequest represents the request body for creating a batch.
type OpenAIBatchRequest struct {
	InputFileID        string                     `json:"input_file_id"`
	Endpoint           string                     `json:"endpoint"`
	CompletionWindow   string                     `json:"completion_window"`
	Metadata           map[string]string          `json:"metadata,omitempty"`
	OutputExpiresAfter *schemas.BatchExpiresAfter `json:"output_expires_after,omitempty"`
}

// OpenAIBatchResponse represents an OpenAI batch response.
type OpenAIBatchResponse struct {
	ID               string                    `json:"id"`
	Object           string                    `json:"object"`
	Endpoint         string                    `json:"endpoint"`
	Errors           *schemas.BatchErrors      `json:"errors,omitempty"`
	InputFileID      string                    `json:"input_file_id"`
	CompletionWindow string                    `json:"completion_window"`
	Status           string                    `json:"status"`
	OutputFileID     *string                   `json:"output_file_id,omitempty"`
	ErrorFileID      *string                   `json:"error_file_id,omitempty"`
	CreatedAt        int64                     `json:"created_at"`
	InProgressAt     *int64                    `json:"in_progress_at,omitempty"`
	ExpiresAt        *int64                    `json:"expires_at,omitempty"`
	FinalizingAt     *int64                    `json:"finalizing_at,omitempty"`
	CompletedAt      *int64                    `json:"completed_at,omitempty"`
	FailedAt         *int64                    `json:"failed_at,omitempty"`
	ExpiredAt        *int64                    `json:"expired_at,omitempty"`
	CancellingAt     *int64                    `json:"cancelling_at,omitempty"`
	CancelledAt      *int64                    `json:"cancelled_at,omitempty"`
	RequestCounts    *OpenAIBatchRequestCounts `json:"request_counts,omitempty"`
	Metadata         map[string]string         `json:"metadata,omitempty"`
}

// OpenAIBatchRequestCounts represents the request counts for a batch.
type OpenAIBatchRequestCounts struct {
	Total     int `json:"total"`
	Completed int `json:"completed"`
	Failed    int `json:"failed"`
}

// OpenAIBatchListResponse represents the response from listing batches.
type OpenAIBatchListResponse struct {
	Object  string                `json:"object"`
	Data    []OpenAIBatchResponse `json:"data"`
	FirstID *string               `json:"first_id,omitempty"`
	LastID  *string               `json:"last_id,omitempty"`
	HasMore bool                  `json:"has_more"`
}

// ToBifrostBatchStatus converts OpenAI status to Bifrost status.
func ToBifrostBatchStatus(status string) schemas.BatchStatus {
	switch status {
	case "validating":
		return schemas.BatchStatusValidating
	case "failed":
		return schemas.BatchStatusFailed
	case "in_progress":
		return schemas.BatchStatusInProgress
	case "finalizing":
		return schemas.BatchStatusFinalizing
	case "completed":
		return schemas.BatchStatusCompleted
	case "expired":
		return schemas.BatchStatusExpired
	case "cancelling":
		return schemas.BatchStatusCancelling
	case "cancelled":
		return schemas.BatchStatusCancelled
	default:
		return schemas.BatchStatus(status)
	}
}

// ToBifrostBatchCreateResponse converts OpenAI batch response to Bifrost batch response.
func (r *OpenAIBatchResponse) ToBifrostBatchCreateResponse(latency time.Duration, sendBackRawRequest bool, sendBackRawResponse bool, rawRequest interface{}, rawResponse interface{}) *schemas.BifrostBatchCreateResponse {
	resp := &schemas.BifrostBatchCreateResponse{
		ID:               r.ID,
		Object:           r.Object,
		Endpoint:         r.Endpoint,
		InputFileID:      r.InputFileID,
		CompletionWindow: r.CompletionWindow,
		Status:           ToBifrostBatchStatus(r.Status),
		Metadata:         r.Metadata,
		CreatedAt:        r.CreatedAt,
		OutputFileID:     r.OutputFileID,
		ErrorFileID:      r.ErrorFileID,
		ExtraFields: schemas.BifrostResponseExtraFields{
			Latency: latency.Milliseconds(),
		},
	}

	if sendBackRawRequest {
		resp.ExtraFields.RawRequest = rawRequest
	}

	if r.ExpiresAt != nil {
		resp.ExpiresAt = r.ExpiresAt
	}

	if r.RequestCounts != nil {
		resp.RequestCounts = schemas.BatchRequestCounts{
			Total:     r.RequestCounts.Total,
			Completed: r.RequestCounts.Completed,
			Failed:    r.RequestCounts.Failed,
		}
	}

	if sendBackRawResponse {
		resp.ExtraFields.RawResponse = rawResponse
	}

	return resp
}

// ToBifrostBatchRetrieveResponse converts OpenAI batch response to Bifrost batch retrieve response.
func (r *OpenAIBatchResponse) ToBifrostBatchRetrieveResponse(latency time.Duration, sendBackRawRequest bool, sendBackRawResponse bool, rawRequest interface{}, rawResponse interface{}) *schemas.BifrostBatchRetrieveResponse {
	resp := &schemas.BifrostBatchRetrieveResponse{
		ID:               r.ID,
		Object:           r.Object,
		Endpoint:         r.Endpoint,
		InputFileID:      r.InputFileID,
		CompletionWindow: r.CompletionWindow,
		Status:           ToBifrostBatchStatus(r.Status),
		Metadata:         r.Metadata,
		CreatedAt:        r.CreatedAt,
		InProgressAt:     r.InProgressAt,
		FinalizingAt:     r.FinalizingAt,
		CompletedAt:      r.CompletedAt,
		FailedAt:         r.FailedAt,
		ExpiredAt:        r.ExpiredAt,
		CancellingAt:     r.CancellingAt,
		CancelledAt:      r.CancelledAt,
		OutputFileID:     r.OutputFileID,
		ErrorFileID:      r.ErrorFileID,
		Errors:           r.Errors,
		ExtraFields: schemas.BifrostResponseExtraFields{
			Latency: latency.Milliseconds(),
		},
	}

	if sendBackRawRequest {
		resp.ExtraFields.RawRequest = rawRequest
	}

	if r.ExpiresAt != nil {
		resp.ExpiresAt = r.ExpiresAt
	}

	if r.RequestCounts != nil {
		resp.RequestCounts = schemas.BatchRequestCounts{
			Total:     r.RequestCounts.Total,
			Completed: r.RequestCounts.Completed,
			Failed:    r.RequestCounts.Failed,
		}
	}

	if sendBackRawResponse {
		resp.ExtraFields.RawResponse = rawResponse
	}

	return resp
}