Skip to content

Commit 90206c6

Browse files
heavy-dclaude
andcommitted
feat(sdk): enumerate provider models of every modality in the model catalog
The SDK model catalog (/api/sdk/v1/models) built its list from getAllModels, which only enumerates language models from configured providers. Image, TTS, music, ASR, video, and embedding models appeared only when they happened to be in RECOMMENDED_MODELS — an SDK client asking for image_model compatibility got exactly one entry (GPT Image 2) while the web editor's picker, which goes through availableForKind, showed the full provider lists. Add collectProviderCatalogModels: one pass per configured provider calling each getAvailable*Models once, no task filtering, so text_to_image and image_to_image capable models are both included. Per-list failures degrade to an empty list instead of dropping the provider. The catalog service merges these into the existing gathering behind a 60s per-user TTL cache. Only the remote provider enumeration is cached — local download state (HF cache scan, download manager) stays fresh on every call, so a finished download still flips to ready_local immediately. Worker-scoped catalogs are unchanged and never enumerate providers. getAllModels and availableForKind are untouched; existing consumers keep their exact behavior. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1 parent 3ad78c4 commit 90206c6

3 files changed

Lines changed: 199 additions & 1 deletion

File tree

packages/websocket/src/sdk/sdk-model-catalog-service.ts

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import {
99
type SdkV1ModelCatalogQuery
1010
} from "@nodetool-ai/protocol/api-schemas/sdk-models-v1.js";
1111
import {
12+
collectProviderCatalogModels,
1213
getAllModels,
1314
getAvailableProviderIds
1415
} from "../trpc/routers/models.js";
@@ -161,6 +162,32 @@ export function projectSdkModelCatalog(
161162
});
162163
}
163164

