diff --git a/internal/catalog/catalog.go b/internal/catalog/catalog.go index d28d3467..1d5be24f 100644 --- a/internal/catalog/catalog.go +++ b/internal/catalog/catalog.go @@ -556,6 +556,19 @@ func DrawsID(id string) bool { return false } +// ImagesAPI is whether a model draws on an images API (/images/generations) +// rather than answering in chat with pictures: gpt-image, dall-e, imagen, +// flux, seedream… — not gemini-*-image or gpt-5-image, which chat. +func ImagesAPI(id string) bool { + id = strings.ToLower(id) + for _, w := range []string{"gpt-image", "chatgpt-image", "dall-e", "imagen", "imagine", "qwen-image", "wanx", "wan2", "seedream", "cogview", "flux", "stable-diffusion", "sdxl", "kolors", "hidream"} { + if strings.Contains(id, w) { + return true + } + } + return false +} + // Drawers are the models of one models.dev provider that make images, // newest first. func Drawers(id string) []Model { diff --git a/internal/gateway/draw.go b/internal/gateway/draw.go index 05ecef2d..f179e7c5 100644 --- a/internal/gateway/draw.go +++ b/internal/gateway/draw.go @@ -435,15 +435,12 @@ func isGoogle(p provider.Provider) bool { // viaFor is where model is best asked to draw at p. func viaFor(p provider.Provider, model string) drawVia { - m := strings.ToLower(model) - onImages := slices.ContainsFunc([]string{"gpt-image", "chatgpt-image", "dall-e", "imagen", "imagine", "qwen-image", "wanx", "wan2", "seedream", "cogview", "flux", "stable-diffusion", "sdxl", "kolors", "hidream"}, - func(w string) bool { return strings.Contains(m, w) }) switch { - case isGoogle(p) && strings.Contains(m, "gemini"): + case isGoogle(p) && strings.Contains(strings.ToLower(model), "gemini"): return viaGemini case provider.HostOf(p.Base(provider.Chat)) == "openrouter.ai": return viaChat - case onImages: + case catalog.ImagesAPI(model): return viaImages } return viaChat diff --git a/internal/gui/providers.go b/internal/gui/providers.go index 5b93dc39..24ac53d1 100644 --- a/internal/gui/providers.go +++ b/internal/gui/providers.go @@ -583,7 +583,8 @@ func providerRoutes(mux *http.ServeMux, w Windows) { return } if len(req.Test) > 0 { - ctx, cancel := context.WithTimeout(r.Context(), 90*time.Second) + // an image model's test draws a picture, which takes longer + ctx, cancel := context.WithTimeout(r.Context(), 3*time.Minute) defer cancel() writeJSON(rw, map[string]any{"results": p.TestModels(ctx, req.Test)}) return diff --git a/internal/provider/test.go b/internal/provider/test.go index 9ad1a39d..2946aeeb 100644 --- a/internal/provider/test.go +++ b/internal/provider/test.go @@ -45,7 +45,7 @@ func (p Provider) Test(ctx context.Context) []Result { continue } url, body := tiny(q, proto, model) - out = append(out, probe(ctx, q, proto, url, q.Prepare([]byte(body)), model)) + out = append(out, probe(ctx, q, proto, url, q.Prepare([]byte(body)), model, testWait)) } return out } @@ -72,6 +72,32 @@ func tinyBody(q Provider, proto Protocol, model string) (url, body string) { return "", "" } +// testWait is how long a probe waits for an answer; drawWait, for a +// picture, which takes an images API longer than a word takes chat. +const ( + testWait = 20 * time.Second + drawWait = 2 * time.Minute +) + +// drawsOnImages is whether model is asked on p's images API, as the +// gateway draws with it there: chat would be turned away, or answered in +// no shape a chat test knows. OpenRouter draws everything in chat. +func (p Provider) drawsOnImages(model string) bool { + return p.Account == nil && p.Chat != "" && HostOf(p.Chat) != "openrouter.ai" && catalog.ImagesAPI(model) +} + +// tinyDrawing is the smallest images request for model: one picture of +// next to nothing, at the lowest quality gpt-image offers, at the size +// the vendor draws by default (the smallest one takes differs by model). +func tinyDrawing(q Provider, model string) (url, body string) { + req := map[string]any{"model": model, "prompt": "a dot", "n": 1} + if m := strings.ToLower(model); strings.Contains(m, "gpt-image") || strings.Contains(m, "chatgpt-image") { + req["quality"] = "low" + } + b, _ := json.Marshal(req) + return strings.TrimRight(q.Chat, "/") + "/images/generations", string(b) +} + // TestModels sends each of models the smallest request, a few at a time, // on the endpoint it's served on and with a key that sees it: whether each // answers, not only whether the vendor does. Results are in models' order. @@ -103,9 +129,13 @@ func (p Provider) testOne(ctx context.Context, model string) Result { return Result{Model: model, Error: "this provider can't be sent a test request"} } // the endpoint the vendor's list says serves it, else Anthropic's for a - // Claude model, else the one it prefers + // Claude model, else the one it prefers; an image model draws on the + // chat endpoint's images API proto := protos[0] - if apis := p.APIs(model); apis != nil { + draws := p.drawsOnImages(model) && slices.Contains(protos, Chat) + if draws { + proto = Chat + } else if apis := p.APIs(model); apis != nil { if i := slices.IndexFunc(protos, func(pr Protocol) bool { return slices.Contains(apis, pr) }); i >= 0 { proto = protos[i] } @@ -129,8 +159,21 @@ func (p Provider) testOne(ctx context.Context, model string) Result { if !ok { return Result{Protocol: proto, Model: model, Error: "no key that's on sees this model"} } + if draws { + // as the gateway does, one the images API doesn't serve is tried + // in chat, and only its answer said when that fails too + url, body := tinyDrawing(q, model) + r := probe(ctx, q, proto, url, []byte(body), model, drawWait) + if !r.OK && (r.Status == 404 || r.Status == 405) { + url, body := tiny(q, proto, model) + if c := probe(ctx, q, proto, url, q.Prepare([]byte(body)), model, testWait); c.OK { + return c + } + } + return r + } url, body := tiny(q, proto, model) - return probe(ctx, q, proto, url, q.Prepare([]byte(body)), model) + return probe(ctx, q, proto, url, q.Prepare([]byte(body)), model, testWait) } // keyFor is p using the first key on that works with proto: one made for @@ -152,7 +195,8 @@ func (p Provider) keyFor(proto Protocol) (Provider, bool) { // testModel is the model a probe of proto's endpoint asks for: the first // exposed one q's key sees, else the first it sees at all — preferring a -// Claude model on the Anthropic endpoint. +// Claude model on the Anthropic endpoint, and one that chats to one that +// draws on an images API. func (p Provider) testModel(q Provider, proto Protocol) string { k := q.first() var pools [][]catalog.Model @@ -162,6 +206,9 @@ func (p Provider) testModel(q Provider, proto Protocol) string { pools = append(pools, p.Available()) for _, want := range []func(string) bool{ func(id string) bool { + if p.drawsOnImages(id) { + return false + } if apis := p.APIs(id); apis != nil { return slices.Contains(apis, proto) } @@ -204,13 +251,13 @@ func AuthHeaders(p Provider, proto Protocol) map[string]string { return map[string]string{"Authorization": "Bearer " + p.Key} } -func probe(ctx context.Context, p Provider, proto Protocol, url string, body []byte, model string) Result { +func probe(ctx context.Context, p Provider, proto Protocol, url string, body []byte, model string, wait time.Duration) Result { r := Result{Protocol: proto, Model: model} if model == "" { r.Error = "no model to try: expose one, or refresh the model list" return r } - ctx, cancel := context.WithTimeout(ctx, 20*time.Second) + ctx, cancel := context.WithTimeout(ctx, wait) defer cancel() req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body)) if err != nil { diff --git a/internal/provider/test_test.go b/internal/provider/test_test.go index dfb8c5f3..d569288d 100644 --- a/internal/provider/test_test.go +++ b/internal/provider/test_test.go @@ -6,6 +6,7 @@ import ( "net/http" "net/http/httptest" "path/filepath" + "strings" "sync" "testing" ) @@ -106,3 +107,64 @@ func TestTestModels(t *testing.T) { t.Fatalf("seen %v", seen) } } + +// An image model is tested as the gateway draws with it, on the images +// API with one small picture — a chat request it would turn away — while a +// chat model is still asked in chat. The vendor-wide Test picks the chat one. +func TestTestModelsDraws(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + t.Setenv("XDG_CONFIG_HOME", filepath.Join(home, ".config")) + t.Setenv("XDG_CACHE_HOME", filepath.Join(home, ".cache")) + var mu sync.Mutex + seen := map[string]string{} // model → path it was asked on + bodies := map[string]map[string]any{} + srv := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodGet { + rw.WriteHeader(http.StatusNotFound) + return + } + var in map[string]any + _ = json.NewDecoder(r.Body).Decode(&in) + model, _ := in["model"].(string) + mu.Lock() + seen[model] = r.URL.Path + bodies[model] = in + mu.Unlock() + draws := strings.Contains(model, "image") || strings.Contains(model, "flux") + switch { + case r.URL.Path == "/v1/images/generations" && draws: + _, _ = rw.Write([]byte(`{"data":[{"url":"https://img.example/1.png"}]}`)) + case r.URL.Path == "/v1/chat/completions" && !draws: + _, _ = rw.Write([]byte(`{"choices":[{"message":{"content":"hi"}}]}`)) + default: + rw.WriteHeader(http.StatusBadRequest) + _, _ = rw.Write([]byte(`{"error":{"message":"wrong endpoint for ` + model + `"}}`)) + } + })) + defer srv.Close() + p := Provider{ID: "custom", Name: "Custom", Chat: srv.URL + "/v1", Key: "sk-x", + Models: []string{"gpt-image-1", "flux-kontext-pro", "gpt-6-sol"}} + rs := p.TestModels(context.Background(), []string{"gpt-image-1", "flux-kontext-pro", "gpt-6-sol"}) + for _, r := range rs { + if !r.OK { + t.Fatalf("results %+v", rs) + } + } + if seen["gpt-image-1"] != "/v1/images/generations" || seen["flux-kontext-pro"] != "/v1/images/generations" || seen["gpt-6-sol"] != "/v1/chat/completions" { + t.Fatalf("seen %v", seen) + } + if b := bodies["gpt-image-1"]; b["n"] != float64(1) || b["quality"] != "low" || b["prompt"] == "" { + t.Fatalf("gpt-image body %v", b) + } + if b := bodies["flux-kontext-pro"]; b["quality"] != nil || b["messages"] != nil { + t.Fatalf("flux body %v", b) + } + // the vendor-wide Test asks the chat model, though an image one is first + mu.Lock() + clear(seen) + mu.Unlock() + if rs := p.Test(context.Background()); len(rs) != 1 || !rs[0].OK || rs[0].Model != "gpt-6-sol" { + t.Fatalf("Test %+v", rs) + } +}