From dad2dc11586780f772d6b71efb89001f84a29b3a Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 6 Aug 2026 14:40:16 -0700 Subject: [PATCH] fix(ui): restore playground model filtering by endpoint Bring back the prior Chat model dropdown filter (including chat models on responses/anthropic/interactions and image models on image_edits), and map mode realtime so the realtime endpoint only lists compatible models --- .../components/chat_ui/ChatUI.test.tsx | 68 ++++++++++++++ .../playground/components/chat_ui/ChatUI.tsx | 9 +- .../components/chat_ui/EndpointUtils.test.tsx | 90 ++++++++++++------- .../components/chat_ui/EndpointUtils.tsx | 35 +++++--- .../chat_ui/mode_endpoint_mapping.tsx | 3 +- 5 files changed, 163 insertions(+), 42 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx index 7c2fc47bd63..dd1dfbc3389 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx @@ -117,6 +117,74 @@ describe("ChatUI", () => { }); }); + it("shows only endpoint-compatible models when chat endpoint is selected", async () => { + (fetchModelsModule.fetchAvailableModels as ReturnType).mockResolvedValueOnce([ + { model_group: "ChatModel", mode: "chat" }, + { model_group: "SpeechModel", mode: "audio_speech" }, + { model_group: "ImageModel", mode: "image_generation" }, + { model_group: "ResponsesModel", mode: "responses" }, + { model_group: "RealtimeModel", mode: "realtime" }, + { model_group: "NoModeModel" }, + ]); + + render( + , + ); + + await waitFor(() => { + expect(screen.getByText("Test Key")).toBeInTheDocument(); + }); + + await selectComboboxOption("Select an endpoint", "/v1/chat/completions"); + await openComboboxByPlaceholder("Select a Model"); + + await waitFor(() => { + expect(screen.getAllByText("ChatModel").length).toBeGreaterThan(0); + expect(screen.getAllByText("NoModeModel").length).toBeGreaterThan(0); + expect(screen.queryByText("SpeechModel")).toBeNull(); + expect(screen.queryByText("ImageModel")).toBeNull(); + expect(screen.queryByText("ResponsesModel")).toBeNull(); + expect(screen.queryByText("RealtimeModel")).toBeNull(); + }); + }); + + it("shows only realtime models when realtime endpoint is selected", async () => { + (fetchModelsModule.fetchAvailableModels as ReturnType).mockResolvedValueOnce([ + { model_group: "ChatModel", mode: "chat" }, + { model_group: "RealtimeModel", mode: "realtime" }, + { model_group: "NoModeModel" }, + ]); + + render( + , + ); + + await waitFor(() => { + expect(screen.getByText("Test Key")).toBeInTheDocument(); + }); + + await selectComboboxOption("Select an endpoint", "/v1/realtime"); + await openComboboxByPlaceholder("Select a Model"); + + await waitFor(() => { + expect(screen.getAllByText("RealtimeModel").length).toBeGreaterThan(0); + expect(screen.getAllByText("NoModeModel").length).toBeGreaterThan(0); + expect(screen.queryByText("ChatModel")).toBeNull(); + }); + }); + it("should show 'Enter custom model' option in model selector", async () => { render( = ({ }; const supportsStreamingToggle = endpointType === EndpointType.CHAT || endpointType === EndpointType.RESPONSES; + const modelsForEndpoint = useMemo( + () => filterModelsForEndpoint(modelInfo, endpointType as EndpointType), + [modelInfo, endpointType], + ); let modelEmptyText = "No models available for this key"; if (modelLoadError) { modelEmptyText = "Unable to load models for this key"; } else if (apiKeySource === "custom" && !apiKey.trim()) { modelEmptyText = "Enter a Virtual Key to load models"; + } else if (modelInfo.length > 0 && modelsForEndpoint.length === 0) { + modelEmptyText = "No models available for this endpoint"; } const inputPlaceholder = @@ -1405,7 +1412,7 @@ const ChatUI: React.FC = ({ onValueChange={onModelChange} options={[ { value: "custom", label: "Enter custom model" }, - ...modelInfo.map((model) => ({ + ...modelsForEndpoint.map((model) => ({ value: model.model_group, label: model.model_group, sublabel: model.mode ? `Mode: ${model.mode}` : undefined, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointUtils.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointUtils.test.tsx index 2bb1395428b..8e2a160cacc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointUtils.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointUtils.test.tsx @@ -1,37 +1,16 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import type { ModelGroup } from "@/components/llm_calls/fetch_models"; -import { determineEndpointType } from "./EndpointUtils"; +import { determineEndpointType, filterModelsForEndpoint, isModelCompatibleWithEndpoint } from "./EndpointUtils"; import { EndpointType } from "@/components/chat_ui/mode_endpoint_mapping"; -// Mock the getEndpointType function -vi.mock("@/components/chat_ui/mode_endpoint_mapping", () => ({ - EndpointType: { - IMAGE: "image", - VIDEO: "video", - CHAT: "chat", - RESPONSES: "responses", - IMAGE_EDITS: "image_edits", - ANTHROPIC_MESSAGES: "anthropic_messages", - EMBEDDINGS: "embeddings", - SPEECH: "speech", - TRANSCRIPTION: "transcription", - A2A_AGENTS: "a2a_agents", - }, - getEndpointType: vi.fn(), - ModelMode: { - AUDIO_SPEECH: "audio_speech", - AUDIO_TRANSCRIPTION: "audio_transcription", - IMAGE_GENERATION: "image_generation", - VIDEO_GENERATION: "video_generation", - CHAT: "chat", - RESPONSES: "responses", - IMAGE_EDITS: "image_edits", - ANTHROPIC_MESSAGES: "anthropic_messages", - EMBEDDING: "embedding", - }, -})); +vi.mock("@/components/chat_ui/mode_endpoint_mapping", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + getEndpointType: vi.fn(actual.getEndpointType), + }; +}); -// Import the mocked function import { getEndpointType } from "@/components/chat_ui/mode_endpoint_mapping"; describe("determineEndpointType", () => { @@ -210,10 +189,61 @@ describe("determineEndpointType", () => { vi.mocked(getEndpointType).mockReturnValue(EndpointType.CHAT); - // Test with different case - should not match const result = determineEndpointType("gpt-3.5-turbo", mockModelInfo); expect(getEndpointType).not.toHaveBeenCalled(); expect(result).toBe(EndpointType.CHAT); }); }); + +describe("isModelCompatibleWithEndpoint / filterModelsForEndpoint", () => { + beforeEach(() => { + vi.mocked(getEndpointType).mockImplementation((mode: string) => { + const map: Record = { + chat: EndpointType.CHAT, + responses: EndpointType.RESPONSES, + image_generation: EndpointType.IMAGE, + image_edits: EndpointType.IMAGE_EDITS, + audio_speech: EndpointType.SPEECH, + realtime: EndpointType.REALTIME, + embedding: EndpointType.EMBEDDINGS, + }; + return map[mode] ?? EndpointType.CHAT; + }); + }); + + it("keeps models with no mode for every endpoint", () => { + const model: ModelGroup = { model_group: "custom-proxy-model" }; + expect(isModelCompatibleWithEndpoint(model, EndpointType.CHAT)).toBe(true); + expect(isModelCompatibleWithEndpoint(model, EndpointType.REALTIME)).toBe(true); + expect(isModelCompatibleWithEndpoint(model, EndpointType.SPEECH)).toBe(true); + }); + + it("keeps chat models for responses, anthropic messages, and interactions", () => { + const chatModel: ModelGroup = { model_group: "gpt-4o", mode: "chat" }; + expect(isModelCompatibleWithEndpoint(chatModel, EndpointType.RESPONSES)).toBe(true); + expect(isModelCompatibleWithEndpoint(chatModel, EndpointType.ANTHROPIC_MESSAGES)).toBe(true); + expect(isModelCompatibleWithEndpoint(chatModel, EndpointType.INTERACTIONS)).toBe(true); + expect(isModelCompatibleWithEndpoint(chatModel, EndpointType.SPEECH)).toBe(false); + }); + + it("keeps image models for image_edits", () => { + const imageModel: ModelGroup = { model_group: "dall-e-3", mode: "image_generation" }; + expect(isModelCompatibleWithEndpoint(imageModel, EndpointType.IMAGE_EDITS)).toBe(true); + expect(isModelCompatibleWithEndpoint(imageModel, EndpointType.IMAGE)).toBe(true); + expect(isModelCompatibleWithEndpoint(imageModel, EndpointType.CHAT)).toBe(false); + }); + + it("keeps only realtime models for the realtime endpoint", () => { + const models: ModelGroup[] = [ + { model_group: "gpt-4o", mode: "chat" }, + { model_group: "gpt-realtime", mode: "realtime" }, + { model_group: "no-mode" }, + ]; + + expect(filterModelsForEndpoint(models, EndpointType.REALTIME).map((m) => m.model_group)).toEqual([ + "gpt-realtime", + "no-mode", + ]); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointUtils.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointUtils.tsx index 84579610a41..5b1a504be43 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointUtils.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointUtils.tsx @@ -1,22 +1,37 @@ import { ModelGroup } from "@/components/llm_calls/fetch_models"; import { EndpointType, getEndpointType } from "@/components/chat_ui/mode_endpoint_mapping"; -/** - * Determines the appropriate endpoint type based on the selected model - * - * @param selectedModel - The model identifier string - * @param modelInfo - Array of model information - * @returns The appropriate endpoint type - */ export const determineEndpointType = (selectedModel: string, modelInfo: ModelGroup[]): EndpointType => { - // Find the model information for the selected model const selectedModelInfo = modelInfo.find((option) => option.model_group === selectedModel); - // If model info is found and it has a mode, determine the endpoint type if (selectedModelInfo?.mode) { return getEndpointType(selectedModelInfo.mode); } - // Default to chat endpoint if no match is found return EndpointType.CHAT; }; + +export const isModelCompatibleWithEndpoint = (model: ModelGroup, endpointType: EndpointType): boolean => { + if (!model.mode) { + return true; + } + + const optionEndpoint = getEndpointType(model.mode); + + if ( + endpointType === EndpointType.RESPONSES || + endpointType === EndpointType.ANTHROPIC_MESSAGES || + endpointType === EndpointType.INTERACTIONS + ) { + return optionEndpoint === endpointType || optionEndpoint === EndpointType.CHAT; + } + + if (endpointType === EndpointType.IMAGE_EDITS) { + return optionEndpoint === endpointType || optionEndpoint === EndpointType.IMAGE; + } + + return optionEndpoint === endpointType; +}; + +export const filterModelsForEndpoint = (models: ModelGroup[], endpointType: EndpointType): ModelGroup[] => + models.filter((model) => isModelCompatibleWithEndpoint(model, endpointType)); diff --git a/ui/litellm-dashboard/src/components/chat_ui/mode_endpoint_mapping.tsx b/ui/litellm-dashboard/src/components/chat_ui/mode_endpoint_mapping.tsx index cd216d87d7b..765167b0a2c 100644 --- a/ui/litellm-dashboard/src/components/chat_ui/mode_endpoint_mapping.tsx +++ b/ui/litellm-dashboard/src/components/chat_ui/mode_endpoint_mapping.tsx @@ -11,7 +11,7 @@ export enum ModelMode { IMAGE_EDITS = "image_edits", ANTHROPIC_MESSAGES = "anthropic_messages", EMBEDDING = "embedding", - // add additional modes as needed + REALTIME = "realtime", } // Define an enum for the endpoint types your UI calls @@ -42,6 +42,7 @@ export const litellmModeMapping: Record = { [ModelMode.AUDIO_SPEECH]: EndpointType.SPEECH, [ModelMode.AUDIO_TRANSCRIPTION]: EndpointType.TRANSCRIPTION, [ModelMode.EMBEDDING]: EndpointType.EMBEDDINGS, + [ModelMode.REALTIME]: EndpointType.REALTIME, }; export const getEndpointType = (mode: string): EndpointType => {