wiseflow/test/get_info_test.py

# -*- coding: utf-8 -*-
import os, sys
import json
import asyncio
import time
from datetime import datetime

# 将core目录添加到Python路径
core_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', 'core')
sys.path.append(core_path)

# 现在可以直接导入模块，因为core目录已经在Python路径中
from scrapers import *
from agents.get_info import pre_process

from utils.general_utils import is_chinese
from agents.get_info import get_author_and_publish_date, get_info, get_more_related_urls
from agents.get_info_prompts import *

benchmark_model = 'Qwen/Qwen2.5-72B-Instruct'
models = ['Qwen/Qwen2.5-7B-Instruct', 'Qwen/Qwen2.5-14B-Instruct',  'Qwen/Qwen2.5-32B-Instruct', 'deepseek-ai/DeepSeek-V2.5']

async def main(sample: dict, include_ap: bool, prompts: list, focus_dict: dict, record_file: str):
    link_dict, links_parts, contents = sample['link_dict'], sample['links_part'], sample['contents']
    get_link_sys_prompt, get_link_suffix_prompt, get_info_sys_prompt, get_info_suffix_prompt = prompts

    for model in [benchmark_model] + models:
        links_texts = []
        for _parts in links_parts:
            links_texts.extend(_parts.split('\n\n'))
        contents = sample['contents'].copy()

        print(f"running {model} ...")
        start_time = time.time()
        if include_ap:
            author, publish_date = await get_author_and_publish_date(contents[0], model, test_mode=True)
            get_ap_time = time.time() - start_time
            print(f"get author and publish date time: {get_ap_time}")
        else:
            author, publish_date = '', ''
            get_ap_time = 0

        start_time = time.time()
        more_url = await get_more_related_urls(links_texts, link_dict, [get_link_sys_prompt, get_link_suffix_prompt, model], test_mode=True)
        get_more_url_time = time.time() - start_time
        print(f"get more related urls time: {get_more_url_time}")

        start_time = time.time()
        infos = await get_info(contents, link_dict, [get_info_sys_prompt, get_info_suffix_prompt, model], focus_dict, author, publish_date, test_mode=True)
        get_info_time = time.time() - start_time
        print(f"get info time: {get_info_time}")

        if model == benchmark_model:
            benchmark_result = more_url.copy()
            diff = f'benchmark: {len(benchmark_result)} results'
        else:
            missing_in_cache = len(benchmark_result - more_url)  # benchmark中有但cache中没有的
            extra_in_cache = len(more_url - benchmark_result)    # cache中有但benchmark中没有的
            total_diff = missing_in_cache + extra_in_cache
            diff = f'差异{total_diff}个(遗漏{missing_in_cache}个,多出{extra_in_cache}个)'

        related_urls_to_record = '\n'.join(more_url)
        infos_to_record = [f"{fi['tag']}: {fi['content']}" for fi in infos]
        infos_to_record = '\n'.join(infos_to_record)
        with open(record_file, 'a') as f:
            f.write(f"model: {model}\n")
            if include_ap:
                f.write(f"get author and publish date time: {get_ap_time}\n")
                f.write(f"author: {author}\n")
                f.write(f"publish date: {publish_date}\n")
            f.write(f"get more related urls time: {get_more_url_time}\n")
            f.write(f"diff from benchmark: {diff}\n")
            f.write(f"get info time: {get_info_time}\n")
            f.write(f"related urls: \n{related_urls_to_record}\n")
            f.write(f"final result: \n{infos_to_record}\n")
            f.write('\n\n')
        print('\n\n')

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--sample_dir', '-D', type=str, default='')
    parser.add_argument('--include_ap', '-I', type=bool, default=False)
    args = parser.parse_args()

    sample_dir = args.sample_dir
    include_ap = args.include_ap
    if not os.path.exists(os.path.join(sample_dir, 'focus_point.json')):
        raise ValueError(f'{sample_dir} focus_point.json not found')
    
    focus_points = json.load(open(os.path.join(sample_dir, 'focus_point.json'), 'r'))
    focus_statement = ''
    for item in focus_points:
        tag = item["focuspoint"]
        expl = item["explanation"]
        focus_statement = f"{focus_statement}//{tag}//\n"
        if expl:
            if is_chinese(expl):
                focus_statement = f"{focus_statement}解释：{expl}\n"
            else:
                focus_statement = f"{focus_statement}Explanation: {expl}\n"
    
    focus_dict = {item["focuspoint"]: item["focuspoint"] for item in focus_points}
    date_stamp = datetime.now().strftime('%Y-%m-%d')
    if is_chinese(focus_statement):
        get_link_sys_prompt = get_link_system.replace('{focus_statement}', focus_statement)
        get_link_sys_prompt = f"今天的日期是{date_stamp}，{get_link_sys_prompt}"
        get_link_suffix_prompt = get_link_suffix
        get_info_sys_prompt = get_info_system.replace('{focus_statement}', focus_statement)
        get_info_sys_prompt = f"今天的日期是{date_stamp}，{get_info_sys_prompt}"
        get_info_suffix_prompt = get_info_suffix
    else:
        get_link_sys_prompt = get_link_system_en.replace('{focus_statement}', focus_statement)
        get_link_sys_prompt = f"today is {date_stamp}, {get_link_sys_prompt}"
        get_link_suffix_prompt = get_link_suffix_en
        get_info_sys_prompt = get_info_system_en.replace('{focus_statement}', focus_statement)
        get_info_sys_prompt = f"today is {date_stamp}, {get_info_sys_prompt}"
        get_info_suffix_prompt = get_info_suffix_en

    prompts = [get_link_sys_prompt, get_link_suffix_prompt, get_info_sys_prompt, get_info_suffix_prompt]

    time_stamp = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime())
    record_file = os.path.join(sample_dir, f'record-{time_stamp}.txt')
    with open(record_file, 'w') as f:
        f.write(f"focus statement: \n{focus_statement}\n\n")

    for file in os.listdir(sample_dir):
        if not file.endswith('_processed.json'):
            continue
        sample = json.load(open(os.path.join(sample_dir, file), 'r'))
        if 'links_part' not in sample or 'link_dict' not in sample or 'contents' not in sample:
            print(f'{file} not valid sample, skip')
            continue
        with open(record_file, 'a') as f:
            f.write(f"raw materials: {file}\n\n")
        print(f'start testing {file}')
        asyncio.run(main(sample, include_ap, prompts, focus_dict, record_file))