mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-24 05:45:11 +00:00
f6ae25939e
* feat(guardrails): retry generic_guardrail_api checks lost after connect with num_retries * fix(guardrails): close each guardrail retry connection before the next attempt and cap num_retries at 5 A failed retry's fresh client stayed open until every later retry finished, so a check that kept timing out held one idle client per retry. Each retry now closes its own client as soon as it ends, and num_retries is capped at 5 so the retry chain stays shallow * refactor(guardrails): type the generic_guardrail_api inbound header sanitizer as a mapping Keeps the type check gate at its merge-base count after the num_retries constructor params added unknown kwargs errors in the Qohash subclass * chore(guardrails): regenerate OpenAPI snapshot and UI types for num_retries, allow the retry client The guardrail retry opens its own client on purpose to get a new connection and closes it after the attempt, so the per-request async client check allowlists generic_guardrail_api.py
121 lines
4.7 KiB
Python
121 lines
4.7 KiB
Python
import ast
|
|
import os
|
|
|
|
ALLOWED_FILES = [
|
|
# Lens data traffic owns one pool per app lifespan, isolated from model traffic and closed on shutdown.
|
|
"../../litellm/tracing/remote.py",
|
|
"./litellm/tracing/remote.py",
|
|
# Only a guardrail check retry opens a client, on purpose to get a new connection, and closes it right after.
|
|
"../../litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py",
|
|
"./litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py",
|
|
# local files
|
|
"../../litellm/__init__.py",
|
|
"../../litellm/llms/custom_httpx/http_handler.py",
|
|
"../../litellm/router_utils/client_initalization_utils.py",
|
|
"../../litellm/llms/custom_httpx/http_handler.py",
|
|
"../../litellm/llms/huggingface_restapi.py",
|
|
"../../litellm/llms/base.py",
|
|
"../../litellm/llms/custom_httpx/httpx_handler.py",
|
|
"../../litellm/llms/openai/common_utils.py",
|
|
"../../litellm/experimental_mcp_client/client.py",
|
|
# when running on ci/cd
|
|
"./litellm/__init__.py",
|
|
"./litellm/llms/custom_httpx/http_handler.py",
|
|
"./litellm/router_utils/client_initalization_utils.py",
|
|
"./litellm/llms/custom_httpx/http_handler.py",
|
|
"./litellm/llms/huggingface_restapi.py",
|
|
"./litellm/llms/base.py",
|
|
"./litellm/llms/custom_httpx/httpx_handler.py",
|
|
"./litellm/llms/openai/common_utils.py",
|
|
"./litellm/experimental_mcp_client/client.py",
|
|
]
|
|
|
|
warning_msg = "this is a serious violation that can impact latency. Creating Async clients per request can add +500ms per request"
|
|
|
|
|
|
def check_for_async_http_handler(file_path):
|
|
"""
|
|
Checks if AsyncHttpHandler is instantiated in the given file.
|
|
Returns a list of line numbers where AsyncHttpHandler is used.
|
|
"""
|
|
print("..checking file=", file_path)
|
|
if file_path in ALLOWED_FILES:
|
|
return []
|
|
with open(file_path, "r") as file:
|
|
try:
|
|
tree = ast.parse(file.read())
|
|
except SyntaxError:
|
|
print(f"Warning: Syntax error in file {file_path}")
|
|
return []
|
|
|
|
violations = []
|
|
target_names = [
|
|
"AsyncHttpHandler",
|
|
"AsyncHTTPHandler",
|
|
"AsyncClient",
|
|
"httpx.AsyncClient",
|
|
] # Add variations here
|
|
for node in ast.walk(tree):
|
|
if isinstance(node, ast.Call):
|
|
if isinstance(node.func, ast.Name) and node.func.id.lower() in [
|
|
name.lower() for name in target_names
|
|
]:
|
|
raise ValueError(
|
|
f"found violation in file {file_path} line: {node.lineno}. Please use `get_async_httpx_client` instead. {warning_msg}"
|
|
)
|
|
# Check for attribute calls like httpx.AsyncClient()
|
|
elif isinstance(node.func, ast.Attribute):
|
|
full_name = ""
|
|
current = node.func
|
|
while isinstance(current, ast.Attribute):
|
|
full_name = "." + current.attr + full_name
|
|
current = current.value
|
|
if isinstance(current, ast.Name):
|
|
full_name = current.id + full_name
|
|
if full_name.lower() in [name.lower() for name in target_names]:
|
|
raise ValueError(
|
|
f"found violation in file {file_path} line: {node.lineno}. Please use `get_async_httpx_client` instead. {warning_msg}"
|
|
)
|
|
return violations
|
|
|
|
|
|
def scan_directory_for_async_handler(base_dir):
|
|
"""
|
|
Scans all Python files in the directory tree for AsyncHttpHandler usage.
|
|
Returns a dict of files and line numbers where violations were found.
|
|
"""
|
|
violations = {}
|
|
|
|
for root, _, files in os.walk(base_dir):
|
|
for file in files:
|
|
if file.endswith(".py"):
|
|
file_path = os.path.join(root, file)
|
|
file_violations = check_for_async_http_handler(file_path)
|
|
if file_violations:
|
|
violations[file_path] = file_violations
|
|
|
|
return violations
|
|
|
|
|
|
def test_no_async_http_handler_usage():
|
|
"""
|
|
Test to ensure AsyncHttpHandler is not used anywhere in the codebase.
|
|
"""
|
|
base_dir = "./litellm" # Adjust this path as needed
|
|
|
|
# base_dir = "../../litellm" # LOCAL TESTING
|
|
violations = scan_directory_for_async_handler(base_dir)
|
|
|
|
if violations:
|
|
violation_messages = []
|
|
for file_path, line_numbers in violations.items():
|
|
violation_messages.append(
|
|
f"Found AsyncHttpHandler in {file_path} at lines: {line_numbers}"
|
|
)
|
|
raise AssertionError(
|
|
"AsyncHttpHandler usage detected:\n" + "\n".join(violation_messages)
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_no_async_http_handler_usage()
|