vibe-proxy/backend/internal/runtime/executor/xai_executor_media.go
2026-08-24 00:10:41 +02:00

150 lines
5.3 KiB
Go

package executor
import (
"bytes"
"context"
"io"
"net/http"
"net/url"
"strings"
xaiauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/xai"
"github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps"
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
log "github.com/sirupsen/logrus"
"github.com/tidwall/gjson"
)
func (e *XAIExecutor) executeImages(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, endpointPath string) (resp cliproxyexecutor.Response, err error) {
model := strings.TrimSpace(gjson.GetBytes(req.Payload, "model").String())
if model == "" {
model = strings.TrimSpace(req.Model)
}
reporter := helps.NewExecutorUsageReporter(ctx, e, model, auth)
defer reporter.TrackFailure(ctx, &err)
token, baseURL := xaiCreds(auth)
if baseURL == "" {
baseURL = xaiauth.DefaultAPIBaseURL
}
logXAIResolvedBaseURL(ctx, baseURL)
if endpointPath == "" {
endpointPath = xaiDefaultImageEndpointPath
}
payload := normalizeXAIImageRefs(req.Payload)
url := strings.TrimSuffix(baseURL, "/") + endpointPath
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(payload))
if err != nil {
return resp, err
}
applyXAIHeaders(httpReq, auth, token, false, "", opts.Headers)
e.recordXAIRequest(ctx, auth, url, httpReq.Header.Clone(), payload)
httpClient := helps.NewProxyAwareHTTPClient(ctx, e.cfg, auth, 0)
httpClient = reporter.TrackHTTPClient(httpClient)
httpResp, err := httpClient.Do(httpReq)
if err != nil {
helps.RecordAPIResponseError(ctx, e.cfg, err)
return resp, err
}
defer func() {
if errClose := httpResp.Body.Close(); errClose != nil {
log.Errorf("xai executor: close response body error: %v", errClose)
}
}()
helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
data, err := io.ReadAll(httpResp.Body)
if err != nil {
helps.RecordAPIResponseError(ctx, e.cfg, err)
return resp, err
}
helps.AppendAPIResponseChunk(ctx, e.cfg, data)
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data))
err = xaiStatusErr(httpResp.StatusCode, data)
return resp, err
}
reporter.EnsurePublished(ctx)
return cliproxyexecutor.Response{Payload: data, Headers: httpResp.Header.Clone()}, nil
}
func (e *XAIExecutor) executeVideos(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) {
model := strings.TrimSpace(gjson.GetBytes(req.Payload, "model").String())
if model == "" {
model = strings.TrimSpace(req.Model)
}
reporter := helps.NewExecutorUsageReporter(ctx, e, model, auth)
defer reporter.TrackFailure(ctx, &err)
token, baseURL := xaiCreds(auth)
if baseURL == "" {
baseURL = xaiauth.DefaultAPIBaseURL
}
logXAIResolvedBaseURL(ctx, baseURL)
payload := normalizeXAIImageRefs(req.Payload)
method := http.MethodPost
endpointPath := xaiVideosGenerationsPath
var body io.Reader = bytes.NewReader(payload)
switch path := xaiVideoEndpointPath(opts); path {
case xaiVideosGenerationsPath, xaiVideosEditsPath, xaiVideosExtensionsPath:
endpointPath = path
default:
if requestID := strings.TrimSpace(gjson.GetBytes(payload, "request_id").String()); requestID != "" {
method = http.MethodGet
endpointPath = xaiVideosPath + "/" + url.PathEscape(requestID)
body = nil
}
}
requestURL := strings.TrimSuffix(baseURL, "/") + endpointPath
httpReq, err := http.NewRequestWithContext(ctx, method, requestURL, body)
if err != nil {
return resp, err
}
applyXAIHeaders(httpReq, auth, token, false, "", opts.Headers)
if method == http.MethodPost {
key := xaiMetadataString(opts.Metadata, xaiIdempotencyKeyMetaKey)
if key == "" && opts.Headers != nil {
key = strings.TrimSpace(opts.Headers.Get("x-idempotency-key"))
}
if key != "" {
httpReq.Header.Set("x-idempotency-key", key)
}
}
e.recordXAIRequest(ctx, auth, requestURL, httpReq.Header.Clone(), payload)
httpClient := helps.NewProxyAwareHTTPClient(ctx, e.cfg, auth, 0)
httpClient = reporter.TrackHTTPClient(httpClient)
httpResp, err := httpClient.Do(httpReq)
if err != nil {
helps.RecordAPIResponseError(ctx, e.cfg, err)
return resp, err
}
defer func() {
if errClose := httpResp.Body.Close(); errClose != nil {
log.Errorf("xai executor: close response body error: %v", errClose)
}
}()
helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
data, err := io.ReadAll(httpResp.Body)
if err != nil {
helps.RecordAPIResponseError(ctx, e.cfg, err)
return resp, err
}
helps.AppendAPIResponseChunk(ctx, e.cfg, data)
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data))
return resp, xaiStatusErr(httpResp.StatusCode, data)
}
reporter.EnsurePublished(ctx)
return cliproxyexecutor.Response{Payload: data, Headers: httpResp.Header.Clone()}, nil
}