GPT4comment/utils/read_file.py

import openpyxl

from utils.openai_api import get_tokens_length


def read_xlsx_file(data_source, file_name):
    workbook = openpyxl.load_workbook(file_name)
    sheet = workbook.active
    data = []
    if data_source == "七麦":
        for row in sheet.iter_rows(min_row=2, values_only=True):
            data.append({
                '发表时间': row[1],
                '作者': row[2],
                '评级': row[3],
                '标题': row[4],
                '内容': row[5]
            })

        return data[2:]
    else:
        for row in sheet.iter_rows(min_row=1, values_only=True):
            data.append({
                '发表时间': row[0],
                '作者': row[1],
                '评级': row[2],
                '内容': row[3]
            })

        return data[2:]


def process_data(data_source,
                 max_length,
                 data,
                 message,
                 length,
                 count_now=0,
                 reply=""):
    temp = []
    result_message = ""
    data = data[count_now:]

    while len(data) > 0:
        temp_message = message
        ele = data.pop(0)
        if data_source == "七麦":
            temp_string = "({}, {})".format(ele["标题"], ele["内容"])
        else:
            temp_string = "{}".format(ele["内容"])
        temp.append(temp_string)
        temp_message = temp_message.format(
            reply_before=reply,
            comment_num=count_now,
            comment_num_start=count_now + 1,
            comment_num_end=count_now + len(temp),
            all_num=length,
            data_string="[{}]".format(", ".join(temp)))

        if len(temp_message) > max_length:
            break
        result_message = temp_message
    return result_message, len(temp)