tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-04 20:29:02 +00:00
Code Issues Packages Projects Releases Wiki Activity

style: run black formatter on files from main merge

Browse Source
This commit is contained in:
Ishaan Jaffer
2026-04-17 13:02:59 -07:00
parent ab44df202b
commit e8461b5b97
1364 changed files with 47031 additions and 34359 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/test_tool_allowlist_script.py
+53 -13
View File
@@ -24,14 +24,42 @@ def test_extraction():
from litellm.proxy.guardrails.tool_name_extraction import extract_request_tool_names
cases = [
("OpenAI chat tools", "/v1/chat/completions", {"tools": [{"type": "function", "function": {"name": "get_weather"}}]}),
("OpenAI chat functions", "/v1/chat/completions", {"functions": [{"name": "run_sql"}]}),
("OpenAI responses function", "/v1/responses", {"tools": [{"type": "function", "name": "get_current_weather"}]}),
("OpenAI responses MCP", "/v1/responses", {"tools": [{"type": "mcp", "server_label": "dmcp"}]}),
("Anthropic", "/v1/messages", {"tools": [{"name": "get_weather"}, {"name": "run_sql"}]}),
("Google generateContent", "/generate_content", {"tools": [{"functionDeclarations": [{"name": "schedule_meeting"}]}]}),
(
"OpenAI chat tools",
"/v1/chat/completions",
{"tools": [{"type": "function", "function": {"name": "get_weather"}}]},
),
(
"OpenAI chat functions",
"/v1/chat/completions",
{"functions": [{"name": "run_sql"}]},
),
(
"OpenAI responses function",
"/v1/responses",
{"tools": [{"type": "function", "name": "get_current_weather"}]},
),
(
"OpenAI responses MCP",
"/v1/responses",
{"tools": [{"type": "mcp", "server_label": "dmcp"}]},
),
(
"Anthropic",
"/v1/messages",
{"tools": [{"name": "get_weather"}, {"name": "run_sql"}]},
),
(
"Google generateContent",
"/generate_content",
{"tools": [{"functionDeclarations": [{"name": "schedule_meeting"}]}]},
),
("MCP call_tool", "/mcp/call_tool", {"name": "my_tool", "arguments": {}}),
("Non-tool route", "/v1/embeddings", {"tools": [{"type": "function", "function": {"name": "x"}}]}),
(
"Non-tool route",
"/v1/embeddings",
{"tools": [{"type": "function", "function": {"name": "x"}}]},
),
]
print("=== extract_request_tool_names(route, data) ===\n")
for label, route, data in cases:
@@ -60,7 +88,9 @@ async def test_check_tools_allowlist():
# No allowlist -> pass
await check_tools_allowlist(
request_body={"tools": [{"type": "function", "function": {"name": "get_weather"}}]},
request_body={
"tools": [{"type": "function", "function": {"name": "get_weather"}}]
},
valid_token=token(),
team_object=None,
route="/v1/chat/completions",
@@ -69,7 +99,9 @@ async def test_check_tools_allowlist():
# Allowed tool -> pass
await check_tools_allowlist(
request_body={"tools": [{"type": "function", "function": {"name": "get_weather"}}]},
request_body={
"tools": [{"type": "function", "function": {"name": "get_weather"}}]
},
valid_token=token(metadata={"allowed_tools": ["get_weather"]}),
team_object=None,
route="/v1/chat/completions",
@@ -79,7 +111,9 @@ async def test_check_tools_allowlist():
# Disallowed tool -> raise
try:
await check_tools_allowlist(
request_body={"tools": [{"type": "function", "function": {"name": "get_weather"}}]},
request_body={
"tools": [{"type": "function", "function": {"name": "get_weather"}}]
},
valid_token=token(metadata={"allowed_tools": ["other_tool"]}),
team_object=None,
route="/v1/chat/completions",
@@ -87,7 +121,9 @@ async def test_check_tools_allowlist():
print(" DISALLOWED: expected ProxyException")
except ProxyException as e:
if e.type == ProxyErrorTypes.tool_access_denied:
print(" allowed_tools=['other_tool'], body has get_weather: PASS (raised tool_access_denied)")
print(
" allowed_tools=['other_tool'], body has get_weather: PASS (raised tool_access_denied)"
)
else:
print(f" Unexpected ProxyException type: {e.type}")
except Exception as e:
@@ -95,7 +131,9 @@ async def test_check_tools_allowlist():
# Team allowlist when key empty
await check_tools_allowlist(
request_body={"tools": [{"type": "function", "function": {"name": "get_weather"}}]},
request_body={
"tools": [{"type": "function", "function": {"name": "get_weather"}}]
},
valid_token=token(team_metadata={"allowed_tools": ["get_weather"]}),
team_object=None,
route="/v1/chat/completions",
@@ -109,7 +147,9 @@ def main():
test_extraction()
asyncio.run(test_check_tools_allowlist())
print("Done. For full unit tests run:")
print(" poetry run pytest tests/test_litellm/proxy/test_tools_allowlist_enforcement.py -v")
print(
" poetry run pytest tests/test_litellm/proxy/test_tools_allowlist_enforcement.py -v"
)
if __name__ == "__main__":