scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

feat: add selectable mcp servers to the playground

Browse source
This commit is contained in:
Yuta Saito 2026-01-02 11:47:17 +09:00
parent d2f2fe1be5
commit 9815b00deb
7 changed files with 236 additions and 181 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
litellm/responses/mcp/litellm_proxy_mcp_handler.py
View file

@ -205,7 +205,15 @@ class LiteLLM_Proxy_MCP_Handler:
else:
tool_name = getattr(mcp_tool, "name", None)
if tool_name and tool_name in allowed_tool_names:
if not tool_name:
continue
if tool_name in allowed_tool_names:
filtered_tools.append(mcp_tool)
continue
unprefixed_name, _ = split_server_prefix_from_name(tool_name)
if unprefixed_name in allowed_tool_names:
filtered_tools.append(mcp_tool)
return filtered_tools

253
ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx
View file

@ -45,8 +45,8 @@ import { makeOpenAIAudioSpeechRequest } from "../llm_calls/audio_speech";
import { makeOpenAIAudioTranscriptionRequest } from "../llm_calls/audio_transcriptions";
import { makeOpenAIChatCompletionRequest } from "../llm_calls/chat_completion";
import { makeOpenAIEmbeddingsRequest } from "../llm_calls/embeddings_api";
import type { MCPTool } from "../llm_calls/fetch_mcp_tools";
import { fetchAvailableMCPTools } from "../llm_calls/fetch_mcp_tools";
import { listMCPTools, fetchMCPServers } from "../../networking";
import { MCPServer } from "../../mcp_tools/types";
import { fetchAvailableModels, ModelGroup } from "../llm_calls/fetch_models";
import { makeOpenAIImageEditsRequest } from "../llm_calls/image_edits";
import { makeOpenAIImageGenerationRequest } from "../llm_calls/image_generation";
@ -97,20 +97,27 @@ const ChatUI: React.FC<ChatUIProps> = ({
disabledPersonalKeyCreation,
proxySettings,
}) => {
const [isMCPToolsModalVisible, setIsMCPToolsModalVisible] = useState(false);
const [mcpTools, setMCPTools] = useState<MCPTool[]>([]);
const [selectedMCPTools, setSelectedMCPTools] = useState<string[]>(() => {
const saved = sessionStorage.getItem("selectedMCPTools");
const [mcpServers, setMCPServers] = useState<MCPServer[]>([]);
const [selectedMCPServers, setSelectedMCPServers] = useState<string[]>(() => {
const saved = sessionStorage.getItem("selectedMCPServers");
try {
const parsed = saved ? JSON.parse(saved) : [];
// Convert from single string to array if needed for backward compatibility
return Array.isArray(parsed) ? parsed : parsed ? [parsed] : [];
return saved ? JSON.parse(saved) : [];
} catch (error) {
console.error("Error parsing selectedMCPTools from sessionStorage", error);
console.error("Error parsing selectedMCPServers from sessionStorage", error);
return [];
}
});
const [isLoadingMCPTools, setIsLoadingMCPTools] = useState(false);
const [isLoadingMCPServers, setIsLoadingMCPServers] = useState(false);
const [serverToolsMap, setServerToolsMap] = useState<Record<string, any[]>>({});
const [mcpServerToolRestrictions, setMCPServerToolRestrictions] = useState<Record<string, string[]>>(() => {
const saved = sessionStorage.getItem("mcpServerToolRestrictions");
try {
return saved ? JSON.parse(saved) : {};
} catch (error) {
console.error("Error parsing mcpServerToolRestrictions from sessionStorage", error);
return {};
}
});
const [apiKeySource, setApiKeySource] = useState<"session" | "custom">(() => {
const saved = sessionStorage.getItem("apiKeySource");
if (saved) {
@ -211,27 +218,38 @@ const ChatUI: React.FC<ChatUIProps> = ({
const chatEndRef = useRef<HTMLDivElement>(null);
// Fetch MCP tools
const loadMCPTools = async () => {
// Fetch MCP servers
const loadMCPServers = async () => {
const userApiKey = apiKeySource === "session" ? accessToken : apiKey;
if (!userApiKey) return;
setIsLoadingMCPTools(true);
setIsLoadingMCPServers(true);
try {
const tools = await fetchAvailableMCPTools(userApiKey);
setMCPTools(tools);
const servers = await fetchMCPServers(userApiKey);
setMCPServers(Array.isArray(servers) ? servers : servers.data || []);
} catch (error) {
console.error("Error fetching MCP tools:", error);
console.error("Error fetching MCP servers:", error);
} finally {
setIsLoadingMCPTools(false);
setIsLoadingMCPServers(false);
}
};
useEffect(() => {
if (isMCPToolsModalVisible) {
loadMCPTools();
// Fetch tools for a specific server
const loadServerTools = async (serverId: string) => {
const userApiKey = apiKeySource === "session" ? accessToken : apiKey;
if (!userApiKey || serverToolsMap[serverId]) return;
try {
const response = await listMCPTools(userApiKey, serverId);
setServerToolsMap(prev => ({
...prev,
[serverId]: response.tools || []
}));
} catch (error) {
console.error(`Error fetching tools for server ${serverId}:`, error);
}
}, [isMCPToolsModalVisible, accessToken, apiKey, apiKeySource]);
};
useEffect(() => {
if (isGetCodeModalVisible) {
@ -244,7 +262,9 @@ const ChatUI: React.FC<ChatUIProps> = ({
selectedTags,
selectedVectorStores,
selectedGuardrails,
selectedMCPTools,
selectedMCPServers,
mcpServers,
mcpServerToolRestrictions,
endpointType,
selectedModel,
selectedSdk,
@ -264,7 +284,9 @@ const ChatUI: React.FC<ChatUIProps> = ({
selectedTags,
selectedVectorStores,
selectedGuardrails,
selectedMCPTools,
selectedMCPServers,
mcpServers,
mcpServerToolRestrictions,
endpointType,
selectedModel,
proxySettings,
@ -287,8 +309,10 @@ const ChatUI: React.FC<ChatUIProps> = ({
sessionStorage.setItem("selectedTags", JSON.stringify(selectedTags));
sessionStorage.setItem("selectedVectorStores", JSON.stringify(selectedVectorStores));
sessionStorage.setItem("selectedGuardrails", JSON.stringify(selectedGuardrails));
sessionStorage.setItem("selectedMCPTools", JSON.stringify(selectedMCPTools));
sessionStorage.setItem("selectedMCPServers", JSON.stringify(selectedMCPServers));
sessionStorage.setItem("mcpServerToolRestrictions", JSON.stringify(mcpServerToolRestrictions));
sessionStorage.setItem("selectedVoice", selectedVoice);
sessionStorage.removeItem("selectedMCPTools"); // Clean up old key
if (selectedModel) {
sessionStorage.setItem("selectedModel", selectedModel);
@ -318,7 +342,8 @@ const ChatUI: React.FC<ChatUIProps> = ({
messageTraceId,
responsesSessionId,
useApiSessionManagement,
selectedMCPTools,
selectedMCPServers,
mcpServerToolRestrictions,
selectedVoice,
]);
@ -355,7 +380,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
};
loadModels();
loadMCPTools();
loadMCPServers();
}, [accessToken, userID, userRole, apiKeySource, apiKey, token]);
// Fetch agents when A2A endpoint is selected
@ -869,12 +894,14 @@ const ChatUI: React.FC<ChatUIProps> = ({
traceId,
selectedVectorStores.length > 0 ? selectedVectorStores : undefined,
selectedGuardrails.length > 0 ? selectedGuardrails : undefined,
selectedMCPTools,
selectedMCPServers,
updateChatImageUI,
updateSearchResults,
useAdvancedParams ? temperature : undefined,
useAdvancedParams ? maxTokens : undefined,
updateTotalLatency,
mcpServers,
mcpServerToolRestrictions,
);
} else if (endpointType === EndpointType.IMAGE) {
// For image generation
@ -940,12 +967,14 @@ const ChatUI: React.FC<ChatUIProps> = ({
traceId,
selectedVectorStores.length > 0 ? selectedVectorStores : undefined,
selectedGuardrails.length > 0 ? selectedGuardrails : undefined,
selectedMCPTools, // Pass the selected tools array
selectedMCPServers, // Pass the selected servers array
useApiSessionManagement ? responsesSessionId : null, // Only pass session ID if API mode is enabled
handleResponseId, // Pass callback to capture new response ID
handleMCPEvent, // Pass MCP event handler
codeInterpreter.enabled, // Enable Code Interpreter tool
codeInterpreter.setResult, // Handle code interpreter output
mcpServers,
mcpServerToolRestrictions,
);
} else if (endpointType === EndpointType.ANTHROPIC_MESSAGES) {
const apiChatHistory = [
@ -968,7 +997,9 @@ const ChatUI: React.FC<ChatUIProps> = ({
traceId,
selectedVectorStores.length > 0 ? selectedVectorStores : undefined,
selectedGuardrails.length > 0 ? selectedGuardrails : undefined,
selectedMCPTools, // Pass the selected tools array
selectedMCPServers, // Pass the selected servers array
mcpServers,
mcpServerToolRestrictions,
);
} else if (endpointType === EndpointType.EMBEDDINGS) {
await makeOpenAIEmbeddingsRequest(
@ -1329,13 +1360,13 @@ const ChatUI: React.FC<ChatUIProps> = ({
/>
</div>
{/* MCP Tool Selection */}
{/* MCP Server Selection */}
<div>
<Text className="font-medium block mb-2 text-gray-700 flex items-center">
<ToolOutlined className="mr-2" /> MCP Tool
<ToolOutlined className="mr-2" /> MCP Servers
<Tooltip
className="ml-1"
title="Select MCP tools to use in your conversation."
title="Select MCP servers to use in your conversation."
>
<InfoCircleOutlined />
</Tooltip>
@ -1343,30 +1374,103 @@ const ChatUI: React.FC<ChatUIProps> = ({
<Select
mode="multiple"
style={{ width: "100%" }}
placeholder="Select MCP tools"
value={selectedMCPTools}
onChange={(value) => setSelectedMCPTools(value)}
loading={isLoadingMCPTools}
className="mb-4"
placeholder="Select MCP servers"
value={selectedMCPServers}
onChange={(value) => {
if (value.includes("__all__")) {
setSelectedMCPServers(["__all__"]);
setMCPServerToolRestrictions({});
} else {
setSelectedMCPServers(value);
// Clean up tool restrictions for removed servers
setMCPServerToolRestrictions(prev => {
const updated = { ...prev };
Object.keys(updated).forEach(serverId => {
if (!value.includes(serverId)) delete updated[serverId];
});
return updated;
});
// Load tools for newly selected servers
value.forEach(serverId => {
if (!serverToolsMap[serverId]) {
loadServerTools(serverId);
}
});
}
}}
loading={isLoadingMCPServers}
className="mb-2"
allowClear
optionLabelProp="label"
disabled={!MCP_SUPPORTED_ENDPOINTS.has(endpointType as EndpointType)}
maxTagCount="responsive"
>
{Array.isArray(mcpTools) &&
mcpTools.map((tool) => (
<Select.Option
key={tool.name}
value={tool.name}
label={<div className="font-medium">{tool.name}</div>}
>
<div className="flex flex-col py-1">
<span className="font-medium">{tool.name}</span>
<span className="text-xs text-gray-500 mt-1">{tool.description}</span>
</div>
</Select.Option>
))}
{/* All MCP Servers option */}
<Select.Option key="__all__" value="__all__" label="All MCP Servers">
<div className="flex flex-col py-1">
<span className="font-medium">All MCP Servers</span>
<span className="text-xs text-gray-500 mt-1">Use all available MCP servers</span>
</div>
</Select.Option>
{/* Individual servers */}
{mcpServers.map((server) => (
<Select.Option
key={server.server_id}
value={server.server_id}
label={server.alias || server.server_name || server.server_id}
disabled={selectedMCPServers.includes("__all__")}
>
<div className="flex flex-col py-1">
<span className="font-medium">
{server.alias || server.server_name || server.server_id}
</span>
{server.description && (
<span className="text-xs text-gray-500 mt-1">{server.description}</span>
)}
</div>
</Select.Option>
))}
</Select>
{/* Tool restrictions UI (optional) */}
{selectedMCPServers.length > 0 &&
!selectedMCPServers.includes("__all__") &&
MCP_SUPPORTED_ENDPOINTS.has(endpointType as EndpointType) && (
<div className="mt-3 space-y-2">
{selectedMCPServers.map(serverId => {
const server = mcpServers.find(s => s.server_id === serverId);
const tools = serverToolsMap[serverId] || [];
if (tools.length === 0) return null;
return (
<div key={serverId} className="border rounded p-2">
<Text className="text-xs text-gray-600 mb-1">
Limit tools for {server?.alias || server?.server_name || serverId}:
</Text>
<Select
mode="multiple"
size="small"
style={{ width: "100%" }}
placeholder="All tools (default)"
value={mcpServerToolRestrictions[serverId] || []}
onChange={(selectedTools) => {
setMCPServerToolRestrictions(prev => ({
...prev,
[serverId]: selectedTools
}));
}}
options={tools.map(tool => ({
value: tool.name,
label: tool.name
}))}
maxTagCount={2}
/>
</div>
);
})}
</div>
)}
</div>
<div>
@ -2041,53 +2145,6 @@ const ChatUI: React.FC<ChatUIProps> = ({
{generatedCode}
</SyntaxHighlighter>
</Modal>
{apiKeySource === "custom" && (
<Modal
title="Select MCP Tool"
visible={isMCPToolsModalVisible}
onCancel={() => setIsMCPToolsModalVisible(false)}
onOk={() => {
setIsMCPToolsModalVisible(false);
NotificationsManager.success("MCP tool selection updated");
}}
width={800}
>
{isLoadingMCPTools ? (
<div className="flex justify-center items-center py-8">
<Spin indicator={<LoadingOutlined style={{ fontSize: 24 }} spin />} />
</div>
) : (
<div className="space-y-4">
<Text className="text-gray-600 block mb-4">
Select the MCP tools you want to use in your conversation.
</Text>
<Select
mode="multiple"
style={{ width: "100%" }}
placeholder="Select MCP tools"
value={selectedMCPTools}
onChange={(value) => setSelectedMCPTools(value)}
optionLabelProp="label"
allowClear
maxTagCount="responsive"
>
{mcpTools.map((tool) => (
<Select.Option
key={tool.name}
value={tool.name}
label={<div className="font-medium">{tool.name}</div>}
>
<div className="flex flex-col py-1">
<span className="font-medium">{tool.name}</span>
<span className="text-xs text-gray-500 mt-1">{tool.description}</span>
</div>
</Select.Option>
))}
</Select>
</div>
)}
</Modal>
)}
</div>
);
};

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

@ -1,5 +1,6 @@
import { MessageType } from "./types";
import { EndpointType } from "./mode_endpoint_mapping";
import { MCPServer } from "../../mcp_tools/types";
interface CodeGenMetadata {
tags?: string[];
@ -16,7 +17,9 @@ interface GenerateCodeParams {
selectedTags: string[];
selectedVectorStores: string[];
selectedGuardrails: string[];
selectedMCPTools: string[];
selectedMCPServers: string[];
mcpServers?: MCPServer[];
mcpServerToolRestrictions?: Record<string, string[]>;
selectedVoice?: string;
endpointType: string;
selectedModel: string | undefined;
@ -37,7 +40,9 @@ export const generateCodeSnippet = (params: GenerateCodeParams): string => {
selectedTags,
selectedVectorStores,
selectedGuardrails,
selectedMCPTools,
selectedMCPServers,
mcpServers,
mcpServerToolRestrictions,
selectedVoice,
endpointType,
selectedModel,

23
ui/litellm-dashboard/src/components/playground/llm_calls/anthropic_messages.tsx
View file

@ -17,7 +17,6 @@ export async function makeAnthropicMessagesRequest(
traceId?: string,
vector_store_ids?: string[],
guardrails?: string[],
selectedMCPTools?: string[],
) {
if (!accessToken) {
throw new Error("Virtual Key is required");
@ -47,23 +46,6 @@ export async function makeAnthropicMessagesRequest(
const startTime = Date.now();
let firstTokenReceived = false;
// Format MCP tools if selected
const tools =
selectedMCPTools && selectedMCPTools.length > 0
? [
{
type: "mcp",
server_label: "litellm",
server_url: `${proxyBaseUrl}/mcp`,
require_approval: "never",
allowed_tools: selectedMCPTools,
headers: {
"x-litellm-api-key": `Bearer ${accessToken}`,
},
},
]
: undefined;
const requestBody: any = {
model: selectedModel,
messages: messages.map((m) => ({ role: m.role, content: m.content })),
@ -75,11 +57,6 @@ export async function makeAnthropicMessagesRequest(
if (vector_store_ids) requestBody.vector_store_ids = vector_store_ids;
if (guardrails) requestBody.guardrails = guardrails;
if (tools) {
requestBody.tools = tools;
requestBody.tool_choice = "auto";
}
// Use the streaming helper method for cleaner async iteration
// @ts-ignore - The SDK types might not include all litellm-specific parameters
const stream = client.messages.stream(requestBody, { signal });

53
ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx
View file

@ -3,6 +3,7 @@ import { ChatCompletionMessageParam } from "openai/resources/chat/completions";
import { TokenUsage } from "../chat_ui/ResponseMetrics";
import { VectorStoreSearchResponse } from "../chat_ui/types";
import { getProxyBaseUrl } from "@/components/networking";
import { MCPServer } from "../../mcp_tools/types";
export async function makeOpenAIChatCompletionRequest(
chatHistory: { role: string; content: string | any[] }[],
@ -17,12 +18,14 @@ export async function makeOpenAIChatCompletionRequest(
traceId?: string,
vector_store_ids?: string[],
guardrails?: string[],
selectedMCPTools?: string[],
selectedMCPServers?: string[],
onImageGenerated?: (imageUrl: string, model?: string) => void,
onSearchResults?: (searchResults: VectorStoreSearchResponse[]) => void,
temperature?: number,
max_tokens?: number,
onTotalLatency?: (latency: number) => void,
mcpServers?: MCPServer[],
mcpServerToolRestrictions?: Record<string, string[]>,
) {
// base url should be the current base_url
const isLocal = process.env.NODE_ENV === "development";
@ -53,22 +56,36 @@ export async function makeOpenAIChatCompletionRequest(
let fullResponseContent = "";
let fullReasoningContent = "";
// Format MCP tools if selected
const tools =
selectedMCPTools && selectedMCPTools.length > 0
? [
{
type: "mcp",
server_label: "litellm",
server_url: 'litellm_proxy/mcp',
require_approval: "never",
allowed_tools: selectedMCPTools,
headers: {
"x-litellm-api-key": `Bearer ${accessToken}`,
},
},
]
: undefined;
// Build tools array
const tools: any[] = [];
// Add MCP servers if selected
if (selectedMCPServers && selectedMCPServers.length > 0) {
if (selectedMCPServers.includes("__all__")) {
// All MCP Servers selected
tools.push({
type: "mcp",
server_label: "litellm",
server_url: "litellm_proxy/mcp",
require_approval: "never",
});
} else {
// Individual servers selected - create one entry per server
selectedMCPServers.forEach((serverId) => {
const server = mcpServers?.find((s) => s.server_id === serverId);
const serverName = server?.alias || server?.server_name || serverId;
const allowedTools = mcpServerToolRestrictions?.[serverId] || [];
tools.push({
type: "mcp",
server_label: "litellm",
server_url: `litellm_proxy/mcp/${serverName}`,
require_approval: "never",
...(allowedTools.length > 0 ? { allowed_tools: allowedTools } : {}),
});
});
}
}
// @ts-ignore
const response = await client.chat.completions.create(
@ -82,7 +99,7 @@ export async function makeOpenAIChatCompletionRequest(
messages: chatHistory as ChatCompletionMessageParam[],
...(vector_store_ids ? { vector_store_ids } : {}),
...(guardrails ? { guardrails } : {}),
...(tools ? { tools, tool_choice: "auto" } : {}),
...(tools.length > 0 ? { tools, tool_choice: "auto" } : {}),
...(temperature !== undefined ? { temperature } : {}),
...(max_tokens !== undefined ? { max_tokens } : {}),
},

29
ui/litellm-dashboard/src/components/playground/llm_calls/fetch_mcp_tools.tsx
View file

@ -1,29 +0,0 @@
import { mcpToolsCall } from "../../networking";
export interface MCPTool {
name: string;
description: string;
title: string | null;
inputSchema: {
type: string;
properties: Record<string, any>;
required: string[];
};
outputSchema: any;
annotations: any;
_meta: any;
}
interface MCPToolsResponse {
tools: MCPTool[];
}
export async function fetchAvailableMCPTools(accessToken: string): Promise<MCPTool[]> {
try {
const data = (await mcpToolsCall(accessToken)) as MCPToolsResponse;
return data.tools || [];
} catch (error) {
console.error("Error fetching MCP tools:", error);
return [];
}
}

40
ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.tsx
View file

@ -4,6 +4,7 @@ import { TokenUsage } from "../chat_ui/ResponseMetrics";
import { getProxyBaseUrl } from "@/components/networking";
import NotificationManager from "@/components/molecules/notifications_manager";
import { MCPEvent } from "../chat_ui/MCPEventsDisplay";
import { MCPServer } from "../../mcp_tools/types";
import {
CodeInterpreterResult,
CodeInterpreterState,
@ -26,12 +27,14 @@ export async function makeOpenAIResponsesRequest(
traceId?: string,
vector_store_ids?: string[],
guardrails?: string[],
selectedMCPTools?: string[],
selectedMCPServers?: string[],
previousResponseId?: string | null,
onResponseId?: (responseId: string) => void,
onMCPEvent?: (event: MCPEvent) => void,
codeInterpreterEnabled?: boolean,
onCodeInterpreterResult?: (result: CodeInterpreterResult) => void,
mcpServers?: MCPServer[],
mcpServerToolRestrictions?: Record<string, string[]>,
) {
if (!accessToken) {
throw new Error("Virtual Key is required");
@ -86,15 +89,32 @@ export async function makeOpenAIResponsesRequest(
// Build tools array
const tools: any[] = [];
// Add MCP tools if selected
if (selectedMCPTools && selectedMCPTools.length > 0) {
tools.push({
type: "mcp",
server_label: "litellm",
server_url: `litellm_proxy/mcp`,
require_approval: "never",
allowed_tools: selectedMCPTools,
});
// Add MCP servers if selected
if (selectedMCPServers && selectedMCPServers.length > 0) {
if (selectedMCPServers.includes("__all__")) {
// All MCP Servers selected
tools.push({
type: "mcp",
server_label: "litellm",
server_url: "litellm_proxy/mcp",
require_approval: "never",
});
} else {
// Individual servers selected - create one entry per server
selectedMCPServers.forEach((serverId) => {
const server = mcpServers?.find((s) => s.server_id === serverId);
const serverName = server?.alias || server?.server_name || serverId;
const allowedTools = mcpServerToolRestrictions?.[serverId] || [];
tools.push({
type: "mcp",
server_label: "litellm",
server_url: `litellm_proxy/mcp/${serverName}`,
require_approval: "never",
...(allowedTools.length > 0 ? { allowed_tools: allowedTools } : {}),
});
});
}
}
// Add code_interpreter tool if enabled (OpenAI auto-creates container)