feat(app): more examples, better parsing and error handling

app.py

@@ -1,21 +1,24 @@
+import copy
 import json
 import os
 import zipfile
 
-import torch
 import gradio as gr
 import spaces
+import torch
 from peft import PeftConfig, PeftModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 from schema_to_sql import dd_to_sql
 from utils import (
+    MAX_NEW_TOKENS,
+    TEMPERATURE,
     create_summary_tables,
-    get_example_ai_model_output,
+    get_example_ai_model_output_many,
+    get_example_ai_model_output_simple,
     get_prompt_with_files_uploaded,
 )
-
-from utils import MAX_NEW_TOKENS, TEMPERATURE
+from parsing import try_parsing_actual_model_output
 
 LOCAL_DIR = "tsvs"
 ZIP_PATH = "tsvs.zip"
@@ -98,6 +101,11 @@ def run_llm_inference(model_prompt):
         output_data_model = tokenizer.decode(outputs[0][prompt_length:])
         output_data_model = output_data_model.split("<|eot_id|>")[0]
         print(output_data_model)
+
+        parsed_output_data_model = try_parsing_actual_model_output(output_data_model)
+        if "error" not in parsed_output_data_model:
+            output_data_model = copy.deepcopy(parsed_output_data_model)
+
         # Test output for JSON schema validity
         try:
             json.loads(output_data_model)
@@ -141,8 +149,16 @@ def gen_output_from_files_uploaded(filepaths: list[str] = None):
     return model_response, sql, nodes_df, properties_df
 
 
-def gen_output_from_example():
-    model_response = get_example_ai_model_output()
+def gen_output_from_example_simple():
+    model_response = get_example_ai_model_output_simple()
+    model_response_json = json.loads(model_response)
+    sql, validation = dd_to_sql(model_response_json)
+
+    return model_response, sql
+
+
+def gen_output_from_example_many():
+    model_response = get_example_ai_model_output_many()
     model_response_json = json.loads(model_response)
     sql, validation = dd_to_sql(model_response_json)
 
@@ -167,14 +183,12 @@ with gr.Blocks() as demo:
 
     gr.Markdown("## (Optional) Get Sample TSV(s) to Upload")
 
-    gr.Markdown("### Example 1: A single TSV")
+    gr.Markdown("### Example 1: A Simple TSV")
     download_btn = gr.DownloadButton(
-        label="Download Single TSV", value="sample_metadata.tsv"
+        label="Download Simple TSV", value="sample_metadata.tsv"
     )
     gr.Markdown("### Example 2: Many TSVs in a single .zip file.")
-    download_btn = gr.DownloadButton(
-        label="Download All Sample TSVs as .zip", value=zip_tsvs
-    )
+    download_btn = gr.DownloadButton(label="Download Many TSVs as .zip", value=zip_tsvs)
     gr.Markdown("You need to extract the .zip if you want to use them.")
 
     gr.Markdown("## Upload TSVs With Headers (No Data Rows Required)")
@@ -222,10 +236,16 @@ with gr.Blocks() as demo:
             outputs=[json_out, sql_out],
         )
 
-    gr.Markdown("Run out of FreeGPU or having issues? Try the example output!")
-    demo_btn = gr.Button("Manually Load Example Output from Previous Run")
+    gr.Markdown("Run out of FreeGPU or having issues? Try the example outputs!")
+    demo_btn2 = gr.Button("Manually Load 'Simple' Example Output from Previous Run")
+    demo_btn2.click(
+        fn=gen_output_from_example_simple,
+        outputs=[json_out, sql_out],
+    )
+
+    demo_btn = gr.Button("Manually Load 'Many' Example Output from Previous Run")
     demo_btn.click(
-        fn=gen_output_from_example,
+        fn=gen_output_from_example_many,
         outputs=[json_out, sql_out],
     )
 

models.py

