coze_znjd/commom/utils.py

# -*- coding:utf-8 -*-
"""
@author: isaacqyang
@time: 2023/12/28
@desc: 
"""

import json
import os
import time
from urllib.parse import unquote

import lark_oapi as lark
import tos
from docx import Document
from lark_oapi.api.drive.v1 import CreateExportTaskRequest, ExportTask, CreateExportTaskResponse, GetExportTaskRequest, \
    GetExportTaskResponse, DownloadExportTaskRequest, DownloadExportTaskResponse
from tos import HttpMethodType

from config import BaseConfig


def f_upload_file(save_path) -> str:
    ak = BaseConfig.cos_access_key_id
    sk = BaseConfig.cos_secret_access_key
    endpoint = BaseConfig.endpoint
    region = BaseConfig.region
    bucket_name = BaseConfig.bucket_name

    try:
        # 创建 TosClientV2 对象，对桶和对象的操作都通过 TosClientV2 实现
        client = tos.TosClientV2(ak, sk, endpoint, region)
        object_key = os.path.basename(save_path)
        client.put_object_from_file(bucket_name, object_key, save_path)
        pre_signed_url_output = client.pre_signed_url(HttpMethodType.Http_Method_Get, bucket_name, object_key)
        return pre_signed_url_output.signed_url

    except tos.exceptions.TosClientError as e:
        # 操作失败，捕获客户端异常，一般情况为非法请求参数或网络异常
        print('fail with client error, message:{}, cause: {}'.format(e.message, e.cause))
    except tos.exceptions.TosServerError as e:
        # 操作失败，捕获服务端异常，可从返回信息中获取详细错误信息
        print('fail with server error, code: {}'.format(e.code))
        # request id 可定位具体问题，强烈建议日志中保存
        print('error with request id: {}'.format(e.request_id))
        print('error with message: {}'.format(e.message))
        print('error with http code: {}'.format(e.status_code))
        print('error with ec: {}'.format(e.ec))
        print('error with request url: {}'.format(e.request_url))
    except Exception as e:
        print('fail with unknown error: {}'.format(e))


def f_doc_export(token: str) -> str:
    # 飞书在线文档转word
    app_id = BaseConfig.app_id
    app_secret = BaseConfig.app_secret
    word_save_dir = BaseConfig.word_save_dir

    client = lark.Client.builder() \
        .app_id(app_id) \
        .app_secret(app_secret) \
        .log_level(lark.LogLevel.DEBUG) \
        .build()

    # 构造请求对象
    request1: CreateExportTaskRequest = CreateExportTaskRequest.builder() \
        .request_body(ExportTask.builder()
                      .file_extension("docx")
                      .token(token)
                      .type("docx")
                      .build()) \
        .build()

    # 发起请求
    response1: CreateExportTaskResponse = client.drive.v1.export_task.create(request1)

    # 处理失败返回
    if not response1.success():
        lark.logger.error(
            f"client.drive.v1.export_task.create failed, code: {response1.code}, msg: {response1.msg}, log_id: {response1.get_log_id()}, resp: \n{json.dumps(json.loads(response1.raw.content), indent=4, ensure_ascii=False)}")
        return

    # 处理业务结果
    lark.logger.info(lark.JSON.marshal(response1.data, indent=4))
    ticket = response1.data.ticket

    time.sleep(5)
    # 构造请求对象
    request2: GetExportTaskRequest = GetExportTaskRequest.builder() \
        .ticket(ticket) \
        .token(token) \
        .build()

    # 发起请求
    response2: GetExportTaskResponse = client.drive.v1.export_task.get(request2)

    # 处理失败返回
    if not response2.success():
        lark.logger.error(
            f"client.drive.v1.export_task.get failed, code: {response2.code}, msg: {response2.msg}, log_id: {response2.get_log_id()}, resp: \n{json.dumps(json.loads(response2.raw.content), indent=4, ensure_ascii=False)}")
        return

    # 处理业务结果
    lark.logger.info(lark.JSON.marshal(response2.data, indent=4))
    file_token = response2.data.result.file_token

    # 构造请求对象
    request3: DownloadExportTaskRequest = DownloadExportTaskRequest.builder() \
        .file_token(file_token) \
        .build()

    # 发起请求
    response3: DownloadExportTaskResponse = client.drive.v1.export_task.download(request3)

    # 处理失败返回
    if not response3.success():
        lark.logger.error(
            f"client.drive.v1.export_task.download failed, code: {response3.code}, msg: {response3.msg}, log_id: {response3.get_log_id()}")
        return

    # 处理业务结果
    file_name = unquote(response3.file_name)
    save_path = os.path.join(word_save_dir, file_name)
    with open(save_path, "wb") as f:
        f.write(response3.file.read())

    # 操作word
    doc = Document(save_path)
    placeholder = ""
    for paragraph in doc.paragraphs:
        if not placeholder in paragraph.text:
            continue
        # 清除占位符
        for run in paragraph.runs:
            run.text = run.text.replace(placeholder, "")
        table = doc.add_table(rows=1, cols=[])
        paragraph._element.addnext(table._element)
    doc.save(save_path)

    word_download_url = f_upload_file(save_path)

    return word_download_url


if __name__ == "__main__":
    f_doc_export('YKNBdbs10oA3pCxTdnAczcvOnxc')
    # f_upload_file("/root/project/coze_znjd/大模型企业调查报告/1.docx")