This commit is contained in:
mrT23
2024-02-18 07:56:14 +02:00
parent 481a4fe7a1
commit c6cb0524b4

View File

@ -320,7 +320,7 @@ def _fix_key_value(key: str, value: str):
def load_yaml(response_text: str, keys_fix_yaml: List[str] = []) -> dict:
response_text = response_text.removeprefix('```yaml').rstrip('`').rstrip(':\n')
response_text = response_text.removeprefix('```yaml').rstrip('`')
try:
data = yaml.safe_load(response_text)
except Exception as e:
@ -361,9 +361,9 @@ def try_fix_yaml(response_text: str, keys_fix_yaml: List[str] = []) -> dict:
pass
# third fallback - try to remove leading and trailing curly brackets
response_text_copy = response_text.strip().rstrip().removeprefix('{').removesuffix('}')
response_text_copy = response_text.strip().rstrip().removeprefix('{').removesuffix('}').rstrip(':\n')
try:
data = yaml.safe_load(response_text_copy,)
data = yaml.safe_load(response_text_copy)
get_logger().info(f"Successfully parsed AI prediction after removing curly brackets")
return data
except:
@ -374,7 +374,7 @@ def try_fix_yaml(response_text: str, keys_fix_yaml: List[str] = []) -> dict:
for i in range(1, len(response_text_lines)):
response_text_lines_tmp = '\n'.join(response_text_lines[:-i])
try:
data = yaml.safe_load(response_text_lines_tmp,)
data = yaml.safe_load(response_text_lines_tmp)
get_logger().info(f"Successfully parsed AI prediction after removing {i} lines")
return data
except: