fix: bridge grok composer image inputs

This commit is contained in:
Heatherm Huang
2026-07-07 10:55:07 +08:00
parent db9d92ba8f
commit c34db70a88
7 changed files with 419 additions and 12 deletions
+2
View File
@@ -12,6 +12,7 @@ type Model struct {
var defaultModels = []Model{
{ID: "grok-4.3", Object: "model", OwnedBy: "xai", DisplayName: "Grok 4.3"},
{ID: "grok-build-0.1", Object: "model", OwnedBy: "xai", DisplayName: "Grok Build 0.1"},
{ID: "grok-composer-2.5-fast", Object: "model", OwnedBy: "xai", DisplayName: "Grok Composer 2.5 Fast"},
{ID: "grok-4.20-0309-reasoning", Object: "model", OwnedBy: "xai", DisplayName: "Grok 4.20 Reasoning"},
{ID: "grok-4.20-0309-non-reasoning", Object: "model", OwnedBy: "xai", DisplayName: "Grok 4.20 Non Reasoning"},
{ID: "grok-4.20-multi-agent-0309", Object: "model", OwnedBy: "xai", DisplayName: "Grok 4.20 Multi Agent"},
@@ -46,6 +47,7 @@ func DefaultModelMapping() map[string]string {
mapping["grok"] = "grok-4.3"
mapping["grok-latest"] = "grok-4.3"
mapping["grok-build"] = "grok-build-0.1"
mapping["grok-composer"] = "grok-composer-2.5-fast"
mapping["grok-4.20-reasoning"] = "grok-4.20-0309-reasoning"
mapping["grok-4.20-non-reasoning"] = "grok-4.20-0309-non-reasoning"
return mapping
+1
View File
@@ -210,6 +210,7 @@ func TestDefaultModelMappingIncludesGrokAliases(t *testing.T) {
require.Equal(t, "grok-4.3", mapping["grok"])
require.Equal(t, "grok-4.3", mapping["grok-latest"])
require.Equal(t, "grok-build-0.1", mapping["grok-build"])
require.Equal(t, "grok-composer-2.5-fast", mapping["grok-composer"])
require.Equal(t, "grok-4.20-0309-reasoning", mapping["grok-4.20-reasoning"])
require.Equal(t, "grok-4.20-0309-non-reasoning", mapping["grok-4.20-non-reasoning"])
require.Equal(t, "grok-4.20-multi-agent-0309", mapping["grok-4.20-multi-agent-0309"])
@@ -105,6 +105,33 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions(
return nil, policyErr
}
upstreamBody = updatedBody
// Grok Composer does not accept image_url parts directly, but Grok Build
// can describe the images first. Bridge only this exact failure mode.
token, tokenKind, err := s.GetAccessToken(ctx, account)
if err != nil {
return nil, err
}
if strings.TrimSpace(token) == "" {
return nil, fmt.Errorf("account %d missing %s credential", account.ID, tokenKind)
}
var bridgeUsage OpenAIUsage
if account.Platform == PlatformGrok {
bridgedBody, usage, bridged, bridgeErr := s.bridgeGrokComposerImageInputs(ctx, c, account, upstreamBody, token)
if bridgeErr != nil {
var failoverErr *UpstreamFailoverError
if !errors.As(bridgeErr, &failoverErr) && c != nil && c.Writer != nil && !c.Writer.Written() {
writeChatCompletionsError(c, http.StatusBadGateway, "upstream_error", bridgeErr.Error())
}
return nil, bridgeErr
}
if bridged {
upstreamBody = bridgedBody
addOpenAIUsage(&bridgeUsage, usage)
}
}
if clientStream {
var usageErr error
upstreamBody, usageErr = ensureOpenAIChatStreamUsage(upstreamBody)
@@ -122,14 +149,6 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions(
)
// 5. Build upstream request
token, tokenKind, err := s.GetAccessToken(ctx, account)
if err != nil {
return nil, err
}
if strings.TrimSpace(token) == "" {
return nil, fmt.Errorf("account %d missing %s credential", account.ID, tokenKind)
}
targetURL, err := s.rawChatCompletionsURL(account)
if err != nil {
return nil, err
@@ -245,10 +264,17 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions(
}
// 8. Forward response
var result *OpenAIForwardResult
var forwardErr error
if clientStream {
return s.streamRawChatCompletions(c, resp, account, originalModel, billingModel, upstreamModel, reasoningEffort, serviceTier, startTime, len(body))
result, forwardErr = s.streamRawChatCompletions(c, resp, account, originalModel, billingModel, upstreamModel, reasoningEffort, serviceTier, startTime, len(body))
} else {
result, forwardErr = s.bufferRawChatCompletions(c, resp, originalModel, billingModel, upstreamModel, reasoningEffort, serviceTier, startTime)
}
return s.bufferRawChatCompletions(c, resp, originalModel, billingModel, upstreamModel, reasoningEffort, serviceTier, startTime)
if result != nil {
addOpenAIUsage(&result.Usage, bridgeUsage)
}
return result, forwardErr
}
func (s *OpenAIGatewayService) rawChatCompletionsURL(account *Account) (string, error) {
@@ -10,12 +10,18 @@ import (
"strings"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/apicompat"
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
"github.com/gin-gonic/gin"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
)
const (
grokComposerImageBridgeVisionModel = "grok-build-0.1"
grokComposerImageBridgeMaxOutputTokens = 512
)
func (s *OpenAIGatewayService) forwardGrokResponses(
ctx context.Context,
c *gin.Context,
@@ -309,6 +315,303 @@ func shouldDropGrokToolChoice(toolChoice gjson.Result, tools []json.RawMessage)
return false
}
func (s *OpenAIGatewayService) bridgeGrokComposerImageInputs(
ctx context.Context,
c *gin.Context,
account *Account,
body []byte,
token string,
) ([]byte, OpenAIUsage, bool, error) {
if !shouldBridgeGrokComposerImageInputs(body) {
return body, OpenAIUsage{}, false, nil
}
var reqBody map[string]any
if err := json.Unmarshal(body, &reqBody); err != nil {
return body, OpenAIUsage{}, false, fmt.Errorf("parse grok composer image bridge request: %w", err)
}
imageURLs := collectGrokComposerImageURLs(reqBody)
if len(imageURLs) == 0 {
return body, OpenAIUsage{}, false, nil
}
descriptions := make([]string, 0, len(imageURLs))
var bridgeUsage OpenAIUsage
for index, imageURL := range imageURLs {
description, usage, err := s.describeGrokComposerImage(ctx, c, account, token, imageURL, index+1)
if err != nil {
return body, bridgeUsage, false, err
}
descriptions = append(descriptions, description)
addOpenAIUsage(&bridgeUsage, usage)
}
if !rewriteGrokComposerImagesAsText(reqBody, descriptions) {
return body, bridgeUsage, false, nil
}
bridgedBody, err := marshalOpenAIUpstreamJSON(reqBody)
if err != nil {
return body, bridgeUsage, false, fmt.Errorf("serialize grok composer image bridge request: %w", err)
}
return bridgedBody, bridgeUsage, true, nil
}
func shouldBridgeGrokComposerImageInputs(body []byte) bool {
if len(body) == 0 || !isGrokComposerModel(gjson.GetBytes(body, "model").String()) {
return false
}
messages := gjson.GetBytes(body, "messages")
if !messages.Exists() {
return false
}
return openAIJSONValueMayContainImageInput(messages)
}
func isGrokComposerModel(model string) bool {
model = strings.TrimSpace(strings.ToLower(model))
if model == "" {
return false
}
if strings.Contains(model, "/") {
parts := strings.Split(model, "/")
model = strings.TrimSpace(parts[len(parts)-1])
}
return strings.Contains(model, "composer")
}
func collectGrokComposerImageURLs(reqBody map[string]any) []string {
messages, ok := reqBody["messages"].([]any)
if !ok {
return nil
}
var imageURLs []string
for _, msg := range messages {
msgMap, ok := msg.(map[string]any)
if !ok {
continue
}
parts, ok := msgMap["content"].([]any)
if !ok {
continue
}
for _, part := range parts {
if imageURL := grokComposerImageURLFromPart(part); imageURL != "" {
imageURLs = append(imageURLs, imageURL)
}
}
}
return imageURLs
}
func grokComposerImageURLFromPart(part any) string {
partMap, ok := part.(map[string]any)
if !ok {
return ""
}
if strings.TrimSpace(strings.ToLower(fmt.Sprint(partMap["type"]))) != "image_url" {
return ""
}
switch imageURL := partMap["image_url"].(type) {
case string:
return normalizeGrokComposerImageURL(imageURL)
case map[string]any:
raw, _ := imageURL["url"].(string)
return normalizeGrokComposerImageURL(raw)
default:
return ""
}
}
func normalizeGrokComposerImageURL(raw string) string {
trimmed := strings.TrimSpace(raw)
if trimmed == "" || isEmptyBase64DataURI(trimmed) {
return ""
}
return trimmed
}
func (s *OpenAIGatewayService) describeGrokComposerImage(
ctx context.Context,
c *gin.Context,
account *Account,
token string,
imageURL string,
index int,
) (string, OpenAIUsage, error) {
body, err := buildGrokComposerImageDescriptionBody(imageURL, index)
if err != nil {
return "", OpenAIUsage{}, err
}
upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx)
upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, body, token)
releaseUpstreamCtx()
if err != nil {
return "", OpenAIUsage{}, fmt.Errorf("build grok composer image bridge request: %w", err)
}
proxyURL := ""
if account.ProxyID != nil && account.Proxy != nil {
proxyURL = account.Proxy.URL()
}
resp, err := s.httpUpstream.Do(upstreamReq, proxyURL, account.ID, account.Concurrency)
if err != nil {
return "", OpenAIUsage{}, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, false)
}
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode >= 400 {
respBody := s.readUpstreamErrorBody(resp)
s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode))
upstreamMsg := sanitizeUpstreamErrorMessage(extractUpstreamErrorMessage(respBody))
if upstreamMsg == "" {
upstreamMsg = fmt.Sprintf("xAI image bridge upstream returned status %d", resp.StatusCode)
}
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
Platform: account.Platform,
AccountID: account.ID,
AccountName: account.Name,
UpstreamStatusCode: resp.StatusCode,
UpstreamRequestID: firstNonEmpty(resp.Header.Get("x-request-id"), resp.Header.Get("xai-request-id")),
Kind: "failover",
Message: upstreamMsg,
})
s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
if s.shouldFailoverUpstreamError(resp.StatusCode) {
return "", OpenAIUsage{}, &UpstreamFailoverError{
StatusCode: resp.StatusCode,
ResponseBody: respBody,
RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode),
}
}
return "", OpenAIUsage{}, fmt.Errorf("grok composer image bridge upstream error: %s", upstreamMsg)
}
s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode))
respBody, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, nil)
if err != nil {
return "", OpenAIUsage{}, fmt.Errorf("read grok composer image bridge response: %w", err)
}
var parsed apicompat.ResponsesResponse
if err := json.Unmarshal(respBody, &parsed); err != nil {
return "", OpenAIUsage{}, fmt.Errorf("parse grok composer image bridge response: %w", err)
}
description := strings.TrimSpace(grokResponsesOutputText(&parsed))
if description == "" {
return "", copyOpenAIUsageFromResponsesUsage(parsed.Usage), fmt.Errorf("grok composer image bridge returned empty description")
}
return description, copyOpenAIUsageFromResponsesUsage(parsed.Usage), nil
}
func buildGrokComposerImageDescriptionBody(imageURL string, index int) ([]byte, error) {
prompt := fmt.Sprintf("Describe image %d in concise, factual text for a downstream coding/composer model. Include visible text, UI elements, diagrams, errors, and spatial relationships. Do not mention that you are an image analysis bridge.", index)
req := map[string]any{
"model": grokComposerImageBridgeVisionModel,
"stream": false,
"store": false,
"max_output_tokens": grokComposerImageBridgeMaxOutputTokens,
"input": []any{
map[string]any{
"type": "message",
"role": "user",
"content": []any{
map[string]any{"type": "input_text", "text": prompt},
map[string]any{"type": "input_image", "image_url": imageURL},
},
},
},
}
return marshalOpenAIUpstreamJSON(req)
}
func grokResponsesOutputText(resp *apicompat.ResponsesResponse) string {
if resp == nil {
return ""
}
var parts []string
for _, output := range resp.Output {
for _, content := range output.Content {
if content.Type == "output_text" || content.Type == "text" || content.Type == "input_text" {
if text := strings.TrimSpace(content.Text); text != "" {
parts = append(parts, text)
}
}
}
}
return strings.Join(parts, "\n\n")
}
func rewriteGrokComposerImagesAsText(reqBody map[string]any, descriptions []string) bool {
messages, ok := reqBody["messages"].([]any)
if !ok {
return false
}
imageIndex := 0
changed := false
for _, msg := range messages {
msgMap, ok := msg.(map[string]any)
if !ok {
continue
}
parts, ok := msgMap["content"].([]any)
if !ok {
continue
}
var textParts []string
messageChanged := false
for _, part := range parts {
if imageURL := grokComposerImageURLFromPart(part); imageURL != "" {
if imageIndex < len(descriptions) {
textParts = append(textParts, fmt.Sprintf("Image %d description: %s", imageIndex+1, strings.TrimSpace(descriptions[imageIndex])))
}
imageIndex++
messageChanged = true
continue
}
if text := grokComposerTextFromPart(part); text != "" {
textParts = append(textParts, text)
}
}
if messageChanged {
msgMap["content"] = strings.Join(textParts, "\n\n")
changed = true
}
}
return changed
}
func grokComposerTextFromPart(part any) string {
partMap, ok := part.(map[string]any)
if !ok {
return ""
}
partType := strings.TrimSpace(strings.ToLower(fmt.Sprint(partMap["type"])))
switch partType {
case "text", "input_text":
text, _ := partMap["text"].(string)
return strings.TrimSpace(text)
default:
return ""
}
}
func addOpenAIUsage(dst *OpenAIUsage, usage OpenAIUsage) {
if dst == nil {
return
}
dst.InputTokens += usage.InputTokens
dst.ImageInputTokens += usage.ImageInputTokens
dst.OutputTokens += usage.OutputTokens
dst.CacheCreationInputTokens += usage.CacheCreationInputTokens
dst.CacheReadInputTokens += usage.CacheReadInputTokens
dst.ImageOutputTokens += usage.ImageOutputTokens
}
func buildGrokResponsesRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token string) (*http.Request, error) {
targetURL, err := xai.BuildResponsesURL(account.GetGrokBaseURL())
if err != nil {
@@ -651,6 +651,76 @@ func TestForwardAsChatCompletionsForGrokStreamingUsesRawXAIChatCompletions(t *te
require.NotNil(t, repo.updates[53][grokQuotaSnapshotExtraKey])
}
func TestForwardAsChatCompletionsForGrokComposerBridgesImageInput(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{"model":"grok-composer-2.5-fast","messages":[{"role":"system","content":"You are concise."},{"role":"user","content":[{"type":"text","text":"What is shown?"},{"type":"image_url","image_url":{"url":"data:image/png;base64,QUJD"}}]}],"stream":false}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
c.Request.Header.Set("Content-Type", "application/json")
account := &Account{
ID: 55,
Name: "grok",
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Concurrency: 1,
Credentials: map[string]any{
"access_token": "access-token",
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
"base_url": xai.DefaultCLIBaseURL,
},
}
repo := &grokQuotaAccountRepo{
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
accountsByID: map[int64]*Account{55: account},
},
}
upstream := &httpUpstreamRecorder{responses: []*http.Response{
{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}, "xai-request-id": []string{"vision-req"}},
Body: io.NopCloser(strings.NewReader(`{"id":"resp_vision","object":"response","model":"grok-build-0.1","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"A small diagram with ABC letters."}]}],"usage":{"input_tokens":11,"output_tokens":7,"total_tokens":18}}`)),
},
{
StatusCode: http.StatusOK,
Header: http.Header{
"Content-Type": []string{"application/json"},
"X-Request-Id": []string{"composer-req"},
"X-Ratelimit-Limit-Requests": []string{"10"},
"X-Ratelimit-Remaining-Requests": []string{"9"},
"X-Ratelimit-Limit-Tokens": []string{"1000"},
"X-Ratelimit-Remaining-Tokens": []string{"980"},
},
Body: io.NopCloser(strings.NewReader(`{"id":"chatcmpl_composer","object":"chat.completion","model":"grok-composer-2.5-fast","choices":[{"index":0,"message":{"role":"assistant","content":"It shows ABC."},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":5,"total_tokens":8}}`)),
},
}}
svc := &OpenAIGatewayService{
cfg: rawChatCompletionsTestConfig(),
httpUpstream: upstream,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
accountRepo: repo,
}
result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "")
require.NoError(t, err)
require.NotNil(t, result)
require.Len(t, upstream.requests, 2)
require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.requests[0].URL.String())
require.Equal(t, "grok-build-0.1", gjson.GetBytes(upstream.bodies[0], "model").String())
require.Equal(t, "input_image", gjson.GetBytes(upstream.bodies[0], "input.0.content.1.type").String())
require.Equal(t, xai.DefaultCLIBaseURL+"/chat/completions", upstream.requests[1].URL.String())
require.Equal(t, "grok-composer-2.5-fast", gjson.GetBytes(upstream.bodies[1], "model").String())
require.False(t, strings.Contains(string(upstream.bodies[1]), "image_url"))
require.Contains(t, gjson.GetBytes(upstream.bodies[1], "messages.1.content").String(), "Image 1 description")
require.Contains(t, gjson.GetBytes(upstream.bodies[1], "messages.1.content").String(), "A small diagram with ABC letters.")
require.Equal(t, 14, result.Usage.InputTokens)
require.Equal(t, 12, result.Usage.OutputTokens)
require.Equal(t, "It shows ABC.", gjson.Get(recorder.Body.String(), "choices.0.message.content").String())
require.NotNil(t, repo.updates[55][grokQuotaSnapshotExtraKey])
}
func TestForwardAsAnthropicForGrokUsesXAIResponses(t *testing.T) {
gin.SetMode(gin.TestMode)
+4 -2
View File
@@ -18,9 +18,11 @@ ARG GOSUMDB=sum.golang.google.cn
FROM ${NODE_IMAGE} AS frontend-builder
WORKDIR /app/frontend
ENV NODE_OPTIONS=--max-old-space-size=1536
# Install pnpm
RUN corepack enable && corepack prepare pnpm@latest --activate
# Install pnpm. Keep this pinned to the lockfile-compatible major version so
# Docker builds remain reproducible when pnpm changes config validation rules.
RUN corepack enable && corepack prepare pnpm@9.15.9 --activate
# Install dependencies first (better caching)
COPY frontend/package.json frontend/pnpm-lock.yaml ./
@@ -137,12 +137,14 @@ const metaModels = [
const xaiModels = [
'grok-4.3',
'grok-build-0.1',
'grok-composer-2.5-fast',
'grok-4.20-0309-reasoning',
'grok-4.20-0309-non-reasoning',
'grok-4.20-multi-agent-0309',
'grok',
'grok-latest',
'grok-build',
'grok-composer',
'grok-4.20-reasoning',
'grok-4.20-non-reasoning',
'grok-imagine',
@@ -297,6 +299,7 @@ const grokPresetMappings = [
{ label: 'Grok 4.3', from: 'grok-4.3', to: 'grok-4.3', color: 'bg-slate-100 text-slate-700 hover:bg-slate-200 dark:bg-slate-800/50 dark:text-slate-300' },
{ label: 'Grok Latest', from: 'grok-latest', to: 'grok-4.3', color: 'bg-emerald-100 text-emerald-700 hover:bg-emerald-200 dark:bg-emerald-900/30 dark:text-emerald-400' },
{ label: 'Build 0.1', from: 'grok-build', to: 'grok-build-0.1', color: 'bg-cyan-100 text-cyan-700 hover:bg-cyan-200 dark:bg-cyan-900/30 dark:text-cyan-400' },
{ label: 'Composer 2.5', from: 'grok-composer', to: 'grok-composer-2.5-fast', color: 'bg-teal-100 text-teal-700 hover:bg-teal-200 dark:bg-teal-900/30 dark:text-teal-400' },
{ label: '4.20 Reasoning', from: 'grok-4.20-reasoning', to: 'grok-4.20-0309-reasoning', color: 'bg-indigo-100 text-indigo-700 hover:bg-indigo-200 dark:bg-indigo-900/30 dark:text-indigo-400' },
{ label: '4.20 Non Reasoning', from: 'grok-4.20-non-reasoning', to: 'grok-4.20-0309-non-reasoning', color: 'bg-violet-100 text-violet-700 hover:bg-violet-200 dark:bg-violet-900/30 dark:text-violet-400' },
{ label: 'Imagine Image', from: 'grok-imagine', to: 'grok-imagine-image-quality', color: 'bg-sky-100 text-sky-700 hover:bg-sky-200 dark:bg-sky-900/30 dark:text-sky-400' },