strategy-flow-simulation/strategy_parse.py

# -*- coding: utf-8 -*-
"""
@author: yq
@time: 2024/12/18
@desc: 策略流节点解析
"""
import json
import re
import time
from tqdm import tqdm

import pandas as pd
from PIL import Image
from openpyxl import load_workbook

from commom import f_get_save_path, call_llm, f_file_upload, GeneralException
from enums import ResultCodesEnum
from prompt import f_get_prompt_parse_node, f_get_prompt_parse_flow, f_get_prompt_parse_flow_image


def _f_parse_flow_image(ws, node_list: list):
    image = ws._images[0]
    img = Image.open(image.ref).convert("RGB")
    save_path = f_get_save_path("流程图.png")
    img.save(save_path)
    time.sleep(1)
    file_id = f_file_upload(save_path)
    prompt = f_get_prompt_parse_flow_image(node_list)
    print(prompt)
    prompt = [
        {
            "type": "text",
            "text": prompt
        },
        {
            "type": "image",
            "file_id": file_id
        }
    ]
    prompt = json.dumps(prompt, ensure_ascii=False)
    llm_answer = call_llm(prompt, "object_string")
    print(llm_answer)
    code = re.findall(r"```python\n(.*)\n```", llm_answer, flags=re.DOTALL)[0]
    save_path = f_get_save_path("flow.py")
    with open(save_path, mode="w", encoding="utf8") as f:
        f.write(code)

    save_path = f_get_save_path("__init__.py")
    with open(save_path, mode="w", encoding="utf8") as f:
        f.write("")


def _f_parse_flow(node_list: list, df: pd.DataFrame):
    flow = ""
    for _, row in df.iterrows():
        strategy = row["策略流描述"]
        flow = f"{flow}{strategy}\n"
    flow = flow.strip()

    prompt = f_get_prompt_parse_flow(node_list, flow)
    print(prompt)
    llm_answer = call_llm(prompt)
    print(llm_answer)
    code = re.findall(r"```python\n(.*)\n```", llm_answer, flags=re.DOTALL)[0]
    save_path = f_get_save_path("flow.py")
    with open(save_path, mode="w", encoding="utf8") as f:
        f.write(code)

    save_path = f_get_save_path("__init__.py")
    with open(save_path, mode="w", encoding="utf8") as f:
        f.write("")


def _f_parse_node(df: pd.DataFrame, node_name):
    rules = ""
    for idx, row in df.iterrows():
        var_name = row["变量"]
        var_name = var_name.replace("\n", " ")

        rule_content = row["逻辑"]
        rule_content = rule_content.replace("\n", " ")

        rule_out = row["输出"]

        notes_output = row["输出备注"]
        if notes_output is None or notes_output != notes_output:
            notes_output = ""
        else:
            notes_output = notes_output.replace("\n", " ")
            notes_output = f" 结果备注: {notes_output}"

        notes_input = row["输入备注"]
        if notes_input is None or notes_input != notes_input:
            notes_input = ""
        else:
            notes_input = notes_input.replace("\n", " ")
            notes_input = f" 变量备注: {notes_input}"

        rules = f"{rules}规则{idx + 1}: 变量:{var_name} 逻辑:{rule_content} 输出:{rule_out}{notes_input}{notes_output}\n"

    default_output = list(df["默认输出"])[0]
    if default_output is None or default_output != default_output:
        default_output = ""
    else:
        default_output = str(default_output).replace("\n", " ")
        default_output = f"{default_output}"

    prompt = f_get_prompt_parse_node(node_name, rules, default_output)
    print(prompt)
    llm_answer = call_llm(prompt)
    code = re.findall(r"```python\n(.*)\n```", llm_answer, flags=re.DOTALL)[0]
    func_name = re.findall(r"def (.*)\(data", code)[0]
    save_path = f_get_save_path(f"{func_name}.py")
    print(code)
    with open(save_path, mode="w", encoding="utf8") as f:
        f.write(code)
    return func_name, code


def f_parse_strategy_image(file_path):
    wb = load_workbook(file_path)
    excel = pd.ExcelFile(file_path)
    sheet_names = excel.sheet_names
    if "流程图" not in sheet_names:
        GeneralException(ResultCodesEnum.NOT_FOUND, message=f"sheet【流程图】不存在")
    node_list = []
    for node_name in tqdm(sheet_names):
        if node_name == "流程图":
            continue
        df = excel.parse(sheet_name=node_name)
        func_name, code = _f_parse_node(df, node_name)
        node_list.append((node_name, func_name, code))
    _f_parse_flow_image(wb["流程图"], node_list)
    wb.close()
    excel.close()


def f_parse_strategy(file_path):
    excel = pd.ExcelFile(file_path)
    sheet_names = excel.sheet_names
    if "流程" not in sheet_names:
        GeneralException(ResultCodesEnum.NOT_FOUND, message=f"sheet【流程】不存在")
    node_list = []
    for node_name in tqdm(sheet_names):
        if node_name == "流程":
            continue
        df = excel.parse(sheet_name=node_name)
        func_name, code = _f_parse_node(df, node_name)
        node_list.append((node_name, func_name, code))
    _f_parse_flow(node_list, excel.parse(sheet_name="流程"))
    excel.close()


if __name__ == "__main__":
    f_parse_strategy("./cache/策略节点配置3demo.xlsx")