tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-06 14:24:50 +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
tests/code_coverage_tests/check_get_model_cost_key_performance.py
+98 -56
View File
@@ -15,54 +15,65 @@ def _function_has_on_operations(all_lines, func_name, visited=None):
"""
if visited is None:
visited = set()
# Prevent infinite recursion
if func_name in visited:
return False
visited.add(func_name)
func_start = None
func_end = None
for i, line in enumerate(all_lines):
if func_start is None and f'def {func_name}(' in line:
if func_start is None and f"def {func_name}(" in line:
func_start = i
elif func_start is not None:
# Function ends when we hit next def at module level
if line.strip() and not line.startswith(' ') and not line.startswith('\t') and line.startswith('def '):
if (
line.strip()
and not line.startswith(" ")
and not line.startswith("\t")
and line.startswith("def ")
):
func_end = i
break
if func_start is None or func_end is None:
return False
# Check function body for O(n) patterns
func_lines = all_lines[func_start:func_end]
for line in func_lines:
# Skip comments and docstrings
line_stripped = line.strip()
if line_stripped.startswith('#') or line_stripped.startswith('"""') or line_stripped.startswith("'''"):
if (
line_stripped.startswith("#")
or line_stripped.startswith('"""')
or line_stripped.startswith("'''")
):
continue
# Check for for loops
if re.search(r'\bfor\s+\w+\s+in\s+', line):
if re.search(r"\bfor\s+\w+\s+in\s+", line):
return True
# Check for while loops
if re.search(r'\bwhile\s+', line):
if re.search(r"\bwhile\s+", line):
return True
# Check for comprehensions
if re.search(r'\[.*\s+for\s+.*\s+in\s+', line) or re.search(r'\{.*\s+for\s+.*\s+in\s+', line):
if re.search(r"\[.*\s+for\s+.*\s+in\s+", line) or re.search(
r"\{.*\s+for\s+.*\s+in\s+", line
):
return True
# Recursively check called functions (check all, don't skip any in recursive checks)
func_call_match = re.search(r'\b([a-z_][a-z0-9_]*)\s*\(', line)
func_call_match = re.search(r"\b([a-z_][a-z0-9_]*)\s*\(", line)
if func_call_match:
called_func = func_call_match.group(1)
if called_func.startswith('_'):
if called_func.startswith("_"):
if _function_has_on_operations(all_lines, called_func, visited):
return True
return False
@@ -71,46 +82,51 @@ def check_get_model_cost_key_performance():
Check that _get_model_cost_key doesn't contain O(n) operations.
"""
utils_file = "./litellm/utils.py"
if not os.path.exists(utils_file):
print(f"Warning: File {utils_file} does not exist.")
return []
with open(utils_file, "r", encoding="utf-8") as f:
lines = f.readlines()
# Find the _get_model_cost_key function
func_start = None
func_end = None
for i, line in enumerate(lines):
if func_start is None and 'def _get_model_cost_key(' in line:
if func_start is None and "def _get_model_cost_key(" in line:
func_start = i
elif func_start is not None:
# Function ends when we hit next def at module level (no indentation)
if line.strip() and not line.startswith(' ') and not line.startswith('\t') and line.startswith('def '):
if (
line.strip()
and not line.startswith(" ")
and not line.startswith("\t")
and line.startswith("def ")
):
func_end = i
break
if func_start is None:
print("Warning: Could not find _get_model_cost_key function")
return []
if func_end is None:
func_end = len(lines)
# Extract function body
func_lines = lines[func_start:func_end]
problematic_lines = []
# Track if we're inside a docstring
in_docstring = False
docstring_quote = None
# Check for O(n) patterns
for i, line in enumerate(func_lines, start=func_start + 1):
line_stripped = line.strip()
# Track docstring state (handle both single-line and multi-line docstrings)
if not in_docstring:
if line_stripped.startswith('"""') or line_stripped.startswith("'''"):
@@ -128,72 +144,98 @@ def check_get_model_cost_key_performance():
in_docstring = False
docstring_quote = None
continue # Skip all lines inside docstring
# Skip comments
if line_stripped.startswith('#'):
if line_stripped.startswith("#"):
continue
# Check for for loops
if re.search(r'\bfor\s+\w+\s+in\s+', line):
if re.search(r"\bfor\s+\w+\s+in\s+", line):
# Allow helper function calls (they're conditional)
if not re.search(r'(_rebuild_model_cost_lowercase_map|_handle_stale_map_entry_rebuild|_handle_new_key_with_scan)', line):
if not re.search(
r"(_rebuild_model_cost_lowercase_map|_handle_stale_map_entry_rebuild|_handle_new_key_with_scan)",
line,
):
problematic_lines.append((i, "for loop", line_stripped))
# Check for while loops
if re.search(r'\bwhile\s+', line):
if re.search(r"\bwhile\s+", line):
problematic_lines.append((i, "while loop", line_stripped))
# Check for comprehensions
if re.search(r'\[.*\s+for\s+.*\s+in\s+', line) or re.search(r'\{.*\s+for\s+.*\s+in\s+', line):
if re.search(r"\[.*\s+for\s+.*\s+in\s+", line) or re.search(
r"\{.*\s+for\s+.*\s+in\s+", line
):
problematic_lines.append((i, "comprehension", line_stripped))
# Check for problematic function calls
problematic_funcs = ['enumerate', 'zip', 'map', 'filter', 'sorted', 'any', 'all', 'sum', 'max', 'min']
problematic_funcs = [
"enumerate",
"zip",
"map",
"filter",
"sorted",
"any",
"all",
"sum",
"max",
"min",
]
for func in problematic_funcs:
if re.search(rf'\b{func}\s*\(', line):
if re.search(rf"\b{func}\s*\(", line):
problematic_lines.append((i, f"call to {func}()", line_stripped))
# Check for calls to functions that might have O(n) operations
# Allow known helper functions that are conditional
allowed_helpers = [
'_rebuild_model_cost_lowercase_map',
'_handle_stale_map_entry_rebuild',
'_handle_new_key_with_scan',
"_rebuild_model_cost_lowercase_map",
"_handle_stale_map_entry_rebuild",
"_handle_new_key_with_scan",
]
# Check for function calls (pattern: function_name(...), but not function definitions)
# Skip function definitions (def function_name(...))
if not re.search(r'\bdef\s+', line):
func_call_match = re.search(r'\b([a-z_][a-z0-9_]*)\s*\(', line)
if not re.search(r"\bdef\s+", line):
func_call_match = re.search(r"\b([a-z_][a-z0-9_]*)\s*\(", line)
if func_call_match:
func_name = func_call_match.group(1)
# If it's a call to a function that might have O(n) operations, check it
if func_name not in allowed_helpers and func_name.startswith('_'):
if func_name not in allowed_helpers and func_name.startswith("_"):
# Check if this function has O(n) operations
if _function_has_on_operations(lines, func_name):
problematic_lines.append((i, f"call to {func_name}() which contains O(n) operations", line_stripped))
problematic_lines.append(
(
i,
f"call to {func_name}() which contains O(n) operations",
line_stripped,
)
)
return problematic_lines
def main():
"""Main function to check _get_model_cost_key performance requirements."""
problematic_lines = check_get_model_cost_key_performance()
if problematic_lines:
print("\nERROR: Found O(n) operations in _get_model_cost_key:")
for line_num, operation, context in problematic_lines:
print(f" Line {line_num}: {operation} - {context}")
print("\nWARNING: Only O(1) lookup operations are acceptable in _get_model_cost_key.")
print(
"\nWARNING: Only O(1) lookup operations are acceptable in _get_model_cost_key."
)
print("Any O(n) operations will cause severe CPU overhead.")
raise Exception(
f"Found {len(problematic_lines)} O(n) operation(s) in _get_model_cost_key. "
f"This violates the performance requirement."
)
else:
print("OK: No O(n) operations found in _get_model_cost_key. Performance requirement satisfied.")
print(
"OK: No O(n) operations found in _get_model_cost_key. Performance requirement satisfied."
)
if __name__ == "__main__":
tests/code_coverage_tests/check_licenses.py
+9 -8
View File
@@ -248,9 +248,9 @@ class LicenseChecker:
lock_data = tomllib.load(f)
requirement_lines = list(pyproject["project"].get("dependencies", []))
for extra_reqs in pyproject["project"].get(
"optional-dependencies", {}
).values():
for extra_reqs in (
pyproject["project"].get("optional-dependencies", {}).values()
):
requirement_lines.extend(extra_reqs)
for group_reqs in pyproject.get("dependency-groups", {}).values():
requirement_lines.extend(group_reqs)
@@ -286,7 +286,9 @@ class LicenseChecker:
]
except Exception as e:
source = requirements_file or "pyproject.toml + uv.lock"
raise RuntimeError(f"Error parsing requirements from {source}: {str(e)}") from e
raise RuntimeError(
f"Error parsing requirements from {source}: {str(e)}"
) from e
def check_requirements(self, requirements_file: Optional[Path] = None) -> bool:
"""Check all packages from a requirements file or the default repo deps."""
@@ -303,9 +305,7 @@ class LicenseChecker:
for req in requirements:
try:
version = (
next(iter(req.specifier)).version if req.specifier else None
)
version = next(iter(req.specifier)).version if req.specifier else None
except StopIteration:
version = None
@@ -348,7 +348,8 @@ def main():
unhandled_packages = [
p
for p in (unverified + invalid)
if checker._normalize_package_name(p.name) not in checker.authorized_packages
if checker._normalize_package_name(p.name)
not in checker.authorized_packages
]
if unhandled_packages:
tests/code_coverage_tests/check_spanattributes_value_usage.py
+66 -39
View File
@@ -38,70 +38,85 @@ class SpanAttributesUsageChecker(ast.NodeVisitor):
"""
Checks if SpanAttributes is used without .value when setting attributes in safe_set_attribute calls
and other attribute setting methods in opentelemetry.py.
This is important to ensure consistent enum value access and prevent type errors
when sending data to OpenTelemetry exporters.
"""
def __init__(self, debug=False):
self.violations = []
self.debug = debug
def visit_Call(self, node):
# Check if this is a call to safe_set_attribute or set_attribute
if isinstance(node.func, ast.Attribute) and node.func.attr in ['safe_set_attribute', 'set_attribute']:
if isinstance(node.func, ast.Attribute) and node.func.attr in [
"safe_set_attribute",
"set_attribute",
]:
# Look for the 'key' parameter
for keyword in node.keywords:
if keyword.arg == 'key':
if keyword.arg == "key":
# Check if the value is a SpanAttributes member without .value
if isinstance(keyword.value, ast.Attribute) and \
isinstance(keyword.value.value, ast.Name) and \
keyword.value.value.id == 'SpanAttributes':
if (
isinstance(keyword.value, ast.Attribute)
and isinstance(keyword.value.value, ast.Name)
and keyword.value.value.id == "SpanAttributes"
):
# Get the source code for this attribute
try:
attr_source = ast.unparse(keyword.value)
if not attr_source.endswith('.value'):
if not attr_source.endswith(".value"):
if self.debug:
print(f"AST found violation: {node.lineno}: {attr_source}")
self.violations.append((node.lineno, f"{attr_source} used without .value"))
print(
f"AST found violation: {node.lineno}: {attr_source}"
)
self.violations.append(
(node.lineno, f"{attr_source} used without .value")
)
except AttributeError:
# For Python < 3.9, ast.unparse doesn't exist
# Fallback to our best guess
if keyword.value.attr != 'value' and not hasattr(keyword.value, 'value'):
if keyword.value.attr != "value" and not hasattr(
keyword.value, "value"
):
violation_msg = f"SpanAttributes.{keyword.value.attr} used without .value"
if self.debug:
print(f"AST found violation: {node.lineno}: {violation_msg}")
print(
f"AST found violation: {node.lineno}: {violation_msg}"
)
self.violations.append((node.lineno, violation_msg))
# Continue the visit
self.generic_visit(node)
def check_file(file_path: str, debug: bool = False) -> List[Tuple[int, str]]:
"""
Analyze a Python file to check for SpanAttributes usage without .value
Args:
file_path: Path to the Python file to check
debug: Whether to print debug information
Returns:
List of (line_number, message) tuples identifying violations
"""
with open(file_path, 'r') as file:
with open(file_path, "r") as file:
content = file.read()
# First try AST parsing for accurate code structure analysis
try:
tree = ast.parse(content)
checker = SpanAttributesUsageChecker(debug=debug)
checker.visit(tree)
violations = checker.violations
# Also do a regex check for backup/extra coverage
# This catches cases that might be missed by AST parsing
# Split content into lines for more precise analysis
lines = content.splitlines()
for i, line in enumerate(lines, 1):
# Skip lines that contain ".value" after "SpanAttributes."
# This prevents false positives for correct usage
@@ -109,60 +124,72 @@ def check_file(file_path: str, debug: bool = False) -> List[Tuple[int, str]]:
if debug:
print(f"Line {i} skipped - contains .value: {line.strip()}")
continue
# Pattern: Looking for "key=SpanAttributes.ENUM_NAME" without .value at the end
pattern = r"key\s*=\s*SpanAttributes\.[A-Z_][A-Z0-9_]*(?!\.value)"
match = re.search(pattern, line)
if match:
# Check if this violation was already found by AST
if not any(i == line_num for line_num, _ in violations):
if debug:
print(f"Regex found violation: {i}: {match.group(0)}")
violations.append((i, f"SpanAttributes used without .value: {match.group(0)}"))
violations.append(
(i, f"SpanAttributes used without .value: {match.group(0)}")
)
return violations
except SyntaxError:
print(f"Syntax error in {file_path}")
return []
def main():
"""
Main function to run the SpanAttributes usage check on the OpenTelemetry integration file.
Exits with code 1 if violations are found, 0 otherwise.
"""
parser = argparse.ArgumentParser(description='Check for SpanAttributes used without .value')
parser.add_argument('--debug', action='store_true', help='Enable debug output')
parser = argparse.ArgumentParser(
description="Check for SpanAttributes used without .value"
)
parser.add_argument("--debug", action="store_true", help="Enable debug output")
args = parser.parse_args()
# Path to the OpenTelemetry integration file
target_file = os.path.join("litellm", "integrations", "opentelemetry.py")
if not os.path.exists(target_file):
# Try alternate path for local development
target_file = os.path.join("..", "..", "litellm", "integrations", "opentelemetry.py")
target_file = os.path.join(
"..", "..", "litellm", "integrations", "opentelemetry.py"
)
if not os.path.exists(target_file):
print(f"Error: Could not find file at {target_file}")
exit(1)
violations = check_file(target_file, debug=args.debug)
if violations:
print(f"Found {len(violations)} SpanAttributes without .value in {target_file}:")
print(
f"Found {len(violations)} SpanAttributes without .value in {target_file}:"
)
# Sort violations by line number for better readability
violations.sort(key=lambda x: x[0])
for line, message in violations:
print(f" Line {line}: {message}")
print("\nDirect enum reference can cause errors. Always use .value with SpanAttributes enums.")
print(
"\nDirect enum reference can cause errors. Always use .value with SpanAttributes enums."
)
exit(1)
else:
print(f"All SpanAttributes are used correctly with .value in {target_file}")
exit(0)
if __name__ == "__main__":
main()
main()
tests/code_coverage_tests/check_unsafe_enterprise_import.py
+25 -15
View File
@@ -1,6 +1,7 @@
import ast
import os
class EnterpriseImportFinder(ast.NodeVisitor):
def __init__(self):
self.unsafe_imports = []
@@ -34,26 +35,33 @@ class EnterpriseImportFinder(ast.NodeVisitor):
for name in node.names:
if "litellm_enterprise" in name.name or "enterprise" in name.name:
if not self.in_try_block:
self.unsafe_imports.append({
"file": self.current_file,
"line": node.lineno,
"import": name.name,
"context": "direct import"
})
self.unsafe_imports.append(
{
"file": self.current_file,
"line": node.lineno,
"import": name.name,
"context": "direct import",
}
)
self.generic_visit(node)
def visit_ImportFrom(self, node):
# Check for from litellm_enterprise imports
if node.module and ("litellm_enterprise" in node.module or "enterprise" in node.module):
if node.module and (
"litellm_enterprise" in node.module or "enterprise" in node.module
):
if not self.in_try_block:
self.unsafe_imports.append({
"file": self.current_file,
"line": node.lineno,
"import": f"from {node.module}",
"context": "from import"
})
self.unsafe_imports.append(
{
"file": self.current_file,
"line": node.lineno,
"import": f"from {node.module}",
"context": "from import",
}
)
self.generic_visit(node)
def find_unsafe_enterprise_imports_in_file(file_path):
with open(file_path, "r") as file:
tree = ast.parse(file.read(), filename=file_path)
@@ -62,6 +70,7 @@ def find_unsafe_enterprise_imports_in_file(file_path):
finder.visit(tree)
return finder.unsafe_imports
def find_unsafe_enterprise_imports_in_directory(directory):
unsafe_imports = []
for root, _, files in os.walk(directory):
@@ -73,11 +82,12 @@ def find_unsafe_enterprise_imports_in_directory(directory):
unsafe_imports.extend(imports)
return unsafe_imports
if __name__ == "__main__":
# Check for unsafe enterprise imports in the litellm directory
directory_path = "./litellm"
unsafe_imports = find_unsafe_enterprise_imports_in_directory(directory_path)
if unsafe_imports:
print("🚨 UNSAFE ENTERPRISE IMPORTS FOUND (not in try-except blocks):")
for imp in unsafe_imports:
@@ -86,7 +96,7 @@ if __name__ == "__main__":
print(f"Import: {imp['import']}")
print(f"Context: {imp['context']}")
print("---")
# Raise exception to fail CI/CD
raise Exception(
"🚨 Unsafe enterprise imports found. All enterprise imports must be wrapped in try-except blocks."
tests/code_coverage_tests/code_qa_check_tests.py
+75 -64
View File
@@ -6,7 +6,7 @@ def check_for_litellm_module_deletion(base_dir):
"""
Checks for code patterns that delete litellm modules from sys.modules
in the test_litellm directory.
Specifically looks for patterns like:
for module in list(sys.modules.keys()):
if module.startswith("litellm"):
@@ -14,13 +14,13 @@ def check_for_litellm_module_deletion(base_dir):
"""
problematic_files = []
test_dir = os.path.join(base_dir, "test_litellm")
if not os.path.exists(test_dir):
print(f"Warning: Directory {test_dir} does not exist.")
return []
print(f"Checking directory: {test_dir}")
for root, _, files in os.walk(test_dir):
for file in files:
if file.endswith(".py"):
@@ -31,132 +31,143 @@ def check_for_litellm_module_deletion(base_dir):
except SyntaxError:
print(f"Warning: Syntax error in file {file_path}")
continue
# Check for litellm module deletion patterns
if has_litellm_module_deletion(tree):
relative_path = os.path.relpath(file_path, base_dir)
problematic_files.append(relative_path)
print(f"Found litellm module deletion in: {relative_path}")
return problematic_files
def has_litellm_module_deletion(tree):
"""
Checks if the AST contains patterns that delete litellm modules from sys.modules.
Looks for:
1. Loops over sys.modules.keys()
2. Conditions checking if module startswith "litellm"
3. del sys.modules[module] statements
"""
class LiteLLMDeletionVisitor(ast.NodeVisitor):
def __init__(self):
self.has_sys_modules_loop = False
self.has_litellm_check = False
self.has_del_sys_modules = False
self.current_for_target = None
def visit_For(self, node):
# Check if we're looping over sys.modules.keys()
if (isinstance(node.iter, ast.Call) and
isinstance(node.iter.func, ast.Attribute) and
isinstance(node.iter.func.value, ast.Attribute) and
isinstance(node.iter.func.value.value, ast.Name) and
node.iter.func.value.value.id == "sys" and
node.iter.func.value.attr == "modules" and
node.iter.func.attr == "keys"):
if (
isinstance(node.iter, ast.Call)
and isinstance(node.iter.func, ast.Attribute)
and isinstance(node.iter.func.value, ast.Attribute)
and isinstance(node.iter.func.value.value, ast.Name)
and node.iter.func.value.value.id == "sys"
and node.iter.func.value.attr == "modules"
and node.iter.func.attr == "keys"
):
self.has_sys_modules_loop = True
if isinstance(node.target, ast.Name):
self.current_for_target = node.target.id
# Check the body of the for loop
for stmt in node.body:
self.visit(stmt)
# Also check for list(sys.modules.keys()) pattern
elif (isinstance(node.iter, ast.Call) and
isinstance(node.iter.func, ast.Name) and
node.iter.func.id == "list" and
len(node.iter.args) == 1 and
isinstance(node.iter.args[0], ast.Call) and
isinstance(node.iter.args[0].func, ast.Attribute) and
isinstance(node.iter.args[0].func.value, ast.Attribute) and
isinstance(node.iter.args[0].func.value.value, ast.Name) and
node.iter.args[0].func.value.value.id == "sys" and
node.iter.args[0].func.value.attr == "modules" and
node.iter.args[0].func.attr == "keys"):
elif (
isinstance(node.iter, ast.Call)
and isinstance(node.iter.func, ast.Name)
and node.iter.func.id == "list"
and len(node.iter.args) == 1
and isinstance(node.iter.args[0], ast.Call)
and isinstance(node.iter.args[0].func, ast.Attribute)
and isinstance(node.iter.args[0].func.value, ast.Attribute)
and isinstance(node.iter.args[0].func.value.value, ast.Name)
and node.iter.args[0].func.value.value.id == "sys"
and node.iter.args[0].func.value.attr == "modules"
and node.iter.args[0].func.attr == "keys"
):
self.has_sys_modules_loop = True
if isinstance(node.target, ast.Name):
self.current_for_target = node.target.id
# Check the body of the for loop
for stmt in node.body:
self.visit(stmt)
self.generic_visit(node)
def visit_If(self, node):
# Check for conditions like module.startswith("litellm")
if (isinstance(node.test, ast.Call) and
isinstance(node.test.func, ast.Attribute) and
isinstance(node.test.func.value, ast.Name) and
node.test.func.value.id == self.current_for_target and
node.test.func.attr == "startswith" and
len(node.test.args) == 1 and
isinstance(node.test.args[0], ast.Constant) and
node.test.args[0].value == "litellm"):
if (
isinstance(node.test, ast.Call)
and isinstance(node.test.func, ast.Attribute)
and isinstance(node.test.func.value, ast.Name)
and node.test.func.value.id == self.current_for_target
and node.test.func.attr == "startswith"
and len(node.test.args) == 1
and isinstance(node.test.args[0], ast.Constant)
and node.test.args[0].value == "litellm"
):
self.has_litellm_check = True
# Check the body of the if statement
for stmt in node.body:
self.visit(stmt)
self.generic_visit(node)
def visit_Delete(self, node):
# Check for del sys.modules[module]
for target in node.targets:
if (isinstance(target, ast.Subscript) and
isinstance(target.value, ast.Attribute) and
isinstance(target.value.value, ast.Name) and
target.value.value.id == "sys" and
target.value.attr == "modules" and
isinstance(target.slice, ast.Name) and
target.slice.id == self.current_for_target):
if (
isinstance(target, ast.Subscript)
and isinstance(target.value, ast.Attribute)
and isinstance(target.value.value, ast.Name)
and target.value.value.id == "sys"
and target.value.attr == "modules"
and isinstance(target.slice, ast.Name)
and target.slice.id == self.current_for_target
):
self.has_del_sys_modules = True
self.generic_visit(node)
visitor = LiteLLMDeletionVisitor()
visitor.visit(tree)
return (visitor.has_sys_modules_loop and
visitor.has_litellm_check and
visitor.has_del_sys_modules)
return (
visitor.has_sys_modules_loop
and visitor.has_litellm_check
and visitor.has_del_sys_modules
)
def main():
"""
Main function to check for litellm module deletion patterns in test files.
"""
# local dir
#tests_dir = "../../tests/"
# local dir
# tests_dir = "../../tests/"
# ci/cd dir
tests_dir = "./tests/"
problematic_files = check_for_litellm_module_deletion(tests_dir)
if problematic_files:
print("\nERROR: Found files that delete litellm modules from sys.modules:")
for file_path in problematic_files:
print(f" - {file_path}")
raise Exception(
f"Found {len(problematic_files)} file(s) that delete litellm modules from sys.modules. "
f"This can cause import issues and test failures. Files: {problematic_files}"
tests/code_coverage_tests/enforce_llms_folder_style.py
+1 -1
View File
@@ -28,7 +28,7 @@ ALLOWED_FILES_IN_LLMS_FOLDER = [
"custom_httpx",
"custom_llm",
"deprecated_providers",
"pass_through"
"pass_through",
] + SEARCH_PROVIDERS
tests/code_coverage_tests/info_log_check.py
+190 -112
View File
@@ -8,15 +8,15 @@ class SensitiveLogDetector(ast.NodeVisitor):
"""
Detects logger.info() statements that might log sensitive request/response data.
"""
def __init__(self):
self.violations = []
self.current_file = None
def set_file(self, file_path: str):
"""Set the current file being analyzed"""
self.current_file = file_path
def visit_Call(self, node):
"""Visit function calls to detect logger.info() with sensitive data"""
if self._is_logger_info_call(node):
@@ -28,97 +28,110 @@ class SensitiveLogDetector(ast.NodeVisitor):
"line": node.lineno,
"call": self._get_call_string(node),
"reason": self._get_violation_reason(arg),
"arg": self._get_arg_string(arg)
"arg": self._get_arg_string(arg),
}
self.violations.append(violation)
self.generic_visit(node)
def _is_logger_info_call(self, node) -> bool:
"""Check if this is a logger.info() call"""
if not isinstance(node.func, ast.Attribute):
return False
# Check for various logger patterns:
# logger.info(), verbose_logger.info(), verbose_proxy_logger.info(), etc.
if node.func.attr == "info":
if isinstance(node.func.value, ast.Name):
logger_name = node.func.value.id
return any(pattern in logger_name.lower() for pattern in ["logger", "log"])
return any(
pattern in logger_name.lower() for pattern in ["logger", "log"]
)
return False
def _contains_sensitive_data(self, arg) -> bool:
"""Check if the argument might contain sensitive data"""
# Convert argument to string for analysis
arg_str = self._get_arg_string(arg).lower()
# Skip obvious non-sensitive patterns
non_sensitive_patterns = [
r'^["\'][\w\s\-_:.,!?]*["\']$', # Simple static strings
r'^["\'][^{%]*["\']$', # Strings without format placeholders
]
# Skip common safe phrases that contain sensitive keywords
safe_phrases = [
r'request\s+(completed|finished|started|processing)',
r'response\s+(sent|received|processed)',
r'data\s+(inserted|updated|deleted|saved)\s+into',
r'(successfully|failed)\s+(request|response)',
r'(starting|ending|completed)\s+(request|response)',
r'no\s+(usage\s+)?data\s+found',
r'found\s+\d+.*records',
r'exported\s+\d+.*records',
r"request\s+(completed|finished|started|processing)",
r"response\s+(sent|received|processed)",
r"data\s+(inserted|updated|deleted|saved)\s+into",
r"(successfully|failed)\s+(request|response)",
r"(starting|ending|completed)\s+(request|response)",
r"no\s+(usage\s+)?data\s+found",
r"found\s+\d+.*records",
r"exported\s+\d+.*records",
]
for pattern in non_sensitive_patterns:
if re.search(pattern, arg_str):
# Check if it's a safe phrase first
for safe_pattern in safe_phrases:
if re.search(safe_pattern, arg_str, re.IGNORECASE):
return False
# Then check if the static string mentions sensitive keywords
if not any(keyword in arg_str for keyword in
['request', 'response', 'data', 'body', 'payload', 'token', 'auth', 'credential']):
if not any(
keyword in arg_str
for keyword in [
"request",
"response",
"data",
"body",
"payload",
"token",
"auth",
"credential",
]
):
return False
# Direct variable/attribute patterns that are likely sensitive
sensitive_patterns = [
r'\brequest\b(?!\s*(id|status|method))', # request but not request_id, request_status, request_method
r'\bresponse\b(?!\s*(status|code|time))', # response but not response_status, response_code
r'\bdata\b(?=[\.\[\s]|$)', # data followed by . [ space or end
r'\bbody\b(?=[\.\[\s]|$)',
r'\bpayload\b(?=[\.\[\s]|$)',
r'\bmessages?\b(?=[\.\[\s]|$)',
r'\bcontent\b(?=[\.\[\s]|$)',
r'\binput\b(?=[\.\[\s]|$)',
r'\boutput\b(?=[\.\[\s]|$)',
r'\bargs\b(?=[\.\[\s]|$)',
r'\bkwargs\b(?=[\.\[\s]|$)',
r'\bparams\b(?=[\.\[\s]|$)',
r'\bheaders\b(?=[\.\[\s]|$)',
r'\bapi_key\b',
r'\btoken\b(?!\s*(name|id))', # token but not token_name, token_id
r'\bauth\b(?=[\.\[\s]|$)',
r'\bcredentials?\b'
r"\brequest\b(?!\s*(id|status|method))", # request but not request_id, request_status, request_method
r"\bresponse\b(?!\s*(status|code|time))", # response but not response_status, response_code
r"\bdata\b(?=[\.\[\s]|$)", # data followed by . [ space or end
r"\bbody\b(?=[\.\[\s]|$)",
r"\bpayload\b(?=[\.\[\s]|$)",
r"\bmessages?\b(?=[\.\[\s]|$)",
r"\bcontent\b(?=[\.\[\s]|$)",
r"\binput\b(?=[\.\[\s]|$)",
r"\boutput\b(?=[\.\[\s]|$)",
r"\bargs\b(?=[\.\[\s]|$)",
r"\bkwargs\b(?=[\.\[\s]|$)",
r"\bparams\b(?=[\.\[\s]|$)",
r"\bheaders\b(?=[\.\[\s]|$)",
r"\bapi_key\b",
r"\btoken\b(?!\s*(name|id))", # token but not token_name, token_id
r"\bauth\b(?=[\.\[\s]|$)",
r"\bcredentials?\b",
]
# Check for direct variable references with context
for pattern in sensitive_patterns:
if re.search(pattern, arg_str):
return True
# Check for format strings that might interpolate sensitive data
if self._is_format_string_with_sensitive_data(arg):
return True
# Check for JSON dumps or string formatting of objects
if self._is_object_serialization(arg):
return True
return False
def _is_format_string_with_sensitive_data(self, arg) -> bool:
"""Check if this is a format string that might contain sensitive data"""
# Check for f-strings
@@ -128,13 +141,27 @@ class SensitiveLogDetector(ast.NodeVisitor):
value_str = self._get_arg_string(value.value).lower()
# Check for any sensitive data patterns in f-string interpolations
sensitive_f_string_patterns = [
'request', 'response', 'data', 'body', 'content', 'messages',
'token', 'jwt', 'auth', 'api_key', 'apikey', 'credential',
'secret', 'password', 'passwd'
"request",
"response",
"data",
"body",
"content",
"messages",
"token",
"jwt",
"auth",
"api_key",
"apikey",
"credential",
"secret",
"password",
"passwd",
]
if any(pattern in value_str for pattern in sensitive_f_string_patterns):
if any(
pattern in value_str for pattern in sensitive_f_string_patterns
):
return True
# Check for .format() calls
if isinstance(arg, ast.Call) and isinstance(arg.func, ast.Attribute):
if arg.func.attr == "format":
@@ -143,71 +170,107 @@ class SensitiveLogDetector(ast.NodeVisitor):
if "{}" in base_str or "{" in base_str:
# Check format arguments for sensitive data
sensitive_format_patterns = [
'request', 'response', 'data', 'body', 'content',
'token', 'jwt', 'auth', 'api_key', 'apikey', 'credential',
'secret', 'password', 'passwd'
"request",
"response",
"data",
"body",
"content",
"token",
"jwt",
"auth",
"api_key",
"apikey",
"credential",
"secret",
"password",
"passwd",
]
for format_arg in arg.args:
format_str = self._get_arg_string(format_arg).lower()
if any(pattern in format_str for pattern in sensitive_format_patterns):
if any(
pattern in format_str
for pattern in sensitive_format_patterns
):
return True
return False
def _is_object_serialization(self, arg) -> bool:
"""Check if this is serializing an object that might contain sensitive data"""
arg_str = self._get_arg_string(arg)
# Check for json.dumps() calls
if isinstance(arg, ast.Call):
if (isinstance(arg.func, ast.Attribute) and
arg.func.attr == "dumps" and
isinstance(arg.func.value, ast.Name) and
arg.func.value.id == "json"):
if (
isinstance(arg.func, ast.Attribute)
and arg.func.attr == "dumps"
and isinstance(arg.func.value, ast.Name)
and arg.func.value.id == "json"
):
return True
# Check for str() calls on potentially sensitive objects
if (isinstance(arg.func, ast.Name) and arg.func.id == "str" and
len(arg.args) > 0):
if (
isinstance(arg.func, ast.Name)
and arg.func.id == "str"
and len(arg.args) > 0
):
obj_str = self._get_arg_string(arg.args[0]).lower()
if any(pattern in obj_str for pattern in
['request', 'response', 'data', 'body']):
if any(
pattern in obj_str
for pattern in ["request", "response", "data", "body"]
):
return True
return False
def _get_violation_reason(self, arg) -> str:
"""Get a human-readable reason for the violation"""
arg_str = self._get_arg_string(arg).lower()
if any(pattern in arg_str for pattern in ['jwt', 'token', 'api_key', 'apikey', 'auth', 'credential', 'secret', 'password', 'passwd']):
if any(
pattern in arg_str
for pattern in [
"jwt",
"token",
"api_key",
"apikey",
"auth",
"credential",
"secret",
"password",
"passwd",
]
):
return "Potentially logging authentication/secret data (JWT, token, API key, etc.)"
elif 'request' in arg_str:
elif "request" in arg_str:
return "Potentially logging request data"
elif 'response' in arg_str:
elif "response" in arg_str:
return "Potentially logging response data"
elif any(pattern in arg_str for pattern in ['data', 'body', 'payload', 'content']):
elif any(
pattern in arg_str for pattern in ["data", "body", "payload", "content"]
):
return "Potentially logging sensitive data/body/content"
elif any(pattern in arg_str for pattern in ['messages', 'input', 'output']):
elif any(pattern in arg_str for pattern in ["messages", "input", "output"]):
return "Potentially logging message/input/output data"
else:
return "Potentially logging sensitive data"
def _get_call_string(self, node) -> str:
"""Get string representation of the function call"""
try:
if hasattr(ast, 'unparse'):
if hasattr(ast, "unparse"):
return ast.unparse(node)
else:
# Fallback for older Python versions
return f"{self._get_arg_string(node.func)}(...)"
except:
return "logger.info(...)"
def _get_arg_string(self, arg) -> str:
"""Get string representation of an argument"""
try:
if hasattr(ast, 'unparse'):
if hasattr(ast, "unparse"):
return ast.unparse(arg)
else:
# Fallback for older Python versions
@@ -228,75 +291,88 @@ class SensitiveLogDetector(ast.NodeVisitor):
def check_sensitive_logging(base_dir: str) -> List[Dict[str, Any]]:
"""
Check for logger.info() statements that might log sensitive data.
Args:
base_dir: Base directory to scan (typically the litellm root)
Returns:
List of violations found
"""
detector = SensitiveLogDetector()
all_violations = []
# Directories to scan - only main litellm codebase
scan_dirs = [
"litellm",
"enterprise" # Include enterprise directory if it exists
]
scan_dirs = ["litellm", "enterprise"] # Include enterprise directory if it exists
# Directories to exclude (third-party code, venvs, etc.)
exclude_dirs = {
"venv", "venv313", ".venv", "env", ".env",
"node_modules", "__pycache__", ".git",
"build", "dist", ".tox", "clean_env",
"litellm_env", "myenv", "py313_env",
"venv_sip_bypass", "mypyc_env"
"venv",
"venv313",
".venv",
"env",
".env",
"node_modules",
"__pycache__",
".git",
"build",
"dist",
".tox",
"clean_env",
"litellm_env",
"myenv",
"py313_env",
"venv_sip_bypass",
"mypyc_env",
}
for scan_dir in scan_dirs:
dir_path = os.path.join(base_dir, scan_dir)
if not os.path.exists(dir_path):
print(f"Warning: Directory {dir_path} does not exist, skipping.")
continue
print(f"Scanning directory: {dir_path}")
for root, dirs, files in os.walk(dir_path):
# Skip excluded directories
dirs[:] = [d for d in dirs if d not in exclude_dirs]
# Skip if we're in a virtual environment or third-party directory
relative_root = os.path.relpath(root, base_dir)
if any(excluded in relative_root.split(os.sep) for excluded in exclude_dirs):
if any(
excluded in relative_root.split(os.sep) for excluded in exclude_dirs
):
continue
for file in files:
if file.endswith(".py"):
file_path = os.path.join(root, file)
relative_path = os.path.relpath(file_path, base_dir)
# Skip files that are clearly third-party or generated
if any(excluded in relative_path for excluded in exclude_dirs):
continue
try:
with open(file_path, "r", encoding="utf-8") as f:
content = f.read()
tree = ast.parse(content)
detector.set_file(relative_path)
detector.visit(tree)
except SyntaxError as e:
print(f"Warning: Syntax error in file {relative_path}: {e}")
continue
except UnicodeDecodeError as e:
print(f"Warning: Unicode decode error in file {relative_path}: {e}")
print(
f"Warning: Unicode decode error in file {relative_path}: {e}"
)
continue
except Exception as e:
print(f"Warning: Error processing file {relative_path}: {e}")
continue
return detector.violations
@@ -314,28 +390,30 @@ def main():
# Running in CI/CD
###################
base_dir = "./litellm" # Adjust this path as needed
print(f"Checking for sensitive logging in: {base_dir}")
violations = check_sensitive_logging(base_dir)
if violations:
print(f"\n❌ Found {len(violations)} potential violations:")
print("=" * 80)
for i, violation in enumerate(violations, 1):
print(f"\n{i}. {violation['file']}:{violation['line']}")
print(f" Reason: {violation['reason']}")
print(f" Call: {violation['call']}")
print(f" Argument: {violation['arg']}")
print("\n" + "=" * 80)
print("⚠️ SECURITY WARNING:")
print("These logger.info() statements may log sensitive request/response data.")
print("Consider changing them to logger.debug() or removing sensitive data.")
print("This is critical for PII compliance and security.")
print("Please contact @ishaan-jaff for more details about this check. DO NOT VIOLATE THIS CHECK.")
print(
"Please contact @ishaan-jaff for more details about this check. DO NOT VIOLATE THIS CHECK."
)
return 1 # Exit with error code
else:
print("\n✅ No sensitive logging violations found!")
tests/code_coverage_tests/memory_test.py
+316 -157
View File
@@ -26,89 +26,111 @@ from typing import List, Dict, Any, Optional, Sequence
class Pattern(ABC):
"""Base class for memory violation detection patterns"""
@abstractmethod
def get_pattern_name(self) -> str:
"""Return unique identifier for this violation type"""
pass
@abstractmethod
def visit_assign(self, node: ast.Assign, context: Dict[str, Any]) -> List[Dict[str, Any]]:
def visit_assign(
self, node: ast.Assign, context: Dict[str, Any]
) -> List[Dict[str, Any]]:
"""Detect memory-sensitive operations in assignment. Returns list of {line, var_name, call} dicts."""
pass
@abstractmethod
def check_cleanup(self, operations: List[Dict[str, Any]], function_body: List[ast.stmt],
context: Dict[str, Any]) -> List[Dict[str, Any]]:
def check_cleanup(
self,
operations: List[Dict[str, Any]],
function_body: List[ast.stmt],
context: Dict[str, Any],
) -> List[Dict[str, Any]]:
"""Verify variables are set to None. Returns list of violation dicts."""
pass
class QueueGetPattern(Pattern):
"""Detects queue.get()/get_nowait() operations that aren't cleared"""
def get_pattern_name(self) -> str:
return "queue_reference_not_cleared"
def visit_assign(self, node: ast.Assign, context: Dict[str, Any]) -> List[Dict[str, Any]]:
def visit_assign(
self, node: ast.Assign, context: Dict[str, Any]
) -> List[Dict[str, Any]]:
"""Detect queue.get() or queue.get_nowait() calls where object name contains 'queue'"""
operations = []
if isinstance(node.value, ast.Call):
func = node.value.func
if isinstance(func, ast.Attribute) and func.attr in ("get", "get_nowait"):
obj_name = context["get_attr_string"](func.value)
if "queue" in obj_name.lower() and node.targets and isinstance(node.targets[0], ast.Name):
operations.append({
"line": node.lineno,
"var_name": node.targets[0].id,
"call": context["get_call_string"](node.value),
})
if (
"queue" in obj_name.lower()
and node.targets
and isinstance(node.targets[0], ast.Name)
):
operations.append(
{
"line": node.lineno,
"var_name": node.targets[0].id,
"call": context["get_call_string"](node.value),
}
)
return operations
def check_cleanup(self, operations: List[Dict[str, Any]], function_body: List[ast.stmt],
context: Dict[str, Any]) -> List[Dict[str, Any]]:
def check_cleanup(
self,
operations: List[Dict[str, Any]],
function_body: List[ast.stmt],
context: Dict[str, Any],
) -> List[Dict[str, Any]]:
"""Flag queue variables that aren't set to None"""
violations = []
is_var_set_to_none = context["is_var_set_to_none"]
current_function = context["current_function"]
file_path = context["file_path"]
queue_vars = {op["var_name"]: op["line"] for op in operations}
for var_name, line_num in queue_vars.items():
if not is_var_set_to_none(var_name, function_body):
violations.append({
"line": line_num,
"type": self.get_pattern_name(),
"var_name": var_name,
"function": current_function,
"file_path": file_path,
"message": (
f"Queue variable '{var_name}' in function "
f"'{current_function}' is not set to None after use. "
f"If the runtime is overwhelmed, this can cause OOM (Out of Memory) errors."
),
})
violations.append(
{
"line": line_num,
"type": self.get_pattern_name(),
"var_name": var_name,
"function": current_function,
"file_path": file_path,
"message": (
f"Queue variable '{var_name}' in function "
f"'{current_function}' is not set to None after use. "
f"If the runtime is overwhelmed, this can cause OOM (Out of Memory) errors."
),
}
)
return violations
class UnboundedDataStructurePattern(Pattern):
"""Detects class-level data structures (lists, dicts, sets) that can grow unbounded"""
def get_pattern_name(self) -> str:
return "unbounded_data_structure"
def visit_assign(self, node: ast.Assign, context: Dict[str, Any]) -> List[Dict[str, Any]]:
def visit_assign(
self, node: ast.Assign, context: Dict[str, Any]
) -> List[Dict[str, Any]]:
"""Detect list/dict/set creations that are at class level"""
operations = []
# Check if this is a data structure creation
is_data_structure = False
structure_type = None
if isinstance(node.value, (ast.List, ast.Dict, ast.Set)):
is_data_structure = True
if isinstance(node.value, ast.List):
@@ -128,10 +150,18 @@ class UnboundedDataStructurePattern(Pattern):
# Handle cases like collections.defaultdict(list), collections.deque(), etc.
obj_name = context["get_attr_string"](func.value)
attr_name = func.attr
# Check for collections module data structures
if "collections" in obj_name.lower() or "collections" in str(func.value):
if attr_name in ("deque", "defaultdict", "Counter", "OrderedDict", "ChainMap"):
if "collections" in obj_name.lower() or "collections" in str(
func.value
):
if attr_name in (
"deque",
"defaultdict",
"Counter",
"OrderedDict",
"ChainMap",
):
# For deque, we track it and let size checks determine if it's bounded
# (deque with maxlen parameter is bounded, but we detect that via size checks)
is_data_structure = True
@@ -139,7 +169,11 @@ class UnboundedDataStructurePattern(Pattern):
elif attr_name in ("list", "dict", "set"):
# collections.defaultdict(list) pattern
is_data_structure = True
structure_type = "defaultdict" if "defaultdict" in obj_name.lower() else attr_name
structure_type = (
"defaultdict"
if "defaultdict" in obj_name.lower()
else attr_name
)
# Check for queue.Queue, asyncio.Queue (if unbounded)
elif "queue" in obj_name.lower() or "asyncio" in obj_name.lower():
if attr_name == "Queue":
@@ -153,51 +187,67 @@ class UnboundedDataStructurePattern(Pattern):
is_data_structure = True
structure_type = "queue"
# Direct attribute access like deque(), Counter(), etc.
elif attr_name in ("deque", "defaultdict", "Counter", "OrderedDict", "ChainMap"):
elif attr_name in (
"deque",
"defaultdict",
"Counter",
"OrderedDict",
"ChainMap",
):
is_data_structure = True
structure_type = attr_name
if is_data_structure and node.targets and isinstance(node.targets[0], ast.Name):
scope = context.get("current_scope", "function")
# Only track if it's at class level (not module level)
if scope == "class":
operations.append({
"line": node.lineno,
"var_name": node.targets[0].id,
"structure_type": structure_type,
"scope": scope,
"call": context["get_call_string"](node.value) if isinstance(node.value, ast.Call) else f"{structure_type}()",
})
operations.append(
{
"line": node.lineno,
"var_name": node.targets[0].id,
"structure_type": structure_type,
"scope": scope,
"call": (
context["get_call_string"](node.value)
if isinstance(node.value, ast.Call)
else f"{structure_type}()"
),
}
)
return operations
def check_cleanup(self, operations: List[Dict[str, Any]], function_body: List[ast.stmt],
context: Dict[str, Any]) -> List[Dict[str, Any]]:
def check_cleanup(
self,
operations: List[Dict[str, Any]],
function_body: List[ast.stmt],
context: Dict[str, Any],
) -> List[Dict[str, Any]]:
"""Flag persistent data structures that have add operations without size limits"""
violations = []
current_function = context["current_function"]
current_scope = context.get("current_scope", "function")
file_path = context["file_path"]
get_attr_string = context["get_attr_string"]
# Skip if this is initialization code (module-level, class-level, or __init__ methods)
# Only flag operations in regular methods/functions that can be called during runtime
is_initialization = (
current_scope in ("module", "class") or
current_function in ("__init__", "__new__", "__class_init__") or
current_function is None # Module-level code
current_scope in ("module", "class")
or current_function in ("__init__", "__new__", "__class_init__")
or current_function is None # Module-level code
)
if is_initialization:
return violations # Don't flag initialization code
# Track which variables have add operations and size checks
var_add_operations = {} # var_name -> list of lines with add operations
var_size_checks = {} # var_name -> has size limit check
# Build a set of variable names to check
tracked_vars = {op["var_name"]: op for op in operations}
# Scan body for operations on these variables
for stmt in function_body:
for node in ast.walk(stmt):
@@ -205,38 +255,50 @@ class UnboundedDataStructurePattern(Pattern):
if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute):
attr_name = node.func.attr
obj_name = get_attr_string(node.func.value)
# Check if this is an add operation on one of our tracked variables
for var_name, op in tracked_vars.items():
structure_type = op["structure_type"]
# Match variable name (exact or as attribute)
if obj_name == var_name or obj_name.endswith(f".{var_name}") or obj_name.endswith(f"['{var_name}']"):
if (
obj_name == var_name
or obj_name.endswith(f".{var_name}")
or obj_name.endswith(f"['{var_name}']")
):
# Check for add operations
add_ops = {
"list": ["append", "extend", "insert"],
"dict": ["update", "setdefault"],
"set": ["add", "update"],
"deque": ["append", "appendleft", "extend", "extendleft", "insert"],
"deque": [
"append",
"appendleft",
"extend",
"extendleft",
"insert",
],
"defaultdict": ["update", "setdefault"],
"Counter": ["update"],
"OrderedDict": ["update", "setdefault"],
"ChainMap": ["new_child"],
"queue": ["put", "put_nowait"],
}
if attr_name in add_ops.get(structure_type, []):
if var_name not in var_add_operations:
var_add_operations[var_name] = []
var_add_operations[var_name].append(node.lineno)
# Check for size limit checks (len() calls, maxsize/maxlen attributes)
if (attr_name in ("__len__",) or
"maxsize" in attr_name.lower() or
"max_size" in attr_name.lower() or
attr_name == "maxlen"): # For deque
if (
attr_name in ("__len__",)
or "maxsize" in attr_name.lower()
or "max_size" in attr_name.lower()
or attr_name == "maxlen"
): # For deque
var_size_checks[var_name] = True
# Check for heapq operations on tracked lists (heapq.heappush, heapq.heappop)
if isinstance(node, ast.Call):
func = node.func
@@ -245,67 +307,112 @@ class UnboundedDataStructurePattern(Pattern):
func_obj = get_attr_string(func.value)
func_name = func.attr
# Check if it's a heapq operation
if func_obj == "heapq" and func_name in ("heappush", "heapreplace", "heappushpop"):
if func_obj == "heapq" and func_name in (
"heappush",
"heapreplace",
"heappushpop",
):
# First argument should be our tracked variable
if len(node.args) > 0:
arg_name = get_attr_string(node.args[0])
for var_name, op in tracked_vars.items():
if op["structure_type"] == "list" and (
arg_name == var_name or arg_name.endswith(f".{var_name}")
arg_name == var_name
or arg_name.endswith(f".{var_name}")
):
if var_name not in var_add_operations:
var_add_operations[var_name] = []
var_add_operations[var_name].append(node.lineno)
# Check for dict item assignment: dict[key] = value
if isinstance(node, ast.Assign):
for target in node.targets:
if isinstance(target, ast.Subscript):
target_name = get_attr_string(target.value)
for var_name in tracked_vars:
if target_name == var_name or target_name.endswith(f".{var_name}"):
if target_name == var_name or target_name.endswith(
f".{var_name}"
):
if var_name not in var_add_operations:
var_add_operations[var_name] = []
var_add_operations[var_name].append(node.lineno)
# Check for augmented assignment: list += [...]
if isinstance(node, ast.AugAssign):
target_name = get_attr_string(node.target)
for var_name in tracked_vars:
if target_name == var_name or target_name.endswith(f".{var_name}"):
if target_name == var_name or target_name.endswith(
f".{var_name}"
):
if var_name not in var_add_operations:
var_add_operations[var_name] = []
var_add_operations[var_name].append(node.lineno)
# Check for size comparisons in conditionals
if isinstance(node, (ast.If, ast.While, ast.Assert)):
test = getattr(node, "test", None)
if test:
for comp_node in ast.walk(test):
if isinstance(comp_node, ast.Compare):
left_str = get_attr_string(comp_node.left) if hasattr(comp_node, "left") else ""
left_str = (
get_attr_string(comp_node.left)
if hasattr(comp_node, "left")
else ""
)
# Check for len() calls
if isinstance(comp_node.left, ast.Call):
call_func = comp_node.left.func
if isinstance(call_func, ast.Name) and call_func.id == "len":
if (
isinstance(call_func, ast.Name)
and call_func.id == "len"
):
if len(comp_node.left.args) > 0:
arg_name = get_attr_string(comp_node.left.args[0])
arg_name = get_attr_string(
comp_node.left.args[0]
)
for var_name in tracked_vars:
if arg_name == var_name or arg_name.endswith(f".{var_name}"):
if (
arg_name == var_name
or arg_name.endswith(f".{var_name}")
):
# Check if comparing to a limit
for comparator in comp_node.comparators:
if isinstance(comparator, ast.Constant):
var_size_checks[var_name] = True
elif isinstance(comparator, ast.Name):
for (
comparator
) in comp_node.comparators:
if isinstance(
comparator, ast.Constant
):
var_size_checks[
var_name
] = True
elif isinstance(
comparator, ast.Name
):
# Could be a constant like MAX_SIZE
if "max" in comparator.id.lower() or "limit" in comparator.id.lower():
var_size_checks[var_name] = True
if (
"max"
in comparator.id.lower()
or "limit"
in comparator.id.lower()
):
var_size_checks[
var_name
] = True
# Handle deprecated ast.Num for Python < 3.8
try:
Num = getattr(ast, "Num", None)
if Num and isinstance(comparator, Num):
var_size_checks[var_name] = True
except (AttributeError, TypeError):
Num = getattr(
ast, "Num", None
)
if Num and isinstance(
comparator, Num
):
var_size_checks[
var_name
] = True
except (
AttributeError,
TypeError,
):
pass
# Check for direct variable comparisons
for var_name in tracked_vars:
@@ -320,51 +427,60 @@ class UnboundedDataStructurePattern(Pattern):
var_size_checks[var_name] = True
except (AttributeError, TypeError):
pass
# Flag violations: persistent structures with add operations but no size checks
for op in operations:
var_name = op["var_name"]
structure_type = op["structure_type"]
if var_name in var_add_operations and var_name not in var_size_checks:
violations.append({
"line": op["line"],
"type": self.get_pattern_name(),
"var_name": var_name,
"function": current_function or "class-level",
"file_path": file_path,
"message": (
f"Class-level {structure_type} '{var_name}' "
f"has add operations (lines {var_add_operations[var_name]}) but no size limit checks. "
f"This can lead to unbounded memory growth and OOM errors during runtime."
),
})
violations.append(
{
"line": op["line"],
"type": self.get_pattern_name(),
"var_name": var_name,
"function": current_function or "class-level",
"file_path": file_path,
"message": (
f"Class-level {structure_type} '{var_name}' "
f"has add operations (lines {var_add_operations[var_name]}) but no size limit checks. "
f"This can lead to unbounded memory growth and OOM errors during runtime."
),
}
)
return violations
class MemoryViolationDetector(ast.NodeVisitor):
"""AST visitor that detects memory violations using registered patterns"""
DEFAULT_PATTERNS: List[Pattern] = [QueueGetPattern(), UnboundedDataStructurePattern()]
DEFAULT_PATTERNS: List[Pattern] = [
QueueGetPattern(),
UnboundedDataStructurePattern(),
]
def __init__(self, file_path: str, patterns: Optional[Sequence[Pattern]] = None):
self.file_path = file_path
self.violations: List[Dict[str, Any]] = []
self.current_function: Optional[str] = None
self.current_scope: str = "module" # Track current scope: module, class, function
self.current_scope: str = (
"module" # Track current scope: module, class, function
)
self.patterns = self.DEFAULT_PATTERNS if patterns is None else patterns
self.ast_tree: Optional[ast.Module] = None # Store full AST for module-level checks
self.ast_tree: Optional[ast.Module] = (
None # Store full AST for module-level checks
)
self.pattern_operations: Dict[str, List[Dict[str, Any]]] = {
pattern.get_pattern_name(): [] for pattern in self.patterns
}
# Track class-level operations separately (for checking in functions)
self.class_level_operations: Dict[str, List[Dict[str, Any]]] = {
pattern.get_pattern_name(): [] for pattern in self.patterns
}
self._context = {
"get_call_string": self._get_call_string,
"get_attr_string": self._get_attr_string,
@@ -379,9 +495,9 @@ class MemoryViolationDetector(ast.NodeVisitor):
old_scope = self.current_scope
self.current_scope = "class"
self._context["current_scope"] = "class"
self.generic_visit(node)
self.current_scope = old_scope
self._context["current_scope"] = old_scope
@@ -393,13 +509,13 @@ class MemoryViolationDetector(ast.NodeVisitor):
self.current_scope = "function"
self._context["current_function"] = node.name
self._context["current_scope"] = "function"
for pattern_name in self.pattern_operations:
self.pattern_operations[pattern_name] = []
self.generic_visit(node)
self._check_function_cleanup(node)
self.current_function = old_function
self.current_scope = old_scope
self._context["current_function"] = old_function
@@ -413,13 +529,13 @@ class MemoryViolationDetector(ast.NodeVisitor):
self.current_scope = "function"
self._context["current_function"] = node.name
self._context["current_scope"] = "function"
for pattern_name in self.pattern_operations:
self.pattern_operations[pattern_name] = []
self.generic_visit(node)
self._check_function_cleanup(node)
self.current_function = old_function
self.current_scope = old_scope
self._context["current_function"] = old_function
@@ -435,7 +551,7 @@ class MemoryViolationDetector(ast.NodeVisitor):
for op in operations:
if op.get("scope") == "class":
self.class_level_operations[pattern.get_pattern_name()].append(op)
self.generic_visit(node)
def _check_function_cleanup(self, node):
@@ -445,15 +561,22 @@ class MemoryViolationDetector(ast.NodeVisitor):
if operations:
violations = pattern.check_cleanup(operations, node.body, self._context)
self.violations.extend(violations)
# For UnboundedDataStructurePattern, also check if this function modifies class-level structures
if isinstance(pattern, UnboundedDataStructurePattern):
class_ops = self.class_level_operations[pattern.get_pattern_name()]
if class_ops and self.current_function not in ("__init__", "__new__", "__class_init__", None):
if class_ops and self.current_function not in (
"__init__",
"__new__",
"__class_init__",
None,
):
# Check if this regular function modifies class-level structures
violations = pattern.check_cleanup(class_ops, node.body, self._context)
violations = pattern.check_cleanup(
class_ops, node.body, self._context
)
self.violations.extend(violations)
def _check_module_level_cleanup(self):
"""Check cleanup for module/class level operations"""
# Module-level operations are now checked when visiting functions
@@ -475,20 +598,29 @@ class MemoryViolationDetector(ast.NodeVisitor):
break
if assignment_line:
break
if not assignment_line:
return False
for stmt in body:
for node in ast.walk(stmt):
if isinstance(node, ast.Assign):
for target in node.targets:
if isinstance(target, ast.Name) and target.id == var_name and node.lineno > assignment_line:
if isinstance(node.value, ast.Constant) and node.value.value is None:
if (
isinstance(target, ast.Name)
and target.id == var_name
and node.lineno > assignment_line
):
if (
isinstance(node.value, ast.Constant)
and node.value.value is None
):
return True
try:
NameConstant = getattr(ast, "NameConstant", None)
if NameConstant and isinstance(node.value, NameConstant):
if NameConstant and isinstance(
node.value, NameConstant
):
if getattr(node.value, "value", None) is None:
return True
except (AttributeError, TypeError):
@@ -515,15 +647,17 @@ class MemoryViolationDetector(ast.NodeVisitor):
return str(node)
def check_file_for_memory_violations(file_path: str, patterns: Optional[Sequence[Pattern]] = None) -> List[Dict[str, Any]]:
def check_file_for_memory_violations(
file_path: str, patterns: Optional[Sequence[Pattern]] = None
) -> List[Dict[str, Any]]:
"""Check a single file for memory violations"""
try:
with open(file_path, "r", encoding="utf-8") as f:
content = f.read()
if "test" in file_path.lower() or "__pycache__" in file_path:
return []
tree = ast.parse(content, filename=file_path)
detector = MemoryViolationDetector(file_path, patterns)
detector.ast_tree = tree # Store AST for potential future use
@@ -535,19 +669,34 @@ def check_file_for_memory_violations(file_path: str, patterns: Optional[Sequence
return []
def check_directory_for_memory_violations(directory_path: str, ignore_patterns: Optional[List[str]] = None,
patterns: Optional[Sequence[Pattern]] = None) -> List[Dict[str, Any]]:
def check_directory_for_memory_violations(
directory_path: str,
ignore_patterns: Optional[List[str]] = None,
patterns: Optional[Sequence[Pattern]] = None,
) -> List[Dict[str, Any]]:
"""Recursively scan directory for memory violations"""
if ignore_patterns is None:
ignore_patterns = ["__pycache__", ".pyc", "site-packages", "venv", ".venv", "env", ".env", "node_modules", "tests"]
ignore_patterns = [
"__pycache__",
".pyc",
"site-packages",
"venv",
".venv",
"env",
".env",
"node_modules",
"tests",
]
all_violations = []
for root, _dirs, files in os.walk(directory_path):
if any(pattern in root for pattern in ignore_patterns):
continue
for file in files:
if file.endswith(".py"):
violations = check_file_for_memory_violations(os.path.join(root, file), patterns)
violations = check_file_for_memory_violations(
os.path.join(root, file), patterns
)
all_violations.extend(violations)
return all_violations
@@ -555,16 +704,18 @@ def check_directory_for_memory_violations(directory_path: str, ignore_patterns:
def main():
"""Run memory violation detection on codebase"""
codebase_path = "./litellm"
print("=" * 80)
print("MEMORY VIOLATION DETECTION TEST")
print("=" * 80)
print(f"Scanning: {codebase_path}")
print(f"Active patterns: {', '.join(p.get_pattern_name() for p in MemoryViolationDetector.DEFAULT_PATTERNS)}")
print(
f"Active patterns: {', '.join(p.get_pattern_name() for p in MemoryViolationDetector.DEFAULT_PATTERNS)}"
)
print()
violations = check_directory_for_memory_violations(codebase_path)
if violations:
by_type = {}
for v in violations:
@@ -572,36 +723,44 @@ def main():
if vtype not in by_type:
by_type[vtype] = []
by_type[vtype].append(v)
print("MEMORY VIOLATIONS FOUND:")
print("=" * 80)
total = len(violations)
for vtype, vlist in by_type.items():
print(f"\n{vtype.upper().replace('_', ' ')}: {len(vlist)} violation(s)")
print("-" * 80)
for v in vlist[:10]:
print(f" [VIOLATION] {v['file_path'] if 'file_path' in v else 'unknown'}:{v['line']}")
print(
f" [VIOLATION] {v['file_path'] if 'file_path' in v else 'unknown'}:{v['line']}"
)
print(f" Function: {v['function']}")
print(f" Variable: {v['var_name']}")
print(f" {v['message']}")
print()
if len(vlist) > 10:
print(f" ... and {len(vlist) - 10} more violations of this type")
print("=" * 80)
print(f"TOTAL VIOLATIONS: {total}")
print()
print("RECOMMENDATIONS:")
print(" 1. Set queue variables to None after use: obj = queue.get(); ...; obj = None")
print(
" 1. Set queue variables to None after use: obj = queue.get(); ...; obj = None"
)
print(" 2. Use bounded queues to prevent unbounded accumulation")
print(" 3. Process items faster than they're added, or drain queues periodically")
print(" 4. For class-level data structures (lists, dicts, sets) that are modified at runtime:")
print(
" 3. Process items faster than they're added, or drain queues periodically"
)
print(
" 4. For class-level data structures (lists, dicts, sets) that are modified at runtime:"
)
print(" - Add size limit checks: if len(data) >= MAX_SIZE: ...")
print(" - Implement periodic cleanup or use bounded collections")
print(" - Consider using collections.deque with maxlen for lists")
print("=" * 80)
first_v = violations[0]
raise Exception(
f"Found {total} memory violations! "
tests/code_coverage_tests/test_ban_set_verbose.py
+47 -30
View File
@@ -24,20 +24,34 @@ def find_set_verbose_assignments(file_path):
for target in node.targets:
if isinstance(target, ast.Attribute):
# Check if it's litellm.set_verbose
if (isinstance(target.value, ast.Name) and
target.value.id == "litellm" and
target.attr == "set_verbose"):
if (
isinstance(target.value, ast.Name)
and target.value.id == "litellm"
and target.attr == "set_verbose"
):
# Check if the value being assigned is True
if (isinstance(node.value, ast.Constant) and
node.value.value is True):
if (
isinstance(node.value, ast.Constant)
and node.value.value is True
):
line_num = node.lineno
line_text = content_lines[line_num - 1].strip() if line_num <= len(content_lines) else ""
line_text = (
content_lines[line_num - 1].strip()
if line_num <= len(content_lines)
else ""
)
assignments.append((line_num, line_text))
elif (isinstance(node.value, ast.NameConstant) and
node.value.value is True): # For older Python versions
elif (
isinstance(node.value, ast.NameConstant)
and node.value.value is True
): # For older Python versions
line_num = node.lineno
line_text = content_lines[line_num - 1].strip() if line_num <= len(content_lines) else ""
line_text = (
content_lines[line_num - 1].strip()
if line_num <= len(content_lines)
else ""
)
assignments.append((line_num, line_text))
return assignments
@@ -49,10 +63,7 @@ def scan_litellm_files(base_dir):
Returns a dictionary mapping file paths to lists of assignments.
"""
violations = {}
litellm_dirs = [
"litellm",
"enterprise"
]
litellm_dirs = ["litellm", "enterprise"]
for litellm_dir in litellm_dirs:
dir_path = os.path.join(base_dir, litellm_dir)
@@ -66,7 +77,7 @@ def scan_litellm_files(base_dir):
if file.endswith(".py"):
file_path = os.path.join(root, file)
relative_path = os.path.relpath(file_path, base_dir)
assignments = find_set_verbose_assignments(file_path)
if assignments:
violations[relative_path] = assignments
@@ -79,9 +90,9 @@ def test_no_hardcoded_set_verbose():
Pytest-compatible test function that ensures no hardcoded litellm.set_verbose = True assignments exist.
"""
base_dir = "./" # Adjust path as needed for your setup
violations = scan_litellm_files(base_dir)
if violations:
violation_details = []
total_violations = 0
@@ -89,14 +100,14 @@ def test_no_hardcoded_set_verbose():
for line_num, line_text in assignments:
violation_details.append(f"{file_path}:{line_num} -> {line_text}")
total_violations += 1
error_msg = (
f"Found {total_violations} prohibited litellm.set_verbose = True assignments:\n"
+ "\n".join(violation_details) +
"\n\nREASON: litellm.set_verbose = True should not be hardcoded in production code. "
+ "\n".join(violation_details)
+ "\n\nREASON: litellm.set_verbose = True should not be hardcoded in production code. "
"Instead, use environment variables or configuration files to control verbosity."
)
raise AssertionError(error_msg)
@@ -105,29 +116,35 @@ def main():
Main function that scans for litellm.set_verbose = True assignments and fails if any are found.
"""
base_dir = "./" # Adjust path as needed for your setup
print("Scanning for litellm.set_verbose = True assignments...")
violations = scan_litellm_files(base_dir)
if violations:
print("\n❌ FOUND PROHIBITED litellm.set_verbose = True ASSIGNMENTS:")
print("=" * 60)
total_violations = 0
for file_path, assignments in violations.items():
print(f"\nFile: {file_path}")
for line_num, line_text in assignments:
print(f" Line {line_num}: {line_text}")
total_violations += 1
print(f"\n📊 Total violations found: {total_violations}")
print("\n🚫 REASON: litellm.set_verbose = True should not be hardcoded in production code.")
print(" Instead, use environment variables or configuration files to control verbosity.")
print(
"\n🚫 REASON: litellm.set_verbose = True should not be hardcoded in production code."
)
print(
" Instead, use environment variables or configuration files to control verbosity."
)
print(" Example alternatives:")
print(" - Use LITELLM_LOG=DEBUG environment variable")
print(" - Use litellm.set_verbose = os.getenv('LITELLM_VERBOSE', 'false').lower() == 'true'")
print(
" - Use litellm.set_verbose = os.getenv('LITELLM_VERBOSE', 'false').lower() == 'true'"
)
print(" - Use configuration-based verbosity settings")
raise Exception(
f"Found {total_violations} prohibited litellm.set_verbose = True assignments. "
"Remove these hardcoded verbosity settings and use configuration-based approaches instead."
@@ -137,4 +154,4 @@ def main():
if __name__ == "__main__":
main()
main()
tests/code_coverage_tests/test_chat_completion_imports.py
+19 -13
View File
@@ -8,36 +8,42 @@ from pathlib import Path
def test_chat_completion_no_imports():
"""Test that chat_completion endpoint has no imports in function bodies."""
# Path to the proxy server file
proxy_server_path = Path(__file__).parent.parent.parent / "litellm" / "proxy" / "proxy_server.py"
with open(proxy_server_path, 'r') as f:
proxy_server_path = (
Path(__file__).parent.parent.parent / "litellm" / "proxy" / "proxy_server.py"
)
with open(proxy_server_path, "r") as f:
content = f.read()
# Parse the AST
tree = ast.parse(content)
# Find the chat_completion function
chat_completion_func = None
for node in ast.walk(tree):
if (isinstance(node, ast.AsyncFunctionDef) and node.name == "chat_completion"):
if isinstance(node, ast.AsyncFunctionDef) and node.name == "chat_completion":
chat_completion_func = node
break
assert chat_completion_func is not None, "chat_completion function not found"
# Check for imports inside the function body
import_violations = []
for node in ast.walk(chat_completion_func):
if isinstance(node, (ast.Import, ast.ImportFrom)):
# Get line number
line_num = node.lineno
import_violations.append(line_num)
# Assert no import violations found
if import_violations:
print(f"Found {len(import_violations)} import violations in chat_completion endpoint:")
print(
f"Found {len(import_violations)} import violations in chat_completion endpoint:"
)
for line_num in import_violations:
print(f" - Line {line_num}: Import statement found")
print("\nchat_completion endpoint should not contain imports for optimal performance.")
raise Exception("Import violations found in chat_completion endpoint")
print(
"\nchat_completion endpoint should not contain imports for optimal performance."
)
raise Exception("Import violations found in chat_completion endpoint")
tests/code_coverage_tests/test_proxy_types_import.py
+49 -31
View File
@@ -13,50 +13,64 @@ def test_proxy_types_not_imported():
init_file_path = os.path.join("./litellm", "__init__.py")
if not os.path.exists(init_file_path):
raise Exception(f"Could not find {init_file_path}")
with open(init_file_path, "r") as f:
content = f.read()
lines = content.splitlines() # Get lines for line number reporting
try:
tree = ast.parse(content)
except SyntaxError as e:
raise Exception(f"Could not parse {init_file_path}: {e}")
# Check for direct imports of proxy._types
found_imports = []
for node in ast.walk(tree):
if isinstance(node, ast.Import):
for alias in node.names:
if "proxy._types" in alias.name or "proxy/_types" in alias.name:
line_num = node.lineno
line_content = lines[line_num - 1] if line_num <= len(lines) else "Unknown"
line_content = (
lines[line_num - 1] if line_num <= len(lines) else "Unknown"
)
import_statement = f"import {alias.name}"
found_imports.append({
'type': 'import',
'line': line_num,
'content': line_content.strip(),
'statement': import_statement,
'module': alias.name
})
found_imports.append(
{
"type": "import",
"line": line_num,
"content": line_content.strip(),
"statement": import_statement,
"module": alias.name,
}
)
elif isinstance(node, ast.ImportFrom):
if node.module and ("proxy._types" in node.module or "proxy/_types" in node.module):
if node.module and (
"proxy._types" in node.module or "proxy/_types" in node.module
):
line_num = node.lineno
line_content = lines[line_num - 1] if line_num <= len(lines) else "Unknown"
line_content = (
lines[line_num - 1] if line_num <= len(lines) else "Unknown"
)
import_names = [alias.name for alias in node.names]
import_statement = f"from {node.module} import {', '.join(import_names)}"
found_imports.append({
'type': 'from_import',
'line': line_num,
'content': line_content.strip(),
'statement': import_statement,
'module': node.module
})
import_statement = (
f"from {node.module} import {', '.join(import_names)}"
)
found_imports.append(
{
"type": "from_import",
"line": line_num,
"content": line_content.strip(),
"statement": import_statement,
"module": node.module,
}
)
if found_imports:
print("❌ BAD, this can import time to import litellm. Found direct imports of proxy._types in litellm/__init__.py:")
print(
"❌ BAD, this can import time to import litellm. Found direct imports of proxy._types in litellm/__init__.py:"
)
print("=" * 80)
for imp in found_imports:
print(f"Line {imp['line']}: {imp['content']}")
@@ -65,11 +79,11 @@ def test_proxy_types_not_imported():
print(f" Module: {imp['module']}")
print("-" * 80)
print("To fix this, please conditionally import this TYPE using TYPE_CHECKING")
raise Exception(
f"Found {len(found_imports)} direct import(s) of proxy._types in litellm/__init__.py"
)
print("✓ No direct imports of proxy._types found in litellm/__init__.py")
return True
@@ -80,13 +94,17 @@ def main():
"""
print("=" * 60)
print("Testing litellm import performance")
print("Checking that proxy._types is not directly imported from litellm/__init__.py")
print(
"Checking that proxy._types is not directly imported from litellm/__init__.py"
)
print("=" * 60)
try:
test_proxy_types_not_imported()
print("\n" + "=" * 60)
print("✓ Test passed! proxy._types is not directly imported from litellm/__init__.py")
print(
"✓ Test passed! proxy._types is not directly imported from litellm/__init__.py"
)
print("=" * 60)
except Exception as e:
print(f"\n❌ Test failed: {e}")
@@ -95,4 +113,4 @@ def main():
if __name__ == "__main__":
main()
main()