Update app.py

app.py

@@ -1,11 +1,11 @@
 import json
-import sseclient
 import random
-import aiohttp
 import asyncio
-from flask import Flask, request, Response, stream_with_context
+import aiohttp
+from fastapi import FastAPI, Request, Response
+from fastapi.responses import StreamingResponse
 
-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)}"
@@ -43,15 +43,44 @@ def format_openai_response(content, finish_reason=None):
         }]
     }
 
-@app.route('/hf/v1/chat/completions', methods=['POST'])
-async def chat_completions():
-    data = request.json
+async def parse_sse(response):
+    """
+    Parses Server-Sent Events from an aiohttp response.
+    """
+    event = {}
+    async for line in response.content:
+        line = line.decode('utf-8').strip()
+        if not line:
+            # Empty line indicates end of event
+            if event:
+                yield event
+                event = {}
+            continue
+        if line.startswith(':'):
+            # Comment line, ignore
+            continue
+        if ':' in line:
+            key, value = line.split(':', 1)
+            value = value.lstrip()
+        else:
+            key = line
+            value = ''
+        if key in event:
+            event[key] += '\n' + value
+        else:
+            event[key] = value
+    if event:
+        yield event
+
+@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
-
+    
     model = data.get('model', 'gpt-4o')
     if model.startswith('gpt'):
         endpoint = "openAI"
@@ -74,7 +103,7 @@ async def chat_completions():
         while True:
             conversation = "\n".join([f"{msg['role']}: {msg['content']}" for msg in messages])
             conversation += "\nPlease follow and reply to the user’s recent messages and avoid answers that summarize the conversation history."
-
+            
             payload = {
                 "text": conversation,
                 "endpoint": endpoint,
@@ -82,32 +111,52 @@ async def chat_completions():
             }
 
             async with aiohttp.ClientSession() as session:
-                async with session.post(original_api_url, headers=headers, json=payload) as response:
-                    async for line in response.content:
-                        if line.startswith(b'data: '):
-                            event_data = line.decode().lstrip('data: ')
-                            event = json.loads(event_data)
-                            if event.get("final"):
-                                yield f"data: {json.dumps(format_openai_response('', 'stop'))}\n\n"
-                                return
-                            new_content = event.get('text', '')[len(full_response):]
-                            if new_content:
-                                full_response = event.get('text', '')
-                                yield f"data: {json.dumps(format_openai_response(new_content))}\n\n"
-    
+                async with session.post(original_api_url, headers=headers, json=payload) as resp:
+                    if resp.status != 200:
+                        yield f"data: {json.dumps({'error': 'Failed to connect to upstream server'})}\n\n"
+                        return
+
+                    async for event in parse_sse(resp):
+                        if 'data' in event:
+                            data = event['data']
+                            if data.startswith('{"text":'):
+                                data_json = json.loads(data)
+                                new_content = data_json['text'][len(full_response):]
+                                full_response = data_json['text']
+                                if new_content:
+                                    yield f"data: {json.dumps(format_openai_response(new_content))}\n\n"
+                            elif '"final":true' in data:
+                                final_data = json.loads(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"
         async for chunk in generate():
-            response_data = json.loads(chunk[6:])
-            if 'choices' in response_data and response_data['choices']:
-                delta = response_data['choices'][0].get('delta', {})
-                if 'content' in delta:
-                    full_response += delta['content']
-                if 'finish_reason' in delta:
-                    finish_reason = delta['finish_reason']
+            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']:
+                    delta = response_data['choices'][0].get('delta', {})
+                    if 'content' in delta:
+                        full_response += delta['content']
+                    if 'finish_reason' in delta:
+                        finish_reason = delta['finish_reason']
 
         return {
             "id": "chatcmpl-123",
@@ -128,6 +177,3 @@ async def chat_completions():
                 "total_tokens": 0
             }
         }
-
-if __name__ == '__main__':
-    app.run(debug=True, port=5000)