Update app.py

app.py

@@ -4,7 +4,6 @@ import asyncio
 import aiohttp
 from fastapi import FastAPI, Request, Response
 from fastapi.responses import StreamingResponse
-from aiosseclient import aiosseclient
 
 app = FastAPI()
 
@@ -44,6 +43,20 @@ def format_openai_response(content, finish_reason=None):
         }]
     }
 
+def sse_parser():
+    """Generator function to parse SSE messages."""
+    data = ''
+    while True:
+        line = yield
+        if line == '':
+            if data:
+                yield data
+                data = ''
+        elif line.startswith('data:'):
+            data += line[5:].strip()
+        else:
+            continue
+
 @app.post('/hf/v1/chat/completions')
 async def chat_completions(request: Request):
     data = await request.json()
@@ -82,35 +95,51 @@ async def chat_completions(request: Request):
                 "model": model
             }
 
-            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']
-                        
-                        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"
+            async with aiohttp.ClientSession() as session:
+                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
+
+                    parser = sse_parser()
+                    next(parser)  # Initialize the generator
+
+                    async for line in resp.content:
+                        line = line.decode('utf-8').strip()
+                        if line == '':
+                            continue
+
+                        parser.send(line)
+                        try:
+                            event_data = parser.send(None)
+                            if event_data:
+                                # Process the SSE event
+                                event_json = json.loads(event_data)
+                                if 'text' in event_json:
+                                    new_content = event_json['text'][len(full_response):]
+                                    full_response = event_json['text']
+                                    if new_content:
+                                        yield f"data: {json.dumps(format_openai_response(new_content))}\n\n"
+                                elif '"final":true' in event_data:
+                                    final_data = event_json
+                                    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
+                        except StopIteration:
+                            pass  # No complete event yet
+
+            yield f"data: {json.dumps(format_openai_response('', 'stop'))}\n\n"
+            yield "data: [DONE]\n\n"
 
     if stream:
         return StreamingResponse(generate(), media_type='text/event-stream')