Coverage for open_webui/utils/sanitize.py: 38%

1import re 

2 

3# ANSI escape code pattern - matches all common ANSI sequences 

4# This includes color codes, cursor movement, and other terminal control sequences 

5ANSI_ESCAPE_PATTERN = re.compile(r'\x1b\[[0-9;]*[A-Za-z]|\x1b\([AB]|\x1b[PX^_].*?\x1b\\|\x1b\].*?(?:\x07|\x1b\\)') 

6 

7 

8def strip_ansi_codes(text: str) -> str: 

9 """ 

10 Strip ANSI escape codes from text. 

11 

12 ANSI escape codes can be introduced by LLMs that include terminal 

13 color codes in their output. These codes cause syntax errors when 

14 the code is sent to Jupyter for execution. 

15 

16 Common ANSI codes include: 

17 - Color codes: \x1b[31m (red), \x1b[32m (green), etc. 

18 - Reset codes: \x1b[0m, \x1b[39m 

19 - Cursor movement: \x1b[1A, \x1b[2J, etc. 

20 """ 

21 return ANSI_ESCAPE_PATTERN.sub('', text) 

22 

23 

24def strip_markdown_code_fences(code: str) -> str: 

25 """ 

26 Strip markdown code fences if present. 

27 

28 This is a defensive, non-breaking change — if the code doesn't 

29 contain fences, it passes through unchanged. 

30 

31 Handles patterns like: 

32 - ```python 

33 - ```py 

34 - ``` 

35 """ 

36 code = code.strip() 

37 # Remove opening fence (```python, ```py, ``` etc.) 

38 code = re.sub(r'^```\w*\n?', '', code) 

39 # Remove closing fence 

40 code = re.sub(r'\n?```\s*$', '', code) 

41 return code.strip() 

42 

43 

44def sanitize_code(code: str) -> str: 

45 """ 

46 Sanitize code for execution by applying all necessary cleanup steps. 

47 

48 This is the recommended function to use before sending code to 

49 interpreters like Jupyter or Pyodide. 

50 

51 Steps applied: 

52 1. Strip ANSI escape codes (from LLM output) 

53 2. Strip markdown code fences (if model included them) 

54 """ 

55 code = strip_ansi_codes(code) 

56 code = strip_markdown_code_fences(code) 

57 return code