From fc9aba279ea51ab4e9b71614afed6951dc34f028 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 23 Oct 2025 17:59:29 -0700 Subject: [PATCH] [Feat] UI - Search Tools, allow adding search tools on UI + testing search (#15871) * add LiteLLM_SearchToolsTable * init SearchToolRegistry * fix add SearchToolRegistry * fix add SearchToolRegistry * fix handling search tool management * fix search imports * fix registry * init search tools in memory * fix init tools in mem * fix TypedDict def * add new SCHEMA * bump proxy extras * add LiteLLM_SearchToolsTable_search_tool_name_key * bump extras with migration * fix working CRUD Ops * fix: _init_search_tools_in_db * add UI friendly name for search providers * add ui friendly name for search providers * add providers available * working layout * better layout * clean add search tool * update_router_search_tools * fix remove in memory registry, since router is in mem store * allow testing search tool connection * clean create search tool * add test_search_tool_connection * fix: _init_search_tools_in_db * add searchToolQueryCall * fix icon * clean tester --- .../llms/base_llm/search/transformation.py | 8 + .../llms/dataforseo/search/transformation.py | 4 + litellm/llms/exa_ai/search/transformation.py | 4 + .../llms/google_pse/search/transformation.py | 4 + .../llms/parallel_ai/search/transformation.py | 4 + .../llms/perplexity/search/transformation.py | 4 + litellm/llms/tavily/search/transformation.py | 4 + litellm/proxy/proxy_server.py | 34 +- litellm/proxy/search_endpoints/__init__.py | 8 +- .../search_tool_management.py | 265 ++++++++------ .../search_endpoints/search_tool_registry.py | 79 ---- litellm/router_utils/search_api_router.py | 41 +++ litellm/types/search.py | 5 +- ui/litellm-dashboard/src/app/page.tsx | 3 + .../src/components/leftnav.tsx | 2 + .../src/components/networking.tsx | 254 +++++++++++++ .../search_tools/create_search_tool.tsx | 289 +++++++++++++++ .../src/components/search_tools/index.tsx | 6 + .../search_tools/search_connection_test.tsx | 272 ++++++++++++++ .../search_tools/search_tool_columns.tsx | 94 +++++ .../search_tools/search_tool_tester.tsx | 344 ++++++++++++++++++ .../search_tools/search_tool_view.tsx | 127 +++++++ .../components/search_tools/search_tools.tsx | 292 +++++++++++++++ .../src/components/search_tools/types.tsx | 36 ++ 24 files changed, 1974 insertions(+), 209 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/search_tools/create_search_tool.tsx create mode 100644 ui/litellm-dashboard/src/components/search_tools/index.tsx create mode 100644 ui/litellm-dashboard/src/components/search_tools/search_connection_test.tsx create mode 100644 ui/litellm-dashboard/src/components/search_tools/search_tool_columns.tsx create mode 100644 ui/litellm-dashboard/src/components/search_tools/search_tool_tester.tsx create mode 100644 ui/litellm-dashboard/src/components/search_tools/search_tool_view.tsx create mode 100644 ui/litellm-dashboard/src/components/search_tools/search_tools.tsx create mode 100644 ui/litellm-dashboard/src/components/search_tools/types.tsx diff --git a/litellm/llms/base_llm/search/transformation.py b/litellm/llms/base_llm/search/transformation.py index 8dd66f43641..14941911f17 100644 --- a/litellm/llms/base_llm/search/transformation.py +++ b/litellm/llms/base_llm/search/transformation.py @@ -49,6 +49,14 @@ class BaseSearchConfig: def __init__(self) -> None: pass + @staticmethod + def ui_friendly_name() -> str: + """ + UI-friendly name for the search provider. + Override in provider-specific implementations. + """ + return "Unknown Search Provider" + def get_http_method(self) -> Literal["GET", "POST"]: """ Get HTTP method for search requests. diff --git a/litellm/llms/dataforseo/search/transformation.py b/litellm/llms/dataforseo/search/transformation.py index e245804a00c..86b472f61b8 100644 --- a/litellm/llms/dataforseo/search/transformation.py +++ b/litellm/llms/dataforseo/search/transformation.py @@ -27,6 +27,10 @@ class DataForSEOSearchConfig(BaseSearchConfig): DATAFORSEO_API_BASE = "https://api.dataforseo.com/v3/serp/google/organic/live/advanced" + @staticmethod + def ui_friendly_name() -> str: + return "DataForSEO" + def get_http_method(self) -> Literal["GET", "POST"]: """ DataForSEO uses POST requests with JSON body. diff --git a/litellm/llms/exa_ai/search/transformation.py b/litellm/llms/exa_ai/search/transformation.py index c1c2b932d31..6b51c6cf25d 100644 --- a/litellm/llms/exa_ai/search/transformation.py +++ b/litellm/llms/exa_ai/search/transformation.py @@ -46,6 +46,10 @@ class ExaAISearchRequest(_ExaAISearchRequestRequired, total=False): class ExaAISearchConfig(BaseSearchConfig): EXA_AI_API_BASE = "https://api.exa.ai" + @staticmethod + def ui_friendly_name() -> str: + return "Exa AI" + def validate_environment( self, headers: Dict, diff --git a/litellm/llms/google_pse/search/transformation.py b/litellm/llms/google_pse/search/transformation.py index 30fcbd533e9..c1ba9cfe629 100644 --- a/litellm/llms/google_pse/search/transformation.py +++ b/litellm/llms/google_pse/search/transformation.py @@ -55,6 +55,10 @@ class GooglePSESearchRequest(_GooglePSESearchRequestRequired, total=False): class GooglePSESearchConfig(BaseSearchConfig): GOOGLE_PSE_API_BASE = "https://www.googleapis.com/customsearch/v1" + @staticmethod + def ui_friendly_name() -> str: + return "Google PSE" + def get_http_method(self) -> Literal["GET", "POST"]: """ Google PSE uses GET requests with query parameters. diff --git a/litellm/llms/parallel_ai/search/transformation.py b/litellm/llms/parallel_ai/search/transformation.py index 83e54f66ed2..95919b85c2f 100644 --- a/litellm/llms/parallel_ai/search/transformation.py +++ b/litellm/llms/parallel_ai/search/transformation.py @@ -45,6 +45,10 @@ class ParallelAISearchConfig(BaseSearchConfig): PARALLEL_AI_API_BASE = "https://api.parallel.ai" PARALLEL_HEADER_SEARCH_EXTRACT_VALUE = "search-extract-2025-10-10" + @staticmethod + def ui_friendly_name() -> str: + return "Parallel AI" + def validate_environment( self, headers: Dict, diff --git a/litellm/llms/perplexity/search/transformation.py b/litellm/llms/perplexity/search/transformation.py index 2c4ff3f446c..f1dc0909b4d 100644 --- a/litellm/llms/perplexity/search/transformation.py +++ b/litellm/llms/perplexity/search/transformation.py @@ -33,6 +33,10 @@ class PerplexitySearchRequest(_PerplexitySearchRequestRequired, total=False): class PerplexitySearchConfig(BaseSearchConfig): PERPLEXITY_API_BASE = "https://api.perplexity.ai" + @staticmethod + def ui_friendly_name() -> str: + return "Perplexity" + def validate_environment( self, headers: Dict, diff --git a/litellm/llms/tavily/search/transformation.py b/litellm/llms/tavily/search/transformation.py index e61ed1f0879..7fc33416a0b 100644 --- a/litellm/llms/tavily/search/transformation.py +++ b/litellm/llms/tavily/search/transformation.py @@ -45,6 +45,10 @@ class TavilySearchRequest(_TavilySearchRequestRequired, total=False): class TavilySearchConfig(BaseSearchConfig): TAVILY_API_BASE = "https://api.tavily.com" + @staticmethod + def ui_friendly_name() -> str: + return "Tavily" + def validate_environment( self, headers: Dict, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 591d0e86696..3d892938ab4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3523,15 +3523,35 @@ class ProxyConfig: ) async def _init_search_tools_in_db(self, prisma_client: PrismaClient): - from litellm.proxy.search_endpoints.search_tool_registry import ( - IN_MEMORY_SEARCH_TOOL_HANDLER, - SearchToolRegistry, - ) - from litellm.types.search import SearchTool + """ + Initialize search tools from database into the router on startup. + """ + global llm_router + + from litellm.proxy.search_endpoints.search_tool_registry import SearchToolRegistry + from litellm.router_utils.search_api_router import SearchAPIRouter + try: search_tools = await SearchToolRegistry.get_all_search_tools_from_db(prisma_client=prisma_client) - for search_tool in search_tools: - IN_MEMORY_SEARCH_TOOL_HANDLER.add_search_tool(search_tool=cast(SearchTool, search_tool)) + + verbose_proxy_logger.info( + f"Loading {len(search_tools)} search tool(s) from database into router" + ) + + if llm_router is not None: + # Add search tools to the router + await SearchAPIRouter.update_router_search_tools( + router_instance=llm_router, + search_tools=search_tools + ) + verbose_proxy_logger.info( + f"Successfully loaded {len(search_tools)} search tool(s) into router" + ) + else: + verbose_proxy_logger.debug( + "Router not initialized yet, search tools will be added when router is created" + ) + except Exception as e: verbose_proxy_logger.exception( "litellm.proxy.proxy_server.py::ProxyConfig:_init_search_tools_in_db - {}".format( diff --git a/litellm/proxy/search_endpoints/__init__.py b/litellm/proxy/search_endpoints/__init__.py index 45d9d2f0679..92b88f783ca 100644 --- a/litellm/proxy/search_endpoints/__init__.py +++ b/litellm/proxy/search_endpoints/__init__.py @@ -1,14 +1,8 @@ # litellm/proxy/search_endpoints/__init__.py -from .search_tool_registry import ( - IN_MEMORY_SEARCH_TOOL_HANDLER, - InMemorySearchToolHandler, - SearchToolRegistry, -) +from .search_tool_registry import SearchToolRegistry __all__ = [ "SearchToolRegistry", - "InMemorySearchToolHandler", - "IN_MEMORY_SEARCH_TOOL_HANDLER", ] diff --git a/litellm/proxy/search_endpoints/search_tool_management.py b/litellm/proxy/search_endpoints/search_tool_management.py index 7f4ad58b3b0..e3844b52848 100644 --- a/litellm/proxy/search_endpoints/search_tool_management.py +++ b/litellm/proxy/search_endpoints/search_tool_management.py @@ -2,17 +2,14 @@ CRUD ENDPOINTS FOR SEARCH TOOLS """ from datetime import datetime -from typing import List, Union, cast +from typing import Any, Dict, List, Union, cast from fastapi import APIRouter, Depends, HTTPException from pydantic import BaseModel from litellm._logging import verbose_proxy_logger from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.search_endpoints.search_tool_registry import ( - IN_MEMORY_SEARCH_TOOL_HANDLER, - SearchToolRegistry, -) +from litellm.proxy.search_endpoints.search_tool_registry import SearchToolRegistry from litellm.types.search import ( AvailableSearchProvider, ListSearchToolsResponse, @@ -169,16 +166,10 @@ async def create_search_tool(request: CreateSearchToolRequest): search_tool=request.search_tool, prisma_client=prisma_client ) - # Add to in-memory cache - try: - IN_MEMORY_SEARCH_TOOL_HANDLER.add_search_tool(search_tool=cast(SearchTool, result)) - verbose_proxy_logger.info( - f"Successfully added search tool '{result.get('search_tool_name')}' to in-memory cache" - ) - except Exception as cache_error: - verbose_proxy_logger.warning( - f"Failed to add search tool to in-memory cache: {cache_error}" - ) + verbose_proxy_logger.info( + f"Successfully added search tool '{result.get('search_tool_name')}' to database. " + f"Router will be updated by the cron job." + ) return result except Exception as e: @@ -258,18 +249,10 @@ async def update_search_tool(search_tool_id: str, request: UpdateSearchToolReque prisma_client=prisma_client, ) - # Update in-memory cache - try: - IN_MEMORY_SEARCH_TOOL_HANDLER.update_search_tool( - search_tool_id=search_tool_id, search_tool=cast(SearchTool, result) - ) - verbose_proxy_logger.info( - f"Successfully updated search tool '{result.get('search_tool_name')}' in in-memory cache" - ) - except Exception as cache_error: - verbose_proxy_logger.warning( - f"Failed to update search tool in in-memory cache: {cache_error}" - ) + verbose_proxy_logger.info( + f"Successfully updated search tool '{result.get('search_tool_name')}' in database. " + f"Router will be updated by the cron job." + ) return result except HTTPException as e: @@ -323,18 +306,10 @@ async def delete_search_tool(search_tool_id: str): search_tool_id=search_tool_id, prisma_client=prisma_client ) - # Delete from in-memory cache - try: - IN_MEMORY_SEARCH_TOOL_HANDLER.delete_search_tool( - search_tool_id=search_tool_id - ) - verbose_proxy_logger.info( - f"Successfully removed search tool from in-memory cache" - ) - except Exception as cache_error: - verbose_proxy_logger.warning( - f"Failed to remove search tool from in-memory cache: {cache_error}" - ) + verbose_proxy_logger.info( + f"Successfully deleted search tool from database. " + f"Router will be updated by the cron job." + ) return result except HTTPException as e: @@ -387,12 +362,6 @@ async def get_search_tool_info(search_tool_id: str): search_tool_id=search_tool_id, prisma_client=prisma_client ) - if result is None: - # Try in-memory cache - result = IN_MEMORY_SEARCH_TOOL_HANDLER.get_search_tool_by_id( - search_tool_id=search_tool_id - ) - if result is None: raise HTTPException( status_code=404, @@ -422,6 +391,110 @@ async def get_search_tool_info(search_tool_id: str): raise HTTPException(status_code=500, detail=str(e)) +class TestSearchToolConnectionRequest(BaseModel): + litellm_params: Dict[str, Any] + + +@router.post( + "/search_tools/test_connection", + tags=["Search Tools"], + dependencies=[Depends(user_api_key_auth)], +) +async def test_search_tool_connection(request: TestSearchToolConnectionRequest): + """ + Test connection to a search provider with the given configuration. + + Makes a simple test search query to verify the API key and configuration are valid. + + Example Request: + ```bash + curl -X POST "http://localhost:4000/search_tools/test_connection" \\ + -H "Authorization: Bearer " \\ + -H "Content-Type: application/json" \\ + -d '{ + "litellm_params": { + "search_provider": "perplexity", + "api_key": "sk-..." + } + }' + ``` + + Example Response (Success): + ```json + { + "status": "success", + "message": "Successfully connected to perplexity search provider", + "test_query": "test", + "results_count": 5 + } + ``` + + Example Response (Failure): + ```json + { + "status": "error", + "message": "Authentication failed: Invalid API key", + "error_type": "AuthenticationError" + } + ``` + """ + try: + from litellm.search import asearch + + # Extract params from request + litellm_params = request.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 HTTPException( + status_code=400, + detail="search_provider is required in litellm_params" + ) + + verbose_proxy_logger.debug( + f"Testing connection to search provider: {search_provider}" + ) + + # Make a simple test search query with max_results=1 to minimize cost + test_query = "test" + response = await asearch( + query=test_query, + search_provider=search_provider, + api_key=api_key, + api_base=api_base, + max_results=1, # Minimize results to reduce cost + timeout=10.0, # 10 second timeout for test + ) + + verbose_proxy_logger.info( + f"Successfully tested connection to {search_provider} search provider" + ) + + return { + "status": "success", + "message": f"Successfully connected to {search_provider} search provider", + "test_query": test_query, + "results_count": len(response.results) if response and response.results else 0, + } + + except Exception as e: + error_message = str(e) + error_type = type(e).__name__ + + verbose_proxy_logger.exception( + f"Failed to connect to search provider: {error_message}" + ) + + # Return error details in a structured format + return { + "status": "error", + "message": error_message, + "error_type": error_type, + } + + @router.get( "/search_tools/ui/available_providers", tags=["Search Tools"], @@ -431,7 +504,7 @@ async def get_available_search_providers(): """ Get the list of available search providers with their configuration fields. - This auto-discovers search providers from the SearchProviders enum. + Auto-discovers search providers and their UI-friendly names from transformation configs. Example Request: ```bash @@ -441,80 +514,46 @@ async def get_available_search_providers(): Example Response: ```json - [ - { - "provider": "perplexity", - "display_name": "Perplexity", - "fields": [ - { - "name": "api_key", - "type": "string", - "required": false, - "description": "API key for Perplexity" - }, - { - "name": "api_base", - "type": "string", - "required": false, - "description": "API base URL" - } - ] - } - ] + { + "providers": [ + { + "provider_name": "perplexity", + "ui_friendly_name": "Perplexity" + }, + { + "provider_name": "tavily", + "ui_friendly_name": "Tavily" + } + ] + } ``` """ try: - available_providers: List[AvailableSearchProvider] = [] + from litellm.utils import ProviderConfigManager + + available_providers = [] - # Common fields for all search providers - common_fields = [ - { - "name": "api_key", - "type": "string", - "required": False, - "description": "API key for the search provider", - }, - { - "name": "api_base", - "type": "string", - "required": False, - "description": "Custom API base URL (optional)", - }, - { - "name": "timeout", - "type": "number", - "required": False, - "description": "Request timeout in seconds", - }, - { - "name": "max_retries", - "type": "number", - "required": False, - "description": "Maximum number of retry attempts", - }, - ] - - # Provider display name mapping - provider_display_names = { - SearchProviders.PERPLEXITY: "Perplexity", - SearchProviders.TAVILY: "Tavily", - SearchProviders.PARALLEL_AI: "Parallel AI", - SearchProviders.EXA_AI: "Exa AI", - SearchProviders.GOOGLE_PSE: "Google PSE", - SearchProviders.DATAFORSEO: "DataForSEO", - } - # Auto-discover providers from SearchProviders enum for provider in SearchProviders: - available_providers.append( - AvailableSearchProvider( - provider=provider.value, - display_name=provider_display_names.get(provider, provider.value.title()), - fields=common_fields, + try: + # Get the config class for this provider + config = ProviderConfigManager.get_provider_search_config(provider=provider) + + if config is not None: + # Get the UI-friendly name from the config class + ui_name = config.ui_friendly_name() + + available_providers.append({ + "provider_name": provider.value, + "ui_friendly_name": ui_name, + }) + except Exception as e: + verbose_proxy_logger.debug( + f"Could not get config for search provider {provider.value}: {e}" ) - ) - - return available_providers + continue + + return {"providers": available_providers} except Exception as e: verbose_proxy_logger.exception(f"Error getting available search providers: {e}") raise HTTPException(status_code=500, detail=str(e)) diff --git a/litellm/proxy/search_endpoints/search_tool_registry.py b/litellm/proxy/search_endpoints/search_tool_registry.py index 6f17a8523ae..174f9f6b088 100644 --- a/litellm/proxy/search_endpoints/search_tool_registry.py +++ b/litellm/proxy/search_endpoints/search_tool_registry.py @@ -239,82 +239,3 @@ class SearchToolRegistry: verbose_proxy_logger.exception(f"Error getting search tool from DB: {str(e)}") raise Exception(f"Error getting search tool from DB: {str(e)}") - -class InMemorySearchToolHandler: - """ - Class that handles caching search tools in memory. - """ - - def __init__(self): - self.IN_MEMORY_SEARCH_TOOLS: Dict[str, SearchTool] = {} - """ - Search tool id to SearchTool object mapping - """ - - def add_search_tool(self, search_tool: SearchTool) -> None: - """ - Add a search tool to in-memory cache. - - Args: - search_tool: Search tool configuration - """ - search_tool_id = search_tool.get("search_tool_id") - if search_tool_id: - self.IN_MEMORY_SEARCH_TOOLS[search_tool_id] = search_tool - verbose_proxy_logger.debug( - f"Added search tool '{search_tool.get('search_tool_name')}' to in-memory cache" - ) - - def update_search_tool(self, search_tool_id: str, search_tool: SearchTool) -> None: - """ - Update a search tool in in-memory cache. - - Args: - search_tool_id: ID of search tool to update - search_tool: Updated search tool configuration - """ - self.IN_MEMORY_SEARCH_TOOLS[search_tool_id] = search_tool - verbose_proxy_logger.debug( - f"Updated search tool '{search_tool.get('search_tool_name')}' in in-memory cache" - ) - - def delete_search_tool(self, search_tool_id: str) -> None: - """ - Delete a search tool from in-memory cache. - - Args: - search_tool_id: ID of search tool to delete - """ - self.IN_MEMORY_SEARCH_TOOLS.pop(search_tool_id, None) - verbose_proxy_logger.debug( - f"Deleted search tool with ID '{search_tool_id}' from in-memory cache" - ) - - def list_search_tools(self) -> List[SearchTool]: - """ - List all search tools in in-memory cache. - - Returns: - List of search tool configurations - """ - return list(self.IN_MEMORY_SEARCH_TOOLS.values()) - - def get_search_tool_by_id(self, search_tool_id: str) -> Optional[SearchTool]: - """ - Get a search tool by its ID from in-memory cache. - - Args: - search_tool_id: ID of search tool to retrieve - - Returns: - Search tool configuration or None if not found - """ - return self.IN_MEMORY_SEARCH_TOOLS.get(search_tool_id) - - -######################################################## -# In Memory Search Tool Handler for LiteLLM Proxy -######################################################## -IN_MEMORY_SEARCH_TOOL_HANDLER = InMemorySearchToolHandler() -######################################################## - diff --git a/litellm/router_utils/search_api_router.py b/litellm/router_utils/search_api_router.py index fbad1c39122..247df099462 100644 --- a/litellm/router_utils/search_api_router.py +++ b/litellm/router_utils/search_api_router.py @@ -20,6 +20,47 @@ class SearchAPIRouter: Provides methods for search tool selection, load balancing, and fallback handling. """ + @staticmethod + async def update_router_search_tools(router_instance: Any, search_tools: list): + """ + Update the router with search tools from the database. + + This method is called by a cron job to sync search tools from DB to router. + + Args: + router_instance: The Router instance to update + search_tools: List of search tool configurations from the database + """ + try: + from litellm.types.router import SearchToolTypedDict + + verbose_router_logger.debug(f"Adding {len(search_tools)} search tools to router") + + # Convert search tools to the format expected by the router + router_search_tools: list = [] + for tool in search_tools: + # Create dict that matches SearchToolTypedDict structure + router_search_tool: SearchToolTypedDict = { # type: ignore + "search_tool_id": tool.get("search_tool_id"), + "search_tool_name": tool.get("search_tool_name"), + "litellm_params": tool.get("litellm_params", {}), + "search_tool_info": tool.get("search_tool_info"), + } + router_search_tools.append(router_search_tool) + + # Update the router's search_tools list + router_instance.search_tools = router_search_tools + + verbose_router_logger.info( + f"Successfully updated router with {len(router_search_tools)} search tool(s)" + ) + + except Exception as e: + verbose_router_logger.exception( + f"Error updating router with search tools: {str(e)}" + ) + raise e + @staticmethod def get_matching_search_tools( router_instance: Any, diff --git a/litellm/types/search.py b/litellm/types/search.py index 440a9ea7706..77a4c9ea72c 100644 --- a/litellm/types/search.py +++ b/litellm/types/search.py @@ -69,8 +69,7 @@ class ListSearchToolsResponse(TypedDict): class AvailableSearchProvider(TypedDict): """Information about an available search provider.""" - provider: str - display_name: str - fields: List[dict] + provider_name: str + ui_friendly_name: str diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index 5dc9fe0509a..f09220772eb 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -41,6 +41,7 @@ import { cx } from "@/lib/cva.config"; import useFeatureFlags from "@/hooks/useFeatureFlags"; import SidebarProvider from "@/app/(dashboard)/components/SidebarProvider"; import OldTeams from "@/components/OldTeams"; +import { SearchTools } from "@/components/search_tools"; function getCookie(name: string) { // Safer cookie read + decoding; handles '=' inside values @@ -463,6 +464,8 @@ export default function CreateKeyPage() { /> ) : page == "mcp-servers" ? ( + ) : page == "search-tools" ? ( + ) : page == "tag-management" ? ( ) : page == "vector-stores" ? ( diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index b974cdfaa9f..7a94c069a81 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -18,6 +18,7 @@ import { ToolOutlined, TagsOutlined, BgColorsOutlined, + SearchOutlined, } from "@ant-design/icons"; import { all_admin_roles, rolesWithWriteAccess, internalUserRoles, isAdminRole } from "../utils/roles"; import UsageIndicator from "./usage_indicator"; @@ -110,6 +111,7 @@ const Sidebar: React.FC = ({ accessToken, setPage, userRole, defau icon: , children: [ { key: "18", page: "mcp-servers", label: "MCP Servers", icon: }, + { key: "28", page: "search-tools", label: "Search Tools", icon: }, { key: "21", page: "vector-stores", diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 713cb1d58ce..f4bcfd47e7e 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -5421,6 +5421,226 @@ export const deleteMCPServer = async (accessToken: string, serverId: string) => } }; +// Search Tools API calls +export const fetchSearchTools = async (accessToken: string) => { + try { + const url = proxyBaseUrl ? `${proxyBaseUrl}/search_tools/list` : `/search_tools/list`; + console.log("Fetching search tools from:", url); + + const response = await fetch(url, { + method: HTTP_REQUEST.GET, + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + }); + + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + handleError(errorMessage); + throw new Error(errorMessage); + } + + const data = await response.json(); + console.log("Fetched search tools:", data); + return data; + } catch (error) { + console.error("Failed to fetch search tools:", error); + throw error; + } +}; + +export const fetchSearchToolById = async (accessToken: string, searchToolId: string) => { + try { + const url = proxyBaseUrl ? `${proxyBaseUrl}/search_tools/${searchToolId}` : `/search_tools/${searchToolId}`; + console.log("Fetching search tool by ID from:", url); + + const response = await fetch(url, { + method: HTTP_REQUEST.GET, + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + }); + + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + handleError(errorMessage); + throw new Error(errorMessage); + } + + const data = await response.json(); + console.log("Fetched search tool:", data); + return data; + } catch (error) { + console.error("Failed to fetch search tool:", error); + throw error; + } +}; + +export const createSearchTool = async (accessToken: string, formValues: Record) => { + try { + console.log("Creating search tool with values:", formValues); + const url = proxyBaseUrl ? `${proxyBaseUrl}/search_tools` : `/search_tools`; + + const response = await fetch(url, { + method: HTTP_REQUEST.POST, + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ + search_tool: formValues, + }), + }); + + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + handleError(errorMessage); + throw new Error(errorMessage); + } + + const data = await response.json(); + console.log("Created search tool:", data); + return data; + } catch (error) { + console.error("Failed to create search tool:", error); + throw error; + } +}; + +export const updateSearchTool = async (accessToken: string, searchToolId: string, formValues: Record) => { + try { + console.log("Updating search tool with ID:", searchToolId, "values:", formValues); + const url = proxyBaseUrl ? `${proxyBaseUrl}/search_tools/${searchToolId}` : `/search_tools/${searchToolId}`; + + const response = await fetch(url, { + method: HTTP_REQUEST.PUT, + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ + search_tool: formValues, + }), + }); + + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + handleError(errorMessage); + throw new Error(errorMessage); + } + + const data = await response.json(); + console.log("Updated search tool:", data); + return data; + } catch (error) { + console.error("Failed to update search tool:", error); + throw error; + } +}; + +export const deleteSearchTool = async (accessToken: string, searchToolId: string) => { + try { + const url = (proxyBaseUrl ? `${proxyBaseUrl}` : "") + `/search_tools/${searchToolId}`; + console.log("Deleting search tool:", searchToolId); + + const response = await fetch(url, { + method: HTTP_REQUEST.DELETE, + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + }); + + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + handleError(errorMessage); + throw new Error(errorMessage); + } + + const data = await response.json(); + console.log("Deleted search tool:", data); + return data; + } catch (error) { + console.error("Failed to delete search tool:", error); + throw error; + } +}; + +export const fetchAvailableSearchProviders = async (accessToken: string) => { + try { + const url = proxyBaseUrl + ? `${proxyBaseUrl}/search_tools/ui/available_providers` + : `/search_tools/ui/available_providers`; + console.log("Fetching available search providers from:", url); + + const response = await fetch(url, { + method: HTTP_REQUEST.GET, + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + }); + + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + handleError(errorMessage); + throw new Error(errorMessage); + } + + const data = await response.json(); + console.log("Fetched available search providers:", data); + return data; + } catch (error) { + console.error("Failed to fetch available search providers:", error); + throw error; + } +}; + +export const testSearchToolConnection = async ( + accessToken: string, + litellmParams: Record +) => { + try { + const url = proxyBaseUrl + ? `${proxyBaseUrl}/search_tools/test_connection` + : `/search_tools/test_connection`; + console.log("Testing search tool connection:", url); + + const response = await fetch(url, { + method: HTTP_REQUEST.POST, + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ + litellm_params: litellmParams, + }), + }); + + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + handleError(errorMessage); + throw new Error(errorMessage); + } + + const data = await response.json(); + console.log("Test connection response:", data); + return data; + } catch (error) { + console.error("Failed to test search tool connection:", error); + throw error; + } +}; + export const listMCPTools = async (accessToken: string, serverId: string, authValue?: string, serverAlias?: string) => { try { // Construct base URL @@ -6609,6 +6829,40 @@ export const vectorStoreSearchCall = async ( } }; +export const searchToolQueryCall = async ( + accessToken: string, + searchToolName: string, + query: string, + maxResults?: number, +): Promise => { + try { + const url = `${getProxyBaseUrl()}/v1/search/${searchToolName}`; + const response = await fetch(url, { + method: "POST", + headers: { + Authorization: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ + query: query, + max_results: maxResults || 5, + }), + }); + + if (!response.ok) { + const errorData = await response.text(); + await handleError(errorData); + return null; + } + + const data = await response.json(); + return data; + } catch (error) { + console.error("Error querying search tool:", error); + throw error; + } +}; + export const userAgentAnalyticsCall = async ( accessToken: string, startTime: Date, diff --git a/ui/litellm-dashboard/src/components/search_tools/create_search_tool.tsx b/ui/litellm-dashboard/src/components/search_tools/create_search_tool.tsx new file mode 100644 index 00000000000..1579737e095 --- /dev/null +++ b/ui/litellm-dashboard/src/components/search_tools/create_search_tool.tsx @@ -0,0 +1,289 @@ +import React, { useState } from "react"; +import { Modal, Tooltip, Form, Select, Input, Typography } from "antd"; +import { InfoCircleOutlined } from "@ant-design/icons"; +import { Button, TextInput } from "@tremor/react"; +import { createSearchTool, fetchAvailableSearchProviders } from "../networking"; +import { SearchTool, AvailableSearchProvider } from "./types"; +import { isAdminRole } from "@/utils/roles"; +import NotificationsManager from "../molecules/notifications_manager"; +import { useQuery } from "@tanstack/react-query"; +import SearchConnectionTest from "./search_connection_test"; + +const { TextArea } = Input; + +interface CreateSearchToolProps { + userRole: string; + accessToken: string | null; + onCreateSuccess: (newSearchTool: SearchTool) => void; + isModalVisible: boolean; + setModalVisible: (visible: boolean) => void; +} + +const CreateSearchTool: React.FC = ({ + userRole, + accessToken, + onCreateSuccess, + isModalVisible, + setModalVisible, +}) => { + const [form] = Form.useForm(); + const [isLoading, setIsLoading] = useState(false); + const [formValues, setFormValues] = useState>({}); + const [isTestModalVisible, setIsTestModalVisible] = useState(false); + const [isTestingConnection, setIsTestingConnection] = useState(false); + const [connectionTestId, setConnectionTestId] = useState(""); + + // Fetch available search providers + const { + data: providersResponse, + isLoading: isLoadingProviders, + } = useQuery({ + queryKey: ["searchProviders"], + queryFn: () => { + if (!accessToken) throw new Error("Access Token required"); + return fetchAvailableSearchProviders(accessToken); + }, + enabled: !!accessToken && isModalVisible, + }) as { data: { providers: AvailableSearchProvider[] }; isLoading: boolean }; + + const availableProviders = providersResponse?.providers || []; + + const handleCreate = async (formValues: Record) => { + setIsLoading(true); + try { + // Prepare the payload + const payload = { + search_tool_name: formValues.search_tool_name, + litellm_params: { + search_provider: formValues.search_provider, + api_key: formValues.api_key, + api_base: formValues.api_base, + timeout: formValues.timeout ? parseFloat(formValues.timeout) : undefined, + max_retries: formValues.max_retries ? parseInt(formValues.max_retries) : undefined, + }, + search_tool_info: formValues.description + ? { + description: formValues.description, + } + : undefined, + }; + + console.log(`Creating search tool with payload:`, payload); + + if (accessToken != null) { + const response = await createSearchTool(accessToken, payload); + + NotificationsManager.success("Search tool created successfully"); + form.resetFields(); + setFormValues({}); + setModalVisible(false); + onCreateSuccess(response); + } + } catch (error) { + NotificationsManager.error("Error creating search tool: " + error); + } finally { + setIsLoading(false); + } + }; + + const handleCancel = () => { + form.resetFields(); + setFormValues({}); + setModalVisible(false); + }; + + const handleTestConnection = async () => { + try { + // Validate required fields for testing + await form.validateFields(["search_provider", "api_key"]); + + setIsTestingConnection(true); + // Generate a new test ID (using timestamp for uniqueness) + setConnectionTestId(`test-${Date.now()}`); + // Show the modal with the fresh test + setIsTestModalVisible(true); + } catch (error) { + NotificationsManager.error("Please fill in Search Provider and API Key before testing"); + } + }; + + // Clear formValues when modal closes to reset + React.useEffect(() => { + if (!isModalVisible) { + setFormValues({}); + } + }, [isModalVisible]); + + if (!isAdminRole(userRole)) { + return null; + } + + return ( + + 🔍 +

Add New Search Tool

+ + } + open={isModalVisible} + width={800} + onCancel={handleCancel} + footer={null} + className="top-8" + styles={{ + body: { padding: "24px" }, + header: { padding: "24px 24px 0 24px", border: "none" }, + }} + > +
+
setFormValues(allValues)} + layout="vertical" + className="space-y-6" + > +
+ + Search Tool Name + + + + + } + name="search_tool_name" + rules={[ + { required: true, message: "Please enter a search tool name" }, + { + pattern: /^[a-zA-Z0-9_-]+$/, + message: "Name can only contain letters, numbers, hyphens, and underscores", + }, + ]} + > + + + + + Search Provider + + + + + } + name="search_provider" + rules={[{ required: true, message: "Please select a search provider" }]} + > + + + + + API Key + + + + + } + name="api_key" + rules={[{ required: false, message: "Please enter an API key" }]} + > + + + + Description (Optional)} + name="description" + > +