mirror of
https://github.com/QuantumNous/new-api.git
synced 2026-04-25 20:38:38 +00:00
Merge pull request #1659 from Sh1n3zZ/feat-vertex-veo
feat: vertex veo (#1450)
This commit is contained in:
@@ -12,4 +12,4 @@ var LogSqlType = DatabaseTypeSQLite // Default to SQLite for logging SQL queries
|
|||||||
var UsingMySQL = false
|
var UsingMySQL = false
|
||||||
var UsingClickHouse = false
|
var UsingClickHouse = false
|
||||||
|
|
||||||
var SQLitePath = "one-api.db?_busy_timeout=30000"
|
var SQLitePath = "one-api.db?_busy_timeout=30000"
|
||||||
@@ -178,4 +178,4 @@ func boolToString(b bool) string {
|
|||||||
return "true"
|
return "true"
|
||||||
}
|
}
|
||||||
return "false"
|
return "false"
|
||||||
}
|
}
|
||||||
@@ -94,7 +94,7 @@ func updateVideoSingleTask(ctx context.Context, adaptor channel.TaskAdaptor, cha
|
|||||||
} else if taskResult, err = adaptor.ParseTaskResult(responseBody); err != nil {
|
} else if taskResult, err = adaptor.ParseTaskResult(responseBody); err != nil {
|
||||||
return fmt.Errorf("parseTaskResult failed for task %s: %w", taskId, err)
|
return fmt.Errorf("parseTaskResult failed for task %s: %w", taskId, err)
|
||||||
} else {
|
} else {
|
||||||
task.Data = responseBody
|
task.Data = redactVideoResponseBody(responseBody)
|
||||||
}
|
}
|
||||||
|
|
||||||
now := time.Now().Unix()
|
now := time.Now().Unix()
|
||||||
@@ -117,7 +117,9 @@ func updateVideoSingleTask(ctx context.Context, adaptor channel.TaskAdaptor, cha
|
|||||||
if task.FinishTime == 0 {
|
if task.FinishTime == 0 {
|
||||||
task.FinishTime = now
|
task.FinishTime = now
|
||||||
}
|
}
|
||||||
task.FailReason = taskResult.Url
|
if !(len(taskResult.Url) > 5 && taskResult.Url[:5] == "data:") {
|
||||||
|
task.FailReason = taskResult.Url
|
||||||
|
}
|
||||||
case model.TaskStatusFailure:
|
case model.TaskStatusFailure:
|
||||||
task.Status = model.TaskStatusFailure
|
task.Status = model.TaskStatusFailure
|
||||||
task.Progress = "100%"
|
task.Progress = "100%"
|
||||||
@@ -146,3 +148,37 @@ func updateVideoSingleTask(ctx context.Context, adaptor channel.TaskAdaptor, cha
|
|||||||
|
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func redactVideoResponseBody(body []byte) []byte {
|
||||||
|
var m map[string]any
|
||||||
|
if err := json.Unmarshal(body, &m); err != nil {
|
||||||
|
return body
|
||||||
|
}
|
||||||
|
resp, _ := m["response"].(map[string]any)
|
||||||
|
if resp != nil {
|
||||||
|
delete(resp, "bytesBase64Encoded")
|
||||||
|
if v, ok := resp["video"].(string); ok {
|
||||||
|
resp["video"] = truncateBase64(v)
|
||||||
|
}
|
||||||
|
if vs, ok := resp["videos"].([]any); ok {
|
||||||
|
for i := range vs {
|
||||||
|
if vm, ok := vs[i].(map[string]any); ok {
|
||||||
|
delete(vm, "bytesBase64Encoded")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
b, err := json.Marshal(m)
|
||||||
|
if err != nil {
|
||||||
|
return body
|
||||||
|
}
|
||||||
|
return b
|
||||||
|
}
|
||||||
|
|
||||||
|
func truncateBase64(s string) string {
|
||||||
|
const maxKeep = 256
|
||||||
|
if len(s) <= maxKeep {
|
||||||
|
return s
|
||||||
|
}
|
||||||
|
return s[:maxKeep] + "..."
|
||||||
|
}
|
||||||
|
|||||||
2
main.go
2
main.go
@@ -204,4 +204,4 @@ func InitResources() error {
|
|||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
@@ -166,9 +166,9 @@ func getModelRequest(c *gin.Context) (*ModelRequest, bool, error) {
|
|||||||
c.Set("platform", string(constant.TaskPlatformSuno))
|
c.Set("platform", string(constant.TaskPlatformSuno))
|
||||||
c.Set("relay_mode", relayMode)
|
c.Set("relay_mode", relayMode)
|
||||||
} else if strings.Contains(c.Request.URL.Path, "/v1/video/generations") {
|
} else if strings.Contains(c.Request.URL.Path, "/v1/video/generations") {
|
||||||
err = common.UnmarshalBodyReusable(c, &modelRequest)
|
|
||||||
relayMode := relayconstant.RelayModeUnknown
|
relayMode := relayconstant.RelayModeUnknown
|
||||||
if c.Request.Method == http.MethodPost {
|
if c.Request.Method == http.MethodPost {
|
||||||
|
err = common.UnmarshalBodyReusable(c, &modelRequest)
|
||||||
relayMode = relayconstant.RelayModeVideoSubmit
|
relayMode = relayconstant.RelayModeVideoSubmit
|
||||||
} else if c.Request.Method == http.MethodGet {
|
} else if c.Request.Method == http.MethodGet {
|
||||||
relayMode = relayconstant.RelayModeVideoFetchByID
|
relayMode = relayconstant.RelayModeVideoFetchByID
|
||||||
|
|||||||
344
relay/channel/task/vertex/adaptor.go
Normal file
344
relay/channel/task/vertex/adaptor.go
Normal file
@@ -0,0 +1,344 @@
|
|||||||
|
package vertex
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"encoding/base64"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"regexp"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"github.com/gin-gonic/gin"
|
||||||
|
|
||||||
|
"one-api/common"
|
||||||
|
"one-api/constant"
|
||||||
|
"one-api/dto"
|
||||||
|
"one-api/relay/channel"
|
||||||
|
vertexcore "one-api/relay/channel/vertex"
|
||||||
|
relaycommon "one-api/relay/common"
|
||||||
|
"one-api/service"
|
||||||
|
)
|
||||||
|
|
||||||
|
type requestPayload struct {
|
||||||
|
Instances []map[string]any `json:"instances"`
|
||||||
|
Parameters map[string]any `json:"parameters,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type submitResponse struct {
|
||||||
|
Name string `json:"name"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type operationVideo struct {
|
||||||
|
MimeType string `json:"mimeType"`
|
||||||
|
BytesBase64Encoded string `json:"bytesBase64Encoded"`
|
||||||
|
Encoding string `json:"encoding"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type operationResponse struct {
|
||||||
|
Name string `json:"name"`
|
||||||
|
Done bool `json:"done"`
|
||||||
|
Response struct {
|
||||||
|
Type string `json:"@type"`
|
||||||
|
RaiMediaFilteredCount int `json:"raiMediaFilteredCount"`
|
||||||
|
Videos []operationVideo `json:"videos"`
|
||||||
|
BytesBase64Encoded string `json:"bytesBase64Encoded"`
|
||||||
|
Encoding string `json:"encoding"`
|
||||||
|
Video string `json:"video"`
|
||||||
|
} `json:"response"`
|
||||||
|
Error struct {
|
||||||
|
Message string `json:"message"`
|
||||||
|
} `json:"error"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type TaskAdaptor struct{}
|
||||||
|
|
||||||
|
func (a *TaskAdaptor) Init(info *relaycommon.TaskRelayInfo) {}
|
||||||
|
|
||||||
|
func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycommon.TaskRelayInfo) (taskErr *dto.TaskError) {
|
||||||
|
info.Action = constant.TaskActionTextGenerate
|
||||||
|
|
||||||
|
req := relaycommon.TaskSubmitReq{}
|
||||||
|
if err := common.UnmarshalBodyReusable(c, &req); err != nil {
|
||||||
|
return service.TaskErrorWrapperLocal(err, "invalid_request", http.StatusBadRequest)
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(req.Prompt) == "" {
|
||||||
|
return service.TaskErrorWrapperLocal(fmt.Errorf("prompt is required"), "invalid_request", http.StatusBadRequest)
|
||||||
|
}
|
||||||
|
c.Set("task_request", req)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (a *TaskAdaptor) BuildRequestURL(info *relaycommon.TaskRelayInfo) (string, error) {
|
||||||
|
adc := &vertexcore.Credentials{}
|
||||||
|
if err := json.Unmarshal([]byte(info.ApiKey), adc); err != nil {
|
||||||
|
return "", fmt.Errorf("failed to decode credentials: %w", err)
|
||||||
|
}
|
||||||
|
modelName := info.OriginModelName
|
||||||
|
if v, ok := getRequestModelFromContext(info); ok {
|
||||||
|
modelName = v
|
||||||
|
}
|
||||||
|
if modelName == "" {
|
||||||
|
modelName = "veo-3.0-generate-001"
|
||||||
|
}
|
||||||
|
|
||||||
|
region := vertexcore.GetModelRegion(info.ApiVersion, modelName)
|
||||||
|
if strings.TrimSpace(region) == "" {
|
||||||
|
region = "global"
|
||||||
|
}
|
||||||
|
if region == "global" {
|
||||||
|
return fmt.Sprintf(
|
||||||
|
"https://aiplatform.googleapis.com/v1/projects/%s/locations/global/publishers/google/models/%s:predictLongRunning",
|
||||||
|
adc.ProjectID,
|
||||||
|
modelName,
|
||||||
|
), nil
|
||||||
|
}
|
||||||
|
return fmt.Sprintf(
|
||||||
|
"https://%s-aiplatform.googleapis.com/v1/projects/%s/locations/%s/publishers/google/models/%s:predictLongRunning",
|
||||||
|
region,
|
||||||
|
adc.ProjectID,
|
||||||
|
region,
|
||||||
|
modelName,
|
||||||
|
), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (a *TaskAdaptor) BuildRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.TaskRelayInfo) error {
|
||||||
|
req.Header.Set("Content-Type", "application/json")
|
||||||
|
req.Header.Set("Accept", "application/json")
|
||||||
|
|
||||||
|
adc := &vertexcore.Credentials{}
|
||||||
|
if err := json.Unmarshal([]byte(info.ApiKey), adc); err != nil {
|
||||||
|
return fmt.Errorf("failed to decode credentials: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
token, err := vertexcore.AcquireAccessToken(*adc, info.ChannelSetting.Proxy)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("failed to acquire access token: %w", err)
|
||||||
|
}
|
||||||
|
req.Header.Set("Authorization", "Bearer "+token)
|
||||||
|
req.Header.Set("x-goog-user-project", adc.ProjectID)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, _ *relaycommon.TaskRelayInfo) (io.Reader, error) {
|
||||||
|
v, ok := c.Get("task_request")
|
||||||
|
if !ok {
|
||||||
|
return nil, fmt.Errorf("request not found in context")
|
||||||
|
}
|
||||||
|
req := v.(relaycommon.TaskSubmitReq)
|
||||||
|
|
||||||
|
body := requestPayload{
|
||||||
|
Instances: []map[string]any{{"prompt": req.Prompt}},
|
||||||
|
Parameters: map[string]any{},
|
||||||
|
}
|
||||||
|
if req.Metadata != nil {
|
||||||
|
if v, ok := req.Metadata["storageUri"]; ok {
|
||||||
|
body.Parameters["storageUri"] = v
|
||||||
|
}
|
||||||
|
if v, ok := req.Metadata["sampleCount"]; ok {
|
||||||
|
body.Parameters["sampleCount"] = v
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if _, ok := body.Parameters["sampleCount"]; !ok {
|
||||||
|
body.Parameters["sampleCount"] = 1
|
||||||
|
}
|
||||||
|
|
||||||
|
data, err := json.Marshal(body)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
return bytes.NewReader(data), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (a *TaskAdaptor) DoRequest(c *gin.Context, info *relaycommon.TaskRelayInfo, requestBody io.Reader) (*http.Response, error) {
|
||||||
|
return channel.DoTaskApiRequest(a, c, info, requestBody)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, _ *relaycommon.TaskRelayInfo) (taskID string, taskData []byte, taskErr *dto.TaskError) {
|
||||||
|
responseBody, err := io.ReadAll(resp.Body)
|
||||||
|
if err != nil {
|
||||||
|
return "", nil, service.TaskErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError)
|
||||||
|
}
|
||||||
|
_ = resp.Body.Close()
|
||||||
|
|
||||||
|
var s submitResponse
|
||||||
|
if err := json.Unmarshal(responseBody, &s); err != nil {
|
||||||
|
return "", nil, service.TaskErrorWrapper(err, "unmarshal_response_failed", http.StatusInternalServerError)
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(s.Name) == "" {
|
||||||
|
return "", nil, service.TaskErrorWrapper(fmt.Errorf("missing operation name"), "invalid_response", http.StatusInternalServerError)
|
||||||
|
}
|
||||||
|
localID := encodeLocalTaskID(s.Name)
|
||||||
|
c.JSON(http.StatusOK, gin.H{"task_id": localID})
|
||||||
|
return localID, responseBody, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (a *TaskAdaptor) GetModelList() []string { return []string{"veo-3.0-generate-001"} }
|
||||||
|
func (a *TaskAdaptor) GetChannelName() string { return "vertex" }
|
||||||
|
|
||||||
|
func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any) (*http.Response, error) {
|
||||||
|
taskID, ok := body["task_id"].(string)
|
||||||
|
if !ok {
|
||||||
|
return nil, fmt.Errorf("invalid task_id")
|
||||||
|
}
|
||||||
|
upstreamName, err := decodeLocalTaskID(taskID)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("decode task_id failed: %w", err)
|
||||||
|
}
|
||||||
|
region := extractRegionFromOperationName(upstreamName)
|
||||||
|
if region == "" {
|
||||||
|
region = "us-central1"
|
||||||
|
}
|
||||||
|
project := extractProjectFromOperationName(upstreamName)
|
||||||
|
model := extractModelFromOperationName(upstreamName)
|
||||||
|
if project == "" || model == "" {
|
||||||
|
return nil, fmt.Errorf("cannot extract project/model from operation name")
|
||||||
|
}
|
||||||
|
var url string
|
||||||
|
if region == "global" {
|
||||||
|
url = fmt.Sprintf("https://aiplatform.googleapis.com/v1/projects/%s/locations/global/publishers/google/models/%s:fetchPredictOperation", project, model)
|
||||||
|
} else {
|
||||||
|
url = fmt.Sprintf("https://%s-aiplatform.googleapis.com/v1/projects/%s/locations/%s/publishers/google/models/%s:fetchPredictOperation", region, project, region, model)
|
||||||
|
}
|
||||||
|
payload := map[string]string{"operationName": upstreamName}
|
||||||
|
data, err := json.Marshal(payload)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
adc := &vertexcore.Credentials{}
|
||||||
|
if err := json.Unmarshal([]byte(key), adc); err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to decode credentials: %w", err)
|
||||||
|
}
|
||||||
|
token, err := vertexcore.AcquireAccessToken(*adc, "")
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to acquire access token: %w", err)
|
||||||
|
}
|
||||||
|
req, err := http.NewRequest(http.MethodPost, url, bytes.NewReader(data))
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
req.Header.Set("Content-Type", "application/json")
|
||||||
|
req.Header.Set("Accept", "application/json")
|
||||||
|
req.Header.Set("Authorization", "Bearer "+token)
|
||||||
|
req.Header.Set("x-goog-user-project", adc.ProjectID)
|
||||||
|
return service.GetHttpClient().Do(req)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
|
||||||
|
var op operationResponse
|
||||||
|
if err := json.Unmarshal(respBody, &op); err != nil {
|
||||||
|
return nil, fmt.Errorf("unmarshal operation response failed: %w", err)
|
||||||
|
}
|
||||||
|
ti := &relaycommon.TaskInfo{}
|
||||||
|
if op.Error.Message != "" {
|
||||||
|
ti.Status = "FAILURE"
|
||||||
|
ti.Reason = op.Error.Message
|
||||||
|
ti.Progress = "100%"
|
||||||
|
return ti, nil
|
||||||
|
}
|
||||||
|
if !op.Done {
|
||||||
|
ti.Status = "IN_PROGRESS"
|
||||||
|
ti.Progress = "50%"
|
||||||
|
return ti, nil
|
||||||
|
}
|
||||||
|
ti.Status = "SUCCESS"
|
||||||
|
ti.Progress = "100%"
|
||||||
|
if len(op.Response.Videos) > 0 {
|
||||||
|
v0 := op.Response.Videos[0]
|
||||||
|
if v0.BytesBase64Encoded != "" {
|
||||||
|
mime := strings.TrimSpace(v0.MimeType)
|
||||||
|
if mime == "" {
|
||||||
|
enc := strings.TrimSpace(v0.Encoding)
|
||||||
|
if enc == "" {
|
||||||
|
enc = "mp4"
|
||||||
|
}
|
||||||
|
if strings.Contains(enc, "/") {
|
||||||
|
mime = enc
|
||||||
|
} else {
|
||||||
|
mime = "video/" + enc
|
||||||
|
}
|
||||||
|
}
|
||||||
|
ti.Url = "data:" + mime + ";base64," + v0.BytesBase64Encoded
|
||||||
|
return ti, nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if op.Response.BytesBase64Encoded != "" {
|
||||||
|
enc := strings.TrimSpace(op.Response.Encoding)
|
||||||
|
if enc == "" {
|
||||||
|
enc = "mp4"
|
||||||
|
}
|
||||||
|
mime := enc
|
||||||
|
if !strings.Contains(enc, "/") {
|
||||||
|
mime = "video/" + enc
|
||||||
|
}
|
||||||
|
ti.Url = "data:" + mime + ";base64," + op.Response.BytesBase64Encoded
|
||||||
|
return ti, nil
|
||||||
|
}
|
||||||
|
if op.Response.Video != "" { // some variants use `video` as base64
|
||||||
|
enc := strings.TrimSpace(op.Response.Encoding)
|
||||||
|
if enc == "" {
|
||||||
|
enc = "mp4"
|
||||||
|
}
|
||||||
|
mime := enc
|
||||||
|
if !strings.Contains(enc, "/") {
|
||||||
|
mime = "video/" + enc
|
||||||
|
}
|
||||||
|
ti.Url = "data:" + mime + ";base64," + op.Response.Video
|
||||||
|
return ti, nil
|
||||||
|
}
|
||||||
|
return ti, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func getRequestModelFromContext(info *relaycommon.TaskRelayInfo) (string, bool) {
|
||||||
|
return info.OriginModelName, info.OriginModelName != ""
|
||||||
|
}
|
||||||
|
|
||||||
|
func encodeLocalTaskID(name string) string {
|
||||||
|
return base64.RawURLEncoding.EncodeToString([]byte(name))
|
||||||
|
}
|
||||||
|
|
||||||
|
func decodeLocalTaskID(local string) (string, error) {
|
||||||
|
b, err := base64.RawURLEncoding.DecodeString(local)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
return string(b), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
var regionRe = regexp.MustCompile(`locations/([a-z0-9-]+)/`)
|
||||||
|
|
||||||
|
func extractRegionFromOperationName(name string) string {
|
||||||
|
m := regionRe.FindStringSubmatch(name)
|
||||||
|
if len(m) == 2 {
|
||||||
|
return m[1]
|
||||||
|
}
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
|
||||||
|
var modelRe = regexp.MustCompile(`models/([^/]+)/operations/`)
|
||||||
|
|
||||||
|
func extractModelFromOperationName(name string) string {
|
||||||
|
m := modelRe.FindStringSubmatch(name)
|
||||||
|
if len(m) == 2 {
|
||||||
|
return m[1]
|
||||||
|
}
|
||||||
|
idx := strings.Index(name, "models/")
|
||||||
|
if idx >= 0 {
|
||||||
|
s := name[idx+len("models/"):]
|
||||||
|
if p := strings.Index(s, "/operations/"); p > 0 {
|
||||||
|
return s[:p]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
|
||||||
|
var projectRe = regexp.MustCompile(`projects/([^/]+)/locations/`)
|
||||||
|
|
||||||
|
func extractProjectFromOperationName(name string) string {
|
||||||
|
m := projectRe.FindStringSubmatch(name)
|
||||||
|
if len(m) == 2 {
|
||||||
|
return m[1]
|
||||||
|
}
|
||||||
|
return ""
|
||||||
|
}
|
||||||
@@ -174,6 +174,7 @@ func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *rel
|
|||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
req.Set("Authorization", "Bearer "+accessToken)
|
req.Set("Authorization", "Bearer "+accessToken)
|
||||||
|
req.Set("x-goog-user-project", a.AccountCredentials.ProjectID)
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -12,7 +12,10 @@ func GetModelRegion(other string, localModelName string) string {
|
|||||||
if m[localModelName] != nil {
|
if m[localModelName] != nil {
|
||||||
return m[localModelName].(string)
|
return m[localModelName].(string)
|
||||||
} else {
|
} else {
|
||||||
return m["default"].(string)
|
if v, ok := m["default"]; ok {
|
||||||
|
return v.(string)
|
||||||
|
}
|
||||||
|
return "global"
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
return other
|
return other
|
||||||
|
|||||||
@@ -6,14 +6,15 @@ import (
|
|||||||
"encoding/json"
|
"encoding/json"
|
||||||
"encoding/pem"
|
"encoding/pem"
|
||||||
"errors"
|
"errors"
|
||||||
"github.com/bytedance/gopkg/cache/asynccache"
|
|
||||||
"github.com/golang-jwt/jwt"
|
|
||||||
"net/http"
|
"net/http"
|
||||||
"net/url"
|
"net/url"
|
||||||
relaycommon "one-api/relay/common"
|
relaycommon "one-api/relay/common"
|
||||||
"one-api/service"
|
"one-api/service"
|
||||||
"strings"
|
"strings"
|
||||||
|
|
||||||
|
"github.com/bytedance/gopkg/cache/asynccache"
|
||||||
|
"github.com/golang-jwt/jwt"
|
||||||
|
|
||||||
"fmt"
|
"fmt"
|
||||||
"time"
|
"time"
|
||||||
)
|
)
|
||||||
@@ -137,3 +138,45 @@ func exchangeJwtForAccessToken(signedJWT string, info *relaycommon.RelayInfo) (s
|
|||||||
|
|
||||||
return "", fmt.Errorf("failed to get access token: %v", result)
|
return "", fmt.Errorf("failed to get access token: %v", result)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func AcquireAccessToken(creds Credentials, proxy string) (string, error) {
|
||||||
|
signedJWT, err := createSignedJWT(creds.ClientEmail, creds.PrivateKey)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("failed to create signed JWT: %w", err)
|
||||||
|
}
|
||||||
|
return exchangeJwtForAccessTokenWithProxy(signedJWT, proxy)
|
||||||
|
}
|
||||||
|
|
||||||
|
func exchangeJwtForAccessTokenWithProxy(signedJWT string, proxy string) (string, error) {
|
||||||
|
authURL := "https://www.googleapis.com/oauth2/v4/token"
|
||||||
|
data := url.Values{}
|
||||||
|
data.Set("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer")
|
||||||
|
data.Set("assertion", signedJWT)
|
||||||
|
|
||||||
|
var client *http.Client
|
||||||
|
var err error
|
||||||
|
if proxy != "" {
|
||||||
|
client, err = service.NewProxyHttpClient(proxy)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("new proxy http client failed: %w", err)
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
client = service.GetHttpClient()
|
||||||
|
}
|
||||||
|
|
||||||
|
resp, err := client.PostForm(authURL, data)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
|
||||||
|
var result map[string]interface{}
|
||||||
|
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
|
||||||
|
if accessToken, ok := result["access_token"].(string); ok {
|
||||||
|
return accessToken, nil
|
||||||
|
}
|
||||||
|
return "", fmt.Errorf("failed to get access token: %v", result)
|
||||||
|
}
|
||||||
|
|||||||
@@ -1,7 +1,6 @@
|
|||||||
package relay
|
package relay
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"github.com/gin-gonic/gin"
|
|
||||||
"one-api/constant"
|
"one-api/constant"
|
||||||
"one-api/relay/channel"
|
"one-api/relay/channel"
|
||||||
"one-api/relay/channel/ali"
|
"one-api/relay/channel/ali"
|
||||||
@@ -28,6 +27,7 @@ import (
|
|||||||
taskjimeng "one-api/relay/channel/task/jimeng"
|
taskjimeng "one-api/relay/channel/task/jimeng"
|
||||||
"one-api/relay/channel/task/kling"
|
"one-api/relay/channel/task/kling"
|
||||||
"one-api/relay/channel/task/suno"
|
"one-api/relay/channel/task/suno"
|
||||||
|
taskvertex "one-api/relay/channel/task/vertex"
|
||||||
taskVidu "one-api/relay/channel/task/vidu"
|
taskVidu "one-api/relay/channel/task/vidu"
|
||||||
"one-api/relay/channel/tencent"
|
"one-api/relay/channel/tencent"
|
||||||
"one-api/relay/channel/vertex"
|
"one-api/relay/channel/vertex"
|
||||||
@@ -37,6 +37,8 @@ import (
|
|||||||
"one-api/relay/channel/zhipu"
|
"one-api/relay/channel/zhipu"
|
||||||
"one-api/relay/channel/zhipu_4v"
|
"one-api/relay/channel/zhipu_4v"
|
||||||
"strconv"
|
"strconv"
|
||||||
|
|
||||||
|
"github.com/gin-gonic/gin"
|
||||||
)
|
)
|
||||||
|
|
||||||
func GetAdaptor(apiType int) channel.Adaptor {
|
func GetAdaptor(apiType int) channel.Adaptor {
|
||||||
@@ -126,6 +128,8 @@ func GetTaskAdaptor(platform constant.TaskPlatform) channel.TaskAdaptor {
|
|||||||
return &kling.TaskAdaptor{}
|
return &kling.TaskAdaptor{}
|
||||||
case constant.ChannelTypeJimeng:
|
case constant.ChannelTypeJimeng:
|
||||||
return &taskjimeng.TaskAdaptor{}
|
return &taskjimeng.TaskAdaptor{}
|
||||||
|
case constant.ChannelTypeVertexAi:
|
||||||
|
return &taskvertex.TaskAdaptor{}
|
||||||
case constant.ChannelTypeVidu:
|
case constant.ChannelTypeVidu:
|
||||||
return &taskVidu.TaskAdaptor{}
|
return &taskVidu.TaskAdaptor{}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -15,6 +15,8 @@ import (
|
|||||||
relayconstant "one-api/relay/constant"
|
relayconstant "one-api/relay/constant"
|
||||||
"one-api/service"
|
"one-api/service"
|
||||||
"one-api/setting/ratio_setting"
|
"one-api/setting/ratio_setting"
|
||||||
|
"strconv"
|
||||||
|
"strings"
|
||||||
|
|
||||||
"github.com/gin-gonic/gin"
|
"github.com/gin-gonic/gin"
|
||||||
)
|
)
|
||||||
@@ -33,6 +35,7 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (taskErr *dto.
|
|||||||
platform = GetTaskPlatform(c)
|
platform = GetTaskPlatform(c)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
info.InitChannelMeta(c)
|
||||||
adaptor := GetTaskAdaptor(platform)
|
adaptor := GetTaskAdaptor(platform)
|
||||||
if adaptor == nil {
|
if adaptor == nil {
|
||||||
return service.TaskErrorWrapperLocal(fmt.Errorf("invalid api platform: %s", platform), "invalid_api_platform", http.StatusBadRequest)
|
return service.TaskErrorWrapperLocal(fmt.Errorf("invalid api platform: %s", platform), "invalid_api_platform", http.StatusBadRequest)
|
||||||
@@ -197,6 +200,9 @@ func RelayTaskFetch(c *gin.Context, relayMode int) (taskResp *dto.TaskError) {
|
|||||||
if taskErr != nil {
|
if taskErr != nil {
|
||||||
return taskErr
|
return taskErr
|
||||||
}
|
}
|
||||||
|
if len(respBody) == 0 {
|
||||||
|
respBody = []byte("{\"code\":\"success\",\"data\":null}")
|
||||||
|
}
|
||||||
|
|
||||||
c.Writer.Header().Set("Content-Type", "application/json")
|
c.Writer.Header().Set("Content-Type", "application/json")
|
||||||
_, err := io.Copy(c.Writer, bytes.NewBuffer(respBody))
|
_, err := io.Copy(c.Writer, bytes.NewBuffer(respBody))
|
||||||
@@ -276,10 +282,92 @@ func videoFetchByIDRespBodyBuilder(c *gin.Context) (respBody []byte, taskResp *d
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
respBody, err = json.Marshal(dto.TaskResponse[any]{
|
func() {
|
||||||
Code: "success",
|
channelModel, err2 := model.GetChannelById(originTask.ChannelId, true)
|
||||||
Data: TaskModel2Dto(originTask),
|
if err2 != nil {
|
||||||
})
|
return
|
||||||
|
}
|
||||||
|
if channelModel.Type != constant.ChannelTypeVertexAi {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
baseURL := constant.ChannelBaseURLs[channelModel.Type]
|
||||||
|
if channelModel.GetBaseURL() != "" {
|
||||||
|
baseURL = channelModel.GetBaseURL()
|
||||||
|
}
|
||||||
|
adaptor := GetTaskAdaptor(constant.TaskPlatform(strconv.Itoa(channelModel.Type)))
|
||||||
|
if adaptor == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
resp, err2 := adaptor.FetchTask(baseURL, channelModel.Key, map[string]any{
|
||||||
|
"task_id": originTask.TaskID,
|
||||||
|
"action": originTask.Action,
|
||||||
|
})
|
||||||
|
if err2 != nil || resp == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
body, err2 := io.ReadAll(resp.Body)
|
||||||
|
if err2 != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
ti, err2 := adaptor.ParseTaskResult(body)
|
||||||
|
if err2 == nil && ti != nil {
|
||||||
|
if ti.Status != "" {
|
||||||
|
originTask.Status = model.TaskStatus(ti.Status)
|
||||||
|
}
|
||||||
|
if ti.Progress != "" {
|
||||||
|
originTask.Progress = ti.Progress
|
||||||
|
}
|
||||||
|
if ti.Url != "" {
|
||||||
|
originTask.FailReason = ti.Url
|
||||||
|
}
|
||||||
|
_ = originTask.Update()
|
||||||
|
var raw map[string]any
|
||||||
|
_ = json.Unmarshal(body, &raw)
|
||||||
|
format := "mp4"
|
||||||
|
if respObj, ok := raw["response"].(map[string]any); ok {
|
||||||
|
if vids, ok := respObj["videos"].([]any); ok && len(vids) > 0 {
|
||||||
|
if v0, ok := vids[0].(map[string]any); ok {
|
||||||
|
if mt, ok := v0["mimeType"].(string); ok && mt != "" {
|
||||||
|
if strings.Contains(mt, "mp4") {
|
||||||
|
format = "mp4"
|
||||||
|
} else {
|
||||||
|
format = mt
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
status := "processing"
|
||||||
|
switch originTask.Status {
|
||||||
|
case model.TaskStatusSuccess:
|
||||||
|
status = "succeeded"
|
||||||
|
case model.TaskStatusFailure:
|
||||||
|
status = "failed"
|
||||||
|
case model.TaskStatusQueued, model.TaskStatusSubmitted:
|
||||||
|
status = "queued"
|
||||||
|
}
|
||||||
|
out := map[string]any{
|
||||||
|
"error": nil,
|
||||||
|
"format": format,
|
||||||
|
"metadata": nil,
|
||||||
|
"status": status,
|
||||||
|
"task_id": originTask.TaskID,
|
||||||
|
"url": originTask.FailReason,
|
||||||
|
}
|
||||||
|
respBody, _ = json.Marshal(dto.TaskResponse[any]{
|
||||||
|
Code: "success",
|
||||||
|
Data: out,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
|
||||||
|
if len(respBody) == 0 {
|
||||||
|
respBody, err = json.Marshal(dto.TaskResponse[any]{
|
||||||
|
Code: "success",
|
||||||
|
Data: TaskModel2Dto(originTask),
|
||||||
|
})
|
||||||
|
}
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user