# coding=utf-8
import websocket
import json
import os

MY_OPENAI_API_KEY = os.environ.get('MY_OPENAI_API_KEY')


def on_message(ws, message):
    response = json.loads(message)
    # 处理来自 ChatGPT 的响应
    print(response['choices'][0]['text'])

    # 根据您的条件判断是否需要关闭连接
    if '完成对话的条件':
        ws.close()


def on_error(ws, error):
    print(error)


def on_close(ws):
    print("连接已关闭")


def on_open(ws):
    # 发送初始消息
    ws.send(json.dumps({
        'type': 'input',
        'message': '你好，ChatGPT！'
    }))

    # 发送后续消息（用于完成）
    ws.send(json.dumps({
        'type': 'input',
        'message': '技术问题是...'
    }))


if __name__ == '__main__':
    # 创建 WebSocket 连接和其他初始化代码
    ws = websocket.WebSocketApp(
        'wss://api.openai.com/v1/streaming_chat_completion',
        on_message=on_message,
        on_error=on_error,
        on_close=on_close,
        header={
            'Authorization': 'Bearer YOUR_API_KEY'
        }
    )

    ws.on_open = on_open

    # 运行 WebSocket 连接
    ws.run_forever()