fix(xai): keep image_generation on grok-4.6+ conversation requests

normalizeXAITool still strips Codex hosted image tools on older Grok
conversation models. grok-4.6 and later accept xAI native Imagine
tool, so keep client-supplied image_generation there and rewrite a
forced choice into allowed_tools. grok-4.20-* stays on the old strip
because that product line is not comparable to grok-4.6.

Closes: #5173
This commit is contained in:
程辉
2026-08-22 12:38:30 +08:00
parent 0a14eb70ce
commit dfdf183fcf
2 changed files with 250 additions and 7 deletions

View File

@@ -98,6 +98,7 @@ func (e *XAIExecutor) prepareResponsesRequestTo(ctx context.Context, req cliprox
// configured x_search injection, so no surviving choice references a deleted tool.
body = normalizeXAINamespaceToolChoice(body)
body = normalizeXAIForcedWebSearchToolChoice(body)
body = normalizeXAIForcedImageGenerationToolChoice(body)
body = pruneXAIOrphanedToolChoice(body)
body = normalizeXAIToolChoiceForTools(body)
if e.cfg != nil && e.cfg.XAI.InjectXSearch {
@@ -530,6 +531,91 @@ func preserveXAIResponsesOutputControls(body, source []byte, from sdktranslator.
return body
}
// xaiGrokImageGenerationMinVersion is the first Grok line that accepts xAI's
// native Responses image_generation tool. Older conversation models still
// reject that hosted type, so the executor keeps stripping it there.
var xaiGrokImageGenerationMinVersion = xaiGrokVersion{major: 4, minor: 6}
type xaiGrokVersion struct {
major int
minor int
}
// xaiSupportsNativeImageGeneration reports whether the Grok model accepts
// xAI's native Responses image_generation tool. grok-4.20-* is an older
// product line whose dotted minor is not comparable to grok-4.6.
func xaiSupportsNativeImageGeneration(model string) bool {
name := strings.ToLower(strings.TrimSpace(thinking.ParseSuffix(model).ModelName))
if idx := strings.LastIndex(name, "/"); idx >= 0 {
name = name[idx+1:]
}
if name == "" || !strings.HasPrefix(name, "grok-") {
return false
}
rest := strings.TrimPrefix(name, "grok-")
if rest == "4.20" || strings.HasPrefix(rest, "4.20-") {
return false
}
ver, ok := xaiParseGrokVersionPrefix(rest)
if !ok {
return false
}
return xaiCompareGrokVersion(ver, xaiGrokImageGenerationMinVersion) >= 0
}
func xaiParseGrokVersionPrefix(rest string) (xaiGrokVersion, bool) {
i := 0
for i < len(rest) && rest[i] >= '0' && rest[i] <= '9' {
i++
}
if i == 0 {
return xaiGrokVersion{}, false
}
major, err := strconv.Atoi(rest[:i])
if err != nil {
return xaiGrokVersion{}, false
}
if i == len(rest) || rest[i] != '.' {
return xaiGrokVersion{major: major, minor: -1}, true
}
j := i + 1
for j < len(rest) && rest[j] >= '0' && rest[j] <= '9' {
j++
}
if j == i+1 {
return xaiGrokVersion{major: major, minor: -1}, true
}
minor, err := strconv.Atoi(rest[i+1 : j])
if err != nil {
return xaiGrokVersion{}, false
}
return xaiGrokVersion{major: major, minor: minor}, true
}
func xaiCompareGrokVersion(a, b xaiGrokVersion) int {
if a.major != b.major {
if a.major < b.major {
return -1
}
return 1
}
aMinor := a.minor
if aMinor < 0 {
aMinor = 0
}
bMinor := b.minor
if bMinor < 0 {
bMinor = 0
}
if aMinor < bMinor {
return -1
}
if aMinor > bMinor {
return 1
}
return 0
}
func sanitizeXAIResponsesBody(body []byte, model string) []byte {
// stop is supported by Chat Completions but not by xAI's Responses API.
body, _ = sjson.DeleteBytes(body, "stop")
@@ -590,8 +676,18 @@ func ensureXAINativeXSearchAllowedTools(body []byte) []byte {
// normalizeXAIForcedWebSearchToolChoice rewrites Codex's hosted-tool choice
// into the allowed_tools form accepted by xAI's ModelToolChoice schema.
func normalizeXAIForcedWebSearchToolChoice(body []byte) []byte {
return normalizeXAIForcedHostedToolChoice(body, xaiWebSearchToolType)
}
// normalizeXAIForcedImageGenerationToolChoice rewrites a forced image_generation
// choice into the same allowed_tools form used for web_search.
func normalizeXAIForcedImageGenerationToolChoice(body []byte) []byte {
return normalizeXAIForcedHostedToolChoice(body, xaiImageGenerationToolType)
}
func normalizeXAIForcedHostedToolChoice(body []byte, toolType string) []byte {
choice := gjson.GetBytes(body, "tool_choice")
if !choice.IsObject() || strings.TrimSpace(choice.Get("type").String()) != xaiWebSearchToolType {
if !choice.IsObject() || strings.TrimSpace(choice.Get("type").String()) != toolType {
return body
}
@@ -731,13 +827,14 @@ func normalizeXAITools(body []byte) []byte {
if !gjson.ValidBytes(body) {
return body
}
keepImageGeneration := xaiSupportsNativeImageGeneration(gjson.GetBytes(body, "model").String())
original := body
normalizeAtPath := func(path string) bool {
tools := gjson.GetBytes(body, path)
if !tools.Exists() || !tools.IsArray() {
return true
}
filtered, changed, ok := normalizeXAIToolArray(tools)
filtered, changed, ok := normalizeXAIToolArray(tools, keepImageGeneration)
if !ok {
return false
}
@@ -827,7 +924,7 @@ func promoteXAIAdditionalTools(body []byte) []byte {
return updated
}
func normalizeXAIToolArray(tools gjson.Result) ([]byte, bool, bool) {
func normalizeXAIToolArray(tools gjson.Result, keepImageGeneration bool) ([]byte, bool, bool) {
toolItems := tools.Array()
filtered := make([][]byte, 0, len(toolItems))
changed := false
@@ -838,7 +935,7 @@ func normalizeXAIToolArray(tools gjson.Result) ([]byte, bool, bool) {
namespaceName := tool.Get("name").String()
if namespaceTools := tool.Get("tools"); namespaceTools.IsArray() {
for _, nestedTool := range namespaceTools.Array() {
nestedRaw, nestedChanged, ok := normalizeXAITool(nestedTool, namespaceName)
nestedRaw, nestedChanged, ok := normalizeXAITool(nestedTool, namespaceName, keepImageGeneration)
if !ok {
return nil, false, false
}
@@ -850,7 +947,7 @@ func normalizeXAIToolArray(tools gjson.Result) ([]byte, bool, bool) {
}
continue
}
raw, toolChanged, ok := normalizeXAITool(tool, "")
raw, toolChanged, ok := normalizeXAITool(tool, "", keepImageGeneration)
if !ok {
return nil, false, false
}
@@ -944,10 +1041,13 @@ func normalizeXAINamespaceToolChoice(body []byte) []byte {
return body
}
func normalizeXAITool(tool gjson.Result, namespaceName string) ([]byte, bool, bool) {
func normalizeXAITool(tool gjson.Result, namespaceName string, keepImageGeneration bool) ([]byte, bool, bool) {
toolType := tool.Get("type").String()
changed := false
if toolType == xaiToolSearchType || toolType == xaiImageGenerationToolType {
if toolType == xaiToolSearchType {
return nil, true, true
}
if toolType == xaiImageGenerationToolType && !keepImageGeneration {
return nil, true, true
}
if toolType == xaiCustomToolType && tool.Get("name").String() == "apply_patch" {

View File

@@ -980,6 +980,149 @@ func TestPruneXAIOrphanedToolChoice(t *testing.T) {
}
}
func TestXAISupportsNativeImageGeneration(t *testing.T) {
t.Parallel()
tests := []struct {
model string
want bool
}{
{model: "", want: false},
{model: "grok-4.5", want: false},
{model: "grok-4.3", want: false},
{model: "grok-4", want: false},
{model: "grok-4.20-0309-reasoning", want: false},
{model: "grok-4.20-multi-agent-0309", want: false},
{model: "grok-build-0.1", want: false},
{model: "grok-composer-2.5-fast", want: false},
{model: "grok-3-mini", want: false},
{model: "gpt-5.6", want: false},
{model: "grok-4.6", want: true},
{model: "grok-4.6(high)", want: true},
{model: "xai/grok-4.6", want: true},
{model: "grok-4.7", want: true},
{model: "grok-5", want: true},
{model: "grok-5.0", want: true},
}
for _, tt := range tests {
t.Run(tt.model, func(t *testing.T) {
t.Parallel()
if got := xaiSupportsNativeImageGeneration(tt.model); got != tt.want {
t.Fatalf("xaiSupportsNativeImageGeneration(%q) = %t, want %t", tt.model, got, tt.want)
}
})
}
}
func TestNormalizeXAITools_ImageGenerationByModel(t *testing.T) {
t.Parallel()
tests := []struct {
name string
body []byte
wantKeep bool
wantAction string
}{
{
name: "missing model still strips",
body: []byte(`{"tools":[{"type":"image_generation"},{"type":"web_search"}]}`),
wantKeep: false,
},
{
name: "grok-4.5 strips",
body: []byte(`{"model":"grok-4.5","tools":[{"type":"image_generation"},{"type":"web_search"}]}`),
wantKeep: false,
},
{
name: "grok-4.20 strips despite larger minor",
body: []byte(`{"model":"grok-4.20-0309-reasoning","tools":[{"type":"image_generation"},{"type":"web_search"}]}`),
wantKeep: false,
},
{
name: "grok-4.6 keeps action",
body: []byte(`{"model":"grok-4.6","tools":[{"type":"image_generation","action":"generate"},{"type":"web_search"}]}`),
wantKeep: true,
wantAction: "generate",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
out := normalizeXAITools(tt.body)
tools := gjson.GetBytes(out, "tools").Array()
foundImage := false
foundWebSearch := false
var imageTool gjson.Result
for _, tool := range tools {
switch tool.Get("type").String() {
case "image_generation":
foundImage = true
imageTool = tool
case "web_search":
foundWebSearch = true
}
}
if !foundWebSearch {
t.Fatalf("web_search missing; body=%s", out)
}
if foundImage != tt.wantKeep {
t.Fatalf("image_generation kept=%t, want %t; body=%s", foundImage, tt.wantKeep, out)
}
if tt.wantKeep && tt.wantAction != "" {
if got := imageTool.Get("action").String(); got != tt.wantAction {
t.Fatalf("image_generation.action = %q, want %q; body=%s", got, tt.wantAction, out)
}
}
})
}
}
func TestXAIExecutorPrepareKeepsNativeImageGenerationForGrok46(t *testing.T) {
t.Parallel()
exec := NewXAIExecutor(&config.Config{})
prepared, err := exec.prepareResponsesRequest(context.Background(), cliproxyexecutor.Request{
Model: "grok-4.6",
Payload: []byte(`{
"model":"grok-4.6",
"input":"draw a red circle",
"tools":[{"type":"image_generation","action":"generate"}],
"tool_choice":{"type":"image_generation"}
}`),
}, cliproxyexecutor.Options{
SourceFormat: sdktranslator.FormatOpenAIResponse,
Stream: false,
}, false)
if err != nil {
t.Fatalf("prepareResponsesRequest() error = %v", err)
}
tools := gjson.GetBytes(prepared.body, "tools").Array()
if len(tools) != 1 {
t.Fatalf("tools length = %d, want 1; body=%s", len(tools), prepared.body)
}
if got := tools[0].Get("type").String(); got != "image_generation" {
t.Fatalf("tools.0.type = %q, want image_generation; body=%s", got, prepared.body)
}
if got := tools[0].Get("action").String(); got != "generate" {
t.Fatalf("tools.0.action = %q, want generate; body=%s", got, prepared.body)
}
choice := gjson.GetBytes(prepared.body, "tool_choice")
if got := choice.Get("type").String(); got != "allowed_tools" {
t.Fatalf("tool_choice.type = %q, want allowed_tools; body=%s", got, prepared.body)
}
if got := choice.Get("mode").String(); got != "required" {
t.Fatalf("tool_choice.mode = %q, want required; body=%s", got, prepared.body)
}
allowed := choice.Get("tools").Array()
if len(allowed) != 1 {
t.Fatalf("tool_choice.tools length = %d, want 1; body=%s", len(allowed), prepared.body)
}
if got := allowed[0].Get("type").String(); got != "image_generation" {
t.Fatalf("tool_choice.tools.0.type = %q, want image_generation; body=%s", got, prepared.body)
}
}
func TestXAIExecutorPrepareDropsOrphanedToolChoiceBeforeXSearchInject(t *testing.T) {
t.Parallel()