Update app.py

app.py

@@ -1,10 +1,10 @@
 import json
-import httpx
+import sseclient
+import requests
+from flask import Flask, request, Response, stream_with_context
 import random
-from fastapi import FastAPI, Request, HTTPException
-from sse_starlette.sse import EventSourceResponse
 
-app = FastAPI()
+app = Flask(__name__)
 
 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.post("/hf/v1/chat/completions")
-async def chat_completions(request: Request):
-    data = await request.json()
+@app.route('/hf/v1/chat/completions', methods=['POST'])
+def chat_completions():
+    data = request.json
     messages = data.get('messages', [])
     stream = data.get('stream', False)
     
     if not messages:
-        raise HTTPException(status_code=400, detail="No messages provided")
+        return {"error": "No messages provided"}, 400
     
     model = data.get('model', 'gpt-4o')
     if model.startswith('gpt'):
@@ -59,7 +59,7 @@ async def chat_completions(request: Request):
         endpoint = "claude"
         original_api_url = 'https://chatpro.ai-pro.org/api/ask/claude'
     else:
-        raise HTTPException(status_code=400, detail="Unsupported model")
+        return {"error": "Unsupported model"}, 400
 
     headers = {
         'content-type': 'application/json',
@@ -68,7 +68,7 @@ async def chat_completions(request: Request):
         'user-agent': generate_user_agent()
     }
 
-    async def generate():
+    def generate():
         full_response = ""
         while True:
             conversation = "\n".join([f"{msg['role']}: {msg['content']}" for msg in messages])
@@ -80,35 +80,45 @@ async def chat_completions(request: Request):
                 "model": model
             }
             
-            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"
+            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"
                         
-                        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
+                        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 EventSourceResponse(generate())
+        return Response(stream_with_context(generate()), content_type='text/event-stream')
     else:
         full_response = ""
         finish_reason = "stop"
-        async for chunk in generate():
+        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']:
@@ -137,3 +147,6 @@ async def chat_completions(request: Request):
                 "total_tokens": 0
             }
         }
+
+if __name__ == '__main__':
+    app.run(debug=True, port=5000)