94 lines
3.2 KiB
Python
94 lines
3.2 KiB
Python
import sys
|
|
import os
|
|
|
|
# 添加项目根目录到路径
|
|
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from gpt_api.gpt_agent import *
|
|
from gpt_api.prompt import *
|
|
from gpt_api.task_info import *
|
|
from gpt_api.task_info import get_all_tasks
|
|
from test_gpt_code import *
|
|
|
|
import argparse
|
|
|
|
def generate_code(task_info, las_error=None, message=None):
|
|
"""Generate code for robot task based on task info and previous errors."""
|
|
if message is None:
|
|
message = []
|
|
|
|
# Extract task information
|
|
task_name = task_info['task_name']
|
|
task_description = task_info['task_description']
|
|
current_code = task_info['current_code']
|
|
available_env_function = str(AVAILABLE_ENV_FUNCTOIN)
|
|
function_example = str(FUNCTION_EXAMPLE)
|
|
available_constants = str(AVAILABLE_CONSTANTS)
|
|
|
|
# Generate code based on error status
|
|
if las_error is not None:
|
|
# Handle error case - provide error info to improve generation
|
|
actor_name_keys, actor_data_keys, actor_points_discription = get_actor_keys_and_points_discription(f"gpt_{task_name}")
|
|
Prompt = (
|
|
f"The code is unsuccessful, \nLast Error Message: \n{las_error}\n\n"
|
|
f"Task Discription: \n{task_description}\n\n"
|
|
f"The Actor Points Discription: {actor_points_discription}"
|
|
)
|
|
else:
|
|
# First attempt case - create initial code file
|
|
res = f'''
|
|
from .base_task import Base_task
|
|
from .{task_name} import {task_name}
|
|
from .utils import *
|
|
import sapien
|
|
|
|
class gpt_{task_name}({task_name}):
|
|
def play_once(self):
|
|
pass
|
|
'''
|
|
file_name = f"envs/gpt_{task_name}.py"
|
|
with open(file_name, 'w') as file:
|
|
file.write(res)
|
|
|
|
# Get actor information for prompt
|
|
actor_name_keys, actor_data_keys, actor_points_discription = get_actor_keys_and_points_discription(f"gpt_{task_name}")
|
|
|
|
# Construct full prompt with all necessary information
|
|
Prompt = (
|
|
f"{BASIC_INFO}\n\n"
|
|
f"Task Discription: \n{task_description}\n\n"
|
|
f"Available API: \n{available_env_function}\n\n"
|
|
f"Function Example: \n{function_example}\n\n"
|
|
f"Available Constants: \n{available_constants}\n\n"
|
|
f"The Actor Name List: {actor_name_keys}\n\n"
|
|
f"The Actor Data List: {actor_data_keys}\n\n"
|
|
f"The Actor Points Discription: {actor_points_discription}\n\n"
|
|
f"Current Code:\n{current_code}"
|
|
)
|
|
|
|
# Add prompt to message history
|
|
message.append({"role": "user", "content": Prompt})
|
|
|
|
# Generate code using the model
|
|
res = generate(message, gpt="deepseek", temperature=0.5)
|
|
|
|
# Extract the relevant portion of the generated code
|
|
res = f'''
|
|
from .base_task import Base_task
|
|
from .{task_name} import {task_name}
|
|
from .utils import *
|
|
import sapien
|
|
|
|
class gpt_{task_name}({task_name}):
|
|
''' + res[res.find('def play_once'):res.rfind("```")]
|
|
|
|
# Save generated code to file
|
|
file_name = f"envs/gpt_{task_name}.py"
|
|
with open(file_name, 'w') as file:
|
|
file.write(res)
|
|
|
|
# Test the generated code
|
|
success_rate, error_message, error_count = test_run(f"gpt_{task_name}")
|
|
|
|
return res, success_rate, error_message, error_count
|