@@ -0,0 +1,41 @@
+"""
+This file should contain all relevant data-only information, it should contain minimal to no behavior.
+"""
+
+from typing import Dict, List, Optional
+
+from pydantic import BaseModel, Field, model_validator
+
+
+class Property(BaseModel):
+    name: Optional[str] = None
+    type: Optional[str] = None
+    description: Optional[str] = None
+    model_config = {"extra": "ignore"}
+
+
+def upsert_dict(mapping, kvp):
+    k, v = kvp
+    mapping[k] = v
+    return mapping
+
+
+class DataModel(BaseModel):
+    name: Optional[str] = None
+    description: Optional[str] = None
+    links: Optional[List[str]] = Field(default={})
+    properties: Optional[Dict[str, Property]] = Field(default={})
+    model_config = {"extra": "ignore"}
+
+    @model_validator(mode="before")
+    def _normalize_properties(cls, node_content):
+        props = node_content.get("properties", {})
+        if not isinstance(props, list):
+            return node_content
+        props_to_use = [prop for prop in props if "name" in prop]
+        prop_items = [Property(**prop) for prop in props_to_use]
+        prop_dict = {
+            prop.name: Property(type=prop.type, description=prop.description)
+            for prop in prop_items
+        }
+        return upsert_dict(node_content, ("properties", prop_dict))

parsing.py

