import dataiku

def get_dataset_names(project):
    datasets = project.list_datasets() 
    dataset_names = {"choices":[{"value": dataset['name'], "label": dataset['name']} for dataset in datasets]}
    return dataset_names

def get_dataset_columns(project, dataset_name):
    dataset = project.get_dataset(dataset_name)
    schema = dataset.get_schema()
    return {"choices":[{"value": column['name'], "label": column['name']} for column in schema['columns']]}

def do(payload, config, plugin_config, inputs):
    parameter_name = payload["parameterName"]
    client = dataiku.api_client()
    project = client.get_default_project()

    if parameter_name == "llm_responses_dataset":
        get_dataset_names(project)
    elif parameter_name == "llm_responses_column":
        dataset_name = config.get("llm_responses_dataset")
        return get_dataset_columns(project, dataset_name)