Underground-Digital/Workflow-Engine

1import json2 3from core.llm_generator.output_parser.errors import OutputParserError4 5 6def parse_json_markdown(json_string: str) -> dict:7    # Get json from the backticks/braces8    json_string = json_string.strip()9    starts = ["```json", "```", "``", "`", "{"]10    ends = ["```", "``", "`", "}"]11    end_index = -112    for s in starts:13        start_index = json_string.find(s)14        if start_index != -1:15            if json_string[start_index] != "{":16                start_index += len(s)17            break18    if start_index != -1:19        for e in ends:20            end_index = json_string.rfind(e, start_index)21            if end_index != -1:22                if json_string[end_index] == "}":23                    end_index += 124                break25    if start_index != -1 and end_index != -1 and start_index < end_index:26        extracted_content = json_string[start_index:end_index].strip()27        print("content:", extracted_content, start_index, end_index)28        parsed = json.loads(extracted_content)29    else:30        raise Exception("Could not find JSON block in the output.")31 32    return parsed33 34 35def parse_and_check_json_markdown(text: str, expected_keys: list[str]) -> dict:36    try:37        json_obj = parse_json_markdown(text)38    except json.JSONDecodeError as e:39        raise OutputParserError(f"Got invalid JSON object. Error: {e}")40    for key in expected_keys:41        if key not in json_obj:42            raise OutputParserError(43                f"Got invalid return object. Expected key `{key}` to be present, but got {json_obj}"44            )45    return json_obj46