AutoGPT/scripts/json_utils.py

import re
import json
from config import Config

cfg = Config()


def extract_char_position(error_message: str) -> int:
    """Extract the character position from the JSONDecodeError message.

    Args:
        error_message (str): The error message from the JSONDecodeError
          exception.

    Returns:
        int: The character position.
    """
    import re

    char_pattern = re.compile(r'\(char (\d+)\)')
    if match := char_pattern.search(error_message):
        return int(match[1])
    else:
        raise ValueError("Character position not found in the error message.")


def add_quotes_to_property_names(json_string: str) -> str:
    """
    Add quotes to property names in a JSON string.

    Args:
        json_string (str): The JSON string.

    Returns:
        str: The JSON string with quotes added to property names.
    """

    def replace_func(match):
        return f'"{match.group(1)}":'

    property_name_pattern = re.compile(r'(\w+):')
    corrected_json_string = property_name_pattern.sub(
        replace_func,
        json_string)

    try:
        json.loads(corrected_json_string)
        return corrected_json_string
    except json.JSONDecodeError as e:
        raise e


def balance_braces(json_string: str) -> str:
    """
    Balance the braces in a JSON string.

    Args:
        json_string (str): The JSON string.

    Returns:
        str: The JSON string with braces balanced.
    """

    open_braces_count = json_string.count('{')
    close_braces_count = json_string.count('}')

    while open_braces_count > close_braces_count:
        json_string += '}'
        close_braces_count += 1

    while close_braces_count > open_braces_count:
        json_string = json_string.rstrip('}')
        close_braces_count -= 1

    try:
        json.loads(json_string)
        return json_string
    except json.JSONDecodeError as e:
        pass


def fix_invalid_escape(json_str: str, error_message: str) -> str:
    while error_message.startswith('Invalid \\escape'):
        bad_escape_location = extract_char_position(error_message)
        json_str = json_str[:bad_escape_location] + \
            json_str[bad_escape_location + 1:]
        try:
            json.loads(json_str)
            return json_str
        except json.JSONDecodeError as e:
            if cfg.debug_mode:
                print('json loads error - fix invalid escape', e)
            error_message = str(e)
    return json_str


def correct_json(json_str: str) -> str:
    """
    Correct common JSON errors.

    Args:
        json_str (str): The JSON string.
    """

    try:
        if cfg.debug_mode:
            print("json", json_str)
        json.loads(json_str)
        return json_str
    except json.JSONDecodeError as e:
        if cfg.debug_mode:
            print('json loads error', e)
        error_message = str(e)
        if error_message.startswith('Invalid \\escape'):
            json_str = fix_invalid_escape(json_str, error_message)
        if error_message.startswith('Expecting property name enclosed in double quotes'):
            json_str = add_quotes_to_property_names(json_str)
            try:
                json.loads(json_str)
                return json_str
            except json.JSONDecodeError as e:
                if cfg.debug_mode:
                    print('json loads error - add quotes', e)
                error_message = str(e)
        if balanced_str := balance_braces(json_str):
            return balanced_str
    return json_str