Update utils.py

more robust try_fix_yaml
This commit is contained in:
Xinyu Wu 2025-11-07 15:30:16 +08:00 committed by GitHub
parent 8dbc53271e
commit 3e9b15de15
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -812,14 +812,19 @@ def try_fix_yaml(response_text: str,
pass pass
# second fallback - try to extract only range from first ```yaml to the last ``` # second fallback - try to extract only range from first ```yaml to the last ```
snippet_pattern = r'```yaml([\s\S]*?)```(?=\s*$|")' snippet_pattern = r'```(yaml|yml)?([\s\S]*?)```(?=\s*$|")'
snippet = re.search(snippet_pattern, '\n'.join(response_text_lines_copy)) snippet = re.search(snippet_pattern, '\n'.join(response_text_lines_copy))
if not snippet: if not snippet:
snippet = re.search(snippet_pattern, response_text_original) # before we removed the "```" snippet = re.search(snippet_pattern, response_text_original) # before we removed the "```"
if snippet: if snippet:
snippet_text = snippet.group() snippet_text = snippet.group()
prefix = (
'```yaml'
if snippet_text.startswith('```yaml')
else ('```yml' if snippet_text.startswith('```yml') else '```')
)
try: try:
data = yaml.safe_load(snippet_text.removeprefix('```yaml').rstrip('`')) data = yaml.safe_load(snippet_text.removeprefix(prefix).rstrip('`'))
get_logger().info(f"Successfully parsed AI prediction after extracting yaml snippet") get_logger().info(f"Successfully parsed AI prediction after extracting yaml snippet")
return data return data
except: except: