Skip to content

Commit 387828f

Browse files
Avery WagarAvery Wagar
authored andcommitted
Support user-scoped custom model discovery
Bead: shared-infra-el5k.2.1
1 parent 6d499ba commit 387828f

5 files changed

Lines changed: 66 additions & 6 deletions

File tree

‎packages/api/src/endpoints/config/models.ts‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,10 @@ function headersFingerprint(headers: Record<string, string> | undefined): string
3636
return crypto.createHash('sha256').update(JSON.stringify(ordered)).digest('hex').slice(0, 16);
3737
}
3838

39+
function recordFingerprint(value: Record<string, string> | undefined): string {
40+
return headersFingerprint(value);
41+
}
42+
3943
interface ResolvedEndpoint {
4044
name: string;
4145
endpoint: TEndpoint;
@@ -171,7 +175,7 @@ export function createLoadConfigModels(deps: LoadConfigModelsDeps) {
171175
// endpoints that happen to share the same baseURL+apiKey but configure
172176
// different (potentially user-bound) headers don't reuse each other's
173177
// fetched model list within the same request.
174-
const uniqueKey = `${BASE_URL}__${API_KEY}__${headersFingerprint(endpointHeaders)}`;
178+
const uniqueKey = `${BASE_URL}__${API_KEY}__${headersFingerprint(endpointHeaders)}__${recordFingerprint(models?.fetchQuery)}`;
175179

176180
if (models?.fetch && !apiKeyIsUserProvided && !baseURLIsUserProvided) {
177181
if (!fetchPromisesMap[uniqueKey]) {
@@ -190,6 +194,7 @@ export function createLoadConfigModels(deps: LoadConfigModelsDeps) {
190194
headers: endpointHeaders,
191195
direct: endpoint.directEndpoint,
192196
userIdQuery: models.userIdQuery,
197+
fetchQuery: models.fetchQuery,
193198
tokenKey,
194199
});
195200
}
@@ -233,6 +238,7 @@ export function createLoadConfigModels(deps: LoadConfigModelsDeps) {
233238
headers: baseURLIsUserProvided ? undefined : endpointHeaders,
234239
direct: endpoint.directEndpoint,
235240
userIdQuery: models.userIdQuery,
241+
fetchQuery: models.fetchQuery,
236242
skipCache: true,
237243
/** Fetched with the user's key/URL — always user-scoped */
238244
tokenKey: getTokenConfigKey(endpoint, name, req.user?.id ?? '', tenantId),

‎packages/api/src/endpoints/models.spec.ts‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,21 @@ describe('fetchModels', () => {
111111
);
112112
});
113113

114+
it('adds admin-configured model discovery query parameters', async () => {
115+
const models = await fetchModels({
116+
apiKey: 'testApiKey',
117+
baseURL: 'https://bifrost.test/v1',
118+
fetchQuery: { scope: 'personal' },
119+
name: 'bifrost-loki',
120+
});
121+
122+
expect(models).toEqual(['model-1', 'model-2']);
123+
expect(mockedAxios.get).toHaveBeenCalledWith(
124+
'https://bifrost.test/v1/models?scope=personal',
125+
expect.any(Object),
126+
);
127+
});
128+
114129
it('should pass custom headers to the API request', async () => {
115130
const customHeaders = {
116131
'X-Custom-Header': 'custom-value',

‎packages/api/src/endpoints/models.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,8 @@ export interface FetchModelsParams {
4545
azure?: boolean;
4646
/** Whether to send user ID as query parameter */
4747
userIdQuery?: boolean;
48+
/** Admin-configured query parameters for model discovery. */
49+
fetchQuery?: Record<string, string>;
4850
/** Whether to create token configuration from API response */
4951
createTokenConfig?: boolean;
5052
/** Cache key for token configuration (uses name if omitted) */
@@ -161,6 +163,7 @@ export async function fetchModels({
161163
direct = false,
162164
azure = false,
163165
userIdQuery = false,
166+
fetchQuery,
164167
createTokenConfig = true,
165168
tokenKey,
166169
headers,
@@ -280,6 +283,7 @@ export async function fetchModels({
280283
}
281284

282285
const url = new URL(`${(baseURL ?? '').replace(/\/+$/, '')}${azure ? '' : '/models'}`);
286+
Object.entries(fetchQuery ?? {}).forEach(([key, value]) => url.searchParams.set(key, value));
283287
if (user && userIdQuery) {
284288
url.searchParams.append('user', user);
285289
}

‎packages/data-provider/specs/config-schemas.spec.ts‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -251,6 +251,36 @@ describe('endpointSchema deprecated fields', () => {
251251
});
252252
});
253253

254+
describe('endpointSchema fetched models', () => {
255+
const endpoint = {
256+
name: 'BifrostLoki',
257+
apiKey: 'service-key',
258+
baseURL: 'https://bifrost.example.test/v1',
259+
};
260+
261+
it('accepts a fetch-only endpoint without inventing a static fallback', () => {
262+
const result = endpointSchema.safeParse({
263+
...endpoint,
264+
models: { fetch: true, fetchQuery: { scope: 'personal' } },
265+
});
266+
267+
expect(result.success).toBe(true);
268+
if (result.success) {
269+
expect(result.data.models.default).toEqual([]);
270+
expect(result.data.models.fetchQuery).toEqual({ scope: 'personal' });
271+
}
272+
});
273+
274+
it('rejects an endpoint with neither discovery nor a default model', () => {
275+
const result = endpointSchema.safeParse({
276+
...endpoint,
277+
models: {},
278+
});
279+
280+
expect(result.success).toBe(false);
281+
});
282+
});
283+
254284
describe('endpointSchema addParams validation', () => {
255285
const validEndpoint = {
256286
name: 'CustomEndpoint',

‎packages/data-provider/src/config.ts‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1123,11 +1123,16 @@ export const endpointSchema = baseEndpointSchema.merge(
11231123
* reads can show which key is configured without returning the secret. */
11241124
apiKeyPreview: z.string().optional(),
11251125
baseURL: z.string(),
1126-
models: z.object({
1127-
default: z.array(modelItemSchema).min(1),
1128-
fetch: z.boolean().optional(),
1129-
userIdQuery: z.boolean().optional(),
1130-
}),
1126+
models: z
1127+
.object({
1128+
default: z.array(modelItemSchema).default([]),
1129+
fetch: z.boolean().optional(),
1130+
fetchQuery: z.record(z.string()).optional(),
1131+
userIdQuery: z.boolean().optional(),
1132+
})
1133+
.refine((models) => models.fetch === true || models.default.length > 0, {
1134+
message: 'models requires fetch: true or at least one default model',
1135+
}),
11311136
iconURL: z.string().optional(),
11321137
modelDisplayLabel: z.string().optional(),
11331138
/**

0 commit comments

Comments
 (0)