Function get_sample_data

mycli/packages/special/llm.py:342–368

(
    cur: Cursor,
    dbname: str,
    prompt_field_truncate: int,
    prompt_section_truncate: int,
)

def get_sample_data(
    cur: Cursor,
    dbname: str,
    prompt_field_truncate: int,
    prompt_section_truncate: int,
) -> dict[str, Any]:
    if dbname in SAMPLE_DATA_CACHE:
        return SAMPLE_DATA_CACHE[dbname]
    click.echo("Preparing sample data to feed the LLM")
    tables_query = "SHOW TABLES"
    sample_row_query = "SELECT * FROM `{dbname}`.`{table}` LIMIT 1"
    cur.execute(tables_query)
    sample_data = {}
    for (table_name,) in cur.fetchall():
        try:
            cur.execute(sample_row_query.format(dbname=dbname, table=table_name))
        except Exception:
            continue
        cols = [desc[0] for desc in cur.description]
        row = cur.fetchone()
        if row is None:
            continue
        sample_data[table_name] = list(
            zip(cols, truncate_list_elements(list(row), prompt_field_truncate, prompt_section_truncate), strict=False)
        )
    SAMPLE_DATA_CACHE[dbname] = sample_data
    return sample_data


def sql_using_llm(