Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/search_endpoints/search_tool_management.py: 65%
196 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1"""
2CRUD ENDPOINTS FOR SEARCH TOOLS
3"""
5from collections.abc import Awaitable, Callable
6from datetime import datetime
7from typing import Any, Final, TypeAlias
9from fastapi import APIRouter, Depends, HTTPException
10from pydantic import BaseModel
12from litellm._logging import verbose_proxy_logger
13from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
14from litellm.proxy._types import (
15 LiteLLM_TeamTable,
16 LitellmUserRoles,
17 UserAPIKeyAuth,
18)
19from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
20from litellm.proxy.search_endpoints.search_tool_registry import SearchToolRegistry
21from litellm.types.search import (
22 ListSearchToolsResponse,
23 SearchTool,
24 SearchToolInfoResponse,
25)
26from litellm.types.utils import SearchProviders
28#### SEARCH TOOLS ENDPOINTS ####
30router: Final = APIRouter()
31SEARCH_TOOL_REGISTRY: Final = SearchToolRegistry()
34def _convert_datetime_to_str(value: datetime | str | None) -> str | None:
35 """
36 Convert datetime object to ISO format string.
38 Args:
39 value: datetime object, string, or None
41 Returns:
42 ISO format string or original value if already string or None
43 """
44 if value is None: 44 ↛ 45line 44 didn't jump to line 45 because the condition on line 44 was never true
45 return None
46 if isinstance(value, datetime): 46 ↛ 47line 46 didn't jump to line 47 because the condition on line 46 was never true
47 return value.isoformat()
48 return value
51TeamObjectLookup: TypeAlias = Callable[[str, UserAPIKeyAuth], Awaitable[LiteLLM_TeamTable]]
54async def _refresh_router_search_tools() -> None:
55 """Push the search tools table into this worker's router.
57 Best-effort: the row is already committed, so a refresh failure must not surface as a 500 and
58 push the caller into a retry that creates duplicates.
59 """
60 from litellm.proxy.proxy_server import proxy_config
62 try:
63 await proxy_config.reload_search_tools_from_db()
64 except Exception as e: # noqa: BLE001 # the row is committed; no refresh failure may reach the caller
65 verbose_proxy_logger.exception("Search tool router refresh failed after a management write: %s", e)
68async def _team_object_from_db(team_id: str, user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTable:
69 from litellm.proxy.auth.auth_checks import get_team_object
70 from litellm.proxy.proxy_server import (
71 prisma_client,
72 proxy_logging_obj,
73 user_api_key_cache,
74 )
76 return await get_team_object(
77 team_id=team_id,
78 prisma_client=prisma_client,
79 user_api_key_cache=user_api_key_cache,
80 parent_otel_span=user_api_key_dict.parent_otel_span,
81 proxy_logging_obj=proxy_logging_obj,
82 )
85def _allowlist_team_id(user_api_key_dict: UserAPIKeyAuth) -> str | None:
86 """
87 The team whose object_permission allowlist scopes this caller, or None when there is none.
89 Every Admin UI session key is stamped with UI_SESSION_TOKEN_TEAM_ID, a reserved sentinel that
90 never has a row in LiteLLM_TeamTable (`/team/new` rejects it as a real team id), so looking it
91 up would raise 404 instead of resolving a team. It carries no allowlist of its own, so the
92 caller is scoped by its key-level allowlist alone. Any other team id is looked up for real and
93 a failed lookup still surfaces.
94 """
95 team_id: Final = user_api_key_dict.team_id
96 if not team_id or team_id == UI_SESSION_TOKEN_TEAM_ID:
97 return None
98 return team_id
101async def _filter_visible_search_tools(
102 search_tools: list[SearchToolInfoResponse],
103 user_api_key_dict: UserAPIKeyAuth,
104 lookup_team_object: TeamObjectLookup = _team_object_from_db,
105) -> list[SearchToolInfoResponse]:
106 """
107 Drop search tools the caller is not authorized to invoke, applying the same
108 key/team object_permission allowlists enforced on /search. Admins see all tools.
109 """
110 if user_api_key_dict.user_role in ( 110 ↛ 116line 110 didn't jump to line 116 because the condition on line 110 was always true
111 LitellmUserRoles.PROXY_ADMIN,
112 LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
113 ):
114 return search_tools
116 from litellm.proxy.auth.auth_checks import can_user_view_search_tool
118 allowlist_team_id: Final = _allowlist_team_id(user_api_key_dict)
119 team_object: Final[LiteLLM_TeamTable | None] = (
120 await lookup_team_object(allowlist_team_id, user_api_key_dict) if allowlist_team_id else None
121 )
123 visible: Final[list[SearchToolInfoResponse]] = []
124 for tool in search_tools:
125 tool_name = tool.get("search_tool_name")
126 if tool_name and await can_user_view_search_tool(
127 search_tool_name=tool_name,
128 valid_token=user_api_key_dict,
129 team_object=team_object,
130 ):
131 visible.append(tool)
132 return visible
135@router.get(
136 "/search_tools/list",
137 tags=["Search Tools"],
138 dependencies=[Depends(user_api_key_auth)],
139 response_model=ListSearchToolsResponse,
140)
141async def list_search_tools(
142 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
143):
144 """
145 List all search tools that are available in the database and config file.
147 Example Request:
148 ```bash
149 curl -X GET "http://localhost:4000/search_tools/list" -H "Authorization: Bearer <your_api_key>"
150 ```
152 Example Response:
153 ```json
154 {
155 "search_tools": [
156 {
157 "search_tool_id": "123e4567-e89b-12d3-a456-426614174000",
158 "search_tool_name": "litellm-search",
159 "litellm_params": {
160 "search_provider": "perplexity",
161 "api_key": "sk-***",
162 "api_base": "https://api.perplexity.ai"
163 },
164 "search_tool_info": {
165 "description": "Perplexity search tool"
166 },
167 "created_at": "2023-11-09T12:34:56.789Z",
168 "updated_at": "2023-11-09T12:34:56.789Z",
169 "is_from_config": false
170 },
171 {
172 "search_tool_name": "config-search-tool",
173 "litellm_params": {
174 "search_provider": "tavily",
175 "api_key": "tvly-***"
176 },
177 "is_from_config": true
178 }
179 ]
180 }
181 ```
182 """
183 from litellm.litellm_core_utils.litellm_logging import _get_masked_values
184 from litellm.proxy.proxy_server import prisma_client, proxy_config
186 if prisma_client is None: 186 ↛ 187line 186 didn't jump to line 187 because the condition on line 186 was never true
187 raise HTTPException(status_code=500, detail="Prisma client not initialized")
189 try:
190 search_tools_from_db = await SEARCH_TOOL_REGISTRY.get_all_search_tools_from_db(prisma_client=prisma_client)
192 db_tool_names: Final = {tool.get("search_tool_name") for tool in search_tools_from_db}
194 search_tool_configs: list[SearchToolInfoResponse] = []
196 config_search_tools = []
198 try:
199 config: Final = await proxy_config.get_config()
200 parsed_tools: Final = proxy_config.parse_search_tools(config)
201 if parsed_tools: 201 ↛ 202line 201 didn't jump to line 202 because the condition on line 201 was never true
202 config_search_tools = parsed_tools
203 except Exception as e:
204 verbose_proxy_logger.debug("Could not get config-defined search tools: %s", e)
206 for config_search_tool in config_search_tools: 206 ↛ 207line 206 didn't jump to line 207 because the loop on line 206 never started
207 tool_name = config_search_tool.get("search_tool_name")
208 if tool_name:
209 litellm_params_dict = dict(config_search_tool.get("litellm_params", {}))
210 masked_litellm_params_dict = _get_masked_values(
211 litellm_params_dict,
212 unmasked_length=4,
213 number_of_asterisks=4,
214 )
215 config_tool_info = config_search_tool.get("search_tool_info")
217 search_tool_configs.append(
218 SearchToolInfoResponse(
219 search_tool_id=None,
220 search_tool_name=tool_name,
221 litellm_params=masked_litellm_params_dict,
222 search_tool_info=(dict(config_tool_info) if config_tool_info else None),
223 created_at=None,
224 updated_at=None,
225 is_from_config=True,
226 )
227 )
229 search_tool_configs = [
230 tool for tool in search_tool_configs if tool.get("search_tool_name") not in db_tool_names
231 ]
233 for db_search_tool in search_tools_from_db:
234 litellm_params_dict = dict(db_search_tool.get("litellm_params", {}))
235 masked_litellm_params_dict = _get_masked_values(
236 litellm_params_dict,
237 unmasked_length=4,
238 number_of_asterisks=4,
239 )
241 search_tool_configs.append(
242 SearchToolInfoResponse(
243 search_tool_id=db_search_tool.get("search_tool_id"),
244 search_tool_name=db_search_tool.get("search_tool_name", ""),
245 litellm_params=masked_litellm_params_dict,
246 search_tool_info=db_search_tool.get("search_tool_info"),
247 created_at=_convert_datetime_to_str(db_search_tool.get("created_at")),
248 updated_at=_convert_datetime_to_str(db_search_tool.get("updated_at")),
249 is_from_config=False,
250 )
251 )
253 visible_search_tools: Final = await _filter_visible_search_tools(search_tool_configs, user_api_key_dict)
255 return ListSearchToolsResponse(search_tools=visible_search_tools)
256 except HTTPException:
257 raise
258 except Exception as e:
259 verbose_proxy_logger.exception("Error getting search tools: %s", e)
260 raise HTTPException(status_code=500, detail=str(e))
263class CreateSearchToolRequest(BaseModel):
264 search_tool: SearchTool
267@router.post(
268 "/search_tools",
269 tags=["Search Tools"],
270 dependencies=[Depends(user_api_key_auth)],
271)
272async def create_search_tool(request: CreateSearchToolRequest):
273 """
274 Create a new search tool.
276 Example Request:
277 ```bash
278 curl -X POST "http://localhost:4000/search_tools" \\
279 -H "Authorization: Bearer <your_api_key>" \\
280 -H "Content-Type: application/json" \\
281 -d '{
282 "search_tool": {
283 "search_tool_name": "litellm-search",
284 "litellm_params": {
285 "search_provider": "perplexity",
286 "api_key": "sk-..."
287 },
288 "search_tool_info": {
289 "description": "Perplexity search tool"
290 }
291 }
292 }'
293 ```
295 Example Response:
296 ```json
297 {
298 "search_tool_id": "123e4567-e89b-12d3-a456-426614174000",
299 "search_tool_name": "litellm-search",
300 "litellm_params": {
301 "search_provider": "perplexity",
302 "api_key": "sk-..."
303 },
304 "search_tool_info": {
305 "description": "Perplexity search tool"
306 },
307 "created_at": "2023-11-09T12:34:56.789Z",
308 "updated_at": "2023-11-09T12:34:56.789Z"
309 }
310 ```
311 """
312 from litellm.proxy.proxy_server import prisma_client
314 if prisma_client is None: 314 ↛ 315line 314 didn't jump to line 315 because the condition on line 314 was never true
315 raise HTTPException(status_code=500, detail="Prisma client not initialized")
317 try:
318 result: Final = await SEARCH_TOOL_REGISTRY.add_search_tool_to_db(
319 search_tool=request.search_tool, prisma_client=prisma_client
320 )
322 await _refresh_router_search_tools()
324 verbose_proxy_logger.debug(
325 "Successfully added search tool '%s' to database.",
326 result.get("search_tool_name"),
327 )
329 return result
330 except Exception as e:
331 verbose_proxy_logger.exception("Error adding search tool to db: %s", e)
332 raise HTTPException(status_code=500, detail=str(e))
335class UpdateSearchToolRequest(BaseModel):
336 search_tool: SearchTool
339@router.put(
340 "/search_tools/{search_tool_id}",
341 tags=["Search Tools"],
342 dependencies=[Depends(user_api_key_auth)],
343)
344async def update_search_tool(search_tool_id: str, request: UpdateSearchToolRequest):
345 """
346 Update an existing search tool.
348 Example Request:
349 ```bash
350 curl -X PUT "http://localhost:4000/search_tools/123e4567-e89b-12d3-a456-426614174000" \\
351 -H "Authorization: Bearer <your_api_key>" \\
352 -H "Content-Type: application/json" \\
353 -d '{
354 "search_tool": {
355 "search_tool_name": "updated-search",
356 "litellm_params": {
357 "search_provider": "perplexity",
358 "api_key": "sk-new-key"
359 },
360 "search_tool_info": {
361 "description": "Updated search tool"
362 }
363 }
364 }'
365 ```
367 Example Response:
368 ```json
369 {
370 "search_tool_id": "123e4567-e89b-12d3-a456-426614174000",
371 "search_tool_name": "updated-search",
372 "litellm_params": {
373 "search_provider": "perplexity",
374 "api_key": "sk-new-key"
375 },
376 "search_tool_info": {
377 "description": "Updated search tool"
378 },
379 "created_at": "2023-11-09T12:34:56.789Z",
380 "updated_at": "2023-11-09T13:45:12.345Z"
381 }
382 ```
383 """
384 from litellm.proxy.proxy_server import prisma_client
386 if prisma_client is None: 386 ↛ 387line 386 didn't jump to line 387 because the condition on line 386 was never true
387 raise HTTPException(status_code=500, detail="Prisma client not initialized")
389 try:
390 # Check if search tool exists
391 existing_tool: Final = await SEARCH_TOOL_REGISTRY.get_search_tool_by_id_from_db(
392 search_tool_id=search_tool_id, prisma_client=prisma_client
393 )
395 if existing_tool is None: 395 ↛ 396line 395 didn't jump to line 396 because the condition on line 395 was never true
396 raise HTTPException(
397 status_code=404,
398 detail=f"Search tool with ID {search_tool_id} not found",
399 )
401 result: Final = await SEARCH_TOOL_REGISTRY.update_search_tool_in_db(
402 search_tool_id=search_tool_id,
403 search_tool=request.search_tool,
404 prisma_client=prisma_client,
405 )
407 await _refresh_router_search_tools()
409 verbose_proxy_logger.debug(
410 "Successfully updated search tool '%s' in database.",
411 result.get("search_tool_name"),
412 )
414 return result
415 except HTTPException as e:
416 raise e
417 except Exception as e:
418 verbose_proxy_logger.exception("Error updating search tool: %s", e)
419 raise HTTPException(status_code=500, detail=str(e))
422@router.delete(
423 "/search_tools/{search_tool_id}",
424 tags=["Search Tools"],
425 dependencies=[Depends(user_api_key_auth)],
426)
427async def delete_search_tool(search_tool_id: str):
428 """
429 Delete a search tool.
431 Example Request:
432 ```bash
433 curl -X DELETE "http://localhost:4000/search_tools/123e4567-e89b-12d3-a456-426614174000" \\
434 -H "Authorization: Bearer <your_api_key>"
435 ```
437 Example Response:
438 ```json
439 {
440 "message": "Search tool 123e4567-e89b-12d3-a456-426614174000 deleted successfully",
441 "search_tool_name": "litellm-search"
442 }
443 ```
444 """
445 from litellm.proxy.proxy_server import prisma_client
447 if prisma_client is None: 447 ↛ 448line 447 didn't jump to line 448 because the condition on line 447 was never true
448 raise HTTPException(status_code=500, detail="Prisma client not initialized")
450 try:
451 # Check if search tool exists
452 existing_tool: Final = await SEARCH_TOOL_REGISTRY.get_search_tool_by_id_from_db(
453 search_tool_id=search_tool_id, prisma_client=prisma_client
454 )
456 if existing_tool is None: 456 ↛ 457line 456 didn't jump to line 457 because the condition on line 456 was never true
457 raise HTTPException(
458 status_code=404,
459 detail=f"Search tool with ID {search_tool_id} not found",
460 )
462 result: Final = await SEARCH_TOOL_REGISTRY.delete_search_tool_from_db(
463 search_tool_id=search_tool_id, prisma_client=prisma_client
464 )
466 await _refresh_router_search_tools()
468 verbose_proxy_logger.debug("Successfully deleted search tool from database.")
470 return result
471 except HTTPException as e:
472 raise e
473 except Exception as e:
474 verbose_proxy_logger.exception("Error deleting search tool: %s", e)
475 raise HTTPException(status_code=500, detail=str(e))
478@router.get(
479 "/search_tools/{search_tool_id}",
480 tags=["Search Tools"],
481 dependencies=[Depends(user_api_key_auth)],
482)
483async def get_search_tool_info(search_tool_id: str):
484 """
485 Get detailed information about a specific search tool by ID.
487 Example Request:
488 ```bash
489 curl -X GET "http://localhost:4000/search_tools/123e4567-e89b-12d3-a456-426614174000" \\
490 -H "Authorization: Bearer <your_api_key>"
491 ```
493 Example Response:
494 ```json
495 {
496 "search_tool_id": "123e4567-e89b-12d3-a456-426614174000",
497 "search_tool_name": "litellm-search",
498 "litellm_params": {
499 "search_provider": "perplexity",
500 "api_key": "sk-***"
501 },
502 "search_tool_info": {
503 "description": "Perplexity search tool"
504 },
505 "created_at": "2023-11-09T12:34:56.789Z",
506 "updated_at": "2023-11-09T12:34:56.789Z"
507 }
508 ```
509 """
510 from litellm.litellm_core_utils.litellm_logging import _get_masked_values
511 from litellm.proxy.proxy_server import prisma_client
513 if prisma_client is None: 513 ↛ 514line 513 didn't jump to line 514 because the condition on line 513 was never true
514 raise HTTPException(status_code=500, detail="Prisma client not initialized")
516 try:
517 result: Final = await SEARCH_TOOL_REGISTRY.get_search_tool_by_id_from_db(
518 search_tool_id=search_tool_id, prisma_client=prisma_client
519 )
521 if result is None:
522 raise HTTPException(
523 status_code=404,
524 detail=f"Search tool with ID {search_tool_id} not found",
525 )
527 # Mask sensitive data
528 litellm_params_dict: Final = dict(result.get("litellm_params", {}))
529 masked_litellm_params_dict: Final = _get_masked_values(
530 litellm_params_dict,
531 unmasked_length=4,
532 number_of_asterisks=4,
533 )
535 return SearchToolInfoResponse(
536 search_tool_id=result.get("search_tool_id"),
537 search_tool_name=result.get("search_tool_name", ""),
538 litellm_params=masked_litellm_params_dict,
539 search_tool_info=result.get("search_tool_info"),
540 created_at=_convert_datetime_to_str(result.get("created_at")),
541 updated_at=_convert_datetime_to_str(result.get("updated_at")),
542 is_from_config=False, # This endpoint only returns DB tools
543 )
544 except HTTPException as e:
545 raise e
546 except Exception as e:
547 verbose_proxy_logger.exception("Error getting search tool info: %s", e)
548 raise HTTPException(status_code=500, detail=str(e))
551class TestSearchToolConnectionRequest(BaseModel):
552 litellm_params: dict[str, Any]
555@router.post(
556 "/search_tools/test_connection",
557 tags=["Search Tools"],
558 dependencies=[Depends(user_api_key_auth)],
559)
560async def test_search_tool_connection(request: TestSearchToolConnectionRequest):
561 """
562 Test connection to a search provider with the given configuration.
564 Makes a simple test search query to verify the API key and configuration are valid.
566 Example Request:
567 ```bash
568 curl -X POST "http://localhost:4000/search_tools/test_connection" \\
569 -H "Authorization: Bearer <your_api_key>" \\
570 -H "Content-Type: application/json" \\
571 -d '{
572 "litellm_params": {
573 "search_provider": "perplexity",
574 "api_key": "sk-..."
575 }
576 }'
577 ```
579 Example Response (Success):
580 ```json
581 {
582 "status": "success",
583 "message": "Successfully connected to perplexity search provider",
584 "test_query": "test",
585 "results_count": 5
586 }
587 ```
589 Example Response (Failure):
590 ```json
591 {
592 "status": "error",
593 "message": "Authentication failed: Invalid API key",
594 "error_type": "AuthenticationError"
595 }
596 ```
597 """
598 try:
599 from litellm.search import asearch
601 # Extract params from request
602 litellm_params: Final = request.litellm_params
603 search_provider: Final = litellm_params.get("search_provider")
604 api_key: Final = litellm_params.get("api_key")
605 api_base: Final = litellm_params.get("api_base")
607 if not search_provider: 607 ↛ 610line 607 didn't jump to line 610 because the condition on line 607 was always true
608 raise HTTPException(status_code=400, detail="search_provider is required in litellm_params")
610 verbose_proxy_logger.debug("Testing connection to search provider: %s", search_provider)
612 # Make a simple test search query with max_results=1 to minimize cost
613 test_query: Final = "test"
614 response: Final = await asearch(
615 query=test_query,
616 search_provider=search_provider,
617 api_key=api_key,
618 api_base=api_base,
619 max_results=1, # Minimize results to reduce cost
620 timeout=10.0, # 10 second timeout for test
621 )
623 verbose_proxy_logger.debug("Successfully tested connection to %s search provider", search_provider)
625 return {
626 "status": "success",
627 "message": f"Successfully connected to {search_provider} search provider",
628 "test_query": test_query,
629 "results_count": (len(response.results) if response and response.results else 0),
630 }
632 except Exception as e:
633 error_message: Final = str(e)
634 error_type: Final = type(e).__name__
636 verbose_proxy_logger.exception("Failed to connect to search provider: %s", error_message)
638 # Return error details in a structured format
639 return {
640 "status": "error",
641 "message": error_message,
642 "error_type": error_type,
643 }
646@router.get(
647 "/search_tools/ui/available_providers",
648 tags=["Search Tools"],
649 dependencies=[Depends(user_api_key_auth)],
650)
651async def get_available_search_providers():
652 """
653 Get the list of available search providers with their configuration fields.
655 Auto-discovers search providers and their UI-friendly names from transformation configs.
657 Example Request:
658 ```bash
659 curl -X GET "http://localhost:4000/search_tools/ui/available_providers" \\
660 -H "Authorization: Bearer <your_api_key>"
661 ```
663 Example Response:
664 ```json
665 {
666 "providers": [
667 {
668 "provider_name": "perplexity",
669 "ui_friendly_name": "Perplexity"
670 },
671 {
672 "provider_name": "tavily",
673 "ui_friendly_name": "Tavily"
674 }
675 ]
676 }
677 ```
678 """
679 try:
680 from litellm.utils import ProviderConfigManager
682 available_providers: Final = []
684 # Auto-discover providers from SearchProviders enum
685 for provider in SearchProviders:
686 try:
687 # Get the config class for this provider
688 config = ProviderConfigManager.get_provider_search_config(provider=provider)
690 if config is not None: 690 ↛ 685line 690 didn't jump to line 685 because the condition on line 690 was always true
691 # Get the UI-friendly name from the config class
692 ui_name = config.ui_friendly_name()
694 available_providers.append(
695 {
696 "provider_name": provider.value,
697 "ui_friendly_name": ui_name,
698 }
699 )
700 except Exception as e:
701 verbose_proxy_logger.debug("Could not get config for search provider %s: %s", provider.value, e)
702 continue
704 return {"providers": available_providers}
705 except Exception as e:
706 verbose_proxy_logger.exception("Error getting available search providers: %s", e)
707 raise HTTPException(status_code=500, detail=str(e))