Files
NianAIGC/backend/internal/providers/provider.go

231 lines
6.9 KiB
Go

// Package providers contains bounded protocol adapters for generation services.
package providers
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"net/url"
"strings"
)
type Status string
const (
StatusQueued Status = "queued"
StatusRunning Status = "running"
StatusSucceeded Status = "succeeded"
StatusFailed Status = "failed"
StatusCancelled Status = "cancelled"
StatusExpired Status = "expired"
)
type Request struct {
Capability string `json:"capability"`
Model string `json:"model,omitempty"`
Prompt string `json:"prompt"`
InputURLs []string `json:"inputUrls,omitempty"`
Materials []Material `json:"materials,omitempty"`
Settings map[string]any `json:"settings,omitempty"`
}
type MaterialType string
const (
MaterialImage MaterialType = "image"
MaterialVideo MaterialType = "video"
MaterialAudio MaterialType = "audio"
)
// Material retains the provider-facing type metadata that InputURLs cannot
// express. InputURLs remains supported for existing callers.
type Material struct {
URL string `json:"url"`
Type MaterialType `json:"type"`
Role string `json:"role,omitempty"`
Label string `json:"label,omitempty"`
}
func requestModel(request Request, fallback string) string {
if value := strings.TrimSpace(request.Model); value != "" {
return value
}
return fallback
}
type Result struct {
TaskID string
Status Status
OutputURLs []string
Raw json.RawMessage
ErrorMessage string
Usage map[string]int
}
// HTTPResult is the provider-neutral representation persisted with a Job.
// Output URLs and usage must survive a process restart after the external
// provider has already reached a terminal state.
type HTTPResult struct {
TaskID string `json:"taskId,omitempty"`
Status Status `json:"status"`
OutputURLs []string `json:"outputUrls"`
Raw json.RawMessage `json:"raw,omitempty"`
ErrorMessage string `json:"errorMessage,omitempty"`
Usage map[string]int `json:"usage,omitempty"`
}
func EncodeResult(result Result) (json.RawMessage, error) {
urls := result.OutputURLs
if urls == nil {
urls = []string{}
}
return json.Marshal(HTTPResult{TaskID: result.TaskID, Status: result.Status, OutputURLs: urls, Raw: result.Raw, ErrorMessage: result.ErrorMessage, Usage: result.Usage})
}
type Adapter interface {
Submit(context.Context, Request) (Result, error)
Query(context.Context, string) (Result, error)
}
// ModelQueryAdapter is implemented by providers whose query protocol requires
// the same model identifier that was used when the task was submitted. Callers
// can opt into it without widening the common Adapter contract.
type ModelQueryAdapter interface {
QueryModel(context.Context, string, string) (Result, error)
}
type HTTPClient interface {
Do(*http.Request) (*http.Response, error)
}
type Config struct {
BaseURL, APIKey, Model, AccessKeyID, SecretAccessKey, Region, Service string
MaxResponseBytes int64
}
type ProviderError struct {
Operation string
Status int
}
func (e *ProviderError) Error() string {
if e.Status > 0 {
return fmt.Sprintf("provider %s failed with HTTP %d", e.Operation, e.Status)
}
return "provider " + e.Operation + " failed"
}
type httpAdapter struct {
name string
config Config
client HTTPClient
submitPath func(Request) string
queryPath func(string) string
payload func(Request) any
headers func(*http.Request)
decode func([]byte) Result
}
func (a *httpAdapter) submit(ctx context.Context, input Request) (Result, error) {
body, err := json.Marshal(a.payload(input))
if err != nil {
return Result{}, fmt.Errorf("encode provider request: %w", err)
}
return a.call(ctx, http.MethodPost, a.submitPath(input), body, "submit")
}
func (a *httpAdapter) query(ctx context.Context, id string) (Result, error) {
if strings.TrimSpace(id) == "" {
return Result{}, errors.New("provider task id is required")
}
return a.call(ctx, http.MethodGet, a.queryPath(url.PathEscape(id)), nil, "query")
}
func (a *httpAdapter) call(ctx context.Context, method, path string, body []byte, operation string) (Result, error) {
base, err := url.Parse(strings.TrimRight(a.config.BaseURL, "/"))
if err != nil {
return Result{}, errors.New("invalid provider base URL")
}
rel, err := url.Parse(path)
if err != nil {
return Result{}, errors.New("invalid provider path")
}
base.Path = strings.TrimRight(base.Path, "/") + "/"
rel.Path = strings.TrimLeft(rel.Path, "/")
target := base.ResolveReference(rel)
req, err := http.NewRequestWithContext(ctx, method, target.String(), strings.NewReader(string(body)))
if err != nil {
return Result{}, fmt.Errorf("build provider request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+a.config.APIKey)
if a.headers != nil {
a.headers(req)
}
resp, err := a.client.Do(req)
if err != nil {
return Result{}, &ProviderError{Operation: a.name + " " + operation}
}
defer resp.Body.Close()
limit := a.config.MaxResponseBytes
if limit <= 0 {
limit = 2 << 20
}
raw, err := io.ReadAll(io.LimitReader(resp.Body, limit+1))
if err != nil || int64(len(raw)) > limit {
return Result{}, &ProviderError{Operation: a.name + " " + operation}
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return Result{}, &ProviderError{Operation: a.name + " " + operation, Status: resp.StatusCode}
}
if !json.Valid(raw) {
return Result{}, &ProviderError{Operation: a.name + " " + operation}
}
result := a.decode(raw)
result.Raw = append(json.RawMessage(nil), raw...)
return result, nil
}
func record(raw []byte) map[string]any { var v map[string]any; _ = json.Unmarshal(raw, &v); return v }
func object(v any) map[string]any { x, _ := v.(map[string]any); return x }
func stringValue(values ...any) string {
for _, v := range values {
if s, ok := v.(string); ok && strings.TrimSpace(s) != "" {
return strings.TrimSpace(s)
}
}
return ""
}
func status(v any) Status {
s := strings.ToLower(stringValue(v))
switch s {
case "completed", "complete", "succeeded", "success", "done":
return StatusSucceeded
case "running", "processing", "generating", "in_progress":
return StatusRunning
case "failed", "error", "unknown":
return StatusFailed
case "cancelled", "canceled":
return StatusCancelled
case "expired", "not_found", "timeout":
return StatusExpired
default:
return StatusQueued
}
}
func collectURLs(v any, out *[]string) {
switch x := v.(type) {
case string:
if strings.HasPrefix(x, "http://") || strings.HasPrefix(x, "https://") {
*out = append(*out, x)
}
case []any:
for _, i := range x {
collectURLs(i, out)
}
case map[string]any:
for _, k := range []string{"url", "image_url", "imageUrl", "result_url", "resultUrl", "video_url", "file_url"} {
collectURLs(x[k], out)
}
}
}