aigc-test/gpt_tag.py

import time
import json
import requests
import traceback
from config import set_config
from log import Log

config_ = set_config()
log_ = Log()


def get_tag(prompt):
    retry = 1
    result_content = None
    while retry < 3:
        try:
            log_.info(f"retry = {retry}")
            headers = {
                'Content-Type': 'application/json',
                # 'Authorization': f'Bearer {os.environ["OPENAI_API_KEY"]}',
                'Authorization': f'Bearer {config_.GPT_OPENAI_API_KEY}',
            }
            proxies = config_.PROXIES
            json_data = {
                'model': 'gpt-3.5-turbo',
                'messages': [
                    {
                        'role': 'user',
                        'content': f'{prompt}',
                    },
                ],
            }
            response = requests.post(url=config_.GPT_HOST, headers=headers, json=json_data, proxies=proxies)
            print(response.json())
            print(response.json()['choices'][0]['message']['content'])
            print('\n')
            result_content = response.json()['choices'][0]['message']['content']
            return result_content
        except Exception as e:
            print(e)
            log_.error(traceback.format_exc())
            retry += 1
            time.sleep(30)
            continue
    return result_content


def request_gpt(prompt):
    retry_count = 0
    result_content = None
    while retry_count < config_.RETRY_MAX_COUNT:
        retry_count += 1
        try:
            # response = requests.post(url=config_.GPT_URL, json={'content': prompt, 'auth': config_.GPT_OPENAI_API_KEY})
            response = requests.post(url=config_.GPT_URL, json={'content': prompt})
            # print(response.json())
            # print(response.json()['choices'][0]['message']['content'])
            # print('\n')
            # result_content = response.json()['choices'][0]['message']['content']
            # log_.info(f"response.text: {response.text}")
            res_data = json.loads(response.text)
            if res_data['code'] != 0:
                time.sleep(10)
                continue
            result_content = res_data['data']['choices'][0]['message']['content']
            return result_content
        except Exception:
            time.sleep(10)
            continue
    return result_content