112 lines
3.1 KiB
Go
112 lines
3.1 KiB
Go
package store
|
|
|
|
import (
|
|
"context"
|
|
"strings"
|
|
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
|
)
|
|
|
|
// VolcesCompatibleTaskListFilter mirrors the supported filters of Ark's
|
|
// ListContentsGenerationsTasks API. Task IDs are the gateway's public task
|
|
// IDs, which are the IDs returned by the compatibility create endpoint.
|
|
type VolcesCompatibleTaskListFilter struct {
|
|
CompatibilityMarker string
|
|
Status string
|
|
Model string
|
|
TaskIDs []string
|
|
Page int
|
|
PageSize int
|
|
}
|
|
|
|
// ListVolcesCompatibleTasks returns only video tasks created through a named
|
|
// compatibility surface. Keeping this query separate from ListTasks avoids
|
|
// broadening the ordinary task-list API's filtering semantics.
|
|
func (s *Store) ListVolcesCompatibleTasks(ctx context.Context, user *auth.User, filter VolcesCompatibleTaskListFilter) (TaskListResult, error) {
|
|
page := filter.Page
|
|
if page < 1 {
|
|
page = 1
|
|
}
|
|
if page > 500 {
|
|
page = 500
|
|
}
|
|
pageSize := filter.PageSize
|
|
if pageSize < 1 {
|
|
pageSize = 20
|
|
}
|
|
if pageSize > 500 {
|
|
pageSize = 500
|
|
}
|
|
gatewayUserID := localGatewayUserID(user)
|
|
userID, apiKeyID := "", ""
|
|
if user != nil {
|
|
userID = strings.TrimSpace(user.ID)
|
|
apiKeyID = strings.TrimSpace(user.APIKeyID)
|
|
}
|
|
if gatewayUserID == "" && userID == "" {
|
|
return TaskListResult{}, ErrLocalUserRequired
|
|
}
|
|
taskIDs := make([]string, 0, len(filter.TaskIDs))
|
|
seen := make(map[string]bool, len(filter.TaskIDs))
|
|
for _, taskID := range filter.TaskIDs {
|
|
taskID = strings.TrimSpace(taskID)
|
|
if taskID != "" && !seen[taskID] {
|
|
seen[taskID] = true
|
|
taskIDs = append(taskIDs, taskID)
|
|
}
|
|
}
|
|
args := []any{
|
|
gatewayUserID,
|
|
userID,
|
|
apiKeyID,
|
|
strings.TrimSpace(filter.CompatibilityMarker),
|
|
strings.ToLower(strings.TrimSpace(filter.Status)),
|
|
strings.TrimSpace(filter.Model),
|
|
taskIDs,
|
|
}
|
|
whereSQL := `
|
|
WHERE (
|
|
(
|
|
NULLIF($1, '')::uuid IS NOT NULL
|
|
AND gateway_user_id = NULLIF($1, '')::uuid
|
|
)
|
|
OR (
|
|
NULLIF($1, '')::uuid IS NULL
|
|
AND NULLIF($2, '') IS NOT NULL
|
|
AND user_id = $2
|
|
)
|
|
)
|
|
AND (NULLIF($3, '') IS NULL OR api_key_id = $3)
|
|
AND kind = 'videos.generations'
|
|
AND request->>'_gateway_compatibility' = $4
|
|
AND (NULLIF($5, '') IS NULL OR LOWER(status) = $5)
|
|
AND (NULLIF($6, '') IS NULL OR model = $6 OR resolved_model = $6)
|
|
AND (COALESCE(array_length($7::text[], 1), 0) = 0 OR id::text = ANY($7::text[]))`
|
|
var total int
|
|
if err := s.pool.QueryRow(ctx, `SELECT count(*) FROM gateway_tasks `+whereSQL, args...).Scan(&total); err != nil {
|
|
return TaskListResult{}, err
|
|
}
|
|
rows, err := s.pool.Query(ctx, `
|
|
SELECT `+gatewayTaskColumns+`
|
|
FROM gateway_tasks
|
|
`+whereSQL+`
|
|
ORDER BY created_at DESC
|
|
LIMIT $8 OFFSET $9`, append(args, pageSize, (page-1)*pageSize)...)
|
|
if err != nil {
|
|
return TaskListResult{}, err
|
|
}
|
|
defer rows.Close()
|
|
items := make([]GatewayTask, 0)
|
|
for rows.Next() {
|
|
task, err := scanGatewayTask(rows)
|
|
if err != nil {
|
|
return TaskListResult{}, err
|
|
}
|
|
items = append(items, task)
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return TaskListResult{}, err
|
|
}
|
|
return TaskListResult{Items: items, Total: total, Page: page, PageSize: pageSize}, nil
|
|
}
|