Spaces:
Sleeping
Sleeping
File size: 2,946 Bytes
8496edd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
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)
# print(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)
# print(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
|