mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-12 15:05:01 +00:00
Merge pull request #20371 from BerriAI/litellm_config_search_tools
[Feature] Include Config Defined Search Tools in /search_tools/list
This commit is contained in:
yuneng-jiang
@@ -48,7 +48,7 @@ def _convert_datetime_to_str(value: Union[datetime, str, None]) -> Union[str, No
|
||||
)
|
||||
async def list_search_tools():
|
||||
"""
|
||||
List all search tools that are available in the database.
|
||||
List all search tools that are available in the database and config file.
|
||||
|
||||
Example Request:
|
||||
```bash
|
||||
@@ -71,38 +71,100 @@ async def list_search_tools():
|
||||
"description": "Perplexity search tool"
|
||||
},
|
||||
"created_at": "2023-11-09T12:34:56.789Z",
|
||||
"updated_at": "2023-11-09T12:34:56.789Z"
|
||||
"updated_at": "2023-11-09T12:34:56.789Z",
|
||||
"is_from_config": false
|
||||
},
|
||||
{
|
||||
"search_tool_name": "config-search-tool",
|
||||
"litellm_params": {
|
||||
"search_provider": "tavily",
|
||||
"api_key": "tvly-***"
|
||||
},
|
||||
"is_from_config": true
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
"""
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
from litellm.litellm_core_utils.litellm_logging import _get_masked_values
|
||||
from litellm.proxy.proxy_server import prisma_client, proxy_config
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail="Prisma client not initialized")
|
||||
|
||||
try:
|
||||
search_tools = await SEARCH_TOOL_REGISTRY.get_all_search_tools_from_db(
|
||||
search_tools_from_db = await SEARCH_TOOL_REGISTRY.get_all_search_tools_from_db(
|
||||
prisma_client=prisma_client
|
||||
)
|
||||
|
||||
db_tool_names = {
|
||||
tool.get("search_tool_name") for tool in search_tools_from_db
|
||||
}
|
||||
|
||||
search_tool_configs: List[SearchToolInfoResponse] = []
|
||||
for search_tool in search_tools:
|
||||
|
||||
config_search_tools = []
|
||||
|
||||
try:
|
||||
config = await proxy_config.get_config()
|
||||
parsed_tools = proxy_config.parse_search_tools(config)
|
||||
if parsed_tools:
|
||||
config_search_tools = parsed_tools
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug(
|
||||
f"Could not get config-defined search tools: {e}"
|
||||
)
|
||||
|
||||
for search_tool in config_search_tools:
|
||||
tool_name = search_tool.get("search_tool_name")
|
||||
if tool_name:
|
||||
litellm_params_dict = dict(search_tool.get("litellm_params", {}))
|
||||
masked_litellm_params_dict = _get_masked_values(
|
||||
litellm_params_dict,
|
||||
unmasked_length=4,
|
||||
number_of_asterisks=4,
|
||||
)
|
||||
|
||||
search_tool_configs.append(
|
||||
SearchToolInfoResponse(
|
||||
search_tool_id=None,
|
||||
search_tool_name=tool_name,
|
||||
litellm_params=masked_litellm_params_dict,
|
||||
search_tool_info=search_tool.get("search_tool_info"),
|
||||
created_at=None,
|
||||
updated_at=None,
|
||||
is_from_config=True,
|
||||
)
|
||||
)
|
||||
|
||||
search_tool_configs = [
|
||||
tool for tool in search_tool_configs
|
||||
if tool.get("search_tool_name") not in db_tool_names
|
||||
]
|
||||
|
||||
for search_tool in search_tools_from_db:
|
||||
litellm_params_dict = dict(search_tool.get("litellm_params", {}))
|
||||
masked_litellm_params_dict = _get_masked_values(
|
||||
litellm_params_dict,
|
||||
unmasked_length=4,
|
||||
number_of_asterisks=4,
|
||||
)
|
||||
|
||||
search_tool_configs.append(
|
||||
SearchToolInfoResponse(
|
||||
search_tool_id=search_tool.get("search_tool_id"),
|
||||
search_tool_name=search_tool.get("search_tool_name", ""),
|
||||
litellm_params=dict(search_tool.get("litellm_params", {})),
|
||||
litellm_params=masked_litellm_params_dict,
|
||||
search_tool_info=search_tool.get("search_tool_info"),
|
||||
created_at=_convert_datetime_to_str(search_tool.get("created_at")),
|
||||
updated_at=_convert_datetime_to_str(search_tool.get("updated_at")),
|
||||
is_from_config=False,
|
||||
)
|
||||
)
|
||||
|
||||
return ListSearchToolsResponse(search_tools=search_tool_configs)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception(f"Error getting search tools from db: {e}")
|
||||
verbose_proxy_logger.exception(f"Error getting search tools: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@@ -382,6 +444,7 @@ async def get_search_tool_info(search_tool_id: str):
|
||||
search_tool_info=result.get("search_tool_info"),
|
||||
created_at=_convert_datetime_to_str(result.get("created_at")),
|
||||
updated_at=_convert_datetime_to_str(result.get("updated_at")),
|
||||
is_from_config=False, # This endpoint only returns DB tools
|
||||
)
|
||||
except HTTPException as e:
|
||||
raise e
|
||||
|
||||
@@ -60,6 +60,7 @@ class SearchToolInfoResponse(TypedDict, total=False):
|
||||
search_tool_info: Optional[dict]
|
||||
created_at: Optional[str]
|
||||
updated_at: Optional[str]
|
||||
is_from_config: Optional[bool] # True if this tool is defined in config file, False if from DB
|
||||
|
||||
|
||||
class ListSearchToolsResponse(TypedDict):
|
||||
|
||||
+546
@@ -0,0 +1,546 @@
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../../../../..")
|
||||
) # Adds the parent directory to the system path
|
||||
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
|
||||
# Import proxy_server module first to ensure it's initialized
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
||||
# Now we can safely import app
|
||||
from litellm.proxy.proxy_server import app
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_search_tools_db_only(monkeypatch):
|
||||
"""Test listing search tools when only DB tools exist"""
|
||||
# Mock DB tools
|
||||
db_tools = [
|
||||
{
|
||||
"search_tool_id": "test-id-1",
|
||||
"search_tool_name": "db-tool-1",
|
||||
"litellm_params": {"search_provider": "perplexity", "api_key": "sk-test"},
|
||||
"search_tool_info": {"description": "DB tool 1"},
|
||||
"created_at": datetime(2023, 11, 9, 12, 34, 56),
|
||||
"updated_at": datetime(2023, 11, 9, 13, 45, 12),
|
||||
}
|
||||
]
|
||||
|
||||
# Mock SearchToolRegistry
|
||||
mock_registry = MagicMock()
|
||||
mock_registry.get_all_search_tools_from_db = AsyncMock(return_value=db_tools)
|
||||
with patch(
|
||||
"litellm.proxy.search_endpoints.search_tool_management.SEARCH_TOOL_REGISTRY",
|
||||
mock_registry,
|
||||
):
|
||||
# Mock prisma_client
|
||||
mock_prisma = MagicMock()
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
|
||||
# Mock proxy_config
|
||||
mock_proxy_config = MagicMock()
|
||||
mock_proxy_config.get_config = AsyncMock(return_value={})
|
||||
mock_proxy_config.parse_search_tools = MagicMock(return_value=None)
|
||||
with patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config):
|
||||
# Mock auth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user"
|
||||
)
|
||||
|
||||
try:
|
||||
test_client = TestClient(app)
|
||||
response = test_client.get("/search_tools/list")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "search_tools" in data
|
||||
assert len(data["search_tools"]) == 1
|
||||
|
||||
tool = data["search_tools"][0]
|
||||
assert tool["search_tool_id"] == "test-id-1"
|
||||
assert tool["search_tool_name"] == "db-tool-1"
|
||||
assert tool["is_from_config"] is False
|
||||
# Verify datetime conversion to ISO string
|
||||
assert tool["created_at"] == "2023-11-09T12:34:56"
|
||||
assert tool["updated_at"] == "2023-11-09T13:45:12"
|
||||
# Verify masking of sensitive values
|
||||
assert tool["litellm_params"]["api_key"] != "sk-test"
|
||||
assert "****" in tool["litellm_params"]["api_key"]
|
||||
assert tool["litellm_params"]["search_provider"] == "perplexity"
|
||||
finally:
|
||||
app.dependency_overrides.pop(user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_search_tools_config_only(monkeypatch):
|
||||
"""Test listing search tools when only config tools exist"""
|
||||
# Mock DB tools - empty
|
||||
db_tools = []
|
||||
|
||||
# Mock config tools
|
||||
config_tools = [
|
||||
{
|
||||
"search_tool_name": "config-tool-1",
|
||||
"litellm_params": {"search_provider": "tavily", "api_key": "tvly-secret-key"},
|
||||
"search_tool_info": {"description": "Config tool 1"},
|
||||
}
|
||||
]
|
||||
|
||||
# Mock SearchToolRegistry
|
||||
mock_registry = MagicMock()
|
||||
mock_registry.get_all_search_tools_from_db = AsyncMock(return_value=db_tools)
|
||||
with patch(
|
||||
"litellm.proxy.search_endpoints.search_tool_management.SEARCH_TOOL_REGISTRY",
|
||||
mock_registry,
|
||||
):
|
||||
# Mock prisma_client
|
||||
mock_prisma = MagicMock()
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
|
||||
# Mock proxy_config
|
||||
mock_proxy_config = MagicMock()
|
||||
mock_proxy_config.get_config = AsyncMock(return_value={"search_tools": config_tools})
|
||||
mock_proxy_config.parse_search_tools = MagicMock(return_value=config_tools)
|
||||
with patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config):
|
||||
# Mock auth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user"
|
||||
)
|
||||
|
||||
try:
|
||||
test_client = TestClient(app)
|
||||
response = test_client.get("/search_tools/list")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "search_tools" in data
|
||||
assert len(data["search_tools"]) == 1
|
||||
|
||||
tool = data["search_tools"][0]
|
||||
assert tool["search_tool_name"] == "config-tool-1"
|
||||
assert tool["is_from_config"] is True
|
||||
assert tool["search_tool_id"] is None
|
||||
assert tool["created_at"] is None
|
||||
assert tool["updated_at"] is None
|
||||
# Verify masking
|
||||
assert "tv****ey" in tool["litellm_params"]["api_key"]
|
||||
finally:
|
||||
app.dependency_overrides.pop(user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_search_tools_filters_duplicate_config_tools(monkeypatch):
|
||||
"""
|
||||
Test that config tools with the same name as DB tools are filtered out.
|
||||
This tests the new filtering logic added in lines 139-142.
|
||||
"""
|
||||
# Mock DB tools
|
||||
db_tools = [
|
||||
{
|
||||
"search_tool_id": "db-id-1",
|
||||
"search_tool_name": "existing-tool",
|
||||
"litellm_params": {"search_provider": "perplexity", "api_key": "sk-db"},
|
||||
"search_tool_info": {"description": "DB tool"},
|
||||
"created_at": datetime(2023, 11, 9, 12, 34, 56),
|
||||
"updated_at": datetime(2023, 11, 9, 13, 45, 12),
|
||||
}
|
||||
]
|
||||
|
||||
# Mock config tools - one duplicate, one unique
|
||||
config_tools = [
|
||||
{
|
||||
"search_tool_name": "existing-tool", # Duplicate - should be filtered
|
||||
"litellm_params": {"search_provider": "tavily", "api_key": "tvly-config"},
|
||||
"search_tool_info": {"description": "Config tool - duplicate"},
|
||||
},
|
||||
{
|
||||
"search_tool_name": "unique-config-tool", # Unique - should be included
|
||||
"litellm_params": {"search_provider": "tavily", "api_key": "tvly-unique"},
|
||||
"search_tool_info": {"description": "Config tool - unique"},
|
||||
},
|
||||
]
|
||||
|
||||
# Mock SearchToolRegistry
|
||||
mock_registry = MagicMock()
|
||||
mock_registry.get_all_search_tools_from_db = AsyncMock(return_value=db_tools)
|
||||
with patch(
|
||||
"litellm.proxy.search_endpoints.search_tool_management.SEARCH_TOOL_REGISTRY",
|
||||
mock_registry,
|
||||
):
|
||||
# Mock prisma_client
|
||||
mock_prisma = MagicMock()
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
|
||||
# Mock proxy_config
|
||||
mock_proxy_config = MagicMock()
|
||||
mock_proxy_config.get_config = AsyncMock(return_value={"search_tools": config_tools})
|
||||
mock_proxy_config.parse_search_tools = MagicMock(return_value=config_tools)
|
||||
with patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config):
|
||||
# Mock auth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user"
|
||||
)
|
||||
|
||||
try:
|
||||
test_client = TestClient(app)
|
||||
response = test_client.get("/search_tools/list")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "search_tools" in data
|
||||
# Should have 1 DB tool + 1 unique config tool (duplicate filtered out)
|
||||
assert len(data["search_tools"]) == 2
|
||||
|
||||
# Verify DB tool is present
|
||||
db_tool = next(
|
||||
(t for t in data["search_tools"] if t["search_tool_name"] == "existing-tool"),
|
||||
None,
|
||||
)
|
||||
assert db_tool is not None
|
||||
assert db_tool["is_from_config"] is False
|
||||
assert db_tool["search_tool_id"] == "db-id-1"
|
||||
# Verify masking of sensitive values in DB tool
|
||||
assert db_tool["litellm_params"]["api_key"] != "sk-db"
|
||||
assert "****" in db_tool["litellm_params"]["api_key"]
|
||||
assert db_tool["litellm_params"]["search_provider"] == "perplexity"
|
||||
|
||||
# Verify unique config tool is present
|
||||
config_tool = next(
|
||||
(t for t in data["search_tools"] if t["search_tool_name"] == "unique-config-tool"),
|
||||
None,
|
||||
)
|
||||
assert config_tool is not None
|
||||
assert config_tool["is_from_config"] is True
|
||||
|
||||
# Verify duplicate config tool is NOT present
|
||||
duplicate_tool = next(
|
||||
(
|
||||
t
|
||||
for t in data["search_tools"]
|
||||
if t["search_tool_name"] == "existing-tool" and t["is_from_config"] is True
|
||||
),
|
||||
None,
|
||||
)
|
||||
assert duplicate_tool is None
|
||||
finally:
|
||||
app.dependency_overrides.pop(user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_search_tools_datetime_conversion(monkeypatch):
|
||||
"""
|
||||
Test that datetime objects in DB tools are properly converted to ISO format strings.
|
||||
This tests the new datetime conversion logic using _convert_datetime_to_str.
|
||||
"""
|
||||
# Mock DB tools with datetime objects
|
||||
db_tools = [
|
||||
{
|
||||
"search_tool_id": "test-id-1",
|
||||
"search_tool_name": "datetime-test-tool",
|
||||
"litellm_params": {"search_provider": "perplexity", "api_key": "sk-test"},
|
||||
"search_tool_info": {"description": "Test tool"},
|
||||
"created_at": datetime(2024, 1, 15, 10, 30, 45, 123456),
|
||||
"updated_at": datetime(2024, 1, 16, 14, 20, 30, 789012),
|
||||
},
|
||||
{
|
||||
"search_tool_id": "test-id-2",
|
||||
"search_tool_name": "null-datetime-tool",
|
||||
"litellm_params": {"search_provider": "tavily", "api_key": "tvly-test"},
|
||||
"search_tool_info": None,
|
||||
"created_at": None,
|
||||
"updated_at": None,
|
||||
},
|
||||
{
|
||||
"search_tool_id": "test-id-3",
|
||||
"search_tool_name": "string-datetime-tool",
|
||||
"litellm_params": {"search_provider": "perplexity", "api_key": "sk-test"},
|
||||
"search_tool_info": {"description": "Already string"},
|
||||
"created_at": "2024-01-17T08:15:00", # Already a string
|
||||
"updated_at": "2024-01-18T09:25:00", # Already a string
|
||||
},
|
||||
]
|
||||
|
||||
# Mock SearchToolRegistry
|
||||
mock_registry = MagicMock()
|
||||
mock_registry.get_all_search_tools_from_db = AsyncMock(return_value=db_tools)
|
||||
with patch(
|
||||
"litellm.proxy.search_endpoints.search_tool_management.SEARCH_TOOL_REGISTRY",
|
||||
mock_registry,
|
||||
):
|
||||
# Mock prisma_client
|
||||
mock_prisma = MagicMock()
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
|
||||
# Mock proxy_config
|
||||
mock_proxy_config = MagicMock()
|
||||
mock_proxy_config.get_config = AsyncMock(return_value={})
|
||||
mock_proxy_config.parse_search_tools = MagicMock(return_value=None)
|
||||
with patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config):
|
||||
# Mock auth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user"
|
||||
)
|
||||
|
||||
try:
|
||||
test_client = TestClient(app)
|
||||
response = test_client.get("/search_tools/list")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "search_tools" in data
|
||||
assert len(data["search_tools"]) == 3
|
||||
|
||||
# Test datetime conversion for tool 1
|
||||
tool1 = next(
|
||||
(t for t in data["search_tools"] if t["search_tool_name"] == "datetime-test-tool"),
|
||||
None,
|
||||
)
|
||||
assert tool1 is not None
|
||||
assert isinstance(tool1["created_at"], str)
|
||||
assert tool1["created_at"] == "2024-01-15T10:30:45.123456"
|
||||
assert isinstance(tool1["updated_at"], str)
|
||||
assert tool1["updated_at"] == "2024-01-16T14:20:30.789012"
|
||||
# Verify masking of sensitive values
|
||||
assert tool1["litellm_params"]["api_key"] != "sk-test"
|
||||
assert "****" in tool1["litellm_params"]["api_key"]
|
||||
|
||||
# Test None handling for tool 2
|
||||
tool2 = next(
|
||||
(t for t in data["search_tools"] if t["search_tool_name"] == "null-datetime-tool"),
|
||||
None,
|
||||
)
|
||||
assert tool2 is not None
|
||||
assert tool2["created_at"] is None
|
||||
assert tool2["updated_at"] is None
|
||||
# Verify masking of sensitive values
|
||||
assert tool2["litellm_params"]["api_key"] != "tvly-test"
|
||||
assert "****" in tool2["litellm_params"]["api_key"]
|
||||
|
||||
# Test string passthrough for tool 3
|
||||
tool3 = next(
|
||||
(t for t in data["search_tools"] if t["search_tool_name"] == "string-datetime-tool"),
|
||||
None,
|
||||
)
|
||||
assert tool3 is not None
|
||||
assert tool3["created_at"] == "2024-01-17T08:15:00"
|
||||
assert tool3["updated_at"] == "2024-01-18T09:25:00"
|
||||
# Verify masking of sensitive values
|
||||
assert tool3["litellm_params"]["api_key"] != "sk-test"
|
||||
assert "****" in tool3["litellm_params"]["api_key"]
|
||||
finally:
|
||||
app.dependency_overrides.pop(user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_search_tools_config_error_handling(monkeypatch):
|
||||
"""Test that config errors are handled gracefully"""
|
||||
# Mock DB tools
|
||||
db_tools = [
|
||||
{
|
||||
"search_tool_id": "test-id-1",
|
||||
"search_tool_name": "db-tool-1",
|
||||
"litellm_params": {"search_provider": "perplexity", "api_key": "sk-test"},
|
||||
"search_tool_info": {"description": "DB tool"},
|
||||
"created_at": datetime(2023, 11, 9, 12, 34, 56),
|
||||
"updated_at": datetime(2023, 11, 9, 13, 45, 12),
|
||||
}
|
||||
]
|
||||
|
||||
# Mock SearchToolRegistry
|
||||
mock_registry = MagicMock()
|
||||
mock_registry.get_all_search_tools_from_db = AsyncMock(return_value=db_tools)
|
||||
with patch(
|
||||
"litellm.proxy.search_endpoints.search_tool_management.SEARCH_TOOL_REGISTRY",
|
||||
mock_registry,
|
||||
):
|
||||
# Mock prisma_client
|
||||
mock_prisma = MagicMock()
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
|
||||
# Mock proxy_config to raise an error
|
||||
mock_proxy_config = MagicMock()
|
||||
mock_proxy_config.get_config = AsyncMock(side_effect=Exception("Config error"))
|
||||
with patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config):
|
||||
# Mock auth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user"
|
||||
)
|
||||
|
||||
try:
|
||||
# Should still succeed and return DB tools only
|
||||
response = client.get("/search_tools/list")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "search_tools" in data
|
||||
# Should only have DB tools since config failed
|
||||
assert len(data["search_tools"]) == 1
|
||||
assert data["search_tools"][0]["search_tool_name"] == "db-tool-1"
|
||||
# Verify masking of sensitive values
|
||||
assert data["search_tools"][0]["litellm_params"]["api_key"] != "sk-test"
|
||||
assert "****" in data["search_tools"][0]["litellm_params"]["api_key"]
|
||||
finally:
|
||||
app.dependency_overrides.pop(user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_search_tools_no_prisma_client(monkeypatch):
|
||||
"""Test error handling when prisma_client is None"""
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", None):
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user"
|
||||
)
|
||||
|
||||
try:
|
||||
test_client = TestClient(app)
|
||||
response = test_client.get("/search_tools/list")
|
||||
assert response.status_code == 500
|
||||
data = response.json()
|
||||
assert "Prisma client not initialized" in data["detail"]
|
||||
finally:
|
||||
app.dependency_overrides.pop(user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_search_tools_db_masking_sensitive_values(monkeypatch):
|
||||
"""
|
||||
Test that sensitive values in DB search tools are properly masked.
|
||||
This tests the new masking logic added for database search tools.
|
||||
"""
|
||||
# Mock DB tools with various sensitive fields
|
||||
db_tools = [
|
||||
{
|
||||
"search_tool_id": "test-id-1",
|
||||
"search_tool_name": "perplexity-tool",
|
||||
"litellm_params": {
|
||||
"search_provider": "perplexity",
|
||||
"api_key": "pplx-sk-1234567890abcdef",
|
||||
"api_base": "https://api.perplexity.ai",
|
||||
},
|
||||
"search_tool_info": {"description": "Perplexity tool"},
|
||||
"created_at": datetime(2023, 11, 9, 12, 34, 56),
|
||||
"updated_at": datetime(2023, 11, 9, 13, 45, 12),
|
||||
},
|
||||
{
|
||||
"search_tool_id": "test-id-2",
|
||||
"search_tool_name": "tavily-tool",
|
||||
"litellm_params": {
|
||||
"search_provider": "tavily",
|
||||
"api_key": "tvly-secret-key-12345",
|
||||
"api_base": "https://api.tavily.com",
|
||||
},
|
||||
"search_tool_info": {"description": "Tavily tool"},
|
||||
"created_at": datetime(2023, 11, 9, 12, 34, 56),
|
||||
"updated_at": datetime(2023, 11, 9, 13, 45, 12),
|
||||
},
|
||||
{
|
||||
"search_tool_id": "test-id-3",
|
||||
"search_tool_name": "tool-with-token",
|
||||
"litellm_params": {
|
||||
"search_provider": "custom",
|
||||
"access_token": "token-abcdefghijklmnop",
|
||||
"secret_key": "secret-xyz123",
|
||||
},
|
||||
"search_tool_info": {"description": "Tool with token"},
|
||||
"created_at": datetime(2023, 11, 9, 12, 34, 56),
|
||||
"updated_at": datetime(2023, 11, 9, 13, 45, 12),
|
||||
},
|
||||
{
|
||||
"search_tool_id": "test-id-4",
|
||||
"search_tool_name": "tool-with-non-sensitive",
|
||||
"litellm_params": {
|
||||
"search_provider": "custom",
|
||||
"max_results": 10,
|
||||
"timeout": 30,
|
||||
},
|
||||
"search_tool_info": {"description": "Tool without sensitive fields"},
|
||||
"created_at": datetime(2023, 11, 9, 12, 34, 56),
|
||||
"updated_at": datetime(2023, 11, 9, 13, 45, 12),
|
||||
},
|
||||
]
|
||||
|
||||
# Mock SearchToolRegistry
|
||||
mock_registry = MagicMock()
|
||||
mock_registry.get_all_search_tools_from_db = AsyncMock(return_value=db_tools)
|
||||
with patch(
|
||||
"litellm.proxy.search_endpoints.search_tool_management.SEARCH_TOOL_REGISTRY",
|
||||
mock_registry,
|
||||
):
|
||||
# Mock prisma_client
|
||||
mock_prisma = MagicMock()
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
|
||||
# Mock proxy_config
|
||||
mock_proxy_config = MagicMock()
|
||||
mock_proxy_config.get_config = AsyncMock(return_value={})
|
||||
mock_proxy_config.parse_search_tools = MagicMock(return_value=None)
|
||||
with patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config):
|
||||
# Mock auth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user"
|
||||
)
|
||||
|
||||
try:
|
||||
test_client = TestClient(app)
|
||||
response = test_client.get("/search_tools/list")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "search_tools" in data
|
||||
assert len(data["search_tools"]) == 4
|
||||
|
||||
# Test tool 1: api_key should be masked
|
||||
tool1 = next(
|
||||
(t for t in data["search_tools"] if t["search_tool_name"] == "perplexity-tool"),
|
||||
None,
|
||||
)
|
||||
assert tool1 is not None
|
||||
assert tool1["litellm_params"]["api_key"] != "pplx-sk-1234567890abcdef"
|
||||
assert "****" in tool1["litellm_params"]["api_key"]
|
||||
assert tool1["litellm_params"]["search_provider"] == "perplexity"
|
||||
assert tool1["litellm_params"]["api_base"] == "https://api.perplexity.ai"
|
||||
|
||||
# Test tool 2: api_key should be masked
|
||||
tool2 = next(
|
||||
(t for t in data["search_tools"] if t["search_tool_name"] == "tavily-tool"),
|
||||
None,
|
||||
)
|
||||
assert tool2 is not None
|
||||
assert tool2["litellm_params"]["api_key"] != "tvly-secret-key-12345"
|
||||
assert "****" in tool2["litellm_params"]["api_key"]
|
||||
assert tool2["litellm_params"]["search_provider"] == "tavily"
|
||||
|
||||
# Test tool 3: access_token and secret_key should be masked
|
||||
tool3 = next(
|
||||
(t for t in data["search_tools"] if t["search_tool_name"] == "tool-with-token"),
|
||||
None,
|
||||
)
|
||||
assert tool3 is not None
|
||||
assert tool3["litellm_params"]["access_token"] != "token-abcdefghijklmnop"
|
||||
assert "****" in tool3["litellm_params"]["access_token"]
|
||||
assert tool3["litellm_params"]["secret_key"] != "secret-xyz123"
|
||||
assert "****" in tool3["litellm_params"]["secret_key"]
|
||||
|
||||
# Test tool 4: non-sensitive fields should remain unmasked
|
||||
tool4 = next(
|
||||
(t for t in data["search_tools"] if t["search_tool_name"] == "tool-with-non-sensitive"),
|
||||
None,
|
||||
)
|
||||
assert tool4 is not None
|
||||
assert tool4["litellm_params"]["max_results"] == 10
|
||||
assert tool4["litellm_params"]["timeout"] == 30
|
||||
assert tool4["litellm_params"]["search_provider"] == "custom"
|
||||
finally:
|
||||
app.dependency_overrides.pop(user_api_key_auth, None)
|
||||
Reference in New Issue
Block a user