import requests
import json, os

def get_seed_prompt():
    # 读取seed_tasks.jsonl文件
    prompts = []
    with open('templates/dataset/prompt/seed_tasks.jsonl', 'r', encoding='utf-8') as f:
        for line in f:
            prompt = json.loads(line)
            prompts.append(prompt['instruction'])
    return prompts



def refine_prompt():
    API_KEY = "app-qXdPffm6tnrxrshUTYGq5zzg"
    WORKFLOW_RUN_URL = "http://localhost/v1/chat-messages"  # 或你的 Dify 云端地址
    USER_ID = "test-user-001"

    # 你要测试的输入
    prompttypes = [
        "通用优化",
        "通用优化-带输出格式要求",
        "带建议优化",
        "指令性优化",
        "迭代提示词优化"
    ]

    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json"
    }

    prompts = get_seed_prompt()

    # 创建输出目录
    output_dir = 'output/prompt'
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 生成输出文件名
    output_file = os.path.join(output_dir, f'prompt_optimization.jsonl')

    for prompt in prompts:
        query = "Refine the following prompt to make it more effective and efficient."
        for pt in prompttypes:
            inputs = {
                "prompt": prompt,
                "prompttype": pt
            }
            payload = {
                "inputs": inputs,
                "query": query,
                "response_mode": "blocking",
                "user": USER_ID
            }
            response = requests.post(WORKFLOW_RUN_URL, headers=headers, json=payload)
            print(f"\n==== prompttype: {pt} ====")
            if response.status_code == 200:
                print(json.dumps(response.json(), indent=2, ensure_ascii=False))
            else:
                print("❌ 工作流调用失败:", response.status_code, response.text)
            
            
            # 将原始prompt和优化后的prompt写入文件
            if response.status_code == 200:
                result = response.json()
                optimization_record = {
                    'original_prompt': prompt,
                    'optimized_prompt': result['answer'],
                    'prompt_type': pt
                }
                
                # 以追加模式写入JSONL文件
                with open(output_file, 'a', encoding='utf-8') as f:
                    f.write(json.dumps(optimization_record, ensure_ascii=False) + '\n')

refine_prompt()
        

