mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-08 08:08:37 +00:00
ab78459097
Response transformation crashed with TypeError when counting tokens for multimodal (dict-shaped) embedding inputs, since encoding.encode() was called directly on the dict instead of its text parts. Also removed a dead, equally-buggy duplicate of the same loop in _transform_response. Separately, the embeddings guardrail handler silently skipped multimodal dict inputs instead of scanning them, letting prohibited text/image content bypass the configured guardrail. It now extracts and scans the text parts while leaving image parts untouched.
133 lines
4.2 KiB
Python
133 lines
4.2 KiB
Python
import os
|
|
import re
|
|
import ast
|
|
from pathlib import Path
|
|
|
|
|
|
class DataReplaceVisitor(ast.NodeVisitor):
|
|
"""AST visitor that finds calls to .replace("data:", ...) in the code."""
|
|
|
|
def __init__(self):
|
|
self.issues = []
|
|
self.current_file = None
|
|
|
|
def set_file(self, filename):
|
|
self.current_file = filename
|
|
|
|
def visit_Call(self, node):
|
|
# Check for method calls like x.replace(...)
|
|
if isinstance(node.func, ast.Attribute) and node.func.attr == "replace":
|
|
# Check if first argument is "data:"
|
|
if (
|
|
len(node.args) >= 2
|
|
and isinstance(node.args[0], ast.Constant)
|
|
and isinstance(node.args[0].value, str)
|
|
and "data:" in node.args[0].value
|
|
):
|
|
|
|
self.issues.append(
|
|
{
|
|
"file": self.current_file,
|
|
"line": node.lineno,
|
|
"col": node.col_offset,
|
|
"text": f'Found .replace("data:", ...) at line {node.lineno}',
|
|
}
|
|
)
|
|
|
|
# Continue visiting child nodes
|
|
self.generic_visit(node)
|
|
|
|
|
|
def check_file_with_ast(file_path):
|
|
"""Check a Python file for .replace("data:", ...) using AST parsing."""
|
|
with open(file_path, "r", encoding="utf-8") as f:
|
|
try:
|
|
tree = ast.parse(f.read(), filename=file_path)
|
|
visitor = DataReplaceVisitor()
|
|
visitor.set_file(file_path)
|
|
visitor.visit(tree)
|
|
return visitor.issues
|
|
except SyntaxError:
|
|
return [
|
|
{
|
|
"file": file_path,
|
|
"line": 0,
|
|
"col": 0,
|
|
"text": f"Syntax error in file, could not parse",
|
|
}
|
|
]
|
|
|
|
|
|
def check_file_with_regex(file_path):
|
|
"""Check any file for .replace("data:", ...) using regex."""
|
|
issues = []
|
|
with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
|
|
for i, line in enumerate(f, 1):
|
|
matches = re.finditer(r'\.replace\(\s*[\'"]data:[\'"]', line)
|
|
for match in matches:
|
|
issues.append(
|
|
{
|
|
"file": file_path,
|
|
"line": i,
|
|
"col": match.start(),
|
|
"text": f'Found .replace("data:", ...) at line {i}',
|
|
}
|
|
)
|
|
return issues
|
|
|
|
|
|
def scan_directory(base_dir):
|
|
"""Scan a directory recursively for files containing .replace("data:", ...)."""
|
|
all_issues = []
|
|
|
|
for root, _, files in os.walk(base_dir):
|
|
for file in files:
|
|
print("checking file: ", file)
|
|
file_path = os.path.join(root, file)
|
|
|
|
# Skip directories we don't want to check
|
|
if any(
|
|
d in file_path for d in [".git", "__pycache__", ".venv", "node_modules"]
|
|
):
|
|
continue
|
|
|
|
# For Python files, use AST for more accurate parsing
|
|
if file.endswith(".py"):
|
|
issues = check_file_with_ast(file_path)
|
|
# For other files that might contain code, use regex
|
|
elif file.endswith((".js", ".ts", ".jsx", ".tsx", ".md", ".ipynb")):
|
|
issues = check_file_with_regex(file_path)
|
|
else:
|
|
continue
|
|
|
|
all_issues.extend(issues)
|
|
|
|
return all_issues
|
|
|
|
|
|
def main():
|
|
# Start from the project root directory
|
|
|
|
base_dir = "./litellm"
|
|
|
|
# Local testing
|
|
# base_dir = "../../litellm"
|
|
|
|
print(f"Scanning for .replace('data:', ...) usage in {base_dir}")
|
|
issues = scan_directory(base_dir)
|
|
|
|
if issues:
|
|
print(f"\n⚠️ Found {len(issues)} instances of .replace('data:', ...):")
|
|
for issue in issues:
|
|
print(f"{issue['file']}:{issue['line']} - {issue['text']}")
|
|
|
|
# Fail the test if issues are found
|
|
raise Exception(
|
|
f"Found {len(issues)} instances of .replace('data:', ...) which may be unsafe. Use litellm.CustomStreamWrapper._strip_sse_data_from_chunk instead."
|
|
)
|
|
else:
|
|
print("✅ No instances of .replace('data:', ...) found.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|