Update main.py

main.py

@@ -1,48 +1,43 @@
-from flask import Flask, request, Response, stream_with_context
+from flask import Flask, request, Response
 import requests
 
 app = Flask(__name__)
 
-# 你想代理的网址
-PROXY_TARGET = "https://chatpro.ai-pro.org"
+PROXY_TARGET = "https://chatpro.ai-pro.org/chat"
 
-
-@app.route('/', defaults={'path': ''})
-@app.route('/<path:path>', methods=["GET", "POST", "PUT", "DELETE", "PATCH"])
+@app.route('/')
+def home():
+    global PROXY_TARGET
+    # 获取原始HTTP GET请求的参数
+    resp = requests.get(f"{PROXY_TARGET}{request.full_path}")
+    # 返回从目标网站获取的响应数据
+    excluded_headers = ['content-encoding', 'content-length', 'transfer-encoding', 'connection']
+    headers = [(name, value) for (name, value) in resp.raw.headers.items()
+               if name.lower() not in excluded_headers]
+    response = Response(resp.content, resp.status_code, headers)
+    return response
+
+@app.route('/<path:path>', methods=['GET', 'POST', 'PUT', 'DELETE'])
 def proxy(path):
     global PROXY_TARGET
-
-    # 将原始请求的头部信息转发到目标服务器
-    headers = dict(request.headers)
+    # 构造目标URL
+    url = f"{PROXY_TARGET}/{path}"
+    # 根据原始请求的方法来发送对应的请求到目标服务器
+    if request.method == 'GET':
+        resp = requests.get(url, stream=True)
+    elif request.method == 'POST':
+        resp = requests.post(url, json=request.json, stream=True)
+    elif request.method == 'PUT':
+        resp = requests.put(url, json=request.json, stream=True)
+    elif request.method == 'DELETE':
+        resp = requests.delete(url, stream=True)
     
-    # 移除掉 Flask/Werkzeug 会默认修改的Headers
-    headers.pop('Host', None)
-    headers.pop('Content-Length', None)
-
-    # 使用 requests 发送经过构造的请求到目标服务器
-    resp = requests.request(
-        method=request.method,
-        url=f"{PROXY_TARGET}/{path}",
-        headers=headers,
-        data=request.get_data(),
-        cookies=request.cookies,
-        allow_redirects=False,
-        stream=True
-    )
-
-    # 创建一个生成器, 按块读取响应内容, 不需要一次性全部读入内存
-    def generate():
-        for chunk in resp.iter_content(chunk_size=4096):
-            yield chunk
-
-    # 将目标服务器的响应头转发回客户端
-    response_headers = [(name, value) for (name, value) in resp.raw.headers.items()]
-
-    # 使用 stream_with_context 包装响应内容生成器
-    response_stream = Response(stream_with_context(generate()), resp.status_code, response_headers)
-
-    return response_stream
-
+    # 将从目标网站获取的响应返回给客户端
+    excluded_headers = ['content-encoding', 'content-length', 'transfer-encoding', 'connection']
+    headers = [(name, value) for (name, value) in resp.raw.headers.items()
+               if name.lower() not in excluded_headers]
+    response = Response(resp.content, resp.status_code , headers)
+    return response
 
 if __name__ == '__main__':
-    app.run(debug=True,host='0.0.0.0',port=7860)
+    app.run(debug=True)