Update app.py

app.py

@@ -1,168 +1,168 @@
-from flask import Flask, request, Response, stream_with_context
-import requests
-import json
-import uuid
-import time
-import os
-
-app = Flask(__name__)
-
-MODELS = {
-   'claude-3-5-haiku': 'bedrock-anthropic.claude-3-5-haiku',
-   'claude-3-5-sonnet': 'bedrock-anthropic.claude-3-5-sonnet-v2-0',
-   'nova-lite-v1-0': 'bedrock-amazon.nova-lite-v1-0',
-   'nova-pro-v1-0': 'bedrock-amazon.nova-pro-v1-0',
-   'llama3-1-7b': 'bedrock-meta.llama3-1-8b-instruct-v1',
-   'llama3-1-70b': 'bedrock-meta.llama3-1-70b-instruct-v1',
-   'mistral-small': 'bedrock-mistral.mistral-small-2402-v1-0',
-   'mistral-large': 'bedrock-mistral.mistral-large-2407-v1-0'
-}
-
-def validate_key(key):
-   try:
-       parts = key.split('|||', 3)
-       if len(parts) != 3:
-           return None, None, None
-       return parts[0], parts[2], parts[1]  
-   except:
-       return None, None, None
-
-def create_partyrock_request(openai_req, app_id):
-   return {
-       "messages": [
-           {
-               "role": msg["role"],
-               "content": [{"text": msg["content"]}]
-           }
-           for msg in openai_req['messages']
-       ],
-       "modelName": MODELS.get(openai_req.get('model', 'claude-3-5-haiku')),
-       "context": {"type": "chat-widget", "appId": app_id},
-       "options": {"temperature": 0},
-       "apiVersion": 3
-   }
-
-@app.route('/', methods=['GET'])
-def home():
-   return {"status": "PartyRock API Service Running", "port": 8803}
-
-@app.route('/v1/chat/completions', methods=['POST']) 
-def chat():
-   try:
-       api_key = request.headers.get('Authorization', '').replace('Bearer ', '')
-       app_id, cookie, csrf_token = validate_key(api_key)
-       
-       print(f"App ID: {app_id}")
-       print(f"Cookie length: {len(cookie) if cookie else 'None'}")
-       print(f"CSRF Token: {csrf_token}")
-       
-       if not app_id or not cookie or not csrf_token:
-           return Response("Invalid API key format. Use: appId|||csrf_token|||cookies", status=401)
-
-       headers = {
-           'accept': 'text/event-stream',
-           'accept-language': 'zh-CN,zh;q=0.9,en;q=0.8',
-           'anti-csrftoken-a2z': csrf_token,
-           'content-type': 'application/json',
-           'origin': 'https://partyrock.aws',
-           'referer': f'https://partyrock.aws/u/chatyt/{app_id}',
-           'cookie': cookie,
-           'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
-       }
-
-       response = requests.post(
-           'https://partyrock.aws/stream/getCompletion',
-           headers=headers,
-           json=create_partyrock_request(request.json, app_id),
-           stream=True
-       )
-       
-       if response.status_code != 200:
-           return Response(f"PartyRock API error: {response.text}", status=response.status_code)
-
-       if not request.json.get('stream', False):
-           try:
-               full_content = ""
-               buffer = ""
-               for chunk in response.iter_content(chunk_size=1024, decode_unicode=True):
-                   if chunk:
-                       buffer += chunk
-                       while '\n' in buffer:
-                           line, buffer = buffer.split('\n', 1)
-                           if line.startswith('data: '):
-                               try:
-                                   data = json.loads(line[6:])
-                                   if data["type"] == "text":
-                                       content = data["text"]
-                                       if isinstance(content, str):
-                                           content = content.encode('latin1').decode('utf-8')
-                                       if content and content != " ":
-                                           full_content += content
-                               except:
-                                   continue
-
-               return {
-                   "id": str(uuid.uuid4()),
-                   "object": "chat.completion",
-                   "created": int(time.time()),
-                   "model": request.json.get('model', 'claude-3-haiku'),
-                   "choices": [{
-                       "message": {
-                           "role": "assistant",
-                           "content": full_content
-                       },
-                       "finish_reason": "stop",
-                       "index": 0
-                   }]
-               }
-           except Exception as e:
-               print(f"Error processing response: {str(e)}")
-               return Response("Failed to process response", status=500)
-
-       def generate():
-           try:
-               buffer = ""
-               for chunk in response.iter_content(chunk_size=1024, decode_unicode=True):
-                   if chunk:
-                       buffer += chunk
-                       while '\n' in buffer:
-                           line, buffer = buffer.split('\n', 1)
-                           if line.startswith('data: '):
-                               try:
-                                   data = json.loads(line[6:])
-                                   if data["type"] == "text":
-                                       content = data["text"]
-                                       if isinstance(content, str):
-                                           content = content.encode('latin1').decode('utf-8')
-                                       chunk_resp = {
-                                           "id": str(uuid.uuid4()),
-                                           "object": "chat.completion.chunk",
-                                           "created": int(time.time()),
-                                           "model": request.json.get('model', 'claude-3-haiku'),
-                                           "choices": [{
-                                               "delta": {"content": content},
-                                               "index": 0,
-                                               "finish_reason": None
-                                           }]
-                                       }
-                                       yield f"data: {json.dumps(chunk_resp, ensure_ascii=False)}\n\n"
-                               except:
-                                   continue
-
-               yield f"data: {json.dumps({'choices':[{'delta':{'content':''},'index':0,'finish_reason':'stop'}]})}\n\n"
-               yield "data: [DONE]\n\n"
-           except Exception as e:
-               print(f"Error in generate: {str(e)}")
-               return
-
-       return Response(
-           stream_with_context(generate()),
-           content_type='text/event-stream'
-       )
-
-   except Exception as e:
-       print(f"Error: {str(e)}")
-       return Response(f"Internal server error: {str(e)}", status=500)
-
-if __name__ == '__main__':
-   app.run(host='0.0.0.0', port=8803, debug=True)
+from flask import Flask, request, Response, stream_with_context
+import requests
+import json
+import uuid
+import time
+import os
+
+app = Flask(__name__)
+
+MODELS = {
+   'claude-3-5-haiku': 'bedrock-anthropic.claude-3-5-haiku',
+   'claude-3-5-sonnet': 'bedrock-anthropic.claude-3-5-sonnet-v2-0',
+   'nova-lite-v1-0': 'bedrock-amazon.nova-lite-v1-0',
+   'nova-pro-v1-0': 'bedrock-amazon.nova-pro-v1-0',
+   'llama3-1-7b': 'bedrock-meta.llama3-1-8b-instruct-v1',
+   'llama3-1-70b': 'bedrock-meta.llama3-1-70b-instruct-v1',
+   'mistral-small': 'bedrock-mistral.mistral-small-2402-v1-0',
+   'mistral-large': 'bedrock-mistral.mistral-large-2407-v1-0'
+}
+
+def validate_key(key):
+   try:
+       parts = key.split('|||', 3)
+       if len(parts) != 3:
+           return None, None, None
+       return parts[0], parts[2], parts[1]  
+   except:
+       return None, None, None
+
+def create_partyrock_request(openai_req, app_id):
+   return {
+       "messages": [
+           {
+               "role": msg["role"],
+               "content": [{"text": msg["content"]}]
+           }
+           for msg in openai_req['messages']
+       ],
+       "modelName": MODELS.get(openai_req.get('model', 'claude-3-5-haiku')),
+       "context": {"type": "chat-widget", "appId": app_id},
+       "options": {"temperature": 0},
+       "apiVersion": 3
+   }
+
+@app.route('/', methods=['GET'])
+def home():
+   return {"status": "PartyRock API Service Running", "port": 8803}
+
+@app.route('/v1/chat/completions', methods=['POST']) 
+def chat():
+   try:
+       api_key = request.headers.get('Authorization', '').replace('Bearer ', '')
+       app_id, cookie, csrf_token = validate_key(api_key)
+       
+       print(f"App ID: {app_id}")
+       print(f"Cookie length: {len(cookie) if cookie else 'None'}")
+       print(f"CSRF Token: {csrf_token}")
+       
+       if not app_id or not cookie or not csrf_token:
+           return Response("Invalid API key format. Use: appId|||csrf_token|||cookies", status=401)
+
+       headers = {
+           'accept': 'text/event-stream',
+           'accept-language': 'zh-CN,zh;q=0.9,en;q=0.8',
+           'anti-csrftoken-a2z': csrf_token,
+           'content-type': 'application/json',
+           'origin': 'https://partyrock.aws',
+           'referer': f'https://partyrock.aws/u/chatyt/{app_id}',
+           'cookie': cookie,
+           'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
+       }
+
+       response = requests.post(
+           'https://partyrock.aws/stream/getCompletion',
+           headers=headers,
+           json=create_partyrock_request(request.json, app_id),
+           stream=True
+       )
+       
+       if response.status_code != 200:
+           return Response(f"PartyRock API error: {response.text}", status=response.status_code)
+
+       if not request.json.get('stream', False):
+           try:
+               full_content = ""
+               buffer = ""
+               for chunk in response.iter_content(chunk_size=1024, decode_unicode=True):
+                   if chunk:
+                       buffer += chunk
+                       while '\n' in buffer:
+                           line, buffer = buffer.split('\n', 1)
+                           if line.startswith('data: '):
+                               try:
+                                   data = json.loads(line[6:])
+                                   if data["type"] == "text":
+                                       content = data["text"]
+                                       if isinstance(content, str):
+                                           content = content.encode('latin1').decode('utf-8')
+                                       if content and content != " ":
+                                           full_content += content
+                               except:
+                                   continue
+
+               return {
+                   "id": str(uuid.uuid4()),
+                   "object": "chat.completion",
+                   "created": int(time.time()),
+                   "model": request.json.get('model', 'claude-3-haiku'),
+                   "choices": [{
+                       "message": {
+                           "role": "assistant",
+                           "content": full_content
+                       },
+                       "finish_reason": "stop",
+                       "index": 0
+                   }]
+               }
+           except Exception as e:
+               print(f"Error processing response: {str(e)}")
+               return Response("Failed to process response", status=500)
+
+       def generate():
+           try:
+               buffer = ""
+               for chunk in response.iter_content(chunk_size=1024, decode_unicode=True):
+                   if chunk:
+                       buffer += chunk
+                       while '\n' in buffer:
+                           line, buffer = buffer.split('\n', 1)
+                           if line.startswith('data: '):
+                               try:
+                                   data = json.loads(line[6:])
+                                   if data["type"] == "text":
+                                       content = data["text"]
+                                       if isinstance(content, str):
+                                           content = content.encode('latin1').decode('utf-8')
+                                       chunk_resp = {
+                                           "id": str(uuid.uuid4()),
+                                           "object": "chat.completion.chunk",
+                                           "created": int(time.time()),
+                                           "model": request.json.get('model', 'claude-3-haiku'),
+                                           "choices": [{
+                                               "delta": {"content": content},
+                                               "index": 0,
+                                               "finish_reason": None
+                                           }]
+                                       }
+                                       yield f"data: {json.dumps(chunk_resp, ensure_ascii=False)}\n\n"
+                               except:
+                                   continue
+
+               yield f"data: {json.dumps({'choices':[{'delta':{'content':''},'index':0,'finish_reason':'stop'}]})}\n\n"
+               yield "data: [DONE]\n\n"
+           except Exception as e:
+               print(f"Error in generate: {str(e)}")
+               return
+
+       return Response(
+           stream_with_context(generate()),
+           content_type='text/event-stream'
+       )
+
+   except Exception as e:
+       print(f"Error: {str(e)}")
+       return Response(f"Internal server error: {str(e)}", status=500)
+
+if __name__ == '__main__':
+    app.run(host='0.0.0.0', port=8803, debug=False)  # 关闭 debug 模式