scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

feat(ui): let the playground chat as a team so team-only models show up

Browse source
This commit is contained in:
Devin AI 2026-07-30 04:19:32 +00:00
parent 47f1fb394e
commit 7891e50044
7 changed files with 301 additions and 42 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
ui/litellm-dashboard/eslint-suppressions.json
View file

@ -1137,14 +1137,11 @@
"count": 1
},
"no-nested-ternary": {
"count": 7
"count": 6
},
"no-restricted-imports": {
"count": 2
},
"prefer-const": {
"count": 1
},
"react-hooks/set-state-in-effect": {
"count": 4
},

117
ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx
View file

@ -1,7 +1,11 @@
import { act, fireEvent, render, screen, waitFor } from "@testing-library/react";
import { act, fireEvent, screen, waitFor } from "@testing-library/react";
import { renderWithProviders as render, testQueryClient } from "@/../tests/test-utils";
import { beforeEach, describe, expect, it, vi } from "vitest";
import ChatUI from "./ChatUI";
import * as fetchModelsModule from "@/components/llm_calls/fetch_models";
import * as networking from "@/components/networking";
import { PLAYGROUND_TEAM_KEY_DURATION, playgroundTeamKeyStorageKey } from "../../hooks/usePlaygroundTeamSession";
import { getSecureItem } from "@/utils/secureStorage";
// Mock the fetchAvailableModels function
vi.mock("@/components/llm_calls/fetch_models", () => ({
@ -14,6 +18,8 @@ vi.mock("@/components/networking", () => ({
vectorStoreListCall: vi.fn().mockResolvedValue({ data: [] }),
getGuardrailsList: vi.fn().mockResolvedValue({ data: [] }),
modelHubCall: vi.fn().mockResolvedValue({ data: [] }),
teamListCall: vi.fn().mockResolvedValue([]),
keyCreateCall: vi.fn().mockResolvedValue({ key: "sk-team-key" }),
}));
// Mock scrollIntoView which is not available in jsdom
@ -26,6 +32,7 @@ describe("ChatUI", () => {
// Reset mocks before each test
vi.clearAllMocks();
sessionStorage.clear();
testQueryClient.clear();
// Mock scrollIntoView which is not available in JSDOM
Element.prototype.scrollIntoView = vi.fn();
@ -413,4 +420,112 @@ describe("ChatUI", () => {
}
}
});
describe("team scoped session", () => {
const clickOption = async (label: string) => {
await waitFor(() => {
const option = Array.from(document.querySelectorAll(".ant-select-item-option")).find(
(el) => el.getAttribute("title") === label,
);
expect(option).toBeTruthy();
});
const option = Array.from(document.querySelectorAll(".ant-select-item-option")).find(
(el) => el.getAttribute("title") === label,
);
await act(async () => {
fireEvent.click(option!);
});
};
const selectTeamKeySource = async () => {
const keySourceSelect = screen
.getByText("Virtual Key Source")
.parentElement?.querySelector(".ant-select-selector");
expect(keySourceSelect).toBeTruthy();
await act(async () => {
fireEvent.mouseDown(keySourceSelect!);
});
await clickOption("Team");
};
const selectTeam = async (label: string) => {
const teamSelect = screen.getByTestId("playground-team-select").querySelector(".ant-select-selector");
expect(teamSelect).toBeTruthy();
await act(async () => {
fireEvent.mouseDown(teamSelect!);
});
await clickOption(label);
};
beforeEach(() => {
(networking.teamListCall as any).mockResolvedValue([
{ team_id: "team-alpha", team_alias: "Alpha Team" },
{ team_id: "team-beta", team_alias: "Beta Team" },
]);
(networking.keyCreateCall as any).mockResolvedValue({ key: "sk-team-alpha-key" });
});
const renderChatUI = () =>
render(
<ChatUI
accessToken="ui-session-key"
token="jwt"
userRole="Internal User"
userID="user-1"
disabledPersonalKeyCreation={false}
/>,
);
it("mints a team scoped key and lists the team's models with it", async () => {
renderChatUI();
await waitFor(() => expect(screen.getByText("Test Key")).toBeInTheDocument());
await selectTeamKeySource();
await waitFor(() => expect(networking.teamListCall).toHaveBeenCalledWith("ui-session-key", null, "user-1"));
(fetchModelsModule.fetchAvailableModels as any).mockResolvedValue([{ model_group: "team-only-model" }]);
await selectTeam("Alpha Team");
await waitFor(() =>
expect(networking.keyCreateCall).toHaveBeenCalledWith(
"ui-session-key",
"user-1",
expect.objectContaining({ team_id: "team-alpha", duration: PLAYGROUND_TEAM_KEY_DURATION }),
),
);
await waitFor(() => expect(fetchModelsModule.fetchAvailableModels).toHaveBeenCalledWith("sk-team-alpha-key"));
expect(getSecureItem(playgroundTeamKeyStorageKey("team-alpha"))).toBe("sk-team-alpha-key");
});
it("reuses a cached team key instead of minting another one", async () => {
renderChatUI();
await waitFor(() => expect(screen.getByText("Test Key")).toBeInTheDocument());
await selectTeamKeySource();
await selectTeam("Alpha Team");
await waitFor(() => expect(networking.keyCreateCall).toHaveBeenCalledTimes(1));
await selectTeam("Beta Team");
await waitFor(() => expect(networking.keyCreateCall).toHaveBeenCalledTimes(2));
await selectTeam("Alpha Team");
await waitFor(() => expect(fetchModelsModule.fetchAvailableModels).toHaveBeenCalledWith("sk-team-alpha-key"));
expect(networking.keyCreateCall).toHaveBeenCalledTimes(2);
});
it("surfaces an error and keeps the model list empty when the team key cannot be minted", async () => {
(networking.keyCreateCall as any).mockRejectedValue(new Error("not a member of team"));
renderChatUI();
await waitFor(() => expect(screen.getByText("Test Key")).toBeInTheDocument());
await selectTeamKeySource();
await selectTeam("Alpha Team");
await waitFor(() => expect(screen.getByText("not a member of team")).toBeInTheDocument());
expect(getSecureItem(playgroundTeamKeyStorageKey("team-alpha"))).toBeNull();
});
});
});

104
ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx
View file

@ -23,7 +23,7 @@ import {
} from "@ant-design/icons";
import { Card, Text, TextInput, Title, Button as TremorButton } from "@tremor/react";
import { Button, Input, Modal, Popover, Select, Spin, Tooltip, Typography, Upload } from "antd";
import React, { useEffect, useRef, useState } from "react";
import React, { useEffect, useMemo, useRef, useState } from "react";
import ReactMarkdown from "react-markdown";
import { Prism as SyntaxHighlighter } from "react-syntax-highlighter";
import { coy } from "react-syntax-highlighter/dist/esm/styles/prism";
@ -76,6 +76,7 @@ import RealtimePlayground from "./RealtimePlayground";
import { A2ATaskMetadata, MessageType } from "@/components/chat_ui/types";
import { useCodeInterpreter } from "../../hooks/useCodeInterpreter";
import { useChatHistory } from "../../hooks/useChatHistory";
import { usePlaygroundTeamSession } from "../../hooks/usePlaygroundTeamSession";
import { getSecureItem, setSecureItem } from "@/utils/secureStorage";
import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer";
@ -107,6 +108,8 @@ const MCP_SUPPORTED_ENDPOINTS = new Set<EndpointType>([
const CUSTOM_MODEL_DEBOUNCE_WAIT_MS = 500;
type ApiKeySource = "session" | "team" | "custom";
const ChatUI: React.FC<ChatUIProps> = ({
accessToken,
token,
@ -172,11 +175,11 @@ const ChatUI: React.FC<ChatUIProps> = ({
clearMCPEvents,
} = useChatHistory({ simplified });
// codeql[js/clear-text-storage-of-sensitive-data]
const [apiKeySource, setApiKeySource] = useState<"session" | "custom">(() => {
const [apiKeySource, setApiKeySource] = useState<ApiKeySource>(() => {
const saved = getSecureItem("apiKeySource");
if (saved) {
try {
return JSON.parse(saved) as "session" | "custom";
return JSON.parse(saved) as ApiKeySource;
} catch (error) {
console.error("Error parsing apiKeySource from sessionStorage", error);
}
@ -184,6 +187,9 @@ const ChatUI: React.FC<ChatUIProps> = ({
return disabledPersonalKeyCreation ? "custom" : "session";
});
const [apiKey, setApiKey] = useState<string>(() => getSecureItem("apiKey") || "");
const [selectedTeamId, setSelectedTeamId] = useState<string | null>(
() => sessionStorage.getItem("playgroundTeamId") || null,
);
const [customProxyBaseUrl, setCustomProxyBaseUrl] = useState<string>(
() => sessionStorage.getItem("customProxyBaseUrl") || "",
);
@ -267,9 +273,29 @@ const ChatUI: React.FC<ChatUIProps> = ({
const chatEndRef = useRef<HTMLDivElement>(null);
const {
teams,
teamKey,
isLoadingTeams,
isMintingKey,
error: teamSessionError,
} = usePlaygroundTeamSession({
accessToken,
userID,
teamId: selectedTeamId,
enabled: !simplified && apiKeySource === "team",
});
const effectiveApiKey = useMemo(() => {
if (simplified) return accessToken;
if (apiKeySource === "team") return teamKey;
if (apiKeySource === "custom") return apiKey || null;
return accessToken;
}, [simplified, apiKeySource, accessToken, apiKey, teamKey]);
// Fetch MCP servers and toolsets
const loadMCPServers = async () => {
const userApiKey = apiKeySource === "session" ? accessToken : apiKey;
const userApiKey = effectiveApiKey;
if (!userApiKey) return;
setIsLoadingMCPServers(true);
@ -297,7 +323,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
// Fetch tools for a specific server
const loadServerTools = async (serverId: string) => {
const userApiKey = apiKeySource === "session" ? accessToken : apiKey;
const userApiKey = effectiveApiKey;
if (!userApiKey || serverToolsMap[serverId]) return;
try {
@ -314,9 +340,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
useEffect(() => {
if (isGetCodeModalVisible) {
const code = generateCodeSnippet({
apiKeySource,
accessToken,
apiKey,
apiKey: effectiveApiKey,
inputMessage,
chatHistory,
selectedTags,
@ -337,9 +361,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
}, [
isGetCodeModalVisible,
selectedSdk,
apiKeySource,
accessToken,
apiKey,
effectiveApiKey,
inputMessage,
chatHistory,
selectedTags,
@ -370,6 +392,11 @@ const ChatUI: React.FC<ChatUIProps> = ({
sessionStorage.setItem("mcpServerToolRestrictions", JSON.stringify(mcpServerToolRestrictions));
sessionStorage.setItem("selectedVoice", selectedVoice);
sessionStorage.removeItem("selectedMCPTools"); // Clean up old key
if (selectedTeamId) {
sessionStorage.setItem("playgroundTeamId", selectedTeamId);
} else {
sessionStorage.removeItem("playgroundTeamId");
}
if (!simplified) {
if (selectedModel) {
@ -383,6 +410,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
simplified,
apiKeySource,
apiKey,
selectedTeamId,
selectedModel,
endpointType,
selectedTags,
@ -395,8 +423,9 @@ const ChatUI: React.FC<ChatUIProps> = ({
]);
useEffect(() => {
let userApiKey = apiKeySource === "session" ? accessToken : apiKey;
const userApiKey = effectiveApiKey;
if (!userApiKey || !token || !userRole || !userID) {
setModelInfo([]);
return;
}
@ -426,7 +455,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
loadModels();
}
loadMCPServers();
}, [accessToken, userID, userRole, apiKeySource, apiKey, token, simplified]);
}, [effectiveApiKey, userID, userRole, token, simplified]);
// Load tools when MCP direct mode has a server (or toolset) selected
useEffect(() => {
@ -450,7 +479,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
// Fetch agents when A2A endpoint is selected
useEffect(() => {
const userApiKey = apiKeySource === "session" ? accessToken : apiKey;
const userApiKey = effectiveApiKey;
if (!userApiKey || endpointType !== EndpointType.A2A_AGENTS) {
return;
}
@ -469,7 +498,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
};
loadAgents();
}, [accessToken, apiKeySource, apiKey, endpointType, customProxyBaseUrl, selectedAgent]);
}, [effectiveApiKey, endpointType, customProxyBaseUrl, selectedAgent]);
useEffect(() => {
// Scroll to the bottom of the chat whenever chatHistory updates
@ -650,10 +679,12 @@ const ChatUI: React.FC<ChatUIProps> = ({
return;
}
const effectiveApiKey = simplified ? accessToken : apiKeySource === "session" ? accessToken : apiKey;
if (!effectiveApiKey) {
NotificationsManager.fromBackend("Please provide a Virtual Key or select Current UI Session");
NotificationsManager.fromBackend(
apiKeySource === "team"
? "Please select a team to chat as"
: "Please provide a Virtual Key or select Current UI Session",
);
return;
}
@ -1051,14 +1082,14 @@ const ChatUI: React.FC<ChatUIProps> = ({
<KeyOutlined className="mr-2" /> Virtual Key Source
</Text>
<Select
disabled={disabledPersonalKeyCreation}
value={apiKeySource}
style={{ width: "100%" }}
onChange={(value) => {
setApiKeySource(value as "session" | "custom");
setApiKeySource(value as ApiKeySource);
}}
options={[
{ value: "session", label: "Current UI Session" },
{ value: "session", label: "Current UI Session", disabled: disabledPersonalKeyCreation },
{ value: "team", label: "Team" },
{ value: "custom", label: "Virtual Key" },
]}
className="rounded-md"
@ -1073,6 +1104,31 @@ const ChatUI: React.FC<ChatUIProps> = ({
icon={KeyOutlined}
/>
)}
{apiKeySource === "team" && (
<>
<Select
data-testid="playground-team-select"
className="mt-2 rounded-md"
style={{ width: "100%" }}
placeholder={isLoadingTeams ? "Loading teams..." : "Select a team"}
loading={isLoadingTeams || isMintingKey}
value={selectedTeamId ?? undefined}
onChange={(value: string) => setSelectedTeamId(value)}
showSearch
optionFilterProp="label"
options={teams.map((team) => ({
value: team.team_id,
label: team.team_alias || team.team_id,
}))}
notFoundContent={isLoadingTeams ? <Spin size="small" /> : "You are not a member of any team"}
/>
<Text className="text-xs text-gray-500 mt-1">
Requests use a short lived key scoped to this team, so the model list shows the team&apos;s
models.
</Text>
{teamSessionError && <Text className="text-xs text-red-600 mt-1">{teamSessionError}</Text>}
</>
)}
</div>
<div>
@ -1694,7 +1750,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
{endpointType === EndpointType.RESPONSES && (
<div>
<CodeInterpreterTool
accessToken={apiKeySource === "session" ? accessToken || "" : apiKey}
accessToken={effectiveApiKey || ""}
enabled={codeInterpreter.enabled}
onEnabledChange={codeInterpreter.setEnabled}
selectedContainerId={null}
@ -1711,7 +1767,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
<div className={`flex flex-col bg-white ${simplified ? "flex-1 w-full" : "w-3/4"}`}>
{endpointType === EndpointType.REALTIME ? (
<RealtimePlayground
accessToken={apiKeySource === "session" ? accessToken || "" : apiKey}
accessToken={effectiveApiKey || ""}
selectedModel={selectedModel || ""}
customProxyBaseUrl={customProxyBaseUrl || undefined}
selectedGuardrails={selectedGuardrails.length > 0 ? selectedGuardrails : undefined}
@ -1755,7 +1811,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
endpointType={endpointType as EndpointType}
mcpEvents={mcpEvents}
codeInterpreterResult={codeInterpreter.result}
accessToken={apiKeySource === "session" ? accessToken || "" : apiKey}
accessToken={effectiveApiKey || ""}
/>
</div>
))}

102
ui/litellm-dashboard/src/app/(dashboard)/playground/hooks/usePlaygroundTeamSession.ts Normal file
View file

@ -0,0 +1,102 @@
"use client";
import { useQuery } from "@tanstack/react-query";
import { keyCreateCall, teamListCall } from "@/components/networking";
import { Team } from "@/components/key_team_helpers/key_list";
import { getSecureItem, setSecureItem } from "@/utils/secureStorage";
export const PLAYGROUND_TEAM_KEY_DURATION = "24h";
export const playgroundTeamKeyStorageKey = (teamId: string) => `playgroundTeamKey:${teamId}`;
interface UsePlaygroundTeamSessionArgs {
accessToken: string | null;
userID: string | null;
teamId: string | null;
enabled: boolean;
}
interface PlaygroundTeamSession {
teams: Team[];
teamKey: string | null;
isLoadingTeams: boolean;
isMintingKey: boolean;
error: string | null;
}
const readCachedKey = (teamId: string): string | null => {
try {
return getSecureItem(playgroundTeamKeyStorageKey(teamId));
} catch {
return null;
}
};
const cacheKey = (teamId: string, key: string): void => {
try {
setSecureItem(playgroundTeamKeyStorageKey(teamId), key);
} catch {
return;
}
};
const errorMessage = (error: unknown, fallback: string): string =>
error instanceof Error && error.message ? error.message : fallback;
export const usePlaygroundTeamSession = ({
accessToken,
userID,
teamId,
enabled,
}: UsePlaygroundTeamSessionArgs): PlaygroundTeamSession => {
const teamsQuery = useQuery({
queryKey: ["playgroundTeams", userID],
queryFn: async (): Promise<Team[]> => {
if (!accessToken) return [];
const response: unknown = await teamListCall(accessToken, null, userID);
return Array.isArray(response) ? (response as Team[]) : [];
},
enabled: enabled && !!accessToken,
});
const teamKeyQuery = useQuery({
queryKey: ["playgroundTeamKey", teamId],
queryFn: async (): Promise<string> => {
if (!accessToken || !userID || !teamId) {
throw new Error("Sign in and select a team to chat as that team");
}
const cached = readCachedKey(teamId);
if (cached) {
return cached;
}
const response: { key?: string } = await keyCreateCall(accessToken, userID, {
team_id: teamId,
duration: PLAYGROUND_TEAM_KEY_DURATION,
key_alias: `playground-${teamId}-${Date.now()}`,
});
if (!response?.key) {
throw new Error("Creating a key for this team returned no key");
}
cacheKey(teamId, response.key);
return response.key;
},
enabled: enabled && !!accessToken && !!userID && !!teamId,
retry: false,
staleTime: Infinity,
});
let error: string | null = null;
if (teamsQuery.error) {
error = errorMessage(teamsQuery.error, "Failed to load teams");
} else if (teamKeyQuery.error) {
error = errorMessage(teamKeyQuery.error, "Failed to create a key for this team");
}
return {
teams: teamsQuery.data ?? [],
teamKey: teamKeyQuery.data ?? null,
isLoadingTeams: teamsQuery.isFetching,
isMintingKey: teamKeyQuery.isFetching,
error,
};
};

2
ui/litellm-dashboard/src/components/chat_ui/CodeSnippets.test.tsx
View file

@ -27,8 +27,6 @@ describe("CodeSnippets", () => {
endpointType: EndpointType.EMBEDDINGS,
inputMessage: "Hello, world!",
selectedModel: "text-embedding-3-small",
apiKeySource: "session" as const,
accessToken: "1234567890",
apiKey: "1234567890",
chatHistory: [],
selectedTags: [],

9
ui/litellm-dashboard/src/components/chat_ui/CodeSnippets.tsx
View file

@ -10,9 +10,7 @@ interface CodeGenMetadata {
}
interface GenerateCodeParams {
apiKeySource: "session" | "custom";
accessToken: string | null;
apiKey: string;
apiKey: string | null;
inputMessage: string;
chatHistory: MessageType[];
selectedTags: string[];
@ -34,9 +32,7 @@ interface GenerateCodeParams {
export const generateCodeSnippet = (params: GenerateCodeParams): string => {
const {
apiKeySource,
accessToken,
apiKey,
apiKey: effectiveApiKey,
inputMessage,
chatHistory,
selectedTags,
@ -52,7 +48,6 @@ export const generateCodeSnippet = (params: GenerateCodeParams): string => {
selectedSdk,
proxySettings,
} = params;
const effectiveApiKey = apiKeySource === "session" ? accessToken : apiKey;
// Determine base URL with priority: LITELLM_UI_API_DOC_BASE_URL > PROXY_BASE_URL > window.location.origin
let apiBase = window.location.origin;

4
ui/litellm-dashboard/src/components/public_model_hub.tsx
View file

@ -1020,8 +1020,6 @@ const PublicModelHub: React.FC<PublicModelHubProps> = ({ accessToken, isEmbedded
<pre className="text-sm">
{(() => {
const codeSnippet = generateCodeSnippet({
apiKeySource: "custom",
accessToken: null,
apiKey: "your_api_key",
inputMessage: "Hello, how are you?",
chatHistory: [{ role: "user", content: "Hello, how are you?", isImage: false } as MessageType],
@ -1042,8 +1040,6 @@ const PublicModelHub: React.FC<PublicModelHubProps> = ({ accessToken, isEmbedded
<button
onClick={() => {
const codeSnippet = generateCodeSnippet({
apiKeySource: "custom",
accessToken: null,
apiKey: "your_api_key",
inputMessage: "Hello, how are you?",
chatHistory: [{ role: "user", content: "Hello, how are you?", isImage: false } as MessageType],