@@ -0,0 +1,260 @@
+"""
+This file handles reshaping raw results data into a list of nodes that match
+what we expect the model output to be. That is to say, it handles parsing any raw data.
+"""
+
+import json
+from collections import defaultdict
+from functools import reduce
+from typing import Any, Dict, Hashable, Optional, Tuple, Union
+
+from models import DataModel
+from shared import (
+    TOP_LEVEL_IDENTIFIERS,
+    attempt,
+    get_json_from_model_output,
+    keep_errors,
+    on_fail,
+)
+
+
+def handle_parsing_schema_files(expected_location: str, actual_location: str):
+    raw_reference, raw_generated = read_in_expected_and_actual_json(
+        expected_location, actual_location
+    )
+    read_errors = keep_errors((raw_reference, raw_generated))
+    if len(read_errors) > 0:
+        raise ValueError(f"Could not ingest raw data: {read_errors}")
+    generated_nodes_to_content = try_parsing_actual_model_output(raw_generated)
+    reference_nodes_to_content = derive_nodes_from_actual_json_output(
+        parse_json(raw_reference)
+    )
+    errors = keep_errors((reference_nodes_to_content, generated_nodes_to_content))
+    if len(errors) > 0:
+        raise ValueError(f"Error parsing files: {errors}")
+    return generated_nodes_to_content, reference_nodes_to_content
+
+
+def read_in_expected_and_actual_json(
+    expected_json_location: str, actual_json_location: str
+):
+    unparsed_expected, err = read_json_as_text(expected_json_location)
+    if err:
+        return {"error": f"Could not read in expected file. e: {err}"}, None
+    unparsed_actual, err = read_json_as_text(actual_json_location)
+    if err:
+        return None, {"error": str(err)}
+    return unparsed_expected, unparsed_actual
+
+
+def read_json_as_text(file_path: str) -> Tuple[Optional[str], Optional[Exception]]:
+    """
+    Reads the contents of a file as text without attempting to parse it as JSON.
+    """
+    try:
+        with open(file_path, "r", encoding="utf-8") as file:
+            return file.read(), None
+    except Exception as e:
+        return None, e
+
+
+def try_parsing_actual_model_output(model_output: str):
+    first_parse_json = parse_json(model_output)
+    if isinstance(first_parse_json, list):
+        return {
+            "error": "Could not parse json file. Model output should not be a list."
+        }
+    first_pass_failed = "error" in first_parse_json
+    recovered_json, errors = (
+        get_json_from_model_output(model_output) if first_pass_failed else ({}, 0)
+    )
+    if errors > 0:
+        return {"error": "Could not parse json file, no metrics to calculate"}
+    parsed_json = recovered_json if first_pass_failed else first_parse_json
+    node_derivation_outcome = on_fail(
+        attempt(derive_nodes_from_actual_json_output, (parsed_json,)), []
+    )
+    if not node_derivation_outcome:
+        return {"error": f"Could not derive nodes. Parsed json: {parsed_json}"}
+    return node_derivation_outcome
+
+
+def find_all_nodes(name_and_contents) -> Dict:
+    name, contents = name_and_contents
+    content_contains_nodes = bool(set(contents.keys()) & TOP_LEVEL_IDENTIFIERS)
+    if content_contains_nodes:
+        return dict([name_and_contents])
+    sub_dicts = list(filter(lambda kvp: isinstance(kvp[1], dict), contents.items()))
+    all_sub_nodes = {}
+    for sub_name_and_contents in sub_dicts:
+        sub_nodes = find_all_nodes(sub_name_and_contents)
+        all_sub_nodes.update(sub_nodes)
+    return all_sub_nodes
+
+
+def assign_to_key(key: Hashable):
+    def add_at_key(
+        assignment_mapping: Dict[Hashable, Any], mapping_to_add: Dict[Hashable, Any]
+    ):
+        assignment_id = mapping_to_add[key]
+        assignment_mapping[assignment_id] = mapping_to_add
+        return assignment_mapping
+
+    return add_at_key
+
+
+key_exists = lambda key: lambda mapping: key in mapping
+
+
+def handle_property_correction(all_nodes):
+    has_incorrect_property_shape = lambda node: (
+        "properties" in node[1] and isinstance(node[1].get("properties", {}), list)
+    )
+    nodes_that_need_corrected = dict(
+        filter(has_incorrect_property_shape, all_nodes.items())
+    )
+    nodes_with_corrected_properties = dict(
+        map(correct_properties_for_node, nodes_that_need_corrected.items())
+    )
+    all_corrected_nodes = {**all_nodes, **nodes_with_corrected_properties}
+    return all_corrected_nodes
+
+
+def correct_properties_for_node(node):
+    """Maps node's property names to their actual content."""
+    node_name, node_data = node
+    properties = node_data["properties"]
+    identified_properties = list(filter(key_exists("name"), properties))
+    actual_properties = reduce(assign_to_key("name"), identified_properties, {})
+    node_data["properties"] = actual_properties
+    return node
+
+
+def derive_nodes_from_actual_json_output(json_data: Union[dict, list]):
+    """
+    Find nodes from non-deterministic AI output
+    """
+    if isinstance(json_data, list):
+        return {}
+    all_nodes = flatten_all_nodes(json_data)
+    if json_data.get("nodes", None) is None:
+        return all_nodes
+    nodes_with_properties_corrected = handle_property_correction(all_nodes)
+    return nodes_with_properties_corrected
+
+
+def flatten_all_nodes(json_data) -> Dict[Hashable, Any]:
+    """
+    Model output could have nested nodes, this extracts them.
+    """
+    nodes = json_data.get("nodes", None)
+    if nodes is None:
+        sub_nodes_list = [
+            find_all_nodes((name, contents))
+            for name, contents in json_data.items()
+            if isinstance(contents, dict)
+        ]
+    else:
+        sub_nodes_list = [
+            find_all_nodes((node["name"], node))
+            for node in nodes
+            if isinstance(node, dict) and node.get("name") is not None
+        ]
+    all_nodes = {k: v for sub_nodes in sub_nodes_list for k, v in sub_nodes.items()}
+    return all_nodes
+
+
+def aggregate_desc(acc, node):
+    node_name, node_info = node
+    desc = node_info.get("description", None)
+    if desc is not None and isinstance(desc, str):
+        acc[desc].add(node_name)
+    return acc
+
+
+def reform_links(outer_acc, node):
+    node_name, node_info = node
+    links = node_info.get("links", [])
+    collect_links_to_aggregator = lambda inner_acc, link: upsert_set(
+        inner_acc, (link, node_name)
+    )
+    links_to_node_names = reduce(collect_links_to_aggregator, links, outer_acc)
+    return links_to_node_names
+
+
+def lens(key, default=None):
+    """Simple way to interface with the contents of a dict"""
+    return lambda d: d.get(key, default)
+
+
+def aggregate_properties(outer_acc, node):
+    node_name, node_info = node
+    properties = node_info.get("properties", {})
+    is_list = isinstance(properties, list)
+    property_names = (
+        list(map(lens("name"), properties)) if is_list else list(properties.keys())
+    )
+    aggregate_properties = reduce(
+        lambda inner_acc, prop_name: upsert_set(inner_acc, (prop_name, node_name)),
+        property_names,
+        outer_acc,
+    )
+    return aggregate_properties
+
+
+def conform_node_to_expected_schema(name_to_data_model):
+    name, dm = name_to_data_model
+    conform_result = attempt(DataModel.model_validate, (dm,))
+    model_outcome = (
+        conform_result.model_dump(exclude_none=True, exclude_unset=True)
+        if "error" not in conform_result
+        else {}
+    )
+    errors = conform_result if "error" in conform_result else {}
+    return (name, model_outcome), errors
+
+
+def aggregate_parsed_file(nodes: dict):
+    conformed_nodes_result = [
+        conform_node_to_expected_schema(node) for node in nodes.items()
+    ]
+    conformed_nodes = [node for node, errors in conformed_nodes_result if not errors]
+    aggregated_links = reduce(reform_links, conformed_nodes, defaultdict(set))
+    aggregated_properties = reduce(
+        aggregate_properties, conformed_nodes, defaultdict(set)
+    )
+    aggregated_descriptions = reduce(aggregate_desc, conformed_nodes, defaultdict(set))
+
+    parsed_results = {
+        "node_names": dict(conformed_nodes),
+        "links": aggregated_links,
+        "properties": aggregated_properties,
+        "description": aggregated_descriptions,
+    }
+    return parsed_results
+
+
+def upsert_set(accumulator, kvp):
+    key, value = kvp
+    if isinstance(key, Hashable):
+        accumulator[key].add(value)
+    return accumulator
+
+
+def parse_json(json_string: str) -> Optional[Union[dict, list]]:
+    """
+    Safely parses a JSON string into a Python dictionary or list.
+
+    Args:
+        json_string (str): The JSON string to parse
+
+    Returns:
+        dict/list: Parsed JSON data if successful
+        dict["error"]: If parsing fails, provides error as string in response
+    """
+    try:
+        return json.loads(json_string)
+    except json.JSONDecodeError as e:
+        return {"error": str(e)}
+    except TypeError as e:
+        return {"error": str(e)}

