fix: search tools access for team keys

litellm/proxy/_types.py

@@ -758,6 +758,7 @@ class LiteLLM_ObjectPermissionBase(LiteLLMPydanticObjectBase):
     mcp_access_groups: Optional[List[str]] = None
     mcp_tool_permissions: Optional[Dict[str, List[str]]] = None
     vector_stores: Optional[List[str]] = None
+    search_tools: Optional[List[str]] = None
 
 
 class GenerateRequestBase(LiteLLMPydanticObjectBase):

litellm/router_utils/search_api_router.py

@@ -158,55 +158,110 @@ async def async_search_with_fallbacks_helper(
         """
         Helper function for search API calls - selects a search tool and calls the original function.
         Called by async_function_with_fallbacks for each retry attempt.
-        
+
         Args:
             router_instance: The Router instance
             model: The search tool name (passed as model for compatibility)
             original_generic_function: The original litellm.asearch function
             **kwargs: Search parameters
-            
+
         Returns:
             SearchResponse from the selected search provider
         """
         search_tool_name = model  # model field contains the search_tool_name
-        
+
         try:
+            team_id = kwargs.get("litellm_metadata", {}).get("user_api_key_team_id")
+            if team_id is not None:
+                await SearchAPIRouter._check_team_search_tool_permission(
+                    router_instance=router_instance,
+                    team_id=team_id,
+                    search_tool_name=search_tool_name,
+                )
+
             # Find matching search tools
             matching_tools = SearchAPIRouter.get_matching_search_tools(
                 router_instance=router_instance,
                 search_tool_name=search_tool_name,
             )
-            
+
             # Simple random selection for load balancing across multiple providers with same name
             # For search tools, we use simple random choice since they don't have TPM/RPM constraints
             selected_tool = random.choice(matching_tools)
-            
+
             # Extract search provider and other params from litellm_params
             litellm_params = selected_tool.get("litellm_params", {})
             search_provider = litellm_params.get("search_provider")
             api_key = litellm_params.get("api_key")
             api_base = litellm_params.get("api_base")
-            
+
             if not search_provider:
                 raise ValueError(f"search_provider not found in litellm_params for search tool '{search_tool_name}'")
-            
+
             verbose_router_logger.debug(
                 f"Selected search tool with provider: {search_provider}"
             )
-            
+
             # Call the original search function with the provider config
             response = await original_generic_function(
                 search_provider=search_provider,
                 api_key=api_key,
                 api_base=api_base,
                 **kwargs,
             )
-            
+
             return response
-            
+
         except Exception as e:
             verbose_router_logger.error(
                 f"Error in SearchAPIRouter.async_search_with_fallbacks_helper for {search_tool_name}: {str(e)}"
             )
             raise e
 
+    @staticmethod
+    async def _check_team_search_tool_permission(
+        router_instance: Any,
+        team_id: str,
+        search_tool_name: str,
+    ):
+        try:
+            from litellm.proxy.auth.auth_checks import get_team_object
+            from litellm.proxy.proxy_server import (
+                prisma_client,
+                user_api_key_cache,
+                proxy_logging_obj,
+            )
+
+            team_obj = await get_team_object(
+                team_id=team_id,
+                prisma_client=prisma_client,
+                user_api_key_cache=user_api_key_cache,
+                parent_otel_span=None,
+                proxy_logging_obj=proxy_logging_obj,
+            )
+
+            if team_obj is None:
+                raise Exception(f"Team not found: {team_id}")
+
+            if (
+                team_obj.object_permission is not None
+                and team_obj.object_permission.search_tools is not None
+            ):
+                allowed_search_tools = team_obj.object_permission.search_tools
+                if search_tool_name not in allowed_search_tools:
+                    raise Exception(
+                        f"Team '{team_id}' does not have permission to access search tool '{search_tool_name}'. "
+                        f"Allowed search tools: {allowed_search_tools}"
+                    )
+            else:
+                raise Exception(
+                    f"Team '{team_id}' does not have permission to access search tool '{search_tool_name}'. "
+                    "No search tools are configured for this team."
+                )
+
+        except Exception as e:
+            verbose_router_logger.error(
+                f"Error checking team search tool permission for team {team_id}, tool {search_tool_name}: {str(e)}"
+            )
+            raise e
+

ui/litellm-dashboard/src/components/search_tools/search_tool_selector.tsx

@@ -0,0 +1,68 @@
+import React, { useEffect, useState } from "react";
+import { Select } from "antd";
+import { SearchTool } from "./types";
+import { fetchSearchTools } from "../networking";
+
+interface SearchToolSelectorProps {
+  onChange: (selectedSearchTools: string[]) => void;
+  value?: string[];
+  className?: string;
+  accessToken: string;
+  placeholder?: string;
+  disabled?: boolean;
+}
+
+const SearchToolSelector: React.FC<SearchToolSelectorProps> = ({
+  onChange,
+  value,
+  className,
+  accessToken,
+  placeholder = "Select search tools",
+  disabled = false,
+}) => {
+  const [searchTools, setSearchTools] = useState<SearchTool[]>([]);
+  const [loading, setLoading] = useState(false);
+
+  useEffect(() => {
+    const fetchSearchToolList = async () => {
+      if (!accessToken) return;
+
+      setLoading(true);
+      try {
+        const response = await fetchSearchTools(accessToken);
+        if (response.data) {
+          setSearchTools(response.data);
+        }
+      } catch (error) {
+        console.error("Error fetching search tools:", error);
+      } finally {
+        setLoading(false);
+      }
+    };
+
+    fetchSearchToolList();
+  }, [accessToken]);
+
+  return (
+    <div>
+      <Select
+        mode="multiple"
+        placeholder={placeholder}
+        onChange={onChange}
+        value={value}
+        className={className}
+        disabled={disabled || loading}
+        loading={loading}
+        style={{ width: "100%" }}
+      >
+        {searchTools.map((tool) => (
+          <Select.Option key={tool.search_tool_id} value={tool.search_tool_name}>
+            {tool.search_tool_name}
+          </Select.Option>
+        ))}
+      </Select>
+    </div>
+  );
+};
+
+export default SearchToolSelector;

ui/litellm-dashboard/src/components/team/team_info.tsx

@@ -35,6 +35,7 @@ import ObjectPermissionsView from "../object_permissions_view";
 import VectorStoreSelector from "../vector_store_management/VectorStoreSelector";
 import MCPServerSelector from "../mcp_server_management/MCPServerSelector";
 import MCPToolPermissions from "../mcp_server_management/MCPToolPermissions";
+import SearchToolSelector from "../search_tools/search_tool_selector";
 import { formatNumberWithCommas } from "@/utils/dataUtils";
 import EditLoggingSettings from "./EditLoggingSettings";
 import LoggingSettingsView from "../logging_settings_view";
@@ -92,6 +93,7 @@ export interface TeamData {
       mcp_access_groups?: string[];
       mcp_tool_permissions?: Record<string, string[]>;
       vector_stores: string[];
+      search_tools?: string[];
     };
     team_member_budget_table: {
       max_budget: number;
@@ -682,6 +684,15 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
                     />
                   </Form.Item>
 
+                  <Form.Item label="Search Tools" name="search_tools">
+                    <SearchToolSelector
+                      onChange={(values: string[]) => form.setFieldValue("search_tools", values)}
+                      value={form.getFieldValue("search_tools")}
+                      accessToken={accessToken || ""}
+                      placeholder="Select search tools"
+                    />
+                  </Form.Item>
+
                   <Form.Item label="Allowed Pass Through Routes" name="allowed_passthrough_routes">
                     <PassThroughRoutesSelector
                       onChange={(values: string[]) => form.setFieldValue("allowed_passthrough_routes", values)}