259 lines
6.5 KiB
Go
259 lines
6.5 KiB
Go
package zhipu
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"strconv"
|
|
|
|
"github.com/go-resty/resty/v2"
|
|
)
|
|
|
|
const (
|
|
BatchEndpointV4ChatCompletions = "/v4/chat/completions"
|
|
BatchEndpointV4ImagesGenerations = "/v4/images/generations"
|
|
BatchEndpointV4Embeddings = "/v4/embeddings"
|
|
BatchEndpointV4VideosGenerations = "/v4/videos/generations"
|
|
|
|
BatchCompletionWindow24h = "24h"
|
|
)
|
|
|
|
// BatchRequestCounts represents the counts of the batch requests.
|
|
type BatchRequestCounts struct {
|
|
Total int64 `json:"total"`
|
|
Completed int64 `json:"completed"`
|
|
Failed int64 `json:"failed"`
|
|
}
|
|
|
|
// BatchItem represents a batch item.
|
|
type BatchItem struct {
|
|
ID string `json:"id"`
|
|
Object any `json:"object"`
|
|
Endpoint string `json:"endpoint"`
|
|
InputFileID string `json:"input_file_id"`
|
|
CompletionWindow string `json:"completion_window"`
|
|
Status string `json:"status"`
|
|
OutputFileID string `json:"output_file_id"`
|
|
ErrorFileID string `json:"error_file_id"`
|
|
CreatedAt int64 `json:"created_at"`
|
|
InProgressAt int64 `json:"in_progress_at"`
|
|
ExpiresAt int64 `json:"expires_at"`
|
|
FinalizingAt int64 `json:"finalizing_at"`
|
|
CompletedAt int64 `json:"completed_at"`
|
|
FailedAt int64 `json:"failed_at"`
|
|
ExpiredAt int64 `json:"expired_at"`
|
|
CancellingAt int64 `json:"cancelling_at"`
|
|
CancelledAt int64 `json:"cancelled_at"`
|
|
RequestCounts BatchRequestCounts `json:"request_counts"`
|
|
Metadata json.RawMessage `json:"metadata"`
|
|
}
|
|
|
|
// BatchCreateService is a service to create a batch.
|
|
type BatchCreateService struct {
|
|
client *Client
|
|
|
|
inputFileID string
|
|
endpoint string
|
|
completionWindow string
|
|
metadata any
|
|
}
|
|
|
|
// NewBatchCreateService creates a new BatchCreateService.
|
|
func NewBatchCreateService(client *Client) *BatchCreateService {
|
|
return &BatchCreateService{client: client}
|
|
}
|
|
|
|
// SetInputFileID sets the input file id for the batch.
|
|
func (s *BatchCreateService) SetInputFileID(inputFileID string) *BatchCreateService {
|
|
s.inputFileID = inputFileID
|
|
return s
|
|
}
|
|
|
|
// SetEndpoint sets the endpoint for the batch.
|
|
func (s *BatchCreateService) SetEndpoint(endpoint string) *BatchCreateService {
|
|
s.endpoint = endpoint
|
|
return s
|
|
}
|
|
|
|
// SetCompletionWindow sets the completion window for the batch.
|
|
func (s *BatchCreateService) SetCompletionWindow(window string) *BatchCreateService {
|
|
s.completionWindow = window
|
|
return s
|
|
}
|
|
|
|
// SetMetadata sets the metadata for the batch.
|
|
func (s *BatchCreateService) SetMetadata(metadata any) *BatchCreateService {
|
|
s.metadata = metadata
|
|
return s
|
|
}
|
|
|
|
// Do executes the batch create service.
|
|
func (s *BatchCreateService) Do(ctx context.Context) (res BatchItem, err error) {
|
|
var (
|
|
resp *resty.Response
|
|
apiError APIErrorResponse
|
|
)
|
|
|
|
if resp, err = s.client.request(ctx).
|
|
SetBody(M{
|
|
"input_file_id": s.inputFileID,
|
|
"endpoint": s.endpoint,
|
|
"completion_window": s.completionWindow,
|
|
"metadata": s.metadata,
|
|
}).
|
|
SetResult(&res).
|
|
SetError(&apiError).
|
|
Post("batches"); err != nil {
|
|
return
|
|
}
|
|
|
|
if resp.IsError() {
|
|
err = apiError
|
|
}
|
|
|
|
return
|
|
}
|
|
|
|
// BatchGetService is a service to get a batch.
|
|
type BatchGetService struct {
|
|
client *Client
|
|
batchID string
|
|
}
|
|
|
|
// BatchGetResponse represents the response of the batch get service.
|
|
type BatchGetResponse = BatchItem
|
|
|
|
// NewBatchGetService creates a new BatchGetService.
|
|
func NewBatchGetService(client *Client) *BatchGetService {
|
|
return &BatchGetService{client: client}
|
|
}
|
|
|
|
// SetBatchID sets the batch id for the batch get service.
|
|
func (s *BatchGetService) SetBatchID(batchID string) *BatchGetService {
|
|
s.batchID = batchID
|
|
return s
|
|
}
|
|
|
|
// Do executes the batch get service.
|
|
func (s *BatchGetService) Do(ctx context.Context) (res BatchGetResponse, err error) {
|
|
var (
|
|
resp *resty.Response
|
|
apiError APIErrorResponse
|
|
)
|
|
|
|
if resp, err = s.client.request(ctx).
|
|
SetPathParam("batch_id", s.batchID).
|
|
SetResult(&res).
|
|
SetError(&apiError).
|
|
Get("batches/{batch_id}"); err != nil {
|
|
return
|
|
}
|
|
|
|
if resp.IsError() {
|
|
err = apiError
|
|
}
|
|
|
|
return
|
|
}
|
|
|
|
// BatchCancelService is a service to cancel a batch.
|
|
type BatchCancelService struct {
|
|
client *Client
|
|
batchID string
|
|
}
|
|
|
|
// NewBatchCancelService creates a new BatchCancelService.
|
|
func NewBatchCancelService(client *Client) *BatchCancelService {
|
|
return &BatchCancelService{client: client}
|
|
}
|
|
|
|
// SetBatchID sets the batch id for the batch cancel service.
|
|
func (s *BatchCancelService) SetBatchID(batchID string) *BatchCancelService {
|
|
s.batchID = batchID
|
|
return s
|
|
}
|
|
|
|
// Do executes the batch cancel service.
|
|
func (s *BatchCancelService) Do(ctx context.Context) (err error) {
|
|
var (
|
|
resp *resty.Response
|
|
apiError APIErrorResponse
|
|
)
|
|
|
|
if resp, err = s.client.request(ctx).
|
|
SetPathParam("batch_id", s.batchID).
|
|
SetBody(M{}).
|
|
SetError(&apiError).
|
|
Post("batches/{batch_id}/cancel"); err != nil {
|
|
return
|
|
}
|
|
|
|
if resp.IsError() {
|
|
err = apiError
|
|
}
|
|
|
|
return
|
|
}
|
|
|
|
// BatchListService is a service to list batches.
|
|
type BatchListService struct {
|
|
client *Client
|
|
|
|
after *string
|
|
limit *int
|
|
}
|
|
|
|
// BatchListResponse represents the response of the batch list service.
|
|
type BatchListResponse struct {
|
|
Object string `json:"object"`
|
|
Data []BatchItem `json:"data"`
|
|
FirstID string `json:"first_id"`
|
|
LastID string `json:"last_id"`
|
|
HasMore bool `json:"has_more"`
|
|
}
|
|
|
|
// NewBatchListService creates a new BatchListService.
|
|
func NewBatchListService(client *Client) *BatchListService {
|
|
return &BatchListService{client: client}
|
|
}
|
|
|
|
// SetAfter sets the after cursor for the batch list service.
|
|
func (s *BatchListService) SetAfter(after string) *BatchListService {
|
|
s.after = &after
|
|
return s
|
|
}
|
|
|
|
// SetLimit sets the limit for the batch list service.
|
|
func (s *BatchListService) SetLimit(limit int) *BatchListService {
|
|
s.limit = &limit
|
|
return s
|
|
}
|
|
|
|
// Do executes the batch list service.
|
|
func (s *BatchListService) Do(ctx context.Context) (res BatchListResponse, err error) {
|
|
var (
|
|
resp *resty.Response
|
|
apiError APIErrorResponse
|
|
)
|
|
|
|
req := s.client.request(ctx)
|
|
if s.after != nil {
|
|
req.SetQueryParam("after", *s.after)
|
|
}
|
|
if s.limit != nil {
|
|
req.SetQueryParam("limit", strconv.Itoa(*s.limit))
|
|
}
|
|
|
|
if resp, err = req.
|
|
SetResult(&res).
|
|
SetError(&apiError).
|
|
Get("batches"); err != nil {
|
|
return
|
|
}
|
|
|
|
if resp.IsError() {
|
|
err = apiError
|
|
}
|
|
|
|
return
|
|
}
|