Update app.py

app.py

@@ -1,10 +1,12 @@
 import json
-import sseclient
-import requests
-from flask import Flask, request, Response, stream_with_context
 import random
+import asyncio
+import aiohttp
+from fastapi import FastAPI, Request, Response
+from fastapi.responses import StreamingResponse
+from aiosseclient import aiosseclient
 
-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,9 +44,9 @@ 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)
     
@@ -68,7 +70,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])
@@ -79,46 +81,43 @@ def chat_completions():
                 "endpoint": endpoint,
                 "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 aiosseclient(original_api_url, method='POST', headers=headers, json=payload) as client:
+                async for event in client:
+                    if event.data.startswith('{"text":'):
+                        data = json.loads(event.data)
+                        new_content = data['text'][len(full_response):]
+                        full_response = data['text']
                         
-                        yield f"data: {json.dumps(format_openai_response('', finish_reason))}\n\n"
-                        yield "data: [DONE]\n\n"
-                        return
+                        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  # 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"
+                            
+                            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"
 
     if stream:
-        return Response(stream_with_context(generate()), content_type='text/event-stream')
+        return StreamingResponse(generate(), media_type='text/event-stream')
     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 +146,3 @@ def chat_completions():
                 "total_tokens": 0
             }
         }
-
-if __name__ == '__main__':
-    app.run(debug=True, port=5000)