Skip to content

fix(model-router): hide models not linked to mapped API keys #4255

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
iskhakov wants to merge 1 commit into from
Closed

fix(model-router): hide models not linked to mapped API keys #4255

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 26 additions & 5 deletions platform/backend/src/routes/proxy/model-router-resolver.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ import {
type SupportedProvider,
SupportedProviders,
} from "@shared";
import { ModelModel } from "@/models";
import { LlmProviderApiKeyModelLinkModel, ModelModel } from "@/models";
import { ApiError, type Model } from "@/types";

export type ModelRouterResolution = {
Expand All @@ -15,6 +15,7 @@ export type ModelRouterResolution = {
export async function resolveModelRoute(params: {
requestedModel: string;
allowedProviders?: Set<SupportedProvider>;
allowedApiKeyIds?: string[];
}): Promise<ModelRouterResolution> {
const requestedModel = params.requestedModel.trim();
if (!requestedModel) {
Expand All @@ -38,13 +39,20 @@ export async function resolveModelRoute(params: {
provider: explicit.provider,
});

if (providerMatches.length === 1) {
return toResolution(providerMatches[0], requestedModel);
const accessibleMatches = params.allowedApiKeyIds
? await filterModelsByLinkedApiKeys(
providerMatches,
params.allowedApiKeyIds,
)
: providerMatches;

if (accessibleMatches.length === 1) {
return toResolution(accessibleMatches[0], requestedModel);
}
if (providerMatches.length > 1) {
if (accessibleMatches.length > 1) {
throw new ApiError(
500,
`Ambiguous model resolution: "${requestedModel}" matched ${providerMatches.length} models.`,
`Ambiguous model resolution: "${requestedModel}" matched ${accessibleMatches.length} models.`,
);
}
}
Expand Down Expand Up @@ -94,6 +102,19 @@ function providerSortIndex(provider: SupportedProvider): number {
return index === -1 ? Number.MAX_SAFE_INTEGER : index;
}

async function filterModelsByLinkedApiKeys(
models: Model[],
apiKeyIds: string[],
): Promise<Model[]> {
if (models.length === 0 || apiKeyIds.length === 0) {
return [];
}
const linked =
await LlmProviderApiKeyModelLinkModel.getModelsForApiKeyIds(apiKeyIds);
const linkedIds = new Set(linked.map(({ model }) => model.id));
return models.filter((model) => linkedIds.has(model.id));
}

export function sortRoutableModels(models: Model[]): Model[] {
return [...models].sort((a, b) => {
const providerCompare =
Expand Down
82 changes: 81 additions & 1 deletion platform/backend/src/routes/proxy/routes/model-router.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,12 @@ import {
type ZodTypeProvider,
} from "fastify-type-provider-zod";
import { vi } from "vitest";
import { InteractionModel, ModelModel, VirtualApiKeyModel } from "@/models";
import {
InteractionModel,
LlmProviderApiKeyModelLinkModel,
ModelModel,
VirtualApiKeyModel,
} from "@/models";
import { afterEach, beforeEach, describe, expect, test } from "@/test";
import {
type AnthropicStubOptions,
Expand Down Expand Up @@ -100,6 +105,7 @@ async function createModelRouterVirtualKey(params: {
provider: params.provider,
},
);
await linkAllProviderModelsToApiKey(params.provider, chatApiKey.id);
return VirtualApiKeyModel.create({
chatApiKeyId: chatApiKey.id,
name: `${params.provider} model router virtual key`,
Expand All @@ -113,6 +119,17 @@ async function createModelRouterVirtualKey(params: {
});
}

async function linkAllProviderModelsToApiKey(
provider: SupportedProvider,
apiKeyId: string,
) {
const models = await ModelModel.findAll({ provider });
await LlmProviderApiKeyModelLinkModel.linkModelsToApiKey(
apiKeyId,
models.map((m) => m.id),
);
}

const ROUTABLE_PROVIDER_CASES: Array<{
provider: SupportedProvider;
modelId: string;
Expand Down Expand Up @@ -781,6 +798,7 @@ describe("model router proxy routes", () => {
const chatApiKey = await makeLlmProviderApiKey(organization.id, secret.id, {
provider: "openai",
});
await linkAllProviderModelsToApiKey("openai", chatApiKey.id);
const {
virtualKey: { id: virtualKeyId },
value,
Expand Down Expand Up @@ -852,6 +870,7 @@ describe("model router proxy routes", () => {
name: "Vertex AI",
provider: "gemini",
});
await linkAllProviderModelsToApiKey("gemini", systemKey.id);
const { value } = await VirtualApiKeyModel.create({
chatApiKeyId: systemKey.id,
name: "model-router-gemini-system-vk",
Expand Down Expand Up @@ -1001,6 +1020,8 @@ describe("model router proxy routes", () => {
groqSecret.id,
{ provider: "groq" },
);
await linkAllProviderModelsToApiKey("openai", openaiKey.id);
await linkAllProviderModelsToApiKey("groq", groqKey.id);
const { value } = await VirtualApiKeyModel.create({
chatApiKeyId: openaiKey.id,
name: "model-router-openai-groq-vk",
Expand Down Expand Up @@ -1175,6 +1196,8 @@ describe("model router proxy routes", () => {
groqSecret.id,
{ provider: "groq" },
);
await linkAllProviderModelsToApiKey("openai", openaiKey.id);
await linkAllProviderModelsToApiKey("groq", groqKey.id);
const { value } = await VirtualApiKeyModel.create({
chatApiKeyId: openaiKey.id,
name: "model-router-multi-vk",
Expand Down Expand Up @@ -1344,4 +1367,61 @@ describe("model router proxy routes", () => {
expect(response.statusCode).toBe(400);
expect(response.json().error.message).toContain("not mapped");
});

test("excludes models that exist for the mapped provider but are not linked to the virtual key's chat api key", async ({
makeAgent,
makeOrganization,
makeSecret,
makeLlmProviderApiKey,
}) => {
const app = createFastifyApp();
await app.register(modelRouterProxyRoutes);
await upsertModel({ provider: "groq", modelId: "llama-3.1-8b-instant" });
await upsertModel({ provider: "groq", modelId: "unlinked-model" });
const organization = await makeOrganization();
const { value } = await createModelRouterVirtualKey({
organizationId: organization.id,
provider: "groq",
makeSecret,
makeLlmProviderApiKey,
});
// Add a second model AFTER the virtual key so it's not linked
await upsertModel({ provider: "groq", modelId: "added-after-link" });
const agent = await makeAgent({
organizationId: organization.id,
name: "Unlinked Model Test Agent",
agentType: "llm_proxy",
});

const listResponse = await app.inject({
method: "GET",
url: `/v1/model-router/${agent.id}/models`,
headers: {
authorization: `Bearer ${value}`,
},
});

expect(listResponse.statusCode).toBe(200);
const listedIds = listResponse.json().data.map((m: { id: string }) => m.id);
expect(listedIds).toContain("groq:llama-3.1-8b-instant");
expect(listedIds).toContain("groq:unlinked-model");
expect(listedIds).not.toContain("groq:added-after-link");

const routeResponse = await app.inject({
method: "POST",
url: `/v1/model-router/${agent.id}/chat/completions`,
headers: {
"content-type": "application/json",
authorization: `Bearer ${value}`,
"user-agent": "test-client",
},
payload: {
model: "groq:added-after-link",
messages: [{ role: "user", content: "Hello" }],
},
});

expect(routeResponse.statusCode).toBe(404);
expect(routeResponse.json().error.message).toContain("not available");
});
});
58 changes: 33 additions & 25 deletions platform/backend/src/routes/proxy/routes/model-router.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,12 @@ import type { FastifyReply, FastifyRequest } from "fastify";
import type { FastifyPluginAsyncZod } from "fastify-type-provider-zod";
import { z } from "zod";
import logger from "@/logging";
import { AgentModel, ModelModel, VirtualApiKeyModel } from "@/models";
import {
AgentModel,
LlmProviderApiKeyModelLinkModel,
ModelModel,
VirtualApiKeyModel,
} from "@/models";
import { getSecretValueForLlmProviderApiKey } from "@/secrets-manager";
import type { Agent, LLMProvider } from "@/types";
import {
Expand Down Expand Up @@ -189,11 +194,7 @@ const modelRouterProxyRoutes: FastifyPluginAsyncZod = async (fastify) => {
const auth = await getModelRouterVirtualKeyAuth(request);
const agent = await getDefaultModelRouterAgent();
ensureModelRouterAgentAccess({ agent, auth });
return reply.send(
await listModels({
providers: getMappedProviders(auth),
}),
);
return reply.send(await listModels({ auth }));
},
);

Expand All @@ -215,11 +216,7 @@ const modelRouterProxyRoutes: FastifyPluginAsyncZod = async (fastify) => {
const auth = await getModelRouterVirtualKeyAuth(request);
const agent = await getModelRouterAgent(request.params.agentId);
ensureModelRouterAgentAccess({ agent, auth });
return reply.send(
await listModels({
providers: getMappedProviders(auth),
}),
);
return reply.send(await listModels({ auth }));
},
);

Expand Down Expand Up @@ -326,6 +323,7 @@ async function routeChatCompletion(
const resolution = await resolveModelRoute({
requestedModel: body.model,
allowedProviders: getMappedProviders(auth),
allowedApiKeyIds: getMappedApiKeyIds(auth),
});
const routedBody = {
...body,
Expand Down Expand Up @@ -366,6 +364,7 @@ async function routeResponse(request: FastifyRequest, reply: FastifyReply) {
const resolution = await resolveModelRoute({
requestedModel: chatBody.model,
allowedProviders: getMappedProviders(auth),
allowedApiKeyIds: getMappedApiKeyIds(auth),
});
const routedChatBody = {
...chatBody,
Expand Down Expand Up @@ -518,21 +517,24 @@ function handleModelRouterResponsesProvider(
}
}

async function listModels(params: { providers: Set<SupportedProvider> }) {
const providers = [...params.providers].filter((provider) =>
modelRouterSupportedProviders.has(provider),
);
const allModels = await ModelModel.findAll({ providers });
async function listModels(params: { auth: ModelRouterVirtualKeyAuth }) {
const apiKeyIds = [...params.auth.providerApiKeysByProvider.values()]
.filter((mapping) => modelRouterSupportedProviders.has(mapping.provider))
.map((mapping) => mapping.chatApiKeyId);
const linkedModels =
await LlmProviderApiKeyModelLinkModel.getModelsForApiKeyIds(apiKeyIds);
const chatModels = sortRoutableModels(
allModels.filter((model) => {
if (!ModelModel.supportsTextChat(model)) {
return false;
}
if (!modelRouterSupportedProviders.has(model.provider)) {
return false;
}
return true;
}),
linkedModels
.map(({ model }) => model)
.filter((model) => {
if (!ModelModel.supportsTextChat(model)) {
return false;
}
if (!modelRouterSupportedProviders.has(model.provider)) {
return false;
}
return true;
}),
);

return {
Expand Down Expand Up @@ -641,6 +643,12 @@ function getMappedProviders(
return new Set(auth.providerApiKeysByProvider.keys());
}

function getMappedApiKeyIds(auth: ModelRouterVirtualKeyAuth): string[] {
return [...auth.providerApiKeysByProvider.values()].map(
(mapping) => mapping.chatApiKeyId,
);
}

function isTranslatedModelRouterProvider(
provider: SupportedProvider,
): provider is TranslatedModelRouterProvider {
Expand Down
Loading