11import re
22
3+ # ANSI escape code pattern - matches all common ANSI sequences
4+ # This includes color codes, cursor movement, and other terminal control sequences
5+ ANSI_ESCAPE_PATTERN = re .compile (r'\x1b\[[0-9;]*[A-Za-z]|\x1b\([AB]|\x1b[PX^_].*?\x1b\\|\x1b\].*?(?:\x07|\x1b\\)' )
6+
7+
8+ def 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+
323
424def strip_markdown_code_fences (code : str ) -> str :
525 """
@@ -19,3 +39,20 @@ def strip_markdown_code_fences(code: str) -> str:
1939 # Remove closing fence
2040 code = re .sub (r"\n?```\s*$" , "" , code )
2141 return code .strip ()
42+
43+
44+ def 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
58+
0 commit comments