bnm

main.py

@@ -68,10 +68,10 @@ async def predict_endpoint(payload: PredictRequest):
             try:
                 for chunk in result:
                     yield f"data: {json.dumps(chunk)}\n\n"
-                # Add [END] marker after generator ends
-                yield 'data: "[END]"\n\n'
+                # Return structured JSON to indicate end of stream
+                yield f"data: {json.dumps({'end': True})}\n\n"
             except Exception as e:
-                # Send error in stream
-                yield f'data: "[ERROR] {str(e)}"\n\n'
-        return StreamingResponse(event_stream(), media_type="text/event-stream")
-    return JSONResponse(content=result)
+                # Return structured JSON to indicate error
+                yield f"data: {json.dumps({'error': str(e)})}\n\n"
+    return StreamingResponse(event_stream(), media_type="text/event-stream")
+