From af99aab44a8c546f77b1dd2c8ccedcf2e2cef114 Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Tue, 13 Jan 2026 16:34:08 -0800 Subject: [PATCH] Add code quality check to enforce O(1) operations in _get_model_cost_key - Add check_get_model_cost_key_performance.py to statically analyze _get_model_cost_key - Detects O(n) operations (loops, comprehensions, problematic function calls) - Recursively checks called functions to find nested O(n) operations - Allows conditional O(n) rebuilds in helper functions (_rebuild_model_cost_lowercase_map, _handle_stale_map_entry_rebuild, _handle_new_key_with_scan) --- .../check_get_model_cost_key_performance.py | 178 ++++++++++++++++++ 1 file changed, 178 insertions(+) create mode 100644 tests/code_coverage_tests/check_get_model_cost_key_performance.py diff --git a/tests/code_coverage_tests/check_get_model_cost_key_performance.py b/tests/code_coverage_tests/check_get_model_cost_key_performance.py new file mode 100644 index 00000000000..30457e06f6c --- /dev/null +++ b/tests/code_coverage_tests/check_get_model_cost_key_performance.py @@ -0,0 +1,178 @@ +""" +Code quality check: Ensure _get_model_cost_key only uses O(1) operations. + +Simple pattern-based check for O(n) operations in _get_model_cost_key. +""" + +import re +import os + + +def _function_has_on_operations(all_lines, func_name, visited=None): + """ + Check if a function contains O(n) operations by searching for it in the file. + Recursively checks called functions as well. + """ + 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: + 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 '): + 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("'''"): + continue + + # Check for for loops + if re.search(r'\bfor\s+\w+\s+in\s+', line): + return True + # Check for while loops + 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): + 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) + if func_call_match: + called_func = func_call_match.group(1) + if called_func.startswith('_'): + if _function_has_on_operations(all_lines, called_func, visited): + return True + + return False + + +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: + 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 '): + 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 = [] + + # Check for O(n) patterns + for i, line in enumerate(func_lines, start=func_start + 1): + line_stripped = line.strip() + + # Skip comments and docstrings + 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): + # 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): + problematic_lines.append((i, "for loop", line_stripped)) + + # Check for while loops + 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): + problematic_lines.append((i, "comprehension", line_stripped)) + + # Check for problematic function calls + 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): + 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', + ] + + # 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 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('_'): + # 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)) + + 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("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.") + + +if __name__ == "__main__": + main()