Update app.py

app.py

@@ -1,24 +1,17 @@
-import os
+from flask import Flask, request, jsonify, Response, make_response
+import requests
+import threading
 import time
-import json
-import asyncio
-import aiohttp
-import logging
-from aiohttp import web, ClientTimeout
-from aiohttp.web import StreamResponse
-
-# 设置日志
-log_level = os.getenv('LOG_LEVEL', 'INFO')
-logging.basicConfig(level=getattr(logging, log_level), format='%(asctime)s - %(levelname)s - %(message)s')
-logger = logging.getLogger(__name__)
+import os
+
+app = Flask(__name__)
 
+MODEL = 'claude-3-5-sonnet@20240620'
 PROJECT_ID = os.getenv('PROJECT_ID')
 CLIENT_ID = os.getenv('CLIENT_ID')
 CLIENT_SECRET = os.getenv('CLIENT_SECRET')
 REFRESH_TOKEN = os.getenv('REFRESH_TOKEN')
 API_KEY = os.getenv('API_KEY')
-MODEL = 'claude-3-5-sonnet@20240620'
-
 TOKEN_URL = 'https://www.googleapis.com/oauth2/v4/token'
 
 token_cache = {
@@ -27,85 +20,66 @@ token_cache = {
     'refresh_promise': None
 }
 
-async def get_access_token():
+def get_access_token():
     now = time.time()
 
     if token_cache['access_token'] and now < token_cache['expiry'] - 120:
-        logger.info("Using cached access token")
         return token_cache['access_token']
 
     if token_cache['refresh_promise']:
-        logger.info("Waiting for ongoing token refresh")
-        await token_cache['refresh_promise']
+        token_cache['refresh_promise'].join()
         return token_cache['access_token']
 
-    async def refresh_token():
-        logger.info("Refreshing access token")
-        async with aiohttp.ClientSession() as session:
-            async with session.post(TOKEN_URL, json={
+    def refresh_token():
+        try:
+            response = requests.post(TOKEN_URL, json={
                 'client_id': CLIENT_ID,
                 'client_secret': CLIENT_SECRET,
                 'refresh_token': REFRESH_TOKEN,
                 'grant_type': 'refresh_token'
-            }) as response:
-                data = await response.json()
-                token_cache['access_token'] = data['access_token']
-                token_cache['expiry'] = now + data['expires_in']
-                logger.info("Access token refreshed successfully")
-
-    token_cache['refresh_promise'] = refresh_token()
-    await token_cache['refresh_promise']
-    token_cache['refresh_promise'] = None
+            })
+            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()
     return token_cache['access_token']
 
 def get_location():
     current_seconds = time.localtime().tm_sec
-    location = 'europe-west1' if current_seconds < 30 else 'us-east5'
-    logger.info(f"Selected location: {location}")
-    return location
-
-def construct_api_url(location, model):
-    url = f'https://{location}-aiplatform.googleapis.com/v1/projects/{PROJECT_ID}/locations/{location}/publishers/anthropic/models/{model}:streamRawPredict'
-    logger.info(f"Constructed API URL: {url}")
-    return url
-
-def format_model_name(model):
-    if model == 'claude-3-5-sonnet-20240620':
-        return 'claude-3-5-sonnet@20240620'
-    return model
-
-async def handle_root(request):
-    logger.info("Received request to root path /")
-    return web.Response(text="GCP Vertex Claude API Proxy", status=200)
-    
-async def handle_request(request):
-    logger.info(f"Received {request.method} request to {request.path}")
-    
+    return 'europe-west1' if current_seconds < 30 else 'us-east5'
+
+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():
     if request.method == 'OPTIONS':
-        logger.info("Handling OPTIONS request")
         return handle_options()
 
     api_key = request.headers.get('x-api-key')
     if api_key != API_KEY:
-        logger.warning("Invalid API key")
-        return create_error_response('Your API key does not have permission to use the specified resource.', 403)
-
-    try:
-        request_body = await request.json()
-        logger.info(f"Received request body: {json.dumps(request_body)}")
-    except json.JSONDecodeError:
-        logger.error("Invalid JSON in request body")
-        return create_error_response('Invalid JSON in request body', 400)
-
-    if not request_body:
-        logger.error("Empty request body")
-        return create_error_response('Empty request body', 400)
+        error_response = make_response(jsonify({
+            'type': 'error',
+            'error': {
+                'type': 'permission_error',
+                'message': 'Your API key does not have permission to use the specified resource.'
+            }
+        }), 403)
+        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()
+    access_token = get_access_token()
     location = get_location()
+    api_url = construct_api_url(location)
 
-    model = format_model_name(request_body.get('model', MODEL))
-    logger.info(f"Using model: {model}")
+    request_body = request.json
 
     if 'anthropic_version' in request_body:
         del request_body['anthropic_version']
@@ -114,81 +88,26 @@ async def handle_request(request):
 
     request_body['anthropic_version'] = 'vertex-2023-10-16'
 
-    api_url = construct_api_url(location, model)
-
     headers = {
         'Authorization': f'Bearer {access_token}',
         'Content-Type': 'application/json; charset=utf-8'
     }
 
-    logger.info(f"Sending request to Anthropic API: {api_url}")
-    
-    # 创建一个 StreamResponse 对象
-    response = StreamResponse(status=200)
-    response.headers['Content-Type'] = 'application/json'
-    response.headers['Access-Control-Allow-Origin'] = '*'
-    response.headers['Access-Control-Allow-Methods'] = 'POST, GET, OPTIONS'
-    response.headers['Access-Control-Allow-Headers'] = 'Content-Type, Authorization, x-api-key, anthropic-version, model'
-    
-    await response.prepare(request)
-
-    try:
-        async with aiohttp.ClientSession() as session:
-            async with session.post(api_url, json=request_body, headers=headers, timeout=ClientTimeout(total=30)) as api_response:
-                api_response.raise_for_status()
-                logger.info(f"Received response from Anthropic API. Status: {api_response.status}")
-                logger.info(f"Response headers: {dict(api_response.headers)}")
-                
-                async for chunk in api_response.content.iter_any():
-                    await response.write(chunk)
-                    
-                logger.info("Finished streaming response from Anthropic API")
-    except aiohttp.ClientError as e:
-        logger.error(f"Error communicating with Anthropic API: {str(e)}")
-        await response.write(json.dumps({'error': 'Error communicating with AI service'}).encode())
-    except asyncio.TimeoutError:
-        logger.error("Request to Anthropic API timed out")
-        await response.write(json.dumps({'error': 'Request timed out'}).encode())
-    except Exception as e:
-        logger.error(f"Unexpected error: {str(e)}")
-        await response.write(json.dumps({'error': 'An unexpected error occurred'}).encode())
-    finally:
-        await response.write_eof()
-        logger.info("Finished sending response back to client")
-
-    return response
-
-def create_error_response(message, status_code):
-    logger.error(f"Creating error response: {message} (Status: {status_code})")
-    response = web.Response(
-        text=json.dumps({
-            'type': 'error',
-            'error': {
-                'type': 'request_error',
-                'message': message
-            }
-        }),
-        status=status_code,
-        content_type='application/json'
-    )
-    response.headers['Access-Control-Allow-Origin'] = '*'
-    response.headers['Access-Control-Allow-Methods'] = 'POST, GET, OPTIONS, DELETE, HEAD'
-    response.headers['Access-Control-Allow-Headers'] = 'Content-Type, Authorization, x-api-key, anthropic-version, model'
-    return response
+    response = requests.post(api_url, headers=headers, json=request_body)
+    modified_response = make_response(response.content, response.status_code)
+    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'
+
+    return modified_response
 
 def handle_options():
-    logger.info("Handling OPTIONS request")
     headers = {
         'Access-Control-Allow-Origin': '*',
         'Access-Control-Allow-Methods': 'POST, GET, OPTIONS',
         'Access-Control-Allow-Headers': 'Content-Type, Authorization, x-api-key, anthropic-version, model'
     }
-    return web.Response(status=204, headers=headers)
-
-app = web.Application()
-app.router.add_get('/', handle_root)
-app.router.add_route('*', '/ai/v1/messages', handle_request)
+    return '', 204, headers
 
 if __name__ == '__main__':
-    logger.info("Starting server on port 8080")
-    web.run_app(app, port=8080)
+    app.run(port=8080)