Something went wrong. Try again.
sloproxy
Something went wrong. Try again.
20 kB · 589 lines
Go
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590// Package executor provides runtime execution capabilities for various AI service providers.// This file implements the AI Studio executor that routes requests through a websocket-backed// transport for the AI Studio provider.package executor
import ( "bytes" "context" "encoding/json" "fmt" "io" "net/http" "net/url" "strings"
"github.com/router-for-me/CLIProxyAPI/v7/internal/config" "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" "github.com/router-for-me/CLIProxyAPI/v7/internal/util" "github.com/router-for-me/CLIProxyAPI/v7/internal/wsrelay" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" "github.com/tidwall/gjson" "github.com/tidwall/sjson")
// AIStudioExecutor routes AI Studio requests through a websocket-backed transport.type AIStudioExecutor struct { provider string relay *wsrelay.Manager cfg *config.Config}
// NewAIStudioExecutor creates a new AI Studio executor instance.//// Parameters:// - cfg: The application configuration// - provider: The provider name// - relay: The websocket relay manager//// Returns:// - *AIStudioExecutor: A new AI Studio executor instancefunc NewAIStudioExecutor(cfg *config.Config, provider string, relay *wsrelay.Manager) *AIStudioExecutor { return &AIStudioExecutor{provider: strings.ToLower(provider), relay: relay, cfg: cfg}}
// Identifier returns the executor identifier.func (e *AIStudioExecutor) Identifier() string { return "aistudio" }
// PrepareRequest prepares the HTTP request for execution.func (e *AIStudioExecutor) PrepareRequest(req *http.Request, auth *cliproxyauth.Auth) error { if req == nil { return nil } var attrs map[string]string if auth != nil { attrs = auth.Attributes } util.ApplyCustomHeadersFromAttrs(req, attrs) return nil}
// HttpRequest forwards an arbitrary HTTP request through the websocket relay.func (e *AIStudioExecutor) HttpRequest(ctx context.Context, auth *cliproxyauth.Auth, req *http.Request) (*http.Response, error) { if req == nil { return nil, fmt.Errorf("aistudio executor: request is nil") } if ctx == nil { ctx = req.Context() } if e.relay == nil { return nil, fmt.Errorf("aistudio executor: ws relay is nil") } if auth == nil || auth.ID == "" { return nil, fmt.Errorf("aistudio executor: missing auth") } httpReq := req.WithContext(ctx) if err := e.PrepareRequest(httpReq, auth); err != nil { return nil, err } if httpReq.URL == nil || strings.TrimSpace(httpReq.URL.String()) == "" { return nil, fmt.Errorf("aistudio executor: request URL is empty") }
var body []byte if httpReq.Body != nil { b, errRead := io.ReadAll(httpReq.Body) if errRead != nil { return nil, errRead } body = b httpReq.Body = io.NopCloser(bytes.NewReader(b)) }
wsReq := &wsrelay.HTTPRequest{ Method: httpReq.Method, URL: httpReq.URL.String(), Headers: httpReq.Header.Clone(), Body: body, } wsResp, errRelay := e.relay.NonStream(ctx, auth.ID, wsReq) if errRelay != nil { return nil, errRelay } if wsResp == nil { return nil, fmt.Errorf("aistudio executor: ws response is nil") }
statusText := http.StatusText(wsResp.Status) if statusText == "" { statusText = "Unknown" } resp := &http.Response{ StatusCode: wsResp.Status, Status: fmt.Sprintf("%d %s", wsResp.Status, statusText), Header: wsResp.Headers.Clone(), Body: io.NopCloser(bytes.NewReader(wsResp.Body)), ContentLength: int64(len(wsResp.Body)), Request: httpReq, } return resp, nil}
// Execute performs a non-streaming request to the AI Studio API.func (e *AIStudioExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { if opts.Alt == "responses/compact" { return resp, statusErr{code: http.StatusNotImplemented, msg: "/responses/compact not supported"} } baseModel := thinking.ParseSuffix(req.Model).ModelName reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth) defer reporter.TrackFailure(ctx, &err)
translatedReq, body, err := e.translateRequest(ctx, req, opts, false) if err != nil { return resp, err } reporter.SetTranslatedReasoningEffort(body.payload, body.toFormat.String())
endpoint := e.buildEndpoint(baseModel, body.action, opts.Alt) wsReq := &wsrelay.HTTPRequest{ Method: http.MethodPost, URL: endpoint, Headers: http.Header{"Content-Type": []string{"application/json"}}, Body: body.payload, } var attrs map[string]string if auth != nil { attrs = auth.Attributes } util.ApplyCustomHeadersFromAttrs(&http.Request{Header: wsReq.Headers}, attrs)
var authID, authLabel, authType, authValue string if auth != nil { authID = auth.ID authLabel = auth.Label authType, authValue = auth.AccountInfo() } helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{ URL: endpoint, Method: http.MethodPost, Headers: wsReq.Headers.Clone(), Body: body.payload, Provider: e.Identifier(), AuthID: authID, AuthLabel: authLabel, AuthType: authType, AuthValue: authValue, })
reporter.StartResponseTTFT() wsResp, err := e.relay.NonStream(ctx, authID, wsReq) if err != nil { helps.RecordAPIResponseError(ctx, e.cfg, err) return resp, err } helps.RecordAPIResponseMetadata(ctx, e.cfg, wsResp.Status, wsResp.Headers.Clone()) reporter.StartResponseTTFT() if len(wsResp.Body) > 0 { reporter.MarkFirstResponseByte() helps.AppendAPIResponseChunk(ctx, e.cfg, wsResp.Body) } if wsResp.Status < 200 || wsResp.Status >= 300 { return resp, statusErr{code: wsResp.Status, msg: string(wsResp.Body)} } reporter.Publish(ctx, helps.ParseGeminiUsage(wsResp.Body)) responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) var param any out := sdktranslator.TranslateNonStream(ctx, body.toFormat, responseFormat, req.Model, opts.OriginalRequest, translatedReq, wsResp.Body, ¶m) if responseFormat == sdktranslator.FormatOpenAIResponse { out = helps.EnsureResponsesUsageDetails(out) } resp = cliproxyexecutor.Response{Payload: ensureColonSpacedJSON(out), Headers: wsResp.Headers.Clone()} return resp, nil}
// ExecuteStream performs a streaming request to the AI Studio API.func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (_ *cliproxyexecutor.StreamResult, err error) { if opts.Alt == "responses/compact" { return nil, statusErr{code: http.StatusNotImplemented, msg: "/responses/compact not supported"} } baseModel := thinking.ParseSuffix(req.Model).ModelName reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth) defer reporter.TrackFailure(ctx, &err)
translatedReq, body, err := e.translateRequest(ctx, req, opts, true) if err != nil { return nil, err } reporter.SetTranslatedReasoningEffort(body.payload, body.toFormat.String())
endpoint := e.buildEndpoint(baseModel, body.action, opts.Alt) wsReq := &wsrelay.HTTPRequest{ Method: http.MethodPost, URL: endpoint, Headers: http.Header{"Content-Type": []string{"application/json"}}, Body: body.payload, } var attrs map[string]string if auth != nil { attrs = auth.Attributes } util.ApplyCustomHeadersFromAttrs(&http.Request{Header: wsReq.Headers}, attrs) var authID, authLabel, authType, authValue string if auth != nil { authID = auth.ID authLabel = auth.Label authType, authValue = auth.AccountInfo() } helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{ URL: endpoint, Method: http.MethodPost, Headers: wsReq.Headers.Clone(), Body: body.payload, Provider: e.Identifier(), AuthID: authID, AuthLabel: authLabel, AuthType: authType, AuthValue: authValue, }) reporter.StartResponseTTFT() wsStream, err := e.relay.Stream(ctx, authID, wsReq) if err != nil { helps.RecordAPIResponseError(ctx, e.cfg, err) return nil, err } firstEvent, ok := <-wsStream if !ok { err = fmt.Errorf("wsrelay: stream closed before start") helps.RecordAPIResponseError(ctx, e.cfg, err) return nil, err } if firstEvent.Status > 0 && firstEvent.Status != http.StatusOK { metadataLogged := false if firstEvent.Status > 0 { helps.RecordAPIResponseMetadata(ctx, e.cfg, firstEvent.Status, firstEvent.Headers.Clone()) reporter.StartResponseTTFT() metadataLogged = true } var body bytes.Buffer if len(firstEvent.Payload) > 0 { reporter.MarkFirstResponseByte() helps.AppendAPIResponseChunk(ctx, e.cfg, firstEvent.Payload) body.Write(firstEvent.Payload) } if firstEvent.Type == wsrelay.MessageTypeStreamEnd { return nil, statusErr{code: firstEvent.Status, msg: body.String()} } for event := range wsStream { if event.Err != nil { helps.RecordAPIResponseError(ctx, e.cfg, event.Err) if body.Len() == 0 { body.WriteString(event.Err.Error()) } break } if !metadataLogged && event.Status > 0 { helps.RecordAPIResponseMetadata(ctx, e.cfg, event.Status, event.Headers.Clone()) reporter.StartResponseTTFT() metadataLogged = true } if len(event.Payload) > 0 { reporter.MarkFirstResponseByte() helps.AppendAPIResponseChunk(ctx, e.cfg, event.Payload) body.Write(event.Payload) } if event.Type == wsrelay.MessageTypeStreamEnd { break } } return nil, statusErr{code: firstEvent.Status, msg: body.String()} } out := make(chan cliproxyexecutor.StreamChunk) go func(first wsrelay.StreamEvent) { defer close(out) defer reporter.EnsurePublished(ctx) responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) originalRequest := opts.OriginalRequest if len(originalRequest) == 0 { originalRequest = req.Payload } claudeInputTokens := helps.NewClaudeInputTokenState(opts.SourceFormat, body.toFormat, responseFormat, originalRequest) var param any metadataLogged := false processEvent := func(event wsrelay.StreamEvent) bool { if event.Err != nil { helps.RecordAPIResponseError(ctx, e.cfg, event.Err) reporter.PublishFailure(ctx, event.Err) select { case out <- cliproxyexecutor.StreamChunk{Err: fmt.Errorf("wsrelay: %v", event.Err)}: case <-ctx.Done(): } return false } switch event.Type { case wsrelay.MessageTypeStreamStart: if !metadataLogged && event.Status > 0 { helps.RecordAPIResponseMetadata(ctx, e.cfg, event.Status, event.Headers.Clone()) reporter.StartResponseTTFT() metadataLogged = true } case wsrelay.MessageTypeStreamChunk: if len(event.Payload) > 0 { reporter.MarkFirstResponseByte() helps.AppendAPIResponseChunk(ctx, e.cfg, event.Payload) filtered := helps.FilterSSEUsageMetadata(event.Payload) if detail, ok := helps.ParseGeminiStreamUsage(filtered); ok { reporter.Publish(ctx, detail) } lines := helps.TranslateStreamWithClaudeInputTokens(ctx, body.toFormat, responseFormat, req.Model, opts.OriginalRequest, translatedReq, filtered, ¶m, claudeInputTokens) for i := range lines { select { case out <- cliproxyexecutor.StreamChunk{Payload: ensureColonSpacedJSON(lines[i])}: case <-ctx.Done(): return false } } break } case wsrelay.MessageTypeStreamEnd: return false case wsrelay.MessageTypeHTTPResp: if !metadataLogged && event.Status > 0 { helps.RecordAPIResponseMetadata(ctx, e.cfg, event.Status, event.Headers.Clone()) reporter.StartResponseTTFT() metadataLogged = true } if len(event.Payload) > 0 { reporter.MarkFirstResponseByte() helps.AppendAPIResponseChunk(ctx, e.cfg, event.Payload) } lines := helps.TranslateStreamWithClaudeInputTokens(ctx, body.toFormat, responseFormat, req.Model, opts.OriginalRequest, translatedReq, event.Payload, ¶m, claudeInputTokens) for i := range lines { select { case out <- cliproxyexecutor.StreamChunk{Payload: ensureColonSpacedJSON(lines[i])}: case <-ctx.Done(): return false } } reporter.Publish(ctx, helps.ParseGeminiUsage(event.Payload)) return false case wsrelay.MessageTypeError: helps.RecordAPIResponseError(ctx, e.cfg, event.Err) reporter.PublishFailure(ctx, event.Err) select { case out <- cliproxyexecutor.StreamChunk{Err: fmt.Errorf("wsrelay: %v", event.Err)}: case <-ctx.Done(): } return false } return true } if !processEvent(first) { return } for event := range wsStream { if !processEvent(event) { return } } }(firstEvent) return &cliproxyexecutor.StreamResult{Headers: firstEvent.Headers.Clone(), Chunks: out}, nil}
// CountTokens counts tokens for the given request using the AI Studio API.func (e *AIStudioExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { baseModel := thinking.ParseSuffix(req.Model).ModelName _, body, err := e.translateRequest(ctx, req, opts, false) if err != nil { return cliproxyexecutor.Response{}, err }
body.payload, _ = sjson.DeleteBytes(body.payload, "generationConfig") body.payload, _ = sjson.DeleteBytes(body.payload, "tools") body.payload, _ = sjson.DeleteBytes(body.payload, "safetySettings")
endpoint := e.buildEndpoint(baseModel, "countTokens", "") wsReq := &wsrelay.HTTPRequest{ Method: http.MethodPost, URL: endpoint, Headers: http.Header{"Content-Type": []string{"application/json"}}, Body: body.payload, } var authID, authLabel, authType, authValue string if auth != nil { authID = auth.ID authLabel = auth.Label authType, authValue = auth.AccountInfo() } helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{ URL: endpoint, Method: http.MethodPost, Headers: wsReq.Headers.Clone(), Body: body.payload, Provider: e.Identifier(), AuthID: authID, AuthLabel: authLabel, AuthType: authType, AuthValue: authValue, }) resp, err := e.relay.NonStream(ctx, authID, wsReq) if err != nil { helps.RecordAPIResponseError(ctx, e.cfg, err) return cliproxyexecutor.Response{}, err } helps.RecordAPIResponseMetadata(ctx, e.cfg, resp.Status, resp.Headers.Clone()) if len(resp.Body) > 0 { helps.AppendAPIResponseChunk(ctx, e.cfg, resp.Body) } if resp.Status < 200 || resp.Status >= 300 { return cliproxyexecutor.Response{}, statusErr{code: resp.Status, msg: string(resp.Body)} } totalTokens := gjson.GetBytes(resp.Body, "totalTokens").Int() if totalTokens <= 0 { return cliproxyexecutor.Response{}, fmt.Errorf("wsrelay: totalTokens missing in response") } responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) translated := sdktranslator.TranslateTokenCount(ctx, body.toFormat, responseFormat, totalTokens, resp.Body) return cliproxyexecutor.Response{Payload: translated}, nil}
// Refresh refreshes the authentication credentials (no-op for AI Studio).func (e *AIStudioExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { if refreshed, handled, err := helps.RefreshAuthViaHome(ctx, e.cfg, auth); handled { return refreshed, err } return auth, nil}
type translatedPayload struct { payload []byte action string toFormat sdktranslator.Format}
func (e *AIStudioExecutor) translateRequest(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, stream bool) ([]byte, translatedPayload, error) { baseModel := thinking.ParseSuffix(req.Model).ModelName
from := opts.SourceFormat to := sdktranslator.FromString("gemini") originalPayloadSource := req.Payload if len(opts.OriginalRequest) > 0 { originalPayloadSource = opts.OriginalRequest } originalPayload := originalPayloadSource originalTranslated := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, originalPayload, stream) payload := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, req.Payload, stream) payload, err := helps.ApplyThinkingWithSourcePayload(payload, req.Payload, originalPayloadSource, req.Model, from.String(), to.String(), e.Identifier()) if err != nil { return nil, translatedPayload{}, err } payload = fixGeminiImageAspectRatio(baseModel, payload) requestedModel := helps.PayloadRequestedModel(opts, req.Model) requestPath := helps.PayloadRequestPath(opts) payload = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", payload, originalTranslated, requestedModel, requestPath, opts.Headers) payload, _ = sjson.DeleteBytes(payload, "generationConfig.maxOutputTokens") payload, _ = sjson.DeleteBytes(payload, "generationConfig.responseMimeType") payload, _ = sjson.DeleteBytes(payload, "generationConfig.responseJsonSchema") metadataAction := "generateContent" if req.Metadata != nil { if action, _ := req.Metadata["action"].(string); action == "countTokens" { metadataAction = action } } action := metadataAction if stream && action != "countTokens" { action = "streamGenerateContent" } payload, _ = sjson.DeleteBytes(payload, "session_id") payload = helps.EnsureGeminiLeadingUserContent(payload, "contents") payload = normalizeAIStudioThinkingLevel(payload) return payload, translatedPayload{payload: payload, action: action, toFormat: to}, nil}
// normalizeAIStudioThinkingLevel normalizes thinking levels to Google's canonical uppercase enum values.// The AI Studio upstream validates generationConfig.thinkingConfig.thinkingLevel case-sensitively// and rejects lowercase values with a 400 invalid argument error.func normalizeAIStudioThinkingLevel(payload []byte) []byte { const path = "generationConfig.thinkingConfig.thinkingLevel" level := gjson.GetBytes(payload, path) if level.Type != gjson.String { return payload }
normalized := strings.ToUpper(level.String()) switch normalized { case "MINIMAL", "LOW", "MEDIUM", "HIGH": default: return payload } if normalized == level.String() { return payload }
result, err := sjson.SetBytes(payload, path, normalized) if err != nil { return payload } return result}
func (e *AIStudioExecutor) buildEndpoint(model, action, alt string) string { base := fmt.Sprintf("%s/%s/models/%s:%s", glEndpoint, glAPIVersion, model, action) if action == "streamGenerateContent" { if alt == "" { return base + "?alt=sse" } return base + "?$alt=" + url.QueryEscape(alt) } if alt != "" && action != "countTokens" { return base + "?$alt=" + url.QueryEscape(alt) } return base}
// ensureColonSpacedJSON normalizes JSON objects so that colons are followed by a single space while// keeping the payload otherwise compact. Non-JSON inputs are returned unchanged.func ensureColonSpacedJSON(payload []byte) []byte { trimmed := bytes.TrimSpace(payload) if len(trimmed) == 0 { return payload }
var decoded any if err := json.Unmarshal(trimmed, &decoded); err != nil { return payload }
indented, err := json.MarshalIndent(decoded, "", " ") if err != nil { return payload }
compacted := make([]byte, 0, len(indented)) inString := false skipSpace := false
for i := 0; i < len(indented); i++ { ch := indented[i] if ch == '"' { // A quote is escaped only when preceded by an odd number of consecutive backslashes. // For example: "\\\"" keeps the quote inside the string, but "\\\\" closes the string. backslashes := 0 for j := i - 1; j >= 0 && indented[j] == '\\'; j-- { backslashes++ } if backslashes%2 == 0 { inString = !inString } }
if !inString { if ch == '\n' || ch == '\r' { skipSpace = true continue } if skipSpace { if ch == ' ' || ch == '\t' { continue } skipSpace = false } }
compacted = append(compacted, ch) }
return compacted}