将ChatGLM以SSE模式代理到本地

事情的起因是Zotero的插件Awesome GPT只支持/v1/chat/completions路径访问，而ChatGLM的API是https://open.bigmodel.cn/api/paas/v4/,因此只能出此下策，将端口代理到本地，然后通过localhost访问，代码如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58


'''
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)