From 15b717c3005ddecc50cb2d0b17c7d80365d08145 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 4 Feb 2026 18:38:53 -0800 Subject: [PATCH] [Feat] UI - Add support for MCP Semantic Filtering on UI (#20454) * feat: remove_callbacks_by_type * feat: _init_semantic_filter_settings_in_db * add MCP semantic filtering settings endpoints * v0 - semantic filter UI * fix endpint _init_semantic_filter_settings_in_db * refactor ui * add LITELLM_UI_ALLOW_HEADERS * ui fix * ui fix --- litellm/constants.py | 5 + .../logging_callback_manager.py | 21 ++ litellm/proxy/agent_endpoints/a2a_routing.py | 2 +- litellm/proxy/proxy_server.py | 89 +++++ litellm/proxy/route_llm_request.py | 5 +- .../proxy_setting_endpoints.py | 100 +++++- .../useMCPSemanticFilterSettings.ts | 19 ++ .../useUpdateMCPSemanticFilterSettings.ts | 25 ++ .../MCPSemanticFilterSettings.tsx | 312 ++++++++++++++++++ .../MCPSemanticFilterTestPanel.tsx | 164 +++++++++ .../semanticFilterTestUtils.ts | 95 ++++++ .../src/components/mcp_tools/mcp_servers.tsx | 5 + .../src/components/networking.tsx | 131 ++++++++ 13 files changed, 970 insertions(+), 3 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpSemanticFilterSettings/useMCPSemanticFilterSettings.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpSemanticFilterSettings/useUpdateMCPSemanticFilterSettings.ts create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterSettings.tsx create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterTestPanel.tsx create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/semanticFilterTestUtils.ts diff --git a/litellm/constants.py b/litellm/constants.py index 444e78f8ed4..872ad899f84 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -81,6 +81,11 @@ MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH = int( os.getenv("MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH", 150) ) +LITELLM_UI_ALLOW_HEADERS = [ + "x-litellm-semantic-filter", + "x-litellm-semantic-filter-tools", +] + # Gemini model-specific minimal thinking budget constants DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH = int( os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH", 1) diff --git a/litellm/litellm_core_utils/logging_callback_manager.py b/litellm/litellm_core_utils/logging_callback_manager.py index 4f76a5bad03..435ae078a65 100644 --- a/litellm/litellm_core_utils/logging_callback_manager.py +++ b/litellm/litellm_core_utils/logging_callback_manager.py @@ -114,6 +114,27 @@ class LoggingCallbackManager: for c in remove_list: callback_list.remove(c) + def remove_callbacks_by_type(self, callback_list, callback_type): + """ + Remove all callbacks of a specific type from a callback list. + + Args: + callback_list: The list to remove callbacks from (e.g., litellm.callbacks) + callback_type: The class type to match (e.g., SemanticToolFilterHook) + + Example: + litellm.logging_callback_manager.remove_callbacks_by_type( + litellm.callbacks, SemanticToolFilterHook + ) + """ + if not isinstance(callback_list, list): + return + + remove_list = [c for c in callback_list if isinstance(c, callback_type)] + + for c in remove_list: + callback_list.remove(c) + def _add_string_callback_to_list( self, callback: str, parent_list: List[Union[CustomLogger, Callable, str]] ): diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py index 8e4c705df21..cb277d44ee9 100644 --- a/litellm/proxy/agent_endpoints/a2a_routing.py +++ b/litellm/proxy/agent_endpoints/a2a_routing.py @@ -11,7 +11,7 @@ import litellm from litellm._logging import verbose_proxy_logger -async def route_a2a_agent_request(data: dict, route_type: str) -> Optional[Any]: +def route_a2a_agent_request(data: dict, route_type: str) -> Optional[Any]: """ Route A2A agent requests directly to litellm with injected API base. diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b06071481eb..022626e2992 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -47,6 +47,7 @@ from litellm.constants import ( DEFAULT_SLACK_ALERTING_THRESHOLD, LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS, LITELLM_SETTINGS_SAFE_DB_OVERRIDES, + LITELLM_UI_ALLOW_HEADERS, ) from litellm.litellm_core_utils.litellm_logging import ( _init_custom_logger_compatible_class, @@ -1214,6 +1215,7 @@ app.add_middleware( allow_credentials=True, allow_methods=["*"], allow_headers=["*"], + expose_headers=LITELLM_UI_ALLOW_HEADERS, ) app.add_middleware(PrometheusAuthMiddleware) @@ -3920,6 +3922,93 @@ class ProxyConfig: prisma_client=prisma_client, proxy_config=self ) + if self._should_load_db_object(object_type="semantic_filter_settings"): + await self._init_semantic_filter_settings_in_db( + prisma_client=prisma_client + ) + + async def _init_semantic_filter_settings_in_db(self, prisma_client: PrismaClient): + """ + Initialize MCP semantic filter settings from database. + Called periodically (approximately every 10 seconds) by background task to hot-reload settings across all pods. + """ + import json + + import litellm + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + + try: + # Load litellm_settings from DB + config_record = await prisma_client.db.litellm_config.find_unique( + where={"param_name": "litellm_settings"} + ) + + if config_record is None or config_record.param_value is None: + return + + litellm_settings = config_record.param_value + if isinstance(litellm_settings, str): + litellm_settings = json.loads(litellm_settings) + + mcp_semantic_filter_config = litellm_settings.get( + "mcp_semantic_tool_filter", None + ) + + if mcp_semantic_filter_config is None: + return + + # Check if settings have changed (compare with in-memory state) + if hasattr(self, "_last_semantic_filter_config"): + if self._last_semantic_filter_config == mcp_semantic_filter_config: + # If hook is missing or router isn't built yet, reinitialize anyway + active_hooks = ( + litellm.logging_callback_manager.get_custom_loggers_for_type( + SemanticToolFilterHook + ) + ) + if active_hooks: + for active_hook in active_hooks: + if isinstance(active_hook, SemanticToolFilterHook): + if ( + active_hook.filter is not None + and active_hook.filter.tool_router is not None + ): + verbose_proxy_logger.debug( + "Semantic filter settings unchanged, skipping reinitialization" + ) + return + verbose_proxy_logger.info( + "Semantic filter settings unchanged, but hook is missing or uninitialized. Reinitializing." + ) + + # Remove old hooks using logging callback manager + litellm.logging_callback_manager.remove_callbacks_by_type( + litellm.callbacks, SemanticToolFilterHook + ) + + # Initialize new hook if enabled + if mcp_semantic_filter_config.get("enabled", False): + global llm_router + hook = await SemanticToolFilterHook.initialize_from_config( + config=mcp_semantic_filter_config, + llm_router=llm_router, + ) + if hook: + litellm.logging_callback_manager.add_litellm_callback(hook) + verbose_proxy_logger.info( + "MCP Semantic Filter reinitialized from DB" + ) + else: + verbose_proxy_logger.info("MCP Semantic Filter disabled") + + # Store current config for comparison next time + self._last_semantic_filter_config = mcp_semantic_filter_config.copy() + + except Exception as e: + verbose_proxy_logger.exception( + f"Error initializing semantic filter settings from DB: {e}" + ) + async def _init_sso_settings_in_db(self, prisma_client: PrismaClient): """ Initialize SSO settings from database into the router on startup. diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 39ef5fdd1d5..92fb88f7147 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -332,7 +332,10 @@ async def route_request( route_a2a_agent_request, ) - return await route_a2a_agent_request(data, route_type) + result = route_a2a_agent_request(data, route_type) + if result is not None: + return result + # Fall through to raise exception below if result is None elif user_model is not None: return getattr(litellm, f"{route_type}")(**data) diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 30ec0766dbf..a308d0b2703 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -98,6 +98,40 @@ ALLOWED_UI_SETTINGS_FIELDS = { } +class MCPSemanticFilterSettings(BaseModel): + """Configuration for MCP Semantic Tool Filter""" + + enabled: bool = Field( + default=False, + description="Enable semantic filtering of MCP tools based on query relevance", + ) + + embedding_model: str = Field( + default="text-embedding-3-small", + description="Embedding model to use for semantic similarity (e.g., 'text-embedding-3-small', 'text-embedding-ada-002')", + ) + + top_k: int = Field( + default=10, + description="Number of most relevant tools to return", + ge=1, + le=100, + ) + + similarity_threshold: float = Field( + default=0.3, + description="Minimum similarity score for tool inclusion (0.0 to 1.0, where 1.0 = exact match)", + ge=0.0, + le=1.0, + ) + + +class MCPSemanticFilterSettingsResponse(SettingsResponse): + """Response model for MCP semantic filter settings""" + + pass + + @router.get( "/get/allowed_ips", tags=["Budget & Spend Tracking"], @@ -325,7 +359,7 @@ async def update_default_team_member_budget( async def _update_litellm_setting( - settings: Union[DefaultInternalUserParams, DefaultTeamSSOParams], + settings: Union[DefaultInternalUserParams, DefaultTeamSSOParams, MCPSemanticFilterSettings], settings_key: str, in_memory_var: Any, success_message: str, @@ -769,6 +803,70 @@ async def update_ui_theme_settings(theme_config: UIThemeConfig): } +@router.get( + "/get/mcp_semantic_filter_settings", + tags=["Settings"], + dependencies=[Depends(user_api_key_auth)], + response_model=MCPSemanticFilterSettingsResponse, +) +async def get_mcp_semantic_filter_settings( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Get MCP semantic filter configuration. + Returns current settings for semantic tool filtering. + """ + from litellm.proxy.proxy_server import prisma_client, proxy_config + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": "Database not connected. Please connect a database."}, + ) + + config = await proxy_config.get_config() + + return await _get_settings_with_schema( + settings_key="mcp_semantic_tool_filter", + settings_class=MCPSemanticFilterSettings, + config=config, + ) + + +@router.patch( + "/update/mcp_semantic_filter_settings", + tags=["Settings"], + dependencies=[Depends(user_api_key_auth)], +) +async def update_mcp_semantic_filter_settings( + settings: MCPSemanticFilterSettings, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Update MCP semantic filter settings in database. + Settings will be picked up by all pods within approximately 10 seconds via background polling. + """ + result = await _update_litellm_setting( + settings=settings, + settings_key="mcp_semantic_tool_filter", + in_memory_var=None, + success_message="MCP Semantic Filter settings updated successfully. Changes will be applied across all pods within 10 seconds.", + ) + try: + from litellm.proxy.proxy_server import prisma_client, proxy_config + + if prisma_client is not None: + await proxy_config._init_semantic_filter_settings_in_db( + prisma_client=prisma_client + ) + except Exception as e: + verbose_proxy_logger.warning( + f"Failed to reinitialize MCP semantic filter settings immediately: {e}" + ) + + return result + + @router.get( "/in_product_nudges", tags=["UI Settings"], diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpSemanticFilterSettings/useMCPSemanticFilterSettings.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpSemanticFilterSettings/useMCPSemanticFilterSettings.ts new file mode 100644 index 00000000000..e91f5aa670b --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpSemanticFilterSettings/useMCPSemanticFilterSettings.ts @@ -0,0 +1,19 @@ +import { getMCPSemanticFilterSettings } from "@/components/networking"; +import { useQuery } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; +import useAuthorized from "../useAuthorized"; + +const mcpSemanticFilterSettingsKeys = createQueryKeys( + "mcpSemanticFilterSettings" +); + +export const useMCPSemanticFilterSettings = () => { + const { accessToken } = useAuthorized(); + return useQuery>({ + queryKey: mcpSemanticFilterSettingsKeys.list({}), + queryFn: async () => await getMCPSemanticFilterSettings(accessToken), + enabled: !!accessToken, + staleTime: 60 * 60 * 1000, // 1 hour + gcTime: 60 * 60 * 1000, // 1 hour + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpSemanticFilterSettings/useUpdateMCPSemanticFilterSettings.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpSemanticFilterSettings/useUpdateMCPSemanticFilterSettings.ts new file mode 100644 index 00000000000..2062b4f4c29 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpSemanticFilterSettings/useUpdateMCPSemanticFilterSettings.ts @@ -0,0 +1,25 @@ +import { updateMCPSemanticFilterSettings } from "@/components/networking"; +import { useMutation, useQueryClient } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; + +const mcpSemanticFilterSettingsKeys = createQueryKeys( + "mcpSemanticFilterSettings" +); + +export const useUpdateMCPSemanticFilterSettings = (accessToken: string) => { + const queryClient = useQueryClient(); + + return useMutation({ + mutationFn: async (settings: Record) => { + if (!accessToken) { + throw new Error("Access token is required"); + } + return updateMCPSemanticFilterSettings(accessToken, settings); + }, + onSuccess: () => { + queryClient.invalidateQueries({ + queryKey: mcpSemanticFilterSettingsKeys.all, + }); + }, + }); +}; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterSettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterSettings.tsx new file mode 100644 index 00000000000..6e669f74ab3 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterSettings.tsx @@ -0,0 +1,312 @@ +"use client"; + +import { useMCPSemanticFilterSettings } from "@/app/(dashboard)/hooks/mcpSemanticFilterSettings/useMCPSemanticFilterSettings"; +import { useUpdateMCPSemanticFilterSettings } from "@/app/(dashboard)/hooks/mcpSemanticFilterSettings/useUpdateMCPSemanticFilterSettings"; +import NotificationManager from "@/components/molecules/notifications_manager"; +import { + Alert, + Button, + Card, + Col, + Form, + InputNumber, + Row, + Select, + Skeleton, + Slider, + Space, + Switch, + Typography, + Tooltip, +} from "antd"; +import { QuestionCircleOutlined, CheckCircleOutlined, SaveOutlined } from "@ant-design/icons"; +import { useEffect, useState } from "react"; +import { fetchAvailableModels, ModelGroup } from "@/components/playground/llm_calls/fetch_models"; +import MCPSemanticFilterTestPanel from "./MCPSemanticFilterTestPanel"; +import { getCurlCommand, runSemanticFilterTest, TestResult } from "./semanticFilterTestUtils"; + +interface MCPSemanticFilterSettingsProps { + accessToken: string | null; +} + +export default function MCPSemanticFilterSettings({ accessToken }: MCPSemanticFilterSettingsProps) { + const { data, isLoading, isError, error } = useMCPSemanticFilterSettings(); + const { + mutate: updateSettings, + isPending: isUpdating, + error: updateError, + } = useUpdateMCPSemanticFilterSettings(accessToken || ""); + const [form] = Form.useForm(); + const [saveSuccess, setSaveSuccess] = useState(false); + const [isDirty, setIsDirty] = useState(false); + const [embeddingModels, setEmbeddingModels] = useState([]); + const [loadingModels, setLoadingModels] = useState(true); + + // Test section state + const [testQuery, setTestQuery] = useState(""); + const [testModel, setTestModel] = useState("gpt-4o"); + const [testResult, setTestResult] = useState(null); + const [isTesting, setIsTesting] = useState(false); + const [showCurl, setShowCurl] = useState(false); + + const schema = data?.field_schema; + const values = data?.values ?? {}; + + useEffect(() => { + const loadEmbeddingModels = async () => { + if (!accessToken) return; + try { + setLoadingModels(true); + const models = await fetchAvailableModels(accessToken); + const embeddingOnly = models.filter((model) => model.mode === "embedding"); + setEmbeddingModels(embeddingOnly); + } catch (error) { + console.error("Error fetching embedding models:", error); + } finally { + setLoadingModels(false); + } + }; + + loadEmbeddingModels(); + }, [accessToken]); + + useEffect(() => { + if (values) { + form.setFieldsValue({ + enabled: values.enabled ?? false, + embedding_model: values.embedding_model ?? "text-embedding-3-small", + top_k: values.top_k ?? 10, + similarity_threshold: values.similarity_threshold ?? 0.3, + }); + setIsDirty(false); + } + }, [values, form]); + + const handleSave = async () => { + try { + const formValues = await form.validateFields(); + updateSettings(formValues, { + onSuccess: () => { + setIsDirty(false); + setSaveSuccess(true); + setTimeout(() => setSaveSuccess(false), 3000); + NotificationManager.success( + "Settings updated successfully. Changes will be applied across all pods within 10 seconds." + ); + }, + onError: (error) => { + NotificationManager.fromBackend(error); + }, + }); + } catch (error) { + console.error("Form validation failed:", error); + } + }; + + const handleTest = async () => { + if (!accessToken) { + return; + } + + await runSemanticFilterTest({ + accessToken, + testModel, + testQuery, + setIsTesting, + setTestResult, + }); + }; + + if (!accessToken) { + return ( +
+ Please log in to configure semantic filter settings. +
+ ); + } + + return ( +
+ {isLoading ? ( + + ) : isError ? ( + + ) : ( + <> + + + {saveSuccess && ( + } + showIcon + closable + style={{ marginBottom: 16 }} + /> + )} + + {updateError && ( + + )} + + + {/* Left Column - Settings */} + +
{ + setIsDirty(true); + }} + > + + + Enable Semantic Filtering + + + + + } + valuePropName="checked" + > + + + + + {schema?.properties?.enabled?.description} + + + + + + Embedding Model + + + + + } + > +