165+
// The per-provider model enumeration fans out to remote provider APIs and is
166+
// the slow part of a catalog request. Its result is pure remote data (no
167+
// local download state), so a short TTL cache is safe and keeps repeated SDK
168+
// catalog polls fast.
169+
const PROVIDER_CATALOG_TTL_MS = 60_000;
170+
const providerCatalogCache = new Map<
171+
string,
172+
{ at: number; models: readonly UnifiedModel[] }
173+
>();
174+
175+
async function getCachedProviderCatalogModels(
176+
userId: string
177+
): Promise<readonly UnifiedModel[]> {
178+
const cached = providerCatalogCache.get(userId);
179+
if (cached && Date.now() - cached.at < PROVIDER_CATALOG_TTL_MS) {
180+
return cached.models;
181+
}
182+
const models = await collectProviderCatalogModels(userId);
183+
providerCatalogCache.set(userId, { at: Date.now(), models });
184+
return models;
185+
}
186+
187+
export function clearProviderCatalogCache(): void {
188+
providerCatalogCache.clear();
189+
}
190+
164191
function dedupeCatalogModels(models: readonly UnifiedModel[]): UnifiedModel[] {
165192
const byKey = new Map<string, UnifiedModel>();
166193
for (const model of models) {
@@ -176,8 +203,12 @@ export async function getSdkV1ModelCatalog(args: {
176203
query: SdkV1ModelCatalogQuery;
177204
recommendedModels?: readonly UnifiedModel[];
178205
getWorkerModels?: () => Promise<readonly UnifiedModel[]>;
206+
getProviderCatalogModels?: (
207+
userId: string
208+
) => Promise<readonly UnifiedModel[]>;
179209
}): Promise<SdkV1ModelCatalog> {
180210
let availableModels: readonly UnifiedModel[];
211+
let providerCatalogModels: readonly UnifiedModel[];
181212
let providerIds: readonly string[];
182213
if (args.query.scope === "worker") {
183214
if (!args.getWorkerModels) {
@@ -186,10 +217,14 @@ export async function getSdkV1ModelCatalog(args: {
186217
);
187218
}
188219
availableModels = await args.getWorkerModels();
220+
providerCatalogModels = [];
189221
providerIds = [];
190222
} else {
191-
[availableModels, providerIds] = await Promise.all([
223+
[availableModels, providerCatalogModels, providerIds] = await Promise.all([
192224
getAllModels(args.userId),
225+
(args.getProviderCatalogModels ?? getCachedProviderCatalogModels)(
226+
args.userId
227+
),
193228
getAvailableProviderIds(args.userId)
194229
]);
195230
}
@@ -199,6 +234,7 @@ export async function getSdkV1ModelCatalog(args: {
199234
];
200235
const models = dedupeCatalogModels([
201236
...availableModels,
237+
...providerCatalogModels,
202238
...recommendedModels
203239
]);
204240
const manager =

packages/websocket/src/trpc/routers/models.ts

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -928,6 +928,53 @@ async function collectProviderModelsForKind(
928928
return out;
929929
}
930930

931+
/**
932+
* Every remote model the user's configured providers can enumerate, across all
933+
* non-language modalities, for the SDK model catalog. Unlike
934+
* `collectProviderModelsForKind` this makes one pass per provider (each
935+
* `getAvailable*Models` called once, no task filtering) so text_to_image and
936+
* image_to_image capable models are both included. Language models are not
937+
* collected here — `getAllModels` already enumerates them.
938+
*/
939+
export async function collectProviderCatalogModels(
940+
userId: string
941+
): Promise<UnifiedModel[]> {
942+
const providerIds = await getAvailableProviderIds(userId);
943+
const perProvider = await Promise.all(
944+
providerIds.map((providerId) =>
945+
safeProviderCall(
946+
"catalogModels",
947+
{ provider: providerId, userId },
948+
async () => {
949+
const instance = await instantiateProvider(providerId, userId);
950+
if (!instance) return [];
951+
const collect = (
952+
fetchModels: () => Promise<Parameters<typeof toUnifiedModel>[0][]>,
953+
type: string
954+
) =>
955+
safeProviderCall(
956+
`catalogModels:${type}`,
957+
{ provider: providerId, userId },
958+
async () => (await fetchModels()).map((m) => toUnifiedModel(m, type)),
959+
[] as UnifiedModel[]
960+
);
961+
const lists = await Promise.all([
962+
collect(() => instance.getAvailableImageModels(), "image_model"),
963+
collect(() => instance.getAvailableEmbeddingModels(), "embedding_model"),
964+
collect(() => instance.getAvailableTTSModels(), "tts_model"),
965+
collect(() => instance.getAvailableMusicModels(), "music_model"),
966+
collect(() => instance.getAvailableASRModels(), "asr_model"),
967+
collect(() => instance.getAvailableVideoModels(), "video_model")
968+
]);
969+
return lists.flat();
970+
},
971+
[] as UnifiedModel[]
972+
)
973+
)
974+
);
975+
return perProvider.flat();
976+
}
977+
931978
function curatedForKind(kind: ModelSearchKind): UnifiedModel[] {
932979
const modality = KIND_TO_MODALITY[kind];
933980
// For text_generation/embedding/image/video, RECOMMENDED_MODELS entries are
Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
import { beforeEach, describe, expect, it, vi } from "vitest";
2+
import type { UnifiedModel } from "@nodetool-ai/protocol";
3+
4+
const getAllModels = vi.fn<(userId: string) => Promise<UnifiedModel[]>>();
5+
const getAvailableProviderIds = vi.fn<(userId: string) => Promise<string[]>>();
6+
const collectProviderCatalogModels =
7+
vi.fn<(userId: string) => Promise<UnifiedModel[]>>();
8+
9+
vi.mock("../src/trpc/routers/models.js", () => ({
10+
getAllModels: (userId: string) => getAllModels(userId),
11+
getAvailableProviderIds: (userId: string) => getAvailableProviderIds(userId),
12+
collectProviderCatalogModels: (userId: string) =>
13+
collectProviderCatalogModels(userId)
14+
}));
15+
16+
vi.mock("@nodetool-ai/huggingface", () => ({
17+
getExistingDownloadManager: () => null
18+
}));
19+
20+
import {
21+
clearProviderCatalogCache,
22+
getSdkV1ModelCatalog
23+
} from "../src/sdk/sdk-model-catalog-service.js";
24+
25+
const query = { scope: "local" as const, limit: 200 };
26+
27+
const languageModel: UnifiedModel = {
28+
id: "gpt-test",
29+
name: "GPT Test",
30+
type: "language_model",
31+
provider: "openai"
32+
};
33+
34+
const falImageModel: UnifiedModel = {
35+
id: "fal-ai/flux/schnell",
36+
name: "FLUX.1 Schnell",
37+
type: "image_model",
38+
provider: "fal_ai"
39+
};
40+
41+
const openaiImageModel: UnifiedModel = {
42+
id: "gpt-image-2",
43+
name: "GPT Image 2",
44+
type: "image_model",
45+
provider: "openai"
46+
};
47+
48+
beforeEach(() => {
49+
clearProviderCatalogCache();
50+
getAllModels.mockReset().mockResolvedValue([languageModel]);
51+
getAvailableProviderIds
52+
.mockReset()
53+
.mockResolvedValue(["openai", "fal_ai"]);
54+
collectProviderCatalogModels
55+
.mockReset()
56+
.mockResolvedValue([falImageModel, openaiImageModel]);
57+
});
58+
59+
describe("SDK model catalog provider models", () => {
60+
it("includes provider-enumerated non-language models as ready_remote", async () => {
61+
const catalog = await getSdkV1ModelCatalog({ userId: "alice", query });
62+
63+
const flux = catalog.entries.find(
64+
(entry) => entry.id === "fal-ai/flux/schnell"
65+
);
66+
expect(flux).toMatchObject({
67+
compatibility: "image_model",
68+
availability: "ready_remote",
69+
provider: "fal_ai"
70+
});
71+
});
72+
73+
it("dedupes provider-enumerated models against the recommended list", async () => {
74+
const catalog = await getSdkV1ModelCatalog({ userId: "alice", query });
75+
76+
// gpt-image-2 exists both in RECOMMENDED_MODELS and in the provider
77+
// enumeration; the catalog must carry it once per (type, provider, id).
78+
const gptImage = catalog.entries.filter(
79+
(entry) => entry.id === "gpt-image-2" && entry.provider === "openai"
80+
);
81+
expect(gptImage).toHaveLength(1);
82+
});
83+
84+
it("caches the provider enumeration per user", async () => {
85+
await getSdkV1ModelCatalog({ userId: "alice", query });
86+
await getSdkV1ModelCatalog({ userId: "alice", query });
87+
expect(collectProviderCatalogModels).toHaveBeenCalledTimes(1);
88+
89+
await getSdkV1ModelCatalog({ userId: "bob", query });
90+
expect(collectProviderCatalogModels).toHaveBeenCalledTimes(2);
91+
expect(collectProviderCatalogModels).toHaveBeenLastCalledWith("bob");
92+
});
93+
94+
it("prefers an injected provider-catalog fetcher over the cache", async () => {
95+
const injected = vi.fn().mockResolvedValue([falImageModel]);
96+
await getSdkV1ModelCatalog({
97+
userId: "alice",
98+
query,
99+
getProviderCatalogModels: injected
100+
});
101+
expect(injected).toHaveBeenCalledWith("alice");
102+
expect(collectProviderCatalogModels).not.toHaveBeenCalled();
103+
});
104+
105+
it("never enumerates providers for worker-scoped catalogs", async () => {
106+
const catalog = await getSdkV1ModelCatalog({
107+
userId: "alice",
108+
query: { scope: "worker" as const, limit: 200 },
109+
getWorkerModels: async () => [languageModel]
110+
});
111+
expect(collectProviderCatalogModels).not.toHaveBeenCalled();
112+
expect(getAllModels).not.toHaveBeenCalled();
113+
expect(catalog.scope).toBe("worker");
114+
});
115+
});

0 commit comments

Comments
 (0)