schema_to_sql.py

@@ -99,6 +99,43 @@ def get_foreign_table_and_field(prop_name, node_name):
     return None, None
 
 
+def transform_dd(dd):
+    """
+    Returns transformed DD
+
+    This function takes AI generated DD and ensures all required fields are
+    present in properties and properties are dictionaries.
+
+    Parameters:
+      dd (dict): AI generated DD
+
+    Returns:
+      dd (dict): Transformed DD
+    """
+    for node in dd.get("nodes", []):
+        props = node.get("properties", [])
+        if props and all(isinstance(x, dict) for x in props):
+            prop_names = {p["name"] for p in props}
+        elif props and all(isinstance(x, str) for x in props):
+            prop_names = set(props)
+            # Upgrade to list of dicts
+            props = [
+                {"name": prop, "description": "", "type": "string"} for prop in props
+            ]
+        else:
+            props = []
+            prop_names = set()
+
+        # Ensure each required field is present in properties
+        for req in node.get("required", []):
+            if req not in prop_names:
+                props.append({"name": req, "description": "", "type": "string"})
+                prop_names.add(req)
+
+        node["properties"] = props
+    return dd
+
+
 def generate_create_table(node, table_lookup):
     """
     Returns SQL for the given AI generated node
@@ -189,6 +226,7 @@ def dd_to_sql(dd):
       sql (str): SQL
       validation (str): Validation result
     """
+    dd = transform_dd(dd)
     # Build a lookup for table columns in all nodes
     table_lookup = {}
     for node in dd["nodes"]:

shared.py

@@ -0,0 +1,308 @@
+"""
+Functions used in several different places. This file should not import from any other non-lib files to prevent
+circular dependencies.
+"""
+
+import json
+import logging
+from copy import copy
+from typing import Any, Callable, Dict, Optional, Tuple, Union
+
+TOP_LEVEL_IDENTIFIERS = {"description", "links", "properties"}
+
+
+def get_json_from_model_output(input_generated_json: str):
+    """
+    Parses a string, potentially containing Markdown code fences, into a JSON object.
+
+    This function attempts to extract and parse a JSON object from a string,
+    often the output of a language model. It handles cases where the JSON
+    is enclosed in Markdown code fences (```json ... ``` or ``` ... ```).
+    If the initial parsing fails, it attempts a more robust parsing using
+    `_get_valid_json_from_string` and
+    logs debug messages indicating success or failure. If all attempts fail,
+    it returns an empty dictionary.
+
+    Args:
+        input_generated_json: A string potentially containing a JSON object.
+
+    Returns:
+        A tuple containing:
+            - The parsed JSON object (a dictionary) or an empty dictionary if parsing failed.
+            - An integer representing the number of times parsing failed initially.
+    """
+    originally_invalid_json_count = 0
+
+    generated_json_attempt_1 = copy(input_generated_json)
+    try:
+        code_split = generated_json_attempt_1.split("```")
+        if len(code_split) > 1:
+            generated_json_attempt_1 = json.loads(
+                ("```" + code_split[1]).replace("```json", "")
+            )
+        else:
+            generated_json_attempt_1 = json.loads(
+                generated_json_attempt_1.replace("```json", "").replace("```", "")
+            )
+    except Exception as exc:
+        logging.debug(f"could not parse AI model generated output as JSON. Exc: {exc}.")
+        # originally_invalid_json_count += 1
+        generated_json_attempt_1 = {}
+    some_value_in_attempt_1_is_not_a_dict = check_contents_valid(
+        generated_json_attempt_1
+    )
+    attempt_1_failed = (
+        not bool(generated_json_attempt_1) or some_value_in_attempt_1_is_not_a_dict
+    )
+    generated_json_attempt_2 = copy(input_generated_json) if attempt_1_failed else {}
+    if attempt_1_failed:
+        logging.debug(
+            "Attempting to make output valid to obtain better metrics (this works in limited cases where "
+            "the model output was simply cut off)"
+        )
+        try:
+            code_split = generated_json_attempt_2.split("```")
+            if len(code_split) > 1:
+                generated_json_attempt_2 = json.loads(
+                    _get_valid_json_from_string(
+                        ("```" + code_split[1]).replace("```json", "")
+                    )
+                )
+            else:
+                stripped_output = generated_json_attempt_2.replace(
+                    "```json", ""
+                ).replace("```", "")
+                balance_outcome = attempt(
+                    json.loads, (balance_braces(stripped_output),)
+                )
+                if "error" not in balance_outcome:
+                    generated_json_attempt_2 = balance_outcome
+                else:
+                    generated_json_attempt_2 = json.loads(
+                        _get_valid_json_from_string(stripped_output)
+                    )
+
+            logging.debug(
+                "Success! Reconstructed valid JSON from unparseable model output. Continuing metrics comparison..."
+            )
+        except Exception as exc:
+            logging.debug(
+                "Failed. Setting model output as empty JSON to enable metrics comparison."
+            )
+            generated_json_attempt_2 = {}
+    some_value_in_attempt_2_is_not_a_dict = (
+        attempt_1_failed
+        and isinstance(generated_json_attempt_2, dict)
+        and check_contents_valid(generated_json_attempt_2)
+    )
+    if some_value_in_attempt_1_is_not_a_dict and some_value_in_attempt_2_is_not_a_dict:
+        logging.debug(f"Could not recover model output json, aborting!")
+        originally_invalid_json_count += 1
+    generated_json = (
+        generated_json_attempt_1 if not attempt_1_failed else generated_json_attempt_2
+    )
+    return generated_json, originally_invalid_json_count
+
+
+def check_contents_valid(generated_json_attempt_1: Union[list, dict]):
+    """
+    Checks that the sub nodes are not lists or anything
+
+    Args:
+        generated_json_attempt_1 (Union[list, dict]): data to check
+
+    Returns:
+        truthy based on contents of input
+    """
+    if isinstance(generated_json_attempt_1, list):
+        for item in generated_json_attempt_1:
+            if not isinstance(item, dict):
+                return item
+        return None
+    elif (
+        isinstance(generated_json_attempt_1, dict)
+        and "nodes" in generated_json_attempt_1.keys()
+    ):
+        for item in generated_json_attempt_1.get("nodes", []):
+            if not isinstance(item, dict):
+                return item
+        return None
+    else:
+        for item in generated_json_attempt_1.values():
+            if not isinstance(item, dict):
+                return item
+        return None
+
+
+def _get_valid_json_from_string(s):
+    """
+    Given a JSON string with potentially unclosed strings, arrays, or objects, close those things
+    to hopefully be able to parse as valid JSON
+    """
+    double_quotes = 0
+    single_quotes = 0
+    brackets = []
+
+    for i, c in enumerate(s):
+        if c == '"':
+            double_quotes = 1 - double_quotes  # Toggle between 0 and 1
+        elif c == "'":
+            single_quotes = 1 - single_quotes  # Toggle between 0 and 1
+        elif c in "{[":
+            brackets.append((i, c))
+        elif c in "}]":
+            if double_quotes == 0 and single_quotes == 0:
+                if brackets:
+                    last_opened = brackets.pop()
+                    if (c == "}" and last_opened[1] != "{") or (
+                        c == "]" and last_opened[1] != "["
+                    ):
+                        raise ValueError(
+                            f"Mismatched brackets/quotes found: opened {last_opened[1]} @ {last_opened[0]} "
+                            f"but closed {c} @ {i}"
+                        )
+                else:
+                    # If no matching opening bracket, it's an error, but we can skip this for the task
+                    pass
+
+    # Remove trailing comma if it exists
+    if s.strip().endswith(","):
+        logging.debug("Removing ending ,")
+        s = s.strip().rstrip(",")
+
+    closing_chars = ""
+
+    # Adding closing quotes if there are missing ones
+    if double_quotes > 0:
+        closing_chars += '"'
+    if single_quotes > 0:
+        closing_chars += "'"
+
+    # Add closing brackets for any unmatched opening brackets
+    while brackets:
+        last_opened = brackets.pop()
+        if last_opened[1] == "{":
+            closing_chars += "}"
+        elif last_opened[1] == "[":
+            closing_chars += "]"
+
+    logging.debug(f"closing_chars: {closing_chars}")
+
+    output_string = s + closing_chars
+
+    try:
+        json.loads(output_string)
+    except Exception:
+        logging.debug(
+            "JSON string still fails to be parseable, attempting another modification..."
+        )
+        # it's possible the closing quotes were on a property that didn't have a value, let's
+        # fix that and see if it works
+        new_closing_chars = ""
+        found_first_double_quote = False
+        for char in closing_chars:
+            if not found_first_double_quote and char == '"':
+                # for keys in objects with no value, append an empty value
+                #
+                # For example:
+                # ```
+                # {
+                #    "properties": {
+                #         "annotation
+                # ```
+                new_closing_chars += '": ""'
+            else:
+                new_closing_chars += char
+
+        logging.debug(f"new closing_chars: {new_closing_chars}")
+        output_string = s + new_closing_chars
+
+    return output_string
+
+
+def on_fail(
+    outcome: Union[Any, Dict[str, str]],
+    fallback: Union[Any, Callable] = None,
+):
+    """
+    Allows you to provide a fallback to recover from a failed outcome.
+
+    Args:
+        outcome
+        fallback
+
+    Returns:
+
+    """
+    is_fail = isinstance(outcome, dict) and "error" in outcome
+    is_callable = isinstance(fallback, Callable)
+    if is_fail and is_callable:
+        return fallback(outcome)
+    elif is_fail:
+        return fallback
+    return outcome
+
+
+def attempt(
+    func: Callable,
+    args: Tuple[Any, ...] = (),
+    kwargs: Optional[Dict[str, Any]] = None,
+) -> Union[Any, Dict[str, str]]:
+    """
+    Attempts to execute a function with the provided arguments.
+
+    If the function raises an exception, the exception is caught and returned in a dict.
+    Args:
+        func (Callable): The function to execute.
+        args (Tuple[Any, ...], optional): A tuple of positional arguments for the function.
+        kwargs (Optional[Dict[str, Any]], optional): A dictionary of keyword arguments for the function.
+    Returns:
+        Function result or {"error": <msg>} response
+    """
+    kwargs = kwargs or {}
+    try:
+        return func(*args, **kwargs)
+    except Exception as exc:
+        return {"error": str(exc)}
+
+
+def balance_braces(s: str) -> str:
+    """
+    Primitive function that just tries to add '{}' style braces to try to recover
+    the model string.
+
+    Args:
+        s(str): string to balance braces on.
+
+    Returns:
+        provided string with balanced braces if possible
+    """
+    open_count = s.count("{")
+    close_count = s.count("}")
+
+    if open_count > close_count:
+        s += "}" * (open_count - close_count)
+    elif close_count > open_count:
+        s = "{" * (close_count - open_count) + s
+
+    return s
+
+
+def flatten_list(coll):
+    flattened_data = []
+    for set_list in coll:
+        flattened_data = flattened_data + list(set_list)
+    return flattened_data
+
+
+def keep_errors(collection):
+    """
+    Given a set of outcomes, keeps any that resulted in an error
+
+    Args:
+        collection (Collection): collection of outcomes to filter.
+
+    Returns:
+        All instances of the collection that contain an error response.
+    """
+    return [instance for instance in collection if "error" in (instance or [])]

utils.py

@@ -1,9 +1,10 @@
-import os
 import csv
+import os
+from io import BytesIO
 from string import Template
-import networkx as nx
+
 import matplotlib.pyplot as plt
-from io import BytesIO
+import networkx as nx
 from PIL import Image
 
 MAX_INPUT_TOKEN_LENGTH = 1024
@@ -108,7 +109,227 @@ def create_summary_tables(json_response):
     return node_descriptions, node_property_descriptions
 
 
-def get_example_ai_model_output():
+def get_example_ai_model_output_simple():
+    return """
+{
+    "nodes": [
+      {
+        "name": "project",
+        "description": "Any specifically defined piece of work that is undertaken or attempted to meet a single requirement. (NCIt C47885)",
+        "links": [],
+        "required": [
+          "dbgap_accession_number",
+          "project.id"
+        ],
+        "properties": [
+          {
+            "name": "awg_review",
+            "description": "Indicates that the project is an AWG project.",
+            "type": "boolean"
+          },
+          {
+            "name": "data_citation",
+            "description": "The citation for the published dataset.",
+            "type": "string"
+          },
+          {
+            "name": "data_contributor",
+            "description": "The name of the organization or individual that the contributed dataset belongs to.",
+            "type": "string"
+          },
+          {
+            "name": "data_description",
+            "description": "A brief, free-text description of the data files and associated metadata provided for this dataset.",
+            "type": "string"
+          },
+          {
+            "name": "dbgap_accession_number",
+            "description": "The dbgap accession number provided for the project.",
+            "type": "string"
+          },
+          {
+            "name": "in_review",
+            "description": "Indicates that the project is under review by the submitter. Upload and data modification is disabled.",
+            "type": "boolean"
+          },
+          {
+            "name": "intended_release_date",
+            "description": "Tracks a Project's intended release date.",
+            "type": "string"
+          },
+          {
+            "name": "project.id",
+            "description": "A unique identifier for records in this 'project' table.",
+            "type": "string"
+          },
+          {
+            "name": "protocol_number",
+            "description": "The project's protocol number or similar amount.",
+            "type": "string"
+          },
+          {
+            "name": "releasable",
+            "description": "A project can only be released by the user when `releasable` is true.",
+            "type": "boolean"
+          },
+          {
+            "name": "request_submission",
+            "description": "Indicates that the user has requested submission to the GDC for harmonization.",
+            "type": "boolean"
+          },
+          {
+            "name": "research_design",
+            "description": "A summary of the goals of the research or a general description of the research's relationship to a clinical application.",
+            "type": "string"
+          },
+          {
+            "name": "research_objective",
+            "description": "The general objective of the research; what the researchers hope to discover or determine.",
+            "type": "string"
+          },
+          {
+            "name": "research_setup",
+            "description": "A high level description of the setup used to achieve the research objectives.",
+            "type": "string"
+          }
+        ]
+      },
+      {
+        "name": "dataset",
+        "description": "A set of metadata and associated data file objects originating from single a research study, clinical trial or patient cohort.",
+        "links": [
+          "project"
+        ],
+        "required": [
+          "dataset.id",
+          "project.id"
+        ],
+        "properties": [
+          {
+            "name": "Class_of_Case_Desc",
+            "description": "The text term used to describe the kind of clinical condition that can be defined based on objective criteria or by including all patient information from the case.",
+            "type": "string"
+          },
+          {
+            "name": "data_citation",
+            "description": "The citation for the published dataset.",
+            "type": "string"
+          },
+          {
+            "name": "full_name",
+            "description": "The full name or title of the dataset or publication.",
+            "type": "string"
+          },
+          {
+            "name": "longitudinal",
+            "description": "Indicates whether the dataset has longitudinal or time-series data.",
+            "type": "boolean"
+          },
+          {
+            "name": "project.id",
+            "description": "Unique identifiers for records in the 'project' table that relate via this foreign key to records in this 'dataset' table.",
+            "type": "string"
+          },
+          {
+            "name": "dataset.id",
+            "description": "A unique identifier for records in this 'dataset' table.",
+            "type": "string"
+          }
+        ]
+      },
+      {
+        "name": "subject",
+        "description": "The collection of all data related to a specific subject in the context of a specific experiment.",
+        "links": [
+          "project",
+          "dataset"
+        ],
+        "required": [
+          "dataset.id",
+          "subject.id"
+        ],
+        "properties": [
+          {
+            "name": "date_of_death",
+            "description": "The date of death of the subject in the context of a specific experiment.",
+            "type": "string"
+          },
+          {
+            "name": "dataset.id",
+            "description": "Unique identifiers for records in the 'dataset' table that relate via this foreign key to records in this'subject' table.",
+            "type": "string"
+          },
+          {
+            "name": "subject.id",
+            "description": "A unique identifier for records in this'subject' table.",
+            "type": "string"
+          },
+          {
+            "name": "CancerRegistry_PatientID",
+            "description": "The patient unique id in the case registry.",
+            "type": "string"
+          },
+          {
+            "name": "Ethnicity",
+            "description": "An individual's self-described social and cultural grouping, specifically whether an individual describes themselves as Hispanic or Latino. The provided values are based on the categories defined by the U.S. Office of Management and Business and used by the U.S. Census Bureau.",
+            "type": "string"
+          },
+          {
+            "name": "Last_Name",
+            "description": "The surname(s) of individual(s) in study, in the form used for cultural or ethnic reasons (e.g., Spanish surnames)",
+            "type": "string"
+          },
+          {
+            "name": "Sex_Desc",
+            "description": "The description of the individual's gender.",
+            "type": "string"
+          },
+          {
+            "name": "project.id",
+            "description": "Unique identifiers for records in the 'project' table that relate via this foreign key to records in this'subject' table.",
+            "type": "string"
+          }
+        ]
+      },
+      {
+        "name": "sample",
+        "description": "Any material sample taken from a biological entity for testing, diagnostic, propagation, treatment or research purposes, including a sample obtained from a living organism or taken from the biological object after halting of all its life functions. Biospecimen can contain one or more components including but not limited to cellular molecules, cells, tissues, organs, body fluids, embryos, and body excretory products.",
+        "links": [
+          "subject"
+        ],
+        "required": [
+          "sample.id",
+          "subject.id"
+        ],
+        "properties": [
+          {
+            "name": "body_fluid_code",
+            "description": "The code for the body fluid from which the sample was taken.",
+            "type": "string"
+          },
+          {
+            "name": "procedure_date",
+            "description": "Year the sample was taken for analysis.",
+            "type": "integer"
+          },
+          {
+            "name": "subject.id",
+            "description": "Unique identifiers for records in the'subject' table that relate via this foreign key to records in this'sample' table.",
+            "type": "string"
+          },
+          {
+            "name": "sample.id",
+            "description": "A unique identifier for records in this'sample' table.",
+            "type": "string"
+          }
+        ]
+      }
+    ]
+}
+"""
+
+
+def get_example_ai_model_output_many():
     return """
 {
        "nodes": [