将ChatGLM以SSE模式代理到本地

'''
Author: zhao-leo 18055219130@163.com
Date: 2024-10-09 18:17:52
LastEditTime: 2024-10-09 20:13:55
'''
import requests
from flask import Flask, request, Response
import argparse
import os
import socket

# 创建一个命令行参数解析器 解析一个端口参数
argparser = argparse.ArgumentParser(description='ChatGLM Proxy')
argparser.add_argument('--port', type=int, default=None, help='Proxy server port')

# 查找一个开放的端口
def find_open_port():
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(('', 0))
        return s.getsockname()[1]

# 从环境变量中读取端口，如果未设置则随机选择一个开放的端口
PORT = int(os.getenv('CHATGLM_PROXY_PORT', find_open_port()))

app = Flask(__name__)
TARGET_URL = "https://open.bigmodel.cn/api/paas/v4/chat/completions"  # 替换为你的目标 URL

@app.route('/v1/chat/completions', methods=['POST'])
def proxy_request():
    # 从请求头中获取 Authorization
    auth_header = request.headers.get('Authorization')

    # 从请求体中获取 JSON 数据
    json_data = request.get_json()

    # 构造代理请求
    headers = {
        'Authorization': auth_header,
        'Content-Type': 'application/json'
    }

    # 发送请求到目标 URL
    response = requests.post(TARGET_URL, headers=headers, json=json_data, stream=True)

    # 定义一个生成器，逐块返回响应内容
    def generate():
        for chunk in response.iter_content(chunk_size=1024):
            if chunk:
                yield chunk

    # 返回响应，设置为 SSE 格式
    return Response(generate(), content_type='text/event-stream')

# 优先级 1: 命令行参数 2: 环境变量 3: 默认端口
args = argparser.parse_args()
port = args.port if args.port else PORT
# 启动服务器
app.run(port=port)