Skip to content

Commit f462400

Browse files
heavy-dclaudeCopilot
authored
feat(sdk): enumerate provider models of every modality in the model catalog (#4676)
* 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> * Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> * Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
1 parent 4eb0773 commit f462400

3 files changed

Lines changed: 210 additions & 1 deletion

File tree

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

Lines changed: 48 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,43 @@ 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+
function pruneProviderCatalogCache(now: number): void {
176+
for (const [key, value] of providerCatalogCache) {
177+
if (now - value.at >= PROVIDER_CATALOG_TTL_MS) {
178+
providerCatalogCache.delete(key);
179+
}
180+
}
181+
}
182+
183+
async function getCachedProviderCatalogModels(
184+
userId: string
185+
): Promise<readonly UnifiedModel[]> {
186+
const now = Date.now();
187+
pruneProviderCatalogCache(now);
188+
189+
const cached = providerCatalogCache.get(userId);
190+
if (cached && now - cached.at < PROVIDER_CATALOG_TTL_MS) {
191+
return cached.models;
192+
}
193+
const models = await collectProviderCatalogModels(userId);
194+
providerCatalogCache.set(userId, { at: now, models });
195+
return models;
196+
}
197+
198+
export function clearProviderCatalogCache(): void {
199+
providerCatalogCache.clear();
200+
}
201+
164202
function dedupeCatalogModels(models: readonly UnifiedModel[]): UnifiedModel[] {
165203
const byKey = new Map<string, UnifiedModel>();
166204
for (const model of models) {
@@ -176,8 +214,12 @@ export async function getSdkV1ModelCatalog(args: {
176214
query: SdkV1ModelCatalogQuery;
177215
recommendedModels?: readonly UnifiedModel[];
178216
getWorkerModels?: () => Promise<readonly UnifiedModel[]>;
217+
getProviderCatalogModels?: (
218+
userId: string
219+
) => Promise<readonly UnifiedModel[]>;
179220
}): Promise<SdkV1ModelCatalog> {
180221
let availableModels: readonly UnifiedModel[];
222+
let providerCatalogModels: readonly UnifiedModel[];
181223
let providerIds: readonly string[];
182224
if (args.query.scope === "worker") {
183225
if (!args.getWorkerModels) {
@@ -186,10 +228,14 @@ export async function getSdkV1ModelCatalog(args: {
186228
);
187229
}
188230
availableModels = await args.getWorkerModels();
231+
providerCatalogModels = [];
189232
providerIds = [];
190233
} else {
191-
[availableModels, providerIds] = await Promise.all([
234+
[availableModels, providerCatalogModels, providerIds] = await Promise.all([
192235
getAllModels(args.userId),
236+
(args.getProviderCatalogModels ?? getCachedProviderCatalogModels)(
237+
args.userId
238+
),
193239
getAvailableProviderIds(args.userId)
194240
]);
195241
}
@@ -199,6 +245,7 @@ export async function getSdkV1ModelCatalog(args: {
199245
];
200246
const models = dedupeCatalogModels([
201247
...availableModels,
248+
...providerCatalogModels,
202249
...recommendedModels
203250
]);
204251
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-model types (embeddings, images, TTS, music, ASR, video), for the
934+
* SDK model catalog. Unlike `collectProviderModelsForKind` this makes one pass
935+
* per provider (each `getAvailable*Models` called once, no task filtering) so
936+
* text_to_image and image_to_image capable models are both included. Language
937+
* models are not 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)