scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 8

feat(proxy): add team-level search provider credential resolution

Browse source 
Allow search requests to resolve provider credentials from request metadata, team metadata, and default team settings with clear precedence, and expose this flow in proxy docs/UI with regression tests.

Made-with: Cursor
This commit is contained in:
Sameer Kankute 2026-04-28 16:58:03 +05:30
parent 62920a0cb2
commit 2d2f540480
No known key found for this signature in database
9 changed files with 638 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

92
docs/my-website/docs/proxy/search.md Normal file
View file

@ -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.<provider>`
2. Team DB metadata: `user_api_key_team_metadata.search_provider_config.<provider>`
3. YAML team settings: `default_team_settings[].search_provider_config.<provider>`
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
```

58
docs/my-website/docs/proxy/team_budgets.md Normal file
View file

@ -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;
```

31
litellm/proxy/_types.py
View file

@ -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

110
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -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.<provider>.{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:

10
litellm/proxy/search_endpoints/endpoints.py
View file

@ -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:

127
litellm/router_utils/search_api_router.py
View file

@ -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

133
tests/test_litellm/proxy/search_endpoints/test_team_search_credentials.py Normal file
View file

@ -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)

34
ui/litellm-dashboard/src/components/networking.tsx
View file

@ -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
*

53
ui/litellm-dashboard/src/components/team/TeamInfo.tsx
View file

@ -457,13 +457,26 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
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<string, any> | 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<string, any> | undefined;
if (typeof values.secret_manager_settings === "string") {
const trimmedSecretConfig = values.secret_manager_settings.trim();
@ -513,6 +526,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
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<TeamInfoProps> = ({
: "",
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<TeamInfoProps> = ({
/>
</Form.Item>
<Form.Item
label="Search Provider Configuration"
name="search_provider_config"
tooltip='Team-level provider credentials. Example: {"tavily": {"api_key": "tvly-...", "api_base": "https://api.tavily.com"}}'
rules={[
{
validator: async (_, value) => {
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"));
}
},
},
]}
>
<Input.TextArea rows={8} placeholder='{"tavily":{"api_key":"tvly-...","api_base":"https://api.tavily.com"}}' />
</Form.Item>
<Form.Item
label="Secret Manager Settings"
name="secret_manager_settings"
@ -1614,6 +1654,15 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
</pre>
</div>
)}
{info.metadata?.search_provider_config && (
<div className="pt-4 border-t border-gray-200">
<Text className="font-medium">Search Provider Configuration</Text>
<pre className="mt-2 bg-gray-50 p-3 rounded text-xs overflow-x-auto">
{JSON.stringify(info.metadata.search_provider_config, null, 2)}
</pre>
</div>
)}
</div>
)}
</Card>