dsProject/dsLightRag/KeLing/Kit/KlCommon.py

import os
import json
import logging
import time
from datetime import datetime, timedelta
import requests
import jwt

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
log = logging.getLogger(__name__)


class KlCommon:
    # 获取项目根目录路径
    project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../../..')).replace('\\', '/') + '/dsAi'
    # 拼接相对路径
    base_path = project_root + '/src/main/java/com/dsideal/aiSupport/Util/KeLing/Example/'

    ak = 'c992fd02624d4900a93ca3b6da03d9e9'  # 填写access key
    sk = 'b37f67a00eb44f9bb57e4d530c328e1d'  # 填写secret key

    @staticmethod
    def get_jwt():
        try:
            # 有效时间，此处示例代表当前时间+1800s(30min)
            expired_at = datetime.utcnow() + timedelta(seconds=1800)
            # 开始生效的时间，此处示例代表当前时间-5秒
            not_before = datetime.utcnow() - timedelta(seconds=5)

            # 构建JWT令牌
            payload = {
                'iss': KlCommon.ak,
                'exp': expired_at,
                'nbf': not_before
            }

            # 使用HS256算法签名
            jwt_token = jwt.encode(payload, KlCommon.sk, algorithm='HS256', headers={'alg': 'HS256'})
            return jwt_token
        except Exception as e:
            log.error(f"获取JWT令牌失败: {str(e)}")
            return None

    @staticmethod
    def download_file(file_url, save_file_path):
        """
        从URL下载文件到指定路径

        @param file_url: 文件URL
        @param save_file_path: 保存路径
        @throws Exception: 下载过程中的异常
        """
        try:
            # 发送GET请求下载文件
            response = requests.get(file_url, stream=True)
            response.raise_for_status()  # 检查响应状态码

            # 保存文件
            with open(save_file_path, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    f.write(chunk)

            file_size = os.path.getsize(save_file_path)
            log.info(f"文件下载成功，保存路径: {save_file_path}, 文件大小: {file_size}字节")
        except Exception as e:
            log.error(f"文件下载失败: {str(e)}")
            raise Exception(f"文件下载失败: {str(e)}") from e

    @staticmethod
    def query_task_status(task_id, query_path, log_prefix):
        """
        通用查询任务状态方法

        @param task_id: 任务ID
        @param query_path: 查询路径
        @param log_prefix: 日志前缀，用于区分不同类型的查询
        @return: 任务结果
        @throws Exception: 异常信息
        """
        # 获取JWT令牌
        jwt_token = KlCommon.get_jwt()
        if not jwt_token:
            raise Exception("获取JWT令牌失败")

        # 构建请求URL
        url = f"https://api.klingai.com{query_path}{task_id}"

        # 设置请求头
        headers = {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {jwt_token}'
        }

        # 发送GET请求
        try:
            response = requests.get(url, headers=headers)
            response.raise_for_status()  # 检查响应状态码

            # 解析响应
            response_body = response.text
            response_json = json.loads(response_body)
            log.info(f"{log_prefix}查询任务状态响应：{response_body}")

            # 检查响应状态
            code = response_json.get('code')
            if code != 0:
                message = response_json.get('message', '')
                from .KlErrorCode import KlErrorCode
                solution = KlErrorCode.get_solution_by_code(code)
                error_msg = f"{log_prefix}查询任务状态失败：[{code}] {message} - {solution}"
                raise Exception(error_msg)

            return response_json
        except requests.exceptions.RequestException as e:
            log.error(f"{log_prefix}查询任务状态请求异常: {str(e)}")
            raise Exception(f"{log_prefix}查询任务状态请求异常: {str(e)}") from e
        except json.JSONDecodeError as e:
            log.error(f"{log_prefix}解析响应JSON失败: {str(e)}")
            raise Exception(f"{log_prefix}解析响应JSON失败: {str(e)}") from e
        except Exception as e:
            log.error(f"{log_prefix}查询任务状态发生未知异常: {str(e)}")
            raise