OpenAI JSON Cleaner

json_cleaner.py

def clean_and_parse_json_with_ai(json_string: str, openai_key: str):
    clean_json_fn = {
        "name": "clean_json",
        "description": "Cleans JSON.",
        "parameters": {
            "type": "object",
            "properties": {
                "clean_json": {
                    "type": "string",
                    "description": "The cleaned JSON.",
                },
            },
            "required": ["clean_json"],
        },
    }
    openai.api_key = openai_key
    system_prompt = "You are a helpful assistant."
    messages = [
        {
            "role": "system",
            "content": system_prompt,
        },
        {
            "role": "user",
            "content": f"The bad JSON: {json_string}",
        },
    ]
    openai_response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=messages,
        functions=[clean_json_fn],
        function_call={"name": clean_json_fn["name"]},
    )

    # the arguments come back as text, so those get converted to JSON
    arguments_raw = openai_response.choices[0].message.function_call.arguments
    arguments = json.loads(arguments_raw)
    # Then the cleaned json is also text, which needs to be loaded
    cleaned_json_raw = arguments["clean_json"]
    cleaned_json = json.loads(cleaned_json_raw)  
    return cleaned_json