fix: strip ANSI codes from Python tool output and optimize highlighting
- Add comprehensive ECMA-48 ANSI pattern to strip escape sequences from output - Fix _truncate_line to strip ANSI before length calculation - Cache PythonLexer instance (was creating new one per call) - Memoize token color lookups to avoid repeated parent chain traversal
This commit is contained in:
@@ -14,6 +14,8 @@ from .registry import register_tool_renderer
|
|||||||
MAX_OUTPUT_LINES = 50
|
MAX_OUTPUT_LINES = 50
|
||||||
MAX_LINE_LENGTH = 200
|
MAX_LINE_LENGTH = 200
|
||||||
|
|
||||||
|
ANSI_PATTERN = re.compile(r"\x1b(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~]|\][^\x07]*\x07)")
|
||||||
|
|
||||||
STRIP_PATTERNS = [
|
STRIP_PATTERNS = [
|
||||||
r"\.\.\. \[(stdout|stderr|result|output|error) truncated at \d+k? chars\]",
|
r"\.\.\. \[(stdout|stderr|result|output|error) truncated at \d+k? chars\]",
|
||||||
]
|
]
|
||||||
@@ -25,31 +27,32 @@ def _get_style_colors() -> dict[Any, str]:
|
|||||||
return {token: f"#{style_def['color']}" for token, style_def in style if style_def["color"]}
|
return {token: f"#{style_def['color']}" for token, style_def in style if style_def["color"]}
|
||||||
|
|
||||||
|
|
||||||
|
@cache
|
||||||
|
def _get_lexer() -> PythonLexer:
|
||||||
|
return PythonLexer()
|
||||||
|
|
||||||
|
|
||||||
|
@cache
|
||||||
|
def _get_token_color(token_type: Any) -> str | None:
|
||||||
|
colors = _get_style_colors()
|
||||||
|
while token_type:
|
||||||
|
if token_type in colors:
|
||||||
|
return colors[token_type]
|
||||||
|
token_type = token_type.parent
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
@register_tool_renderer
|
@register_tool_renderer
|
||||||
class PythonRenderer(BaseToolRenderer):
|
class PythonRenderer(BaseToolRenderer):
|
||||||
tool_name: ClassVar[str] = "python_action"
|
tool_name: ClassVar[str] = "python_action"
|
||||||
css_classes: ClassVar[list[str]] = ["tool-call", "python-tool"]
|
css_classes: ClassVar[list[str]] = ["tool-call", "python-tool"]
|
||||||
|
|
||||||
@classmethod
|
|
||||||
def _get_token_color(cls, token_type: Any) -> str | None:
|
|
||||||
colors = _get_style_colors()
|
|
||||||
while token_type:
|
|
||||||
if token_type in colors:
|
|
||||||
return colors[token_type]
|
|
||||||
token_type = token_type.parent
|
|
||||||
return None
|
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def _highlight_python(cls, code: str) -> Text:
|
def _highlight_python(cls, code: str) -> Text:
|
||||||
lexer = PythonLexer()
|
|
||||||
text = Text()
|
text = Text()
|
||||||
|
for token_type, token_value in _get_lexer().get_tokens(code):
|
||||||
for token_type, token_value in lexer.get_tokens(code):
|
if token_value:
|
||||||
if not token_value:
|
text.append(token_value, style=_get_token_color(token_type))
|
||||||
continue
|
|
||||||
color = cls._get_token_color(token_type)
|
|
||||||
text.append(token_value, style=color)
|
|
||||||
|
|
||||||
return text
|
return text
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
@@ -59,11 +62,16 @@ class PythonRenderer(BaseToolRenderer):
|
|||||||
cleaned = re.sub(pattern, "", cleaned)
|
cleaned = re.sub(pattern, "", cleaned)
|
||||||
return cleaned.strip()
|
return cleaned.strip()
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def _strip_ansi(cls, text: str) -> str:
|
||||||
|
return ANSI_PATTERN.sub("", text)
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def _truncate_line(cls, line: str) -> str:
|
def _truncate_line(cls, line: str) -> str:
|
||||||
if len(line) > MAX_LINE_LENGTH:
|
clean_line = cls._strip_ansi(line)
|
||||||
return line[: MAX_LINE_LENGTH - 3] + "..."
|
if len(clean_line) > MAX_LINE_LENGTH:
|
||||||
return line
|
return clean_line[: MAX_LINE_LENGTH - 3] + "..."
|
||||||
|
return clean_line
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def _format_output(cls, output: str) -> Text:
|
def _format_output(cls, output: str) -> Text:
|
||||||
|
|||||||
Reference in New Issue
Block a user