core/input/problem.py

from utils.utils import read_json_file

from prompt.template import PROBLEM_PROMPT
from agent.data_description import DataDescription


def problem_input(problem_path, llm):
    problem = read_json_file(problem_path)
    data_description = problem.get('dataset_description', {})
    ds = DataDescription(llm)
    
    if data_description:
        data_path = problem['dataset_path']
        variable_description = problem['variable_description']
        data_summary = ds.summary(data_description=str(data_description) + '\n' + str(variable_description))
        data_summary = f'Dataset Path:\n{data_path}\n\nData Description:\n{data_summary}'
    else:
        data_summary = ''

    problem['data_summary'] = data_summary
    problem['data_description'] = data_description

    if problem.get('addendum', ''):
        addendum = f"Addendum: \n{problem['addendum']}"
    else:
        addendum = ''

    problem_str = PROBLEM_PROMPT.format(problem_background=problem['background'], problem_requirement=problem['problem_requirement'], addendum=addendum, data_summary=data_summary).strip()
    problem['problem_str'] = problem_str
    return problem_str, problem