stupid-h4er
249 lines
6.8 KiB
Go
249 lines
6.8 KiB
Go
package handler
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"log"
|
|
"mime"
|
|
"mime/multipart"
|
|
"net/http"
|
|
"strings"
|
|
|
|
"github.com/basketikun/infinite-canvas/service"
|
|
)
|
|
|
|
func AIImagesGenerations(w http.ResponseWriter, r *http.Request) {
|
|
proxyAIRequest(w, r, "/images/generations")
|
|
}
|
|
|
|
func AIImagesEdits(w http.ResponseWriter, r *http.Request) {
|
|
proxyAIRequest(w, r, "/images/edits")
|
|
}
|
|
|
|
func AIChatCompletions(w http.ResponseWriter, r *http.Request) {
|
|
proxyAIRequest(w, r, "/chat/completions")
|
|
}
|
|
|
|
func AIVideos(w http.ResponseWriter, r *http.Request) {
|
|
proxyAIRequest(w, r, "/videos")
|
|
}
|
|
|
|
func AIVideo(w http.ResponseWriter, r *http.Request, id string) {
|
|
proxyAIGetRequest(w, r, "/videos/"+id)
|
|
}
|
|
|
|
func AIVideoContent(w http.ResponseWriter, r *http.Request, id string) {
|
|
proxyAIGetRequest(w, r, "/videos/"+id+"/content")
|
|
}
|
|
|
|
func proxyAIGetRequest(w http.ResponseWriter, r *http.Request, path string) {
|
|
modelName := r.URL.Query().Get("model")
|
|
if strings.TrimSpace(modelName) == "" {
|
|
modelName = "grok-imagine-video"
|
|
}
|
|
channel, err := service.SelectModelChannel(modelName)
|
|
if err != nil {
|
|
log.Printf("AI proxy select channel failed: model=%s err=%v", modelName, err)
|
|
Fail(w, "AI 接口请求失败")
|
|
return
|
|
}
|
|
path = resolveAIProxyPath(channel.BaseURL, modelName, path)
|
|
request, err := http.NewRequest(http.MethodGet, service.BuildModelChannelURL(channel, path), nil)
|
|
if err != nil {
|
|
Fail(w, "AI 接口请求失败")
|
|
return
|
|
}
|
|
request.Header.Set("Authorization", "Bearer "+channel.APIKey)
|
|
copyAIResponse(w, request, nil)
|
|
}
|
|
|
|
func proxyAIRequest(w http.ResponseWriter, r *http.Request, path string) {
|
|
body, contentType, modelName, err := readAIRequest(r)
|
|
if err != nil {
|
|
log.Printf("AI proxy request read failed: %v", err)
|
|
Fail(w, "AI 接口请求失败")
|
|
return
|
|
}
|
|
user, ok := service.UserFromContext(r.Context())
|
|
if !ok {
|
|
Fail(w, "未登录或权限不足")
|
|
return
|
|
}
|
|
credits, err := service.ModelCost(modelName)
|
|
if err != nil {
|
|
log.Printf("AI proxy read model cost failed: model=%s err=%v", modelName, err)
|
|
Fail(w, "AI 接口请求失败")
|
|
return
|
|
}
|
|
credits *= readAIRequestCount(body, contentType)
|
|
channel, err := service.SelectModelChannel(modelName)
|
|
if err != nil {
|
|
log.Printf("AI proxy select channel failed: model=%s err=%v", modelName, err)
|
|
Fail(w, "AI 接口请求失败")
|
|
return
|
|
}
|
|
path = resolveAIProxyPath(channel.BaseURL, modelName, path)
|
|
request, err := http.NewRequest(http.MethodPost, service.BuildModelChannelURL(channel, path), bytes.NewReader(body))
|
|
if err != nil {
|
|
log.Printf("AI proxy build request failed: url=%s err=%v", service.BuildModelChannelURL(channel, path), err)
|
|
Fail(w, "AI 接口请求失败")
|
|
return
|
|
}
|
|
request.Header.Set("Authorization", "Bearer "+channel.APIKey)
|
|
if contentType != "" {
|
|
request.Header.Set("Content-Type", contentType)
|
|
}
|
|
if err := service.ConsumeUserCredits(user.ID, modelName, credits, path); err != nil {
|
|
FailError(w, err)
|
|
return
|
|
}
|
|
copyAIResponse(w, request, func() {
|
|
if err := service.RefundUserCredits(user.ID, modelName, credits, path); err != nil {
|
|
log.Printf("AI proxy refund credits failed: user=%s model=%s credits=%d err=%v", user.ID, modelName, credits, err)
|
|
}
|
|
})
|
|
}
|
|
|
|
func copyAIResponse(w http.ResponseWriter, request *http.Request, onFailure func()) {
|
|
response, err := http.DefaultClient.Do(request)
|
|
if err != nil {
|
|
log.Printf("AI proxy request failed: url=%s err=%v", request.URL.String(), err)
|
|
if onFailure != nil {
|
|
onFailure()
|
|
}
|
|
Fail(w, "AI 接口请求失败")
|
|
return
|
|
}
|
|
defer response.Body.Close()
|
|
|
|
if response.StatusCode >= http.StatusBadRequest {
|
|
_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 4096))
|
|
log.Printf("AI upstream error: url=%s status=%d", request.URL.String(), response.StatusCode)
|
|
if onFailure != nil {
|
|
onFailure()
|
|
}
|
|
Fail(w, aiStatusMessage(response.StatusCode))
|
|
return
|
|
}
|
|
|
|
for key, values := range response.Header {
|
|
if strings.EqualFold(key, "Content-Length") {
|
|
continue
|
|
}
|
|
for _, value := range values {
|
|
w.Header().Add(key, value)
|
|
}
|
|
}
|
|
w.WriteHeader(response.StatusCode)
|
|
_, _ = io.Copy(w, response.Body)
|
|
}
|
|
|
|
func readAIRequest(r *http.Request) ([]byte, string, string, error) {
|
|
contentType := r.Header.Get("Content-Type")
|
|
body, err := io.ReadAll(r.Body)
|
|
if err != nil {
|
|
return nil, "", "", err
|
|
}
|
|
modelName := ""
|
|
if strings.HasPrefix(contentType, "multipart/form-data") {
|
|
modelName = readMultipartModel(body, contentType)
|
|
} else {
|
|
var payload struct {
|
|
Model string `json:"model"`
|
|
}
|
|
_ = json.Unmarshal(body, &payload)
|
|
modelName = payload.Model
|
|
}
|
|
if strings.TrimSpace(modelName) == "" {
|
|
return nil, "", "", errMissingModel
|
|
}
|
|
return body, contentType, modelName, nil
|
|
}
|
|
|
|
func readMultipartModel(body []byte, contentType string) string {
|
|
_, params, err := mime.ParseMediaType(contentType)
|
|
if err != nil {
|
|
return ""
|
|
}
|
|
reader := multipart.NewReader(bytes.NewReader(body), params["boundary"])
|
|
form, err := reader.ReadForm(32 << 20)
|
|
if err != nil {
|
|
return ""
|
|
}
|
|
defer form.RemoveAll()
|
|
if values := form.Value["model"]; len(values) > 0 {
|
|
return values[0]
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func readAIRequestCount(body []byte, contentType string) int {
|
|
count := 1
|
|
if strings.HasPrefix(contentType, "multipart/form-data") {
|
|
_, params, err := mime.ParseMediaType(contentType)
|
|
if err != nil {
|
|
return count
|
|
}
|
|
form, err := multipart.NewReader(bytes.NewReader(body), params["boundary"]).ReadForm(32 << 20)
|
|
if err != nil {
|
|
return count
|
|
}
|
|
defer form.RemoveAll()
|
|
if values := form.Value["n"]; len(values) > 0 {
|
|
_, _ = fmt.Sscan(values[0], &count)
|
|
}
|
|
} else {
|
|
var payload struct {
|
|
N int `json:"n"`
|
|
}
|
|
_ = json.Unmarshal(body, &payload)
|
|
count = payload.N
|
|
}
|
|
if count < 1 {
|
|
return 1
|
|
}
|
|
return count
|
|
}
|
|
|
|
var errMissingModel = &aiError{"缺少模型名称"}
|
|
|
|
func resolveAIProxyPath(baseURL string, modelName string, path string) string {
|
|
if !isArkSeedanceVideo(baseURL, modelName) {
|
|
return path
|
|
}
|
|
if path == "/videos" {
|
|
return "/contents/generations/tasks"
|
|
}
|
|
if strings.HasPrefix(path, "/videos/") && !strings.HasSuffix(path, "/content") {
|
|
return "/contents/generations/tasks/" + strings.TrimPrefix(path, "/videos/")
|
|
}
|
|
return path
|
|
}
|
|
|
|
func isArkSeedanceVideo(baseURL string, modelName string) bool {
|
|
base := strings.ToLower(baseURL)
|
|
model := strings.ToLower(modelName)
|
|
return strings.Contains(model, "seedance") || strings.Contains(model, "doubao-seedance") || strings.Contains(base, "/api/plan/v3")
|
|
}
|
|
|
|
func aiStatusMessage(statusCode int) string {
|
|
switch statusCode {
|
|
case http.StatusUnauthorized, http.StatusForbidden:
|
|
return "AI 接口鉴权失败，请检查 API Key、套餐权限或模型权限"
|
|
case http.StatusTooManyRequests:
|
|
return "AI 接口限流或额度不足，请稍后重试或检查额度"
|
|
default:
|
|
return "AI 接口请求失败"
|
|
}
|
|
}
|
|
|
|
type aiError struct {
|
|
message string
|
|
}
|
|
|
|
func (err *aiError) Error() string {
|
|
return err.message
|
|
}
|