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
| import argparse
import requests
from flask import Flask, Response, request
import json
parser = argparse.ArgumentParser(description='Argument Parser with Default Parameters')
parser.add_argument('--source_port', type=int, default=8000, help='Source Port Number')
args, _ = parser.parse_known_args()
print("Source Port:", args.source_port)
source_port = args.source_port
app = Flask(__name__)
sess = requests.Session()
AUTH_TOKEN = "auth_token"
AUTH_HEADER_KEY = "Auth"
GET = "GET"
POST = "POST"
ALLOWED_METHODS = [GET, POST]
@app.route('/', defaults={'path': ''}, methods=ALLOWED_METHODS)
@app.route('/<path:path>', methods=ALLOWED_METHODS)
def chat_completions(path):
url = f'http://localhost:{source_port}{request.full_path}'
request_headers = dict(request.headers)
auth = request_headers.get(AUTH_HEADER_KEY, None)
if auth != AUTH_TOKEN:
return Response(json.dumps({"state": "unauthorized", "message": "Unauthorized access"}), status = 401)
no_auth_headers = {k:v for k,v in request_headers.items() if k !=AUTH_HEADER_KEY}
if request.method == GET:
response = sess.get(url, stream=True, headers=no_auth_headers)
else:
response = sess.post(url, stream=True, headers=no_auth_headers, data=request.data, files=request.files, json=request.json)
def generate():
for chunk in response.iter_content(chunk_size=1024):
yield chunk
return Response(generate(), content_type=response.headers['content-type'], headers = dict(response.headers))
if __name__ == '__main__':
app.run()
|