Update app.py

app.py

@@ -1,10 +1,10 @@
 import json
-import sseclient
-import requests
-from flask import Flask, request, Response, stream_with_context
+import httpx
 import random
+from fastapi import FastAPI, Request, HTTPException
+from sse_starlette.sse import EventSourceResponse
 
-app = Flask(__name__)
+app = FastAPI()
 
 def generate_random_ip():
     return f"{random.randint(1, 255)}.{random.randint(0, 255)}.{random.randint(0, 255)}.{random.randint(0, 255)}"
@@ -42,14 +42,14 @@ def format_openai_response(content, finish_reason=None):
         }]
     }
 
-@app.route('/hf/v1/chat/completions', methods=['POST'])
-def chat_completions():
-    data = request.json
+@app.post("/hf/v1/chat/completions")
+async def chat_completions(request: Request):
+    data = await request.json()
     messages = data.get('messages', [])
     stream = data.get('stream', False)
     
     if not messages:
-        return {"error": "No messages provided"}, 400
+        raise HTTPException(status_code=400, detail="No messages provided")
     
     model = data.get('model', 'gpt-4o')
     if model.startswith('gpt'):
@@ -59,7 +59,7 @@ def chat_completions():
         endpoint = "claude"
         original_api_url = 'https://chatpro.ai-pro.org/api/ask/claude'
     else:
-        return {"error": "Unsupported model"}, 400
+        raise HTTPException(status_code=400, detail="Unsupported model")
 
     headers = {
         'content-type': 'application/json',
@@ -68,7 +68,7 @@ def chat_completions():
         'user-agent': generate_user_agent()
     }
 
-    def generate():
+    async def generate():
         full_response = ""
         while True:
             conversation = "\n".join([f"{msg['role']}: {msg['content']}" for msg in messages])
@@ -80,45 +80,35 @@ def chat_completions():
                 "model": model
             }
             
-            response = requests.post(original_api_url, headers=headers, json=payload, stream=True)
-            client = sseclient.SSEClient(response)
-            
-            for event in client.events():
-                if event.data.startswith('{"text":'):
-                    data = json.loads(event.data)
-                    new_content = data['text'][len(full_response):]
-                    full_response = data['text']
-                    
-                    if new_content:
-                        yield f"data: {json.dumps(format_openai_response(new_content))}\n\n"
-                
-                elif '"final":true' in event.data:
-                    final_data = json.loads(event.data)
-                    response_message = final_data.get('responseMessage', {})
-                    finish_reason = response_message.get('finish_reason', 'stop')
-                    
-                    if finish_reason == 'length':
-                        messages.append({"role": "assistant", "content": full_response})
-                        messages.append({"role": "user", "content": "Please continue your output and do not repeat the previous content"})
-                        break  # Jump out of the current loop and continue with the next request
-                    else:
-                        last_content = response_message.get('text', '')
-                        if last_content and last_content != full_response:
-                            yield f"data: {json.dumps(format_openai_response(last_content[len(full_response):]))}\n\n"
+            async with httpx.AsyncClient() as client:
+                async with client.stream("POST", original_api_url, headers=headers, json=payload) as response:
+                    async for line in response.aiter_lines():
+                        if line.startswith('{"text":'):
+                            data = json.loads(line)
+                            new_content = data['text'][len(full_response):]
+                            full_response = data['text']
+                            
+                            if new_content:
+                                yield f"data: {json.dumps(format_openai_response(new_content))}\n\n"
                         
-                        yield f"data: {json.dumps(format_openai_response('', finish_reason))}\n\n"
-                        yield "data: [DONE]\n\n"
-                        return
-
-        yield f"data: {json.dumps(format_openai_response('', 'stop'))}\n\n"
-        yield "data: [DONE]\n\n"
+                        elif '"final":true' in line:
+                            final_data = json.loads(line)
+                            response_message = final_data.get('responseMessage', {})
+                            finish_reason = response_message.get('finish_reason', 'stop')
+                            last_content = response_message.get('text', '')
+                            if last_content and last_content != full_response:
+                                yield f"data: {json.dumps(format_openai_response(last_content[len(full_response):]))}\n\n"
+                            
+                            yield f"data: {json.dumps(format_openai_response('', finish_reason))}\n\n"
+                            yield "data: [DONE]\n\n"
+                            return
 
     if stream:
-        return Response(stream_with_context(generate()), content_type='text/event-stream')
+        return EventSourceResponse(generate())
     else:
         full_response = ""
         finish_reason = "stop"
-        for chunk in generate():
+        async for chunk in generate():
             if chunk.startswith("data: ") and not chunk.strip() == "data: [DONE]":
                 response_data = json.loads(chunk[6:])
                 if 'choices' in response_data and response_data['choices']:
@@ -147,6 +137,3 @@ def chat_completions():
                 "total_tokens": 0
             }
         }
-
-if __name__ == '__main__':
-    app.run(debug=True, port=5000)