import requests
import pandas as pd
import os,json

def get_research_data():

    # 定义CSV文件路径
    csv_path = 'templates/dataset/research/research_agent_test_data_v2.csv'
    
    # 检查文件是否存在
    if not os.path.exists(csv_path):
        print(f"文件 {csv_path} 不存在")
        return []
    
    try:
        # 读取CSV文件
        df = pd.read_csv(csv_path)
        
        # 将数据转换为字典列表
        research_data = []
        for _, row in df.iterrows():
            research_data.append({
                'research_topic': row.iloc[0],
                'max_loop': 3,
                'user_query': row.iloc[2]
            })
        
        return research_data
        
    except Exception as e:
        print(f"读取CSV文件时出错: {str(e)}")
        return []

API_KEY = 'app-aDHdJldQBFEqHmsD65jHhgcx'
API_URL = 'http://localhost/v1/chat-messages'  # 云端请替换为实际地址

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

def run_research():
    research_data = get_research_data()
    print(f"样本数量: {len(research_data)}")
    for i in range(63,100):
        item = research_data[i]
        print(f"\n样本 {i+1}:")
        print(f"research_topic: {item['research_topic']}")
        print(f"max_loop: {item['max_loop']}")
        print(f"user_query: {item['user_query']}")
        print("-" * 50)

        inputs = {
            "research_topic": item['research_topic'],
            "max_loop": item['max_loop']
        }

        query = item['user_query']

        payload = {
            "inputs": inputs,
            "query": query,
            "response_mode": "blocking",  # 或 "streaming"            
            "user": "abc-123",            # 可选，用户ID
            "conversation_id": ""        # 可选，首次对话留空，后续对话填前一次返回的 conversation_id
        }

        response = requests.post(API_URL, headers=headers, json=payload)
        print(response.text)

        if response.status_code == 200:
            # 确保输出目录存在
            output_dir = 'output/research'
            if not os.path.exists(output_dir):
                os.makedirs(output_dir)
            
            answer = json.loads(response.text)['answer']
            # 将结果写入json文件
            result = {
                "research_topic": item['research_topic'], 
                "user_query": item['user_query'],
                "answer": answer
            }
            
            # 修改文件扩展名为.json
            output_file = os.path.join(output_dir, f'research_log_{i}.json')
            
            with open(output_file, 'w', encoding='utf-8') as f:
                json.dump(result, f, ensure_ascii=False, indent=2)
            print("运行成功，结果已写入JSON文件：", output_file)

run_research()