diff --git a/internal/catalog/public.go b/internal/catalog/public.go new file mode 100644 index 00000000..841662fa --- /dev/null +++ b/internal/catalog/public.go @@ -0,0 +1,96 @@ +package catalog + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/url" + "slices" + "strings" + "time" +) + +// PublicCatalogs are vendors whose /models, asked with a key, leaves out +// the image models the key can use: AIHubMix lists only its Gemini image +// models there, while gpt-image-2 and the rest draw with the same key. Its +// public catalog, by the host of the vendor's base URL, has them. +var PublicCatalogs = map[string]string{ + "aihubmix.com": "https://aihubmix.com/api/v1/models", +} + +// PublicDrawers are the image models the public catalog for base's host +// lists, marked as drawing; nil when there is none or it can't be read. +func PublicDrawers(ctx context.Context, base string) []Model { + u, err := url.Parse(strings.TrimSpace(base)) + if err != nil { + return nil + } + at, ok := PublicCatalogs[strings.TrimPrefix(u.Hostname(), "www.")] + if !ok { + return nil + } + ctx, cancel := context.WithTimeout(ctx, 8*time.Second) + defer cancel() + req, err := http.NewRequestWithContext(ctx, http.MethodGet, at, nil) + if err != nil { + return nil + } + req.Header.Set("Accept", "application/json") + req.Header.Set("User-Agent", "magpie") + res, err := http.DefaultClient.Do(req) + if err != nil { + return nil + } + defer res.Body.Close() + if res.StatusCode != http.StatusOK { + return nil + } + b, _ := io.ReadAll(io.LimitReader(res.Body, 16<<20)) + var list struct { + Data []json.RawMessage `json:"data"` + } + if json.Unmarshal(b, &list) != nil { + return nil + } + var out []Model + for _, raw := range list.Data { + var m struct { + ID string `json:"model_id"` + Name string `json:"model_name"` + Types string `json:"types"` + Out string `json:"output_modalities"` + In string `json:"input_modalities"` + Release string `json:"release_date"` + } + if json.Unmarshal(raw, &m) != nil || m.ID == "" { + continue + } + // what makes images and is asked the way magpie draws: Ideogram's + // V_2, DESCRIBE and UPSCALE are the vendor's own tasks + if !strings.Contains(m.Types, "image_generation") && !strings.Contains(m.Out, "image") || !DrawsID(m.ID) { + continue + } + d := Model{ID: m.ID, Name: m.Name, Draws: true, Released: m.Release} + if d.Name == "" { + d.Name = m.ID + } + if slices.Contains(strings.Split(m.In, ","), "image") { + d.Images, d.ImageInput = true, imageInputOf(true) + } + out = append(out, d) + } + return out +} + +// WithDrawers adds the drawers not already in ms. +func WithDrawers(ms, drawers []Model) []Model { + for _, d := range drawers { + if !slices.ContainsFunc(ms, func(m Model) bool { return m.ID == d.ID }) { + ms = append(ms, d) + } + } + return ms +} + +func imageInputOf(b bool) *bool { return &b } diff --git a/internal/provider/models.go b/internal/provider/models.go index 37d3c0bd..9e14891a 100644 --- a/internal/provider/models.go +++ b/internal/provider/models.go @@ -117,7 +117,7 @@ func (p Provider) fetchOne(ctx context.Context) ([]catalog.Model, string, error) if err != nil { return nil, u, err } - return p.planModels(ms), u, nil + return catalog.WithDrawers(p.planModels(ms), catalog.PublicDrawers(ctx, u)), u, nil } var errs []string for _, proto := range p.Speaks() { @@ -127,7 +127,8 @@ func (p Provider) fetchOne(ctx context.Context) ([]catalog.Model, string, error) if proto != Anthropic { base = p.fixV1(base, at) } - return p.planModels(ms), base, nil + // the image models its list leaves out (AIHubMix's gpt-image-2) + return catalog.WithDrawers(p.planModels(ms), catalog.PublicDrawers(ctx, base)), base, nil } // Chat and Responses at one base say the same thing if !slices.Contains(errs, err.Error()) { diff --git a/internal/provider/public_drawers_test.go b/internal/provider/public_drawers_test.go new file mode 100644 index 00000000..ff7997db --- /dev/null +++ b/internal/provider/public_drawers_test.go @@ -0,0 +1,53 @@ +package provider + +import ( + "context" + "net/http" + "net/http/httptest" + "slices" + "testing" + + "github.com/yetone/magpie/internal/catalog" +) + +// AIHubMix's /models leaves out gpt-image-2, which its key draws with; its +// public catalog has it, and the list kept for the provider gets it. +func TestPublicCatalogsImageModelsJoinTheList(t *testing.T) { + t.Setenv("XDG_CACHE_HOME", t.TempDir()) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.URL.Path == "/api/v1/models" { + if r.Header.Get("Authorization") != "" { + t.Errorf("the public catalog was sent the key") + } + w.Write([]byte(`{"success":true,"data":[ + {"model_id":"gpt-image-2","model_name":"GPT Image 2","types":"image_generation","input_modalities":"text,image","output_modalities":"image","release_date":"2026-04-21","pricing":{"input":5,"output":30}}, + {"model_id":"gemini-2.5-flash-image","model_name":"Gemini","types":"image_generation","output_modalities":"image","release_date":null}, + {"model_id":"V_2","model_name":"Ideogram","types":"image_generation","output_modalities":"image"}, + {"model_id":"gpt-5","model_name":"GPT-5","types":"llm","output_modalities":"text"}]}`)) + return + } + w.Write([]byte(`{"data":[{"id":"gpt-5"},{"id":"gemini-2.5-flash-image"}]}`)) + })) + defer server.Close() + old := catalog.PublicCatalogs + catalog.PublicCatalogs = map[string]string{"127.0.0.1": server.URL + "/api/v1/models"} + t.Cleanup(func() { catalog.PublicCatalogs = old }) + + p := Provider{ID: "hub", Chat: server.URL + "/v1", Key: "k"} + chat, err := p.Fetch(context.Background()) + if err != nil || len(chat) != 1 || chat[0].ID != "gpt-5" { + t.Fatalf("chat models %+v, %v", chat, err) + } + var ids []string + for _, m := range catalog.LiveDrawers("hub") { + ids = append(ids, m.ID) + if m.ID == "gpt-image-2" && (m.Name != "GPT Image 2" || !m.Images || m.Price != nil) { + t.Errorf("gpt-image-2 kept as %+v", m) + } + } + slices.Sort(ids) + if !slices.Equal(ids, []string{"gemini-2.5-flash-image", "gpt-image-2"}) { + t.Fatalf("drawers %v", ids) + } +}