Update app.py

app.py

@@ -1,12 +1,10 @@
-from flask import Flask, request, Response, jsonify
 import os
-import requests
 import time
-import threading
+import json
+import asyncio
+import aiohttp
+from aiohttp import web
 
-app = Flask(__name__)
-
-# 环境变量
 PROJECT_ID = os.getenv('PROJECT_ID')
 CLIENT_ID = os.getenv('CLIENT_ID')
 CLIENT_SECRET = os.getenv('CLIENT_SECRET')
@@ -21,36 +19,31 @@ token_cache = {
     'refresh_promise': None
 }
 
-def get_access_token():
+async def get_access_token():
     now = time.time()
 
-    # 如果 token 仍然有效，直接返回
     if token_cache['access_token'] and now < token_cache['expiry'] - 120:
         return token_cache['access_token']
 
-    # 如果已经有一个刷新操作在进行中，等待它完成
     if token_cache['refresh_promise']:
-        token_cache['refresh_promise'].join()
+        await token_cache['refresh_promise']
         return token_cache['access_token']
 
-    # 开始新的刷新操作
-    def refresh_token():
-        try:
-            response = requests.post(TOKEN_URL, json={
+    async def refresh_token():
+        async with aiohttp.ClientSession() as session:
+            async with session.post(TOKEN_URL, json={
                 'client_id': CLIENT_ID,
                 'client_secret': CLIENT_SECRET,
                 'refresh_token': REFRESH_TOKEN,
                 'grant_type': 'refresh_token'
-            })
-            data = response.json()
-            token_cache['access_token'] = data['access_token']
-            token_cache['expiry'] = now + data['expires_in']
-        finally:
-            token_cache['refresh_promise'] = None
-
-    token_cache['refresh_promise'] = threading.Thread(target=refresh_token)
-    token_cache['refresh_promise'].start()
-    token_cache['refresh_promise'].join()
+            }) as response:
+                data = await response.json()
+                token_cache['access_token'] = data['access_token']
+                token_cache['expiry'] = now + data['expires_in']
+
+    token_cache['refresh_promise'] = refresh_token()
+    await token_cache['refresh_promise']
+    token_cache['refresh_promise'] = None
     return token_cache['access_token']
 
 def get_location():
@@ -60,48 +53,62 @@ def get_location():
 def construct_api_url(location):
     return f'https://{location}-aiplatform.googleapis.com/v1/projects/{PROJECT_ID}/locations/{location}/publishers/anthropic/models/{MODEL}:streamRawPredict'
 
-@app.route('/ai/v1/messages', methods=['POST', 'OPTIONS'])
-def handle_request():
+async def handle_request(request):
     if request.method == 'OPTIONS':
         return handle_options()
 
-    # 检查 x-api-key
     api_key = request.headers.get('x-api-key')
     if api_key != API_KEY:
-        return jsonify({
-            'type': 'error',
-            'error': {
-                'type': 'permission_error',
-                'message': 'Your API key does not have permission to use the specified resource.'
-            }
-        }), 403
-
-    access_token = get_access_token()
+        error_response = web.Response(
+            text=json.dumps({
+                'type': 'error',
+                'error': {
+                    'type': 'permission_error',
+                    'message': 'Your API key does not have permission to use the specified resource.'
+                }
+            }),
+            status=403,
+            content_type='application/json'
+        )
+        error_response.headers['Access-Control-Allow-Origin'] = '*'
+        error_response.headers['Access-Control-Allow-Methods'] = 'POST, GET, OPTIONS, DELETE, HEAD'
+        error_response.headers['Access-Control-Allow-Headers'] = 'Content-Type, Authorization, x-api-key, anthropic-version, model'
+        return error_response
+
+    access_token = await get_access_token()
     location = get_location()
-    model = request.headers.get('model', 'claude-3-5-sonnet@20240620')
-    if model == 'claude-3-5-sonnet-20240620':
-        model = 'claude-3-5-sonnet@20240620'
-
     api_url = construct_api_url(location)
 
-    request_body = request.json
+    request_body = await request.json()
+
     if 'anthropic_version' in request_body:
         del request_body['anthropic_version']
     if 'model' in request_body:
         del request_body['model']
-    request_body['anthropic_version'] = "vertex-2023-10-16"
+
+    request_body['anthropic_version'] = 'vertex-2023-10-16'
 
     headers = {
         'Authorization': f'Bearer {access_token}',
         'Content-Type': 'application/json; charset=utf-8'
     }
 
-    response = requests.post(api_url, headers=headers, json=request_body)
-    return Response(response.content, status=response.status_code, content_type=response.headers['Content-Type'])
+    async with aiohttp.ClientSession() as session:
+        async with session.post(api_url, json=request_body, headers=headers) as response:
+            response_body = await response.read()
+            response_headers = response.headers
+            response_status = response.status
+
+    modified_response = web.Response(
+        body=response_body,
+        status=response_status,
+        headers=response_headers
+    )
+    modified_response.headers['Access-Control-Allow-Origin'] = '*'
+    modified_response.headers['Access-Control-Allow-Methods'] = 'POST, GET, OPTIONS'
+    modified_response.headers['Access-Control-Allow-Headers'] = 'Content-Type, Authorization, x-api-key, anthropic-version, model'
 
-@app.route('/', methods=['GET'])
-def index():
-    return "Vertex Claude API Proxy", 200
+    return modified_response
 
 def handle_options():
     headers = {
@@ -109,7 +116,10 @@ def handle_options():
         'Access-Control-Allow-Methods': 'POST, GET, OPTIONS',
         'Access-Control-Allow-Headers': 'Content-Type, Authorization, x-api-key, anthropic-version, model'
     }
-    return '', 204, headers
+    return web.Response(status=204, headers=headers)
+
+app = web.Application()
+app.router.add_route('*', '/', handle_request)
 
 if __name__ == '__main__':
-    app.run(port=8080)
+    web.run_app(app, port=8080)