agent/checklist.py

from .mini_bench.agent import ChecklistGenerationAgent

def generate_checklist(**data):
    # data: 'intent', 'start_url', 'text_observation'
    agent_config = {
        'model_name': 'WPRM/qwen-3b-ar-reward-cot-mtl-checklist-enhanced',
        'base_url': 'http://165.132.144.84:7701/v1',
        'api_key': 'empty',
        'temperature': 0.7,
        'use_log_probs': True,
        'use_checklist': True,
        'use_multimodal': False,
        'num_generate': 1,
    }
    checklist_generation_agent = ChecklistGenerationAgent(agent_config)
    response_list, cost = checklist_generation_agent.generate_response(data, prompt_type='ours', constraint_str_list=["<think>", "</think>", "<answer>", "</answer>"])
    response = response_list[0]
    return response.split("<answer>")[-1].split("</answer>")[0].strip()