use ServerSentEvent from sse_starlette

Signed-off-by: Ma, Xiangxiang <xiangxiang.ma@intel.com>

Author: maxxgx

application/backend/src/api/endpoints/job_endpoints.py

@@ -5,7 +5,7 @@
 from uuid import UUID
 
 from fastapi import APIRouter, Body, Depends
-from fastapi.responses import StreamingResponse
+from sse_starlette import EventSourceResponse
 
 from api.dependencies import get_job_id, get_job_service
 from api.endpoints import API_PREFIX
@@ -39,6 +39,6 @@ async def submit_train_job(
 async def get_job_logs(
     job_id: Annotated[UUID, Depends(get_job_id)],
     job_service: Annotated[JobService, Depends(get_job_service)],
-) -> StreamingResponse:
+) -> EventSourceResponse:
     """Endpoint to get the logs of a job by its ID"""
-    return StreamingResponse(job_service.stream_logs(job_id=job_id), media_type="text/event-stream")
+    return EventSourceResponse(job_service.stream_logs(job_id=job_id))

application/backend/src/services/job_service.py

@@ -3,10 +3,12 @@
 import asyncio
 import datetime
 import os
+from collections.abc import AsyncGenerator
 from uuid import UUID
 
 import anyio
 from sqlalchemy.exc import IntegrityError
+from sse_starlette import ServerSentEvent
 
 from db import get_async_db_session_ctx
 from exceptions import DuplicateJobException, ResourceNotFoundException
@@ -75,7 +77,7 @@ async def update_job_status(
             await repo.update(job, updates)
 
     @classmethod
-    async def stream_logs(cls, job_id: UUID | str):
+    async def stream_logs(cls, job_id: UUID | str) -> AsyncGenerator[ServerSentEvent, None]:
         from core.logging.utils import get_job_logs_path
 
         log_file = get_job_logs_path(job_id=job_id)
@@ -110,8 +112,5 @@ async def is_job_still_running():
                         continue
                     # No more lines are expected
                     else:
-                        yield "data: DONE\n\n"
                         break
-
-                # Format as an SSE message
-                yield f"data: {line.rstrip()}\n\n"
+                yield ServerSentEvent(data=line.rstrip())

application/backend/tests/unit/endpoints/test_jobs.py

@@ -4,6 +4,7 @@
 
 import pytest
 from fastapi import status
+from sse_starlette import ServerSentEvent
 
 from api.dependencies import get_job_service
 from main import app
@@ -41,8 +42,8 @@ def test_get_job_logs_success(fxt_client, fxt_job_service, fxt_job):
 
     # Mock the stream_logs generator
     async def mock_stream():
-        yield '{"level": "INFO", "message": "Line 1"}\n'
-        yield '{"level": "INFO", "message": "Line 2"}\n'
+        yield ServerSentEvent(data='{"level": "INFO", "message": "Line 1"}')
+        yield ServerSentEvent(data='{"level": "INFO", "message": "Line 2"}')
 
     fxt_job_service.stream_logs.return_value = mock_stream()
 
@@ -52,7 +53,7 @@ async def mock_stream():
     # Verify the streamed content
     content = response.content.decode("utf-8")
     lines = [line for line in content.split("\n") if line]
-    assert len(lines) == 2
+    assert len(lines) == 4  # 2 events + 2 newlines
     assert '"level": "INFO"' in lines[0]
     assert '"message": "Line 1"' in lines[0]
 

application/backend/tests/unit/services/test_job_service.py

@@ -206,7 +206,7 @@ async def consume_stream():
 
     def test_stream_logs_success(self, fxt_job_repository, fxt_job):
         """Test streaming logs successfully from a completed job."""
-        log_lines = ['{"level": "INFO", "message": "Line 1"}\n', '{"level": "INFO", "message": "Line 2"}\n']
+        log_lines = ['{"level": "INFO", "message": "Line 1"}', '{"level": "INFO", "message": "Line 2"}']
 
         # Mock job as completed
         completed_job = fxt_job.model_copy(update={"status": JobStatus.COMPLETED})
@@ -240,7 +240,7 @@ async def mock_anyio_open_file(*args, **kwargs):
             async def consume_stream():
                 result = []
                 async for line in JobService.stream_logs(fxt_job.id):
-                    result.append(line)
+                    result.append(line.data)
                 return result
 
             result = asyncio.run(consume_stream())