| | from typing import List |
| | from pathlib import Path |
| | from .base_agent import BaseAgent |
| | from prompt.template import TASK_DECOMPOSE_PROMPT, TASK_DESCRIPTION_PROMPT, TASK_DECOMPOSE_WO_COO_PROMPT |
| | from utils.utils import read_json_file |
| |
|
| |
|
| | class TaskDecompose(BaseAgent): |
| | def __init__(self, llm, coo=True): |
| | super().__init__(llm) |
| | self.coo = coo |
| | current_file = Path(__file__).resolve() |
| | path = current_file.parent.parent.parent / 'data/actor_data/input/decompose_prompt.json' |
| | self.decomposed_principles = read_json_file(str(path)) |
| |
|
| |
|
| | def decompose(self, modeling_problem: str, problem_analysis: str, modeling_solution: str, problem_type: str, tasknum: int, user_prompt: str=''): |
| | if self.coo: |
| | decomposed_principle = self.decomposed_principles.get(problem_type, self.decomposed_principles['C']) |
| | decomposed_principle = decomposed_principle.get(str(tasknum), decomposed_principle['4']) |
| | prompt = TASK_DECOMPOSE_PROMPT.format(modeling_problem=modeling_problem, problem_analysis=problem_analysis, modeling_solution=modeling_solution, decomposed_principle=decomposed_principle, tasknum=tasknum, user_prompt=user_prompt) |
| | |
| | answer = self.llm.generate(prompt) |
| | tasks = [task.strip() for task in answer.split('---') if task.strip()] |
| | else: |
| | prompt = TASK_DECOMPOSE_WO_COO_PROMPT.format(modeling_problem=modeling_problem, tasknum=tasknum, user_prompt=user_prompt) |
| | |
| | answer = self.llm.generate(prompt) |
| | tasks = [task.strip() for task in answer.split('---') if task.strip()] |
| | return tasks |
| |
|
| | def refine(self, modeling_problem: str, problem_analysis: str, modeling_solution: str, decomposed_subtasks: List[str], task_i: int): |
| | decomposed_subtasks_str = '\n'.join(decomposed_subtasks) |
| | prompt = TASK_DESCRIPTION_PROMPT.format(modeling_problem=modeling_problem, problem_analysis=problem_analysis, modeling_solution=modeling_solution, decomposed_subtasks=decomposed_subtasks_str, task_i=task_i+1) |
| | answer = self.llm.generate(prompt) |
| | return answer |
| |
|
| | def decompose_and_refine(self, modeling_problem: str, problem_analysis: str, modeling_solution: str, decomposed_principle: str, tasknum: int, user_prompt: str=''): |
| | if self.coo: |
| | decomposed_subtasks = self.decompose(modeling_problem, problem_analysis, modeling_solution, decomposed_principle, tasknum, user_prompt) |
| | for task_i in range(len(decomposed_subtasks)): |
| | refined_subtask = self.refine(modeling_problem, problem_analysis, modeling_solution, decomposed_subtasks, task_i) |
| | decomposed_subtasks[task_i] = refined_subtask |
| | else: |
| | decomposed_subtasks = self.decompose(modeling_problem, problem_analysis, modeling_solution, decomposed_principle, tasknum, user_prompt) |
| |
|
| | return decomposed_subtasks |
| |
|
