strategy-flow-simulation/commom/llm_call.py

# -*- coding:utf-8 -*-
"""
@author: isaacqyang
@time: 2023/9/19
@desc: 
"""
import time

import requests

from config import BaseConfig
from .logger import get_logger

logger_error = get_logger("error")


def f_file_upload(file_path: str):
    req_head = {
        "Authorization": BaseConfig.token,
    }
    with open(file_path, 'rb') as file:
        res = requests.post(BaseConfig.file_upload_url, headers=req_head, files={'file': file})
    res_json = res.json()
    print(res_json)
    return res_json['data']['id']


def call_llm(prompt: str, content_type="text"):
    req_head = {
        "Authorization": BaseConfig.token,
        "Content-Type": "application/json",
    }
    req_data = {
        "bot_id": BaseConfig.bot_id,
        "user_id": "test",
        "stream": False,
        "auto_save_history": True,
        "additional_messages": [
            {
                "role": "user",
                "content": prompt,
                "content_type": content_type
            }
        ]
    }

    res_create = requests.post(" https://api.coze.cn/v1/conversation/create",
                               headers=req_head)
    coversition_id = res_create.json()["data"]["id"]
    res_chat = requests.post(f" https://api.coze.cn/v3/chat?conversation_id={coversition_id}",
                             headers=req_head, json=req_data)
    chat_id = res_chat.json()["data"]["id"]
    while True:
        res_retrieve = requests.get(
            f" https://api.coze.cn/v3/chat/retrieve?chat_id={chat_id}&conversation_id={coversition_id}",
            headers=req_head)
        # print(res_retrieve.json()["data"]["status"])
        status = res_retrieve.json()["data"]["status"]
        if status == "completed":
            res_message = requests.get(
                f" https://api.coze.cn/v3/chat/message/list?chat_id={chat_id}&conversation_id={coversition_id}",
                headers=req_head)
            res_json = res_message.json()
            print(res_json)
            data = res_json['data']
            for msg in data:
                if msg["type"] == "answer":
                    return msg["content"]
        time.sleep(1)