diff --git a/docs/my-website/docs/proxy/search.md b/docs/my-website/docs/proxy/search.md new file mode 100644 index 00000000000..65c431eb53e --- /dev/null +++ b/docs/my-website/docs/proxy/search.md @@ -0,0 +1,92 @@ +# Search API + +LiteLLM supports team-aware search provider credentials for providers like Tavily, Perplexity, Brave, Exa, and Serper. + +## Per-team search provider configuration + +Set per-team credentials in team metadata: + +```json +{ + "search_provider_config": { + "tavily": { + "api_key": "tvly-team-a-key", + "api_base": "https://api.tavily.com" + }, + "perplexity": { + "api_key": "pplx-team-a-key" + } + } +} +``` + +Update via API: + +```bash +curl -X POST "http://localhost:4000/team/search_provider_config/update" \ + -H "Authorization: Bearer sk-admin-key" \ + -H "Content-Type: application/json" \ + -d '{ + "team_id": "team-a", + "provider": "tavily", + "api_key": "tvly-team-a-key", + "api_base": "https://api.tavily.com" + }' +``` + +## Request flow and precedence + +Search credentials resolve in this order: + +1. Request metadata: `metadata.search_provider_config.` +2. Team DB metadata: `user_api_key_team_metadata.search_provider_config.` +3. YAML team settings: `default_team_settings[].search_provider_config.` +4. Search tool config: `search_tools[].litellm_params` +5. Provider env fallback (`TAVILY_API_KEY`, etc.) + +## Calling search as an end-user + +The caller only uses their team-bound virtual key. + +```bash +curl -X POST "http://localhost:4000/v1/search" \ + -H "Authorization: Bearer sk-team-a-user-key" \ + -H "Content-Type: application/json" \ + -d '{ + "search_tool_name": "company-search", + "query": "latest AI news", + "max_results": 5 + }' +``` + +or with URL tool name: + +```bash +curl -X POST "http://localhost:4000/v1/search/company-search" \ + -H "Authorization: Bearer sk-team-a-user-key" \ + -H "Content-Type: application/json" \ + -d '{ + "query": "latest AI news", + "max_results": 5 + }' +``` + +## YAML examples + +```yaml +search_tools: + - search_tool_name: company-search + litellm_params: + search_provider: tavily + api_key: os.environ/TAVILY_DEFAULT_API_KEY + +default_team_settings: + - team_id: team-a + search_provider_config: + tavily: + api_key: os.environ/TAVILY_TEAM_A_API_KEY + - team_id: team-b + search_provider_config: + tavily: + api_key: os.environ/TAVILY_TEAM_B_API_KEY +``` diff --git a/docs/my-website/docs/proxy/team_budgets.md b/docs/my-website/docs/proxy/team_budgets.md new file mode 100644 index 00000000000..47f3a832b07 --- /dev/null +++ b/docs/my-website/docs/proxy/team_budgets.md @@ -0,0 +1,58 @@ +# Team Budgets and Search Cost Attribution + +When search requests are made through LiteLLM with a team-bound key, spend is attributed to that team. + +## Cost attribution for search + +Search calls (`search` / `asearch`) are logged with: + +- `metadata.user_api_key_team_id` +- spend rows in `LiteLLM_SpendLogs.team_id` + +This means each team's search usage can be queried independently even when using the same model/provider family. + +## Why per-team search keys matter + +Using one shared Tavily key makes upstream provider billing opaque by team. +With team-specific provider keys: + +- provider-side billing is isolated per team +- LiteLLM spend logs still aggregate by team id +- finance can reconcile provider invoices + LiteLLM spend logs + +## Recommended setup + +1. Issue per-team virtual keys in LiteLLM. +2. Configure `metadata.search_provider_config` per team. +3. Keep a fallback tool-level key only for teams without explicit config. + +## Example team update + +```bash +curl -X POST "http://localhost:4000/team/update" \ + -H "Authorization: Bearer sk-admin-key" \ + -H "Content-Type: application/json" \ + -d '{ + "team_id": "team-research", + "metadata": { + "search_provider_config": { + "tavily": { + "api_key": "tvly-research-key" + }, + "perplexity": { + "api_key": "pplx-research-key" + } + } + } + }' +``` + +## Example spend query + +```sql +SELECT team_id, call_type, SUM(spend) AS total_spend, COUNT(*) AS requests +FROM "LiteLLM_SpendLogs" +WHERE call_type IN ('search', 'asearch') +GROUP BY team_id, call_type +ORDER BY total_spend DESC; +``` diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 92c920ca594..7b92af92f22 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1695,6 +1695,28 @@ class OrgMember(MemberBase): ] +class SearchProviderCredentials(LiteLLMPydanticObjectBase): + """ + Per-team credentials for a search provider. + """ + + api_key: Optional[str] = None + api_base: Optional[str] = None + + +class TeamSearchProviderConfig(LiteLLMPydanticObjectBase): + """ + Structured team-level search provider credentials. + Stored in team metadata under `search_provider_config`. + """ + + tavily: Optional[SearchProviderCredentials] = None + perplexity: Optional[SearchProviderCredentials] = None + brave: Optional[SearchProviderCredentials] = None + exa: Optional[SearchProviderCredentials] = None + serper: Optional[SearchProviderCredentials] = None + + class TeamBase(LiteLLMPydanticObjectBase): team_alias: Optional[str] = None team_id: Optional[str] = None @@ -1703,7 +1725,7 @@ class TeamBase(LiteLLMPydanticObjectBase): members: list = [] members_with_roles: List[Member] = [] team_member_permissions: Optional[List[str]] = None - metadata: Optional[dict] = None + metadata: Optional[dict] = None # may include search_provider_config tpm_limit: Optional[int] = None rpm_limit: Optional[int] = None @@ -1823,6 +1845,13 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): ) +class TeamSearchProviderConfigUpdateRequest(LiteLLMPydanticObjectBase): + team_id: str + provider: str + api_key: Optional[str] = None + api_base: Optional[str] = None + + class ResetTeamBudgetRequest(LiteLLMPydanticObjectBase): """ internal type used to reset the budget on a team diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index f254fea3e7f..abd5239c9a0 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -59,6 +59,7 @@ from litellm.proxy._types import ( TeamMemberUpdateResponse, TeamModelAddRequest, TeamModelDeleteRequest, + TeamSearchProviderConfigUpdateRequest, UpdateTeamRequest, UserAPIKeyAuth, ) @@ -1856,6 +1857,108 @@ async def update_team( # noqa: PLR0915 raise handle_exception_on_proxy(e) +@router.post( + "/team/search_provider_config/update", + tags=["team management"], + dependencies=[Depends(user_api_key_auth)], +) +@management_endpoint_wrapper +async def update_team_search_provider_config( + data: TeamSearchProviderConfigUpdateRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Update per-team search provider credentials in team metadata. + + Stored under: + metadata.search_provider_config..{api_key, api_base} + """ + from litellm.proxy.auth.auth_checks import _cache_team_object + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + provider = data.provider.strip().lower() + if provider == "": + raise HTTPException( + status_code=400, detail={"error": "provider cannot be empty"} + ) + + existing_team_row = await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": data.team_id} + ) + if existing_team_row is None: + raise HTTPException( + status_code=404, + detail={"error": f"Team not found, passed team_id={data.team_id}"}, + ) + + await _verify_team_access( + team_obj=LiteLLM_TeamTable(**existing_team_row.model_dump()), + user_api_key_dict=user_api_key_dict, + ) + + metadata: Dict[str, Any] = {} + if isinstance(existing_team_row.metadata, dict): + metadata = dict(existing_team_row.metadata) + + search_provider_config = metadata.get("search_provider_config") + if not isinstance(search_provider_config, dict): + search_provider_config = {} + + provider_config = search_provider_config.get(provider) + if not isinstance(provider_config, dict): + provider_config = {} + + if data.api_key is not None: + provider_config["api_key"] = data.api_key + if data.api_base is not None: + provider_config["api_base"] = data.api_base + + if provider_config.get("api_key") in (None, "") and provider_config.get( + "api_base" + ) in ( + None, + "", + ): + search_provider_config.pop(provider, None) + else: + search_provider_config[provider] = provider_config + + metadata["search_provider_config"] = search_provider_config + + team_row: Optional[LiteLLM_TeamTable] = ( + await prisma_client.db.litellm_teamtable.update( + where={"team_id": data.team_id}, + data={"metadata": metadata}, + include={"litellm_model_table": True}, # type: ignore + ) + ) + + if team_row is not None and team_row.team_id is not None: + await _cache_team_object( + team_id=team_row.team_id, + team_table=LiteLLM_TeamTableCachedObj(**team_row.model_dump()), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + return { + "message": "Team search provider configuration updated", + "team_id": data.team_id, + "provider": provider, + "search_provider_config": search_provider_config, + } + + def _set_budget_reset_at(data: UpdateTeamRequest, updated_kv: dict) -> None: """Set budget_reset_at in updated_kv if budget_duration is provided.""" if data.budget_duration is not None: @@ -2049,8 +2152,11 @@ async def _process_team_members( # Resolve allowed_models: explicit request value, or fall back to team's default_team_member_models member_allowed_models = data.allowed_models - if member_allowed_models is None and complete_team_data.default_team_member_models: - member_allowed_models = complete_team_data.default_team_member_models + team_default_member_models = getattr( + complete_team_data, "default_team_member_models", None + ) + if member_allowed_models is None and team_default_member_models: + member_allowed_models = team_default_member_models if isinstance(data.member, Member): try: diff --git a/litellm/proxy/search_endpoints/endpoints.py b/litellm/proxy/search_endpoints/endpoints.py index 8bed5b54075..ce2949f707a 100644 --- a/litellm/proxy/search_endpoints/endpoints.py +++ b/litellm/proxy/search_endpoints/endpoints.py @@ -163,6 +163,16 @@ async def search( data["metadata"] = {} data["metadata"]["model_group"] = search_tool_name_value + # Ensure team context is available to search router credential resolution. + # add_litellm_data_to_request() also injects these values, but this keeps + # search endpoint behavior explicit and resilient for direct router paths. + if "metadata" not in data or not isinstance(data.get("metadata"), dict): + data["metadata"] = {} + if getattr(user_api_key_dict, "team_metadata", None) is not None: + data["metadata"]["user_api_key_team_metadata"] = user_api_key_dict.team_metadata + if getattr(user_api_key_dict, "team_id", None) is not None: + data["metadata"]["user_api_key_team_id"] = user_api_key_dict.team_id + # Process request using ProxyBaseLLMRequestProcessing processor = ProxyBaseLLMRequestProcessing(data=data) try: diff --git a/litellm/router_utils/search_api_router.py b/litellm/router_utils/search_api_router.py index a26aa7e71ee..e2e98a65734 100644 --- a/litellm/router_utils/search_api_router.py +++ b/litellm/router_utils/search_api_router.py @@ -8,8 +8,9 @@ import asyncio import random import traceback from functools import partial -from typing import Any, Callable +from typing import Any, Callable, Dict, Optional, Tuple +import litellm from litellm._logging import verbose_router_logger @@ -20,6 +21,96 @@ class SearchAPIRouter: Provides methods for search tool selection, load balancing, and fallback handling. """ + @staticmethod + def _get_team_config_from_default_settings( + team_id: Optional[str], + ) -> Optional[Dict[str, Any]]: + """ + Resolve team config from litellm.default_team_settings. + + This allows search requests to read per-team settings from proxy config + (YAML) similar to completion paths that use ProxyConfig.load_team_config(). + """ + if not team_id: + return None + + default_team_settings = getattr(litellm, "default_team_settings", None) + if not isinstance(default_team_settings, list): + return None + + for team_setting in default_team_settings: + if ( + isinstance(team_setting, dict) + and team_setting.get("team_id") == team_id + ): + return team_setting + return None + + @staticmethod + def _resolve_search_provider_credentials( + *, + search_provider: str, + tool_litellm_params: Dict[str, Any], + request_metadata: Optional[Dict[str, Any]] = None, + team_metadata: Optional[Dict[str, Any]] = None, + team_config: Optional[Dict[str, Any]] = None, + ) -> Tuple[Optional[str], Optional[str]]: + """ + Resolve search provider credentials with precedence: + 1. request metadata.search_provider_config.{provider} + 2. team metadata.search_provider_config.{provider} + 3. default_team_settings.search_provider_config.{provider} + 4. search_tool.litellm_params + 5. env fallback in provider validate_environment() + """ + resolved_api_key: Optional[str] = None + resolved_api_base: Optional[str] = None + + request_provider_config = {} + if isinstance(request_metadata, dict): + search_provider_config = request_metadata.get("search_provider_config") + if isinstance(search_provider_config, dict): + request_provider_config = search_provider_config.get( + search_provider, {} + ) + + team_provider_config = {} + if isinstance(team_metadata, dict): + search_provider_config = team_metadata.get("search_provider_config") + if isinstance(search_provider_config, dict): + team_provider_config = search_provider_config.get(search_provider, {}) + + team_settings_provider_config = {} + if isinstance(team_config, dict): + search_provider_config = team_config.get("search_provider_config") + if isinstance(search_provider_config, dict): + team_settings_provider_config = search_provider_config.get( + search_provider, {} + ) + + if isinstance(request_provider_config, dict): + resolved_api_key = request_provider_config.get("api_key") + resolved_api_base = request_provider_config.get("api_base") + + if resolved_api_key is None and isinstance(team_provider_config, dict): + resolved_api_key = team_provider_config.get("api_key") + if resolved_api_base is None and isinstance(team_provider_config, dict): + resolved_api_base = team_provider_config.get("api_base") + + if resolved_api_key is None and isinstance(team_settings_provider_config, dict): + resolved_api_key = team_settings_provider_config.get("api_key") + if resolved_api_base is None and isinstance( + team_settings_provider_config, dict + ): + resolved_api_base = team_settings_provider_config.get("api_base") + + if resolved_api_key is None: + resolved_api_key = tool_litellm_params.get("api_key") + if resolved_api_base is None: + resolved_api_base = tool_litellm_params.get("api_base") + + return resolved_api_key, resolved_api_base + @staticmethod async def update_router_search_tools(router_instance: Any, search_tools: list): """ @@ -198,16 +289,42 @@ class SearchAPIRouter: # 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}'" ) + request_metadata = kwargs.get("metadata") + litellm_metadata = kwargs.get("litellm_metadata") + if not isinstance(request_metadata, dict) and isinstance( + litellm_metadata, dict + ): + request_metadata = litellm_metadata + + team_metadata = {} + team_id: Optional[str] = None + if isinstance(request_metadata, dict): + _team_metadata = request_metadata.get("user_api_key_team_metadata") + if isinstance(_team_metadata, dict): + team_metadata = _team_metadata + _team_id = request_metadata.get("user_api_key_team_id") + if isinstance(_team_id, str): + team_id = _team_id + + team_config = SearchAPIRouter._get_team_config_from_default_settings( + team_id=team_id + ) + + api_key, api_base = SearchAPIRouter._resolve_search_provider_credentials( + search_provider=search_provider, + tool_litellm_params=litellm_params, + request_metadata=request_metadata, + team_metadata=team_metadata, + team_config=team_config, + ) + verbose_router_logger.debug( - f"Selected search tool with provider: {search_provider}" + f"Selected search tool with provider: {search_provider}, team_id={team_id}" ) # Call the original search function with the provider config diff --git a/tests/test_litellm/proxy/search_endpoints/test_team_search_credentials.py b/tests/test_litellm/proxy/search_endpoints/test_team_search_credentials.py new file mode 100644 index 00000000000..eab167fb75a --- /dev/null +++ b/tests/test_litellm/proxy/search_endpoints/test_team_search_credentials.py @@ -0,0 +1,133 @@ +import os +import sys +from unittest.mock import patch + +import pytest +from fastapi.testclient import TestClient + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.proxy_server import app +from litellm.router_utils.search_api_router import SearchAPIRouter + + +def test_resolve_credentials_team_metadata_overrides_tool_params(): + api_key, api_base = SearchAPIRouter._resolve_search_provider_credentials( + search_provider="tavily", + tool_litellm_params={ + "api_key": "tool-key", + "api_base": "https://tool.example.com", + }, + team_metadata={ + "search_provider_config": { + "tavily": { + "api_key": "team-key", + "api_base": "https://team.example.com", + } + } + }, + ) + assert api_key == "team-key" + assert api_base == "https://team.example.com" + + +def test_resolve_credentials_request_metadata_has_highest_precedence(): + api_key, api_base = SearchAPIRouter._resolve_search_provider_credentials( + search_provider="tavily", + tool_litellm_params={ + "api_key": "tool-key", + "api_base": "https://tool.example.com", + }, + request_metadata={ + "search_provider_config": { + "tavily": { + "api_key": "request-key", + "api_base": "https://request.example.com", + } + } + }, + team_metadata={ + "search_provider_config": { + "tavily": { + "api_key": "team-key", + "api_base": "https://team.example.com", + } + } + }, + ) + assert api_key == "request-key" + assert api_base == "https://request.example.com" + + +def test_resolve_credentials_from_default_team_settings(): + with patch( + "litellm.default_team_settings", + [ + { + "team_id": "team-a", + "search_provider_config": { + "tavily": { + "api_key": "team-settings-key", + "api_base": "https://team-settings.example.com", + } + }, + } + ], + ): + team_config = SearchAPIRouter._get_team_config_from_default_settings("team-a") + api_key, api_base = SearchAPIRouter._resolve_search_provider_credentials( + search_provider="tavily", + tool_litellm_params={}, + team_config=team_config, + ) + assert api_key == "team-settings-key" + assert api_base == "https://team-settings.example.com" + + +@pytest.mark.asyncio +async def test_search_endpoint_injects_team_metadata(): + captured_metadata = {} + + async def _mock_process(self, **kwargs): + nonlocal captured_metadata + captured_metadata = self.data.get("metadata", {}) + return {"object": "search", "results": []} + + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="admin-user", + team_id="team-test", + team_metadata={ + "search_provider_config": { + "tavily": {"api_key": "team-test-key"}, + } + }, + ) + + try: + with patch( + "litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing.base_process_llm_request", + new=_mock_process, + ): + client = TestClient(app) + response = client.post( + "/v1/search", + json={ + "search_tool_name": "tool-a", + "search_provider": "tavily", + "query": "latest ai news", + }, + ) + assert response.status_code == 200 + assert captured_metadata.get("user_api_key_team_id") == "team-test" + assert ( + captured_metadata.get("user_api_key_team_metadata", {}) + .get("search_provider_config", {}) + .get("tavily", {}) + .get("api_key") + == "team-test-key" + ) + finally: + app.dependency_overrides.pop(user_api_key_auth, None) diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 44208904a70..6cd9423c37f 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -3725,6 +3725,40 @@ export const teamUpdateCall = async ( } }; +export const updateTeamSearchProviderConfigCall = async ( + accessToken: string, + formValues: { + team_id: string; + provider: string; + api_key?: string | null; + api_base?: string | null; + }, +) => { + try { + const url = proxyBaseUrl + ? `${proxyBaseUrl}/team/search_provider_config/update` + : `/team/search_provider_config/update`; + const response = await fetch(url, { + method: "POST", + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify(formValues), + }); + + if (!response.ok) { + const errorData = await response.text(); + handleError(errorData); + throw new Error(errorData); + } + return await response.json(); + } catch (error) { + console.error("Failed to update team search provider config:", error); + throw error; + } +}; + /** * Patch update a model * diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index ef98054c080..1a2e97faaaa 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -457,13 +457,26 @@ const TeamInfoView: React.FC = ({ try { const rawMetadata = values.metadata ? JSON.parse(values.metadata) : {}; // Exclude soft_budget_alerting_emails from parsed metadata since it's handled separately - const { soft_budget_alerting_emails, ...rest } = rawMetadata; + const { soft_budget_alerting_emails, search_provider_config, ...rest } = rawMetadata; parsedMetadata = rest; } catch (e) { NotificationsManager.fromBackend("Invalid JSON in metadata field"); return; } + let searchProviderConfig: Record | undefined; + if (typeof values.search_provider_config === "string") { + const trimmedSearchProviderConfig = values.search_provider_config.trim(); + if (trimmedSearchProviderConfig.length > 0) { + try { + searchProviderConfig = JSON.parse(trimmedSearchProviderConfig); + } catch (e) { + NotificationsManager.fromBackend("Invalid JSON in search provider configuration"); + return; + } + } + } + let secretManagerSettings: Record | undefined; if (typeof values.secret_manager_settings === "string") { const trimmedSecretConfig = values.secret_manager_settings.trim(); @@ -513,6 +526,7 @@ const TeamInfoView: React.FC = ({ budget_duration: values.budget_duration, metadata: { ...parsedMetadata, + ...(searchProviderConfig !== undefined ? { search_provider_config: searchProviderConfig } : {}), guardrails: (values.guardrails || []).filter((n: string) => !globalGuardrailNames.has(n)), opted_out_global_guardrails: optedOutGlobalGuardrails, ...(values.logging_settings?.length > 0 ? { logging: values.logging_settings } : {}), @@ -952,11 +966,14 @@ const TeamInfoView: React.FC = ({ : "", metadata: info.metadata ? JSON.stringify( - (({ logging, secret_manager_settings, soft_budget_alerting_emails, model_tpm_limit, model_rpm_limit, ...rest }) => rest)(info.metadata), + (({ logging, secret_manager_settings, soft_budget_alerting_emails, search_provider_config, model_tpm_limit, model_rpm_limit, ...rest }) => rest)(info.metadata), null, 2, ) : "", + search_provider_config: info.metadata?.search_provider_config + ? JSON.stringify(info.metadata.search_provider_config, null, 2) + : "", logging_settings: info.metadata?.logging || [], secret_manager_settings: info.metadata?.secret_manager_settings ? JSON.stringify(info.metadata.secret_manager_settings, null, 2) @@ -1399,6 +1416,29 @@ const TeamInfoView: React.FC = ({ /> + { + if (!value || (typeof value === "string" && value.trim() === "")) { + return Promise.resolve(); + } + try { + JSON.parse(value); + return Promise.resolve(); + } catch (error) { + return Promise.reject(new Error("Please enter valid JSON")); + } + }, + }, + ]} + > + + + = ({ )} + + {info.metadata?.search_provider_config && ( +
+ Search Provider Configuration + 
+                          {JSON.stringify(info.metadata.search_provider_config, null, 2)}
+
+
+ )} )}