feat: add SSE streaming for REST transcription endpoint

- stream=true now returns text/event-stream with per-segment SSE events
- Each segment yields 'data: {json}' followed by 'data: [DONE]'
- Error events streamed as 'data: {"error": ...}'
- Temp files cleaned up in finally block
- 5 new tests in test_server_extended.py (183 total passing)
This commit is contained in:
Aaron Boxer
2026-04-17 10:35:29 -04:00
committed by Aaron Boxer
parent 8bde966c1e
commit 8ac98dceec
2 changed files with 224 additions and 5 deletions
+71 -5
View File
@@ -10,7 +10,8 @@ import tempfile
from typing import Optional, List
from fastapi import FastAPI, UploadFile, Form
from fastapi.middleware.cors import CORSMiddleware
from starlette.responses import PlainTextResponse, JSONResponse
from fastapi.responses import JSONResponse
from starlette.responses import PlainTextResponse, JSONResponse, StreamingResponse
import uvicorn
from faster_whisper import WhisperModel
import torch
@@ -463,6 +464,58 @@ class TranscriptionServer:
wl_metrics.track_connection_closed()
del websocket
def _stream_transcription(self, file, language, prompt, temperature,
timestamp_granularities,
faster_whisper_custom_model_path):
"""Return a StreamingResponse that yields SSE events per segment."""
async def _sse_generator():
tmp_path = None
try:
suffix = os.path.splitext(file.filename)[1] or ".wav"
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
shutil.copyfileobj(file.file, tmp)
tmp_path = tmp.name
device = "cuda" if torch.cuda.is_available() else "cpu"
compute_type = "float16" if device == "cuda" else "int8"
model_name = faster_whisper_custom_model_path or "small"
transcriber = WhisperModel(model_name, device=device, compute_type=compute_type)
segments, info = transcriber.transcribe(
tmp_path,
language=language,
initial_prompt=prompt,
temperature=temperature,
vad_filter=False,
word_timestamps=(timestamp_granularities and "word" in timestamp_granularities),
)
for seg in segments:
seg_dict = {
"id": seg.id,
"start": seg.start,
"end": seg.end,
"text": seg.text.strip(),
}
if timestamp_granularities and "word" in timestamp_granularities:
seg_dict["words"] = [
{"word": w.word, "start": w.start, "end": w.end, "probability": w.probability}
for w in seg.words
]
yield f"data: {json.dumps(seg_dict)}\n\n"
yield "data: [DONE]\n\n"
wl_metrics.track_rest_request(endpoint="transcriptions_stream", status=200)
except Exception as e:
yield f"data: {json.dumps({'error': str(e)})}\n\n"
wl_metrics.track_rest_request(endpoint="transcriptions_stream", status=500)
wl_metrics.track_error("rest_stream")
finally:
if tmp_path and os.path.exists(tmp_path):
os.unlink(tmp_path)
return StreamingResponse(_sse_generator(), media_type="text/event-stream")
def run(self,
host,
port=9090,
@@ -581,10 +634,23 @@ class TranscriptionServer:
hotwords: Optional[str] = Form(default=None),
):
if stream:
wl_metrics.track_rest_request(endpoint="transcriptions", status=400)
return JSONResponse({"error": "Streaming not supported in this backend."}, status_code=400)
if chunking_strategy or known_speaker_names or known_speaker_references:
logging.warning("Diarization/chunking params ignored; not supported.")
return self._stream_transcription(
file, language, prompt, temperature,
timestamp_granularities,
faster_whisper_custom_model_path,
)
ignored_params = []
if chunking_strategy:
ignored_params.append(f"chunking_strategy='{chunking_strategy}'")
if known_speaker_names:
ignored_params.append("known_speaker_names")
if known_speaker_references:
ignored_params.append("known_speaker_references")
if include:
ignored_params.append(f"include={include}")
if ignored_params:
logging.warning(f"Unsupported OpenAI params ignored: {', '.join(ignored_params)}")
supported_formats = ["json", "text", "srt", "verbose_json", "vtt"]
if response_format not in supported_formats: