Update main.py

main.py

@@ -1,85 +1,48 @@
-from fastapi import FastAPI, Request, Response, HTTPException
-from starlette.responses import StreamingResponse
-from starlette.status import HTTP_401_UNAUTHORIZED
+from flask import Flask, request, Response, stream_with_context
 import requests
-import json
-import os
 
-app = FastAPI()
+app = Flask(__name__)
 
-@app.middleware("http")
-async def api_key_auth(request: Request, call_next):
-    apikey = os.getenv('apikey')
-    if (request.method == 'POST') and (request.url.path in ['/yyds/v1/chat/completions']):
-        if not apikey:
-            response = await call_next(request)
-            return response
-        auth_header = request.headers.get('Authorization')
-        if not auth_header or not auth_header.startswith('Bearer '):
-            return Response(content="缺少授权信息", status_code=HTTP_401_UNAUTHORIZED)
-        token = auth_header.split(" ")[1]
-        if token != apikey:
-            return Response(content="授权失败", status_code=HTTP_401_UNAUTHORIZED)
-    response = await call_next(request)
-    return response
+# 你想代理的网址
+PROXY_TARGET = "https://chatpro.ai-pro.org"
 
-# FastAPI uses Pydantic to parse the request body into python object
-@app.post("/yyds/v1/chat/completions")
-@app.options("/yyds/v1/chat/completions")  # This maps to the appropriate URL
-async def main(request: Request):
-    if request.method != 'POST':
-        return Response(None, status_code=204, headers={
-            'Access-Control-Allow-Origin': '*',
-            "Access-Control-Allow-Headers": '*',
-            'Content-Type': 'text/event-stream',
-            'Cache-Control': 'no-cache',
-            'Connection': 'keep-alive'
-        })
 
-    headers = dict(request.headers)
-    # Add or modify headers
-    headers['Content-Type'] = 'application/json'
+@app.route('/', defaults={'path': ''})
+@app.route('/<path:path>', methods=["GET", "POST", "PUT", "DELETE", "PATCH"])
+def proxy(path):
+    global PROXY_TARGET
 
-    url = 'https://multillm.ai-pro.org/api/openai-completion'  # target API address
-    headers ={ "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/123.0.0.0 Safari/537.36", "Content-Type": "application/json"}
-    jsonData = await request.json()
-    jsonData["stream"] = True
-    response = requests.post(url, headers=headers, json=jsonData)  # Send the request
+    # 将原始请求的头部信息转发到目标服务器
+    headers = dict(request.headers)
+    
+    # 移除掉 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
+    )
 
-    if response.status_code != 200:
-        return Response(content='Unable to reach the backend API', status_code=502)
+    # 创建一个生成器, 按块读取响应内容, 不需要一次性全部读入内存
+    def generate():
+        for chunk in resp.iter_content(chunk_size=4096):
+            yield chunk
 
-    response_data = response.json()
-    return StreamingResponse(
-        event_stream(response_data),
-        headers={
-            'Access-Control-Allow-Origin': '*',
-            "Access-Control-Allow-Headers": '*',
-            'Content-Type': 'text/event-stream',
-            'Cache-Control': 'no-cache',
-            'Connection': 'keep-alive'
-        }
-    )
+    # 将目标服务器的响应头转发回客户端
+    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
 
-def event_stream(data):
-    # Simplified eventStream function that does not split the content into chunks
-    output = json.dumps({
-        "id": data['id'],
-        "object": 'chat.completion.chunk',
-        "created": data['created'],
-        "model": data['model'],
-        "system_fingerprint": data['system_fingerprint'],
-        "choices": [{
-            "index": 0,
-            "delta": {"role": 'assistant', "content": data['choices'][0]['message']['content']},
-            "logprobs": None,
-            "finish_reason": data['choices'][0]['finish_reason']
-        }]
-    })
-    yield f'data: {output}\n\n'  # The StreamingResponse expects an iterable
 
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+if __name__ == '__main__':
+    app.run(debug=True,port=7860)