import json

from langchain.schema import OutputParserException


def parse_json_markdown(json_string: str) -> dict:
    # Remove the triple backticks if present
    json_string = json_string.strip()
    start_index = json_string.find("```json")
    end_index = json_string.find("```", start_index + len("```json"))

    if start_index != -1 and end_index != -1:
        extracted_content = json_string[start_index + len("```json"):end_index].strip()

        # Parse the JSON string into a Python dictionary
        parsed = json.loads(extracted_content)
    elif start_index != -1 and end_index == -1 and json_string.endswith("``"):
        end_index = json_string.find("``", start_index + len("```json"))
        extracted_content = json_string[start_index + len("```json"):end_index].strip()

        # Parse the JSON string into a Python dictionary
        parsed = json.loads(extracted_content)
    elif json_string.startswith("{"):
        # Parse the JSON string into a Python dictionary
        parsed = json.loads(json_string)
    else:
        raise Exception("Could not find JSON block in the output.")

    return parsed


def parse_and_check_json_markdown(text: str, expected_keys: list[str]) -> dict:
    try:
        json_obj = parse_json_markdown(text)
    except json.JSONDecodeError as e:
        raise OutputParserException(f"Got invalid JSON object. Error: {e}")
    for key in expected_keys:
        if key not in json_obj:
            raise OutputParserException(
                f"Got invalid return object. Expected key `{key}` "
                f"to be present, but got {json_obj}"
            )
    return json_obj

langchain.schema.OutputParserException
[((526, 555), 'json.loads', 'json.loads', (['extracted_content'], {}), '(extracted_content)\n', (536, 555), False, 'import json\n'), ((871, 900), 'json.loads', 'json.loads', (['extracted_content'], {}), '(extracted_content)\n', (881, 900), False, 'import json\n'), ((1322, 1383), 'langchain.schema.OutputParserException', 'OutputParserException', (['f"""Got invalid JSON object. Error: {e}"""'], {}), "(f'Got invalid JSON object. Error: {e}')\n", (1343, 1383), False, 'from langchain.schema import OutputParserException\n'), ((1464, 1581), 'langchain.schema.OutputParserException', 'OutputParserException', (['f"""Got invalid return object. Expected key `{key}` to be present, but got {json_obj}"""'], {}), "(\n    f'Got invalid return object. Expected key `{key}` to be present, but got {json_obj}'\n    )\n", (1485, 1581), False, 'from langchain.schema import OutputParserException\n'), ((1013, 1036), 'json.loads', 'json.loads', (['json_string'], {}), '(json_string)\n', (1023, 1036), False, 'import json\n')]