coding-alt/AutoGPT

1"""This module contains functions to fix JSON strings using general programmatic approaches, suitable for addressing2common JSON formatting issues."""3from __future__ import annotations4 5import contextlib6import json7import re8from typing import Optional9 10from autogpt.config import Config11from autogpt.json_utils.utilities import extract_char_position12 13CFG = Config()14 15 16def fix_invalid_escape(json_to_load: str, error_message: str) -> str:17    """Fix invalid escape sequences in JSON strings.18 19    Args:20        json_to_load (str): The JSON string.21        error_message (str): The error message from the JSONDecodeError22          exception.23 24    Returns:25        str: The JSON string with invalid escape sequences fixed.26    """27    while error_message.startswith("Invalid \\escape"):28        bad_escape_location = extract_char_position(error_message)29        json_to_load = (30            json_to_load[:bad_escape_location] + json_to_load[bad_escape_location + 1 :]31        )32        try:33            json.loads(json_to_load)34            return json_to_load35        except json.JSONDecodeError as e:36            if CFG.debug_mode:37                print("json loads error - fix invalid escape", e)38            error_message = str(e)39    return json_to_load40 41 42def balance_braces(json_string: str) -> Optional[str]:43    """44    Balance the braces in a JSON string.45 46    Args:47        json_string (str): The JSON string.48 49    Returns:50        str: The JSON string with braces balanced.51    """52 53    open_braces_count = json_string.count("{")54    close_braces_count = json_string.count("}")55 56    while open_braces_count > close_braces_count:57        json_string += "}"58        close_braces_count += 159 60    while close_braces_count > open_braces_count:61        json_string = json_string.rstrip("}")62        close_braces_count -= 163 64    with contextlib.suppress(json.JSONDecodeError):65        json.loads(json_string)66        return json_string67 68 69def add_quotes_to_property_names(json_string: str) -> str:70    """71    Add quotes to property names in a JSON string.72 73    Args:74        json_string (str): The JSON string.75 76    Returns:77        str: The JSON string with quotes added to property names.78    """79 80    def replace_func(match: re.Match) -> str:81        return f'"{match[1]}":'82 83    property_name_pattern = re.compile(r"(\w+):")84    corrected_json_string = property_name_pattern.sub(replace_func, json_string)85 86    try:87        json.loads(corrected_json_string)88        return corrected_json_string89    except json.JSONDecodeError as e:90        raise e91 92 93def correct_json(json_to_load: str) -> str:94    """95    Correct common JSON errors.96    Args:97        json_to_load (str): The JSON string.98    """99 100    try:101        if CFG.debug_mode:102            print("json", json_to_load)103        json.loads(json_to_load)104        return json_to_load105    except json.JSONDecodeError as e:106        if CFG.debug_mode:107            print("json loads error", e)108        error_message = str(e)109        if error_message.startswith("Invalid \\escape"):110            json_to_load = fix_invalid_escape(json_to_load, error_message)111        if error_message.startswith(112            "Expecting property name enclosed in double quotes"113        ):114            json_to_load = add_quotes_to_property_names(json_to_load)115            try:116                json.loads(json_to_load)117                return json_to_load118            except json.JSONDecodeError as e:119                if CFG.debug_mode:120                    print("json loads error - add quotes", e)121                error_message = str(e)122        if balanced_str := balance_braces(json_to_load):123            return balanced_str124    return json_to_load125