diff --git a/README.md b/README.md index 75184fd..fb3fd48 100755 --- a/README.md +++ b/README.md @@ -12,6 +12,7 @@ A unified speech-to-text and translation service that provides a single API for - **Command Line Interface**: Downloadable CLI for communicating with the server for audio processing - **Download and Realtime Transcription** with JSON, SRT, VTT, and text output formats +- **Streaming Audio Input**: Real-time transcription of audio streams via WebSocket - **HTTP API Server**: RESTful API for transcription and translation service - **Docker Support**: Pre-built GPU-enabled containers for easy deployment of the service - **GPU Support**: CUDA, Vulkan, and Metal (macOS) acceleration for local models @@ -170,8 +171,8 @@ flowchart LR end end - CLI --> API - SDK --> API + CLI -->|REST / WebSocket| API + SDK -->|REST / WebSocket| API API --> Orch Orch --> WS WS --> Models diff --git a/doc/API.md b/doc/API.md index 7a6f972..5383673 100644 --- a/doc/API.md +++ b/doc/API.md @@ -168,6 +168,75 @@ Translates audio files to English text. **Response Formats:** Same as transcription (JSON, plain text, SRT, VTT) +### Streaming Transcription (WebSocket) + +```http +GET /api/whisper/stream +``` + +Upgrades the connection to a WebSocket for real-time streaming audio transcription. The client sends raw audio data and receives transcription segments as they are produced, using a sliding-window approach internally. + +**Protocol:** + +1. Connect via WebSocket (`ws://` or `wss://`) +2. Send a JSON configuration message as the first text frame +3. Receive a `stream.ready` event +4. Send binary frames containing 16-bit signed little-endian PCM audio at 16 kHz mono +5. Receive `stream.segment` events as transcription segments are produced +6. Close the connection when done; remaining audio is flushed and a `stream.done` event is sent + +**Configuration Message (first text frame):** + +```json +{ + "model": "ggml-medium-q5_0", + "language": "en", + "temperature": 0.0, + "prompt": "Optional initial prompt", + "translate": false, + "step_ms": 3000, + "length_ms": 10000, + "keep_ms": 200, + "vad_threshold": 0.0 +} +``` + +| Field | Type | Default | Description | +|-------|------|---------|-------------| +| `model` | string | (required) | Model ID to use (must be a local whisper model) | +| `language` | string | `"auto"` | Language code | +| `temperature` | number | `0.0` | Sampling temperature (0–1) | +| `prompt` | string | | Initial prompt to guide transcription | +| `translate` | boolean | `false` | Translate to English | +| `step_ms` | integer | `3000` | How often to run inference on new audio (ms) | +| `length_ms` | integer | `10000` | Maximum window size fed to the model (ms) | +| `keep_ms` | integer | `200` | Overlap retained between consecutive windows (ms) | +| `vad_threshold` | number | `0.0` | RMS energy threshold for voice activity detection (0 = disabled) | + +**Server Events (text frames):** + +```json +{"type": "stream.ready"} +``` + +```json +{"type": "stream.segment", "segment": {"id": 0, "start": 0.0, "end": 2.5, "text": "Hello world."}} +``` + +```json +{"type": "stream.error", "error": "error message"} +``` + +```json +{"type": "stream.done"} +``` + +**Notes:** + +- Only local whisper models are supported for streaming; cloud providers (OpenAI, ElevenLabs) are not available via this endpoint. +- Audio must be 16 kHz, mono, 16-bit signed little-endian PCM. No resampling is performed server-side. +- Each streaming session holds a model context for its duration. With a limited pool size (`--whisper.max-contexts`), long-running streams may prevent other requests from being served. + ## Error Handling The API uses standard HTTP status codes and returns JSON error responses: diff --git a/go.mod b/go.mod index 628800f..f5d0df2 100644 --- a/go.mod +++ b/go.mod @@ -6,6 +6,7 @@ toolchain go1.24.2 require ( github.com/alecthomas/kong v1.13.0 + github.com/coder/websocket v1.8.14 github.com/djthorpe/go-errors v1.0.3 github.com/go-audio/wav v1.1.0 github.com/mutablelogic/go-client v1.3.5 diff --git a/go.sum b/go.sum index 47b2e0f..24a0e1b 100644 --- a/go.sum +++ b/go.sum @@ -20,6 +20,8 @@ github.com/cenkalti/backoff/v5 v5.0.3 h1:ZN+IMa753KfX5hd8vVaMixjnqRZ3y8CuJKRKj1x github.com/cenkalti/backoff/v5 v5.0.3/go.mod h1:rkhZdG3JZukswDf7f0cwqPNk4K0sa+F97BxZthm/crw= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/coder/websocket v1.8.14 h1:9L0p0iKiNOibykf283eHkKUHHrpG7f65OE3BhhO7v9g= +github.com/coder/websocket v1.8.14/go.mod h1:NX3SzP+inril6yawo5CQXx8+fk145lPDC6pumgx0mVg= github.com/containerd/errdefs v1.0.0 h1:tg5yIfIlQIrxYtu9ajqY42W3lpS19XqdxRQeEwYG8PI= github.com/containerd/errdefs v1.0.0/go.mod h1:+YBYIdtsnF4Iw6nWZhJcqGSg/dwvV7tyJ/kCkyJ2k+M= github.com/containerd/errdefs/pkg v0.3.0 h1:9IKJ06FvyNlexW690DXuQNx2KA2cUJXx151Xdx3ZPPE= diff --git a/pkg/httphandler/httphandler.go b/pkg/httphandler/httphandler.go index 8085ccc..7e8c54a 100644 --- a/pkg/httphandler/httphandler.go +++ b/pkg/httphandler/httphandler.go @@ -24,6 +24,7 @@ func RegisterHandlers(router *http.ServeMux, prefix string, manager *pkg.Manager RegisterModelHandlers(router, prefix, manager, middleware) RegisterTranscribeHandlers(router, prefix, manager, middleware) RegisterTranslateHandlers(router, prefix, manager, middleware) + RegisterStreamHandlers(router, prefix, manager, middleware) } /////////////////////////////////////////////////////////////////////////////// diff --git a/pkg/httphandler/stream.go b/pkg/httphandler/stream.go new file mode 100644 index 0000000..ed1738b --- /dev/null +++ b/pkg/httphandler/stream.go @@ -0,0 +1,124 @@ +package httphandler + +import ( + "context" + "encoding/binary" + "net/http" + + // Packages + httpresponse "github.com/mutablelogic/go-server/pkg/httpresponse" + pkg "github.com/mutablelogic/go-whisper/pkg" + schema "github.com/mutablelogic/go-whisper/pkg/schema" + websocket "github.com/coder/websocket" + wsjson "github.com/coder/websocket/wsjson" +) + +/////////////////////////////////////////////////////////////////////////////// +// PUBLIC METHODS + +// RegisterStreamHandlers registers the WebSocket handler for streaming audio transcription +func RegisterStreamHandlers(router *http.ServeMux, prefix string, manager *pkg.Manager, middleware HTTPMiddlewareFuncs) { + // Do NOT wrap with middleware — it may interfere with the WebSocket upgrade + router.HandleFunc(joinPath(prefix, "stream"), func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet { + _ = httpresponse.Error(w, httpresponse.Err(http.StatusMethodNotAllowed), r.Method) + return + } + _ = streamHandle(w, r, manager) + }) +} + +/////////////////////////////////////////////////////////////////////////////// +// PRIVATE METHODS + +// streamHandle upgrades the connection to WebSocket and runs the streaming loop +func streamHandle(w http.ResponseWriter, r *http.Request, manager *pkg.Manager) error { + // Upgrade to WebSocket + conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{}) + if err != nil { + return err + } + + // Use a cancellable context derived from the request context + ctx, cancel := context.WithCancel(r.Context()) + defer cancel() + + // Read first text message as StreamConfig JSON + var cfg schema.StreamConfig + if err := wsjson.Read(ctx, conn, &cfg); err != nil { + conn.Close(websocket.StatusUnsupportedData, "expected StreamConfig JSON as first message") + return err + } + + // Create a new stream session from the manager + session, cleanup, err := manager.NewStreamSession(ctx, &cfg) + if err != nil { + _ = wsjson.Write(ctx, conn, schema.StreamEvent{Type: schema.StreamEventError, Error: err.Error()}) + conn.Close(websocket.StatusInternalError, err.Error()) + return err + } + defer cleanup() + + // Notify client that we are ready + if err := wsjson.Write(ctx, conn, schema.StreamEvent{Type: schema.StreamEventReady}); err != nil { + conn.Close(websocket.StatusInternalError, err.Error()) + return err + } + + // Main loop: read binary PCM frames from client + for { + msgType, data, err := conn.Read(ctx) + if err != nil { + // Client closed or error — run final processing + break + } + + if msgType != websocket.MessageBinary { + // Unexpected message type — skip + continue + } + + // Convert 16-bit signed little-endian PCM to float32 + samples := make([]float32, len(data)/2) + for i := range samples { + raw := binary.LittleEndian.Uint16(data[i*2 : i*2+2]) + samples[i] = float32(int16(raw)) / 32768.0 + } + + session.AddSamples(samples) + + if session.ShouldProcess() { + segments, err := session.Process(ctx) + if err != nil { + _ = wsjson.Write(ctx, conn, schema.StreamEvent{Type: schema.StreamEventError, Error: err.Error()}) + conn.Close(websocket.StatusInternalError, err.Error()) + return err + } + for _, seg := range segments { + if err := wsjson.Write(ctx, conn, schema.StreamEvent{Type: schema.StreamEventSegment, Segment: seg}); err != nil { + conn.Close(websocket.StatusInternalError, err.Error()) + return err + } + } + } + } + + // Process any remaining audio + segments, err := session.ProcessFinal(ctx) + if err != nil { + _ = wsjson.Write(ctx, conn, schema.StreamEvent{Type: schema.StreamEventError, Error: err.Error()}) + conn.Close(websocket.StatusInternalError, err.Error()) + return err + } + for _, seg := range segments { + if err := wsjson.Write(ctx, conn, schema.StreamEvent{Type: schema.StreamEventSegment, Segment: seg}); err != nil { + conn.Close(websocket.StatusInternalError, err.Error()) + return err + } + } + + // Send done event and close cleanly + _ = wsjson.Write(ctx, conn, schema.StreamEvent{Type: schema.StreamEventDone}) + conn.Close(websocket.StatusNormalClosure, "") + return nil +} diff --git a/pkg/manager.go b/pkg/manager.go index a3d81cb..345a17c 100644 --- a/pkg/manager.go +++ b/pkg/manager.go @@ -268,6 +268,56 @@ func (m *Manager) Translate(ctx context.Context, w schema.SegmentWriter, r io.Re } } +// NewStreamSession creates a new streaming transcription session for the +// specified model. The caller MUST call the returned release function when +// the session is finished to return the task to the pool. +// Streaming is only supported for local whisper models. +func (m *Manager) NewStreamSession(ctx context.Context, cfg *schema.StreamConfig) (*whisper.StreamSession, func(), error) { + ctx, endSpan := otel.StartSpan(m.tracer, ctx, "manager.NewStreamSession", + attribute.String("model", cfg.Model), + ) + var err error + defer func() { endSpan(err) }() + + // Resolve the model + model, err := m.GetModel(ctx, cfg.Model) + if err != nil { + return nil, nil, err + } + + // Streaming is only supported for local whisper models + if model.OwnedBy != "whisper" { + err = httpresponse.ErrBadRequest.With("streaming is only supported for local whisper models") + return nil, nil, err + } + + // Acquire a long-lived task from the pool + task, release, err := m.whisper.AcquireTask(model) + if err != nil { + return nil, nil, err + } + + // Configure task parameters from the stream config + if cfg.Language != nil { + task.SetLanguage(*cfg.Language) + } + if cfg.Temperature != nil { + if err = task.SetTemperature(*cfg.Temperature); err != nil { + release() + return nil, nil, err + } + } + if cfg.Prompt != nil { + task.SetPrompt(*cfg.Prompt) + } + if cfg.Translate { + task.SetTranslate(true) + } + + session := whisper.NewStreamSession(task, *cfg) + return session, release, nil +} + /////////////////////////////////////////////////////////////////////////////// // PRIVATE METHODS - WHISPER diff --git a/pkg/schema/stream.go b/pkg/schema/stream.go new file mode 100644 index 0000000..7517a71 --- /dev/null +++ b/pkg/schema/stream.go @@ -0,0 +1,47 @@ +package schema + +////////////////////////////////////////////////////////////////////////////// +// TYPES + +// StreamConfig is the configuration message sent by the client as the first +// WebSocket text message when initiating a streaming transcription session. +type StreamConfig struct { + Model string `json:"model"` + Language *string `json:"language,omitempty"` + Temperature *float64 `json:"temperature,omitempty"` + Prompt *string `json:"prompt,omitempty"` + Translate bool `json:"translate,omitempty"` + StepMs int `json:"step_ms,omitempty"` + LengthMs int `json:"length_ms,omitempty"` + KeepMs int `json:"keep_ms,omitempty"` + VadThreshold float32 `json:"vad_threshold,omitempty"` +} + +// StreamEvent is a JSON message sent back to the client during a streaming +// transcription session. +type StreamEvent struct { + Type string `json:"type"` + Segment *Segment `json:"segment,omitempty"` + Error string `json:"error,omitempty"` +} + +////////////////////////////////////////////////////////////////////////////// +// GLOBALS + +const ( + StreamEventReady = "stream.ready" + StreamEventSegment = "stream.segment" + StreamEventError = "stream.error" + StreamEventDone = "stream.done" +) + +////////////////////////////////////////////////////////////////////////////// +// STRINGIFY + +func (c *StreamConfig) String() string { + return stringify(*c) +} + +func (e *StreamEvent) String() string { + return stringify(*e) +} diff --git a/pkg/whisper/stream.go b/pkg/whisper/stream.go new file mode 100644 index 0000000..c06c8da --- /dev/null +++ b/pkg/whisper/stream.go @@ -0,0 +1,270 @@ +package whisper + +import ( + "context" + "math" + "time" + + // Packages + schema "github.com/mutablelogic/go-whisper/pkg/schema" + whisper "github.com/mutablelogic/go-whisper/sys/whisper" +) + +////////////////////////////////////////////////////////////////////////////// +// GLOBALS + +const ( + // Default sliding-window parameters (milliseconds) + defaultStepMs = 3000 + defaultLengthMs = 10000 + defaultKeepMs = 200 +) + +////////////////////////////////////////////////////////////////////////////// +// TYPES + +// StreamSession implements a sliding-window streaming transcription processor +// that mirrors the approach used in whisper.cpp's examples/stream/stream.cpp. +// +// Audio samples are accumulated in audioBuffer. When enough new samples have +// arrived (>= stepSamples), Process() is called to run whisper_full() on a +// window of up to lengthSamples. After processing, the buffer is trimmed to +// keepSamples of overlap so that speech crossing a window boundary is not lost. +// Segment timestamps are adjusted by windowOffset so they are always expressed +// in absolute time, and duplicate segments (whose end time was already emitted +// in a previous window) are suppressed. +type StreamSession struct { + // task is held for the lifetime of the session. The caller is responsible + // for returning it to the pool (e.g. via Manager.WithModel) once the + // session is finished. + task *Task + + // audioBuffer accumulates raw 16 kHz mono PCM samples. + audioBuffer []float32 + + // stepSamples is the minimum number of new samples that must be present + // before Process() will trigger a whisper_full() run. + stepSamples int + + // lengthSamples is the maximum window size fed to whisper_full(). + lengthSamples int + + // keepSamples is the number of samples retained as overlap when the + // buffer is trimmed after each processing step. + keepSamples int + + // newSamples counts how many samples have been added since the last + // Process() call. + newSamples int + + // lastSegmentEnd is the absolute end timestamp (as time.Duration) of the + // last segment that was emitted. Used to suppress duplicates when the same + // speech appears in consecutive overlapping windows. + lastSegmentEnd time.Duration + + // vadThreshold, if > 0, enables a simple energy gate: Process() will not + // run unless the RMS energy of the newly-arrived samples exceeds this value. + vadThreshold float32 + + // windowOffset is the absolute time position (as time.Duration) of the + // first sample currently in audioBuffer. It is advanced whenever samples + // are trimmed from the front of the buffer. + windowOffset time.Duration +} + +////////////////////////////////////////////////////////////////////////////// +// LIFECYCLE + +// NewStreamSession creates a StreamSession that wraps the given Task. +// The task must already be initialised (via Task.Init / Manager.WithModel) +// and have had CopyParams() called on it so that its params are ready. +// cfg carries the sliding-window parameters; zero values trigger defaults. +func NewStreamSession(task *Task, cfg schema.StreamConfig) *StreamSession { + // Apply defaults for zero values + stepMs := cfg.StepMs + if stepMs <= 0 { + stepMs = defaultStepMs + } + lengthMs := cfg.LengthMs + if lengthMs <= 0 { + lengthMs = defaultLengthMs + } + keepMs := cfg.KeepMs + if keepMs <= 0 { + keepMs = defaultKeepMs + } + + // Clamp: keep must be strictly less than length, and step must be <= length. + if keepMs >= lengthMs { + keepMs = lengthMs / 2 + } + if stepMs > lengthMs { + stepMs = lengthMs + } + + return &StreamSession{ + task: task, + stepSamples: msToSamples(stepMs), + lengthSamples: msToSamples(lengthMs), + keepSamples: msToSamples(keepMs), + vadThreshold: cfg.VadThreshold, + } +} + +////////////////////////////////////////////////////////////////////////////// +// PUBLIC METHODS + +// AddSamples appends raw 16 kHz mono PCM samples to the internal buffer and +// increments the new-sample counter used by ShouldProcess. +func (s *StreamSession) AddSamples(samples []float32) { + s.audioBuffer = append(s.audioBuffer, samples...) + s.newSamples += len(samples) +} + +// ShouldProcess returns true when enough new audio has accumulated to justify +// running a transcription pass. If vadThreshold > 0, the new samples must also +// exceed that RMS energy level (simple voice-activity gate). +func (s *StreamSession) ShouldProcess() bool { + if s.newSamples < s.stepSamples { + return false + } + if s.vadThreshold > 0 { + // TODO: This is a simple RMS energy gate. whisper.cpp has a model-based + // VAD (Silero) via whisper_vad_context / whisper_vad_detect_speech that + // would be far more accurate (distinguishes speech from noise, typing, + // music, etc.), but the Go bindings for the VAD API don't exist yet. + start := len(s.audioBuffer) - s.newSamples + if start < 0 { + start = 0 + } + if rmsEnergy(s.audioBuffer[start:]) < s.vadThreshold { + return false + } + } + return true +} + +// Process runs whisper_full() on the current audio window and returns any new +// segments whose absolute end time is later than the last emitted segment. +// Call this when ShouldProcess() returns true. The audio buffer is trimmed +// and windowOffset is advanced after each call. +func (s *StreamSession) Process(ctx context.Context) ([]*schema.Segment, error) { + return s.process(ctx) +} + +// ProcessFinal is like Process but runs unconditionally on whatever audio +// remains in the buffer. Use this to flush the last few words when the audio +// stream ends. +func (s *StreamSession) ProcessFinal(ctx context.Context) ([]*schema.Segment, error) { + if len(s.audioBuffer) == 0 { + return nil, nil + } + return s.process(ctx) +} + +////////////////////////////////////////////////////////////////////////////// +// PRIVATE METHODS + +// process is the shared implementation for Process and ProcessFinal. +func (s *StreamSession) process(ctx context.Context) ([]*schema.Segment, error) { + // Choose the window: up to the last lengthSamples of the buffer. + window := s.audioBuffer + windowStartOffset := s.windowOffset // absolute time of window[0] + if len(window) > s.lengthSamples { + trimmed := len(window) - s.lengthSamples + window = window[trimmed:] + windowStartOffset = s.windowOffset + samplesToDuration(trimmed) + } + + if len(window) == 0 { + s.newSamples = 0 + return nil, nil + } + + // Run transcription on the window. We pass windowStartOffset as the ts + // argument so that task.Transcribe adds it to the raw whisper timestamps, + // giving us absolute times directly in the returned segments. + var segments []*schema.Segment + err := s.task.Transcribe(ctx, windowStartOffset, window, func(seg *schema.Segment) { + // This callback fires for each new segment as whisper_full progresses. + // We collect them here; deduplication happens below after the full run. + segments = append(segments, seg) + }) + if err != nil { + return nil, err + } + + // If the callback produced no segments, fall back to whatever appendResult + // wrote into task.result (the callback path may not fire for very short + // windows or when context timestamps are disabled). + if len(segments) == 0 { + result := s.task.Result() + if result != nil { + // Collect only the segments that were added during this Transcribe call. + // task.appendResult always appends; we want the ones with Start >= windowStartOffset. + for _, seg := range result.Segments { + if time.Duration(seg.Start) >= windowStartOffset { + segments = append(segments, seg) + } + } + } + } + + // Deduplicate: only emit segments whose end time is strictly after the + // last segment we already returned to the caller. + var newSegments []*schema.Segment + for _, seg := range segments { + if time.Duration(seg.End) > s.lastSegmentEnd { + newSegments = append(newSegments, seg) + } + } + if len(newSegments) > 0 { + s.lastSegmentEnd = time.Duration(newSegments[len(newSegments)-1].End) + } + + // Trim the audio buffer. Keep the last keepSamples as overlap. + s.trimBuffer() + + // Reset the new-sample counter. + s.newSamples = 0 + + return newSegments, nil +} + +// trimBuffer retains at most keepSamples at the end of audioBuffer and +// advances windowOffset to reflect the dropped samples. +func (s *StreamSession) trimBuffer() { + if len(s.audioBuffer) <= s.keepSamples { + // Nothing to trim. + return + } + dropped := len(s.audioBuffer) - s.keepSamples + s.windowOffset += samplesToDuration(dropped) + s.audioBuffer = s.audioBuffer[dropped:] +} + +////////////////////////////////////////////////////////////////////////////// +// HELPERS + +// msToSamples converts a duration in milliseconds to a sample count at 16 kHz. +func msToSamples(ms int) int { + return ms * whisper.SampleRate / 1000 +} + +// samplesToDuration converts a sample count at 16 kHz to a time.Duration. +func samplesToDuration(n int) time.Duration { + return time.Duration(n) * time.Second / time.Duration(whisper.SampleRate) +} + +// rmsEnergy computes the root-mean-square energy of the provided samples. +// Returns 0 for an empty slice. +func rmsEnergy(samples []float32) float32 { + if len(samples) == 0 { + return 0 + } + var sum float64 + for _, s := range samples { + sum += float64(s) * float64(s) + } + return float32(math.Sqrt(sum / float64(len(samples)))) +} diff --git a/pkg/whisper/whisper.go b/pkg/whisper/whisper.go index c340189..03833cd 100644 --- a/pkg/whisper/whisper.go +++ b/pkg/whisper/whisper.go @@ -208,6 +208,34 @@ func (m *Manager) DownloadModel(ctx context.Context, path string, fn func(curByt return m.store.Download(ctx, path, fn) } +// AcquireTask gets a task for the specified model from the pool without +// scoping it to a callback. The caller MUST call the returned release +// function when done (typically via defer) to return the task to the pool. +// This is intended for long-lived sessions such as streaming transcription. +func (m *Manager) AcquireTask(model *schema.Model) (*Task, func(), error) { + if model == nil { + return nil, nil, httpresponse.ErrBadRequest.With("model must be non-nil") + } + + m.RLock() + pool := m.pool + m.RUnlock() + + if pool == nil { + return nil, nil, httpresponse.ErrInternalError.With("pool not initialized") + } + + task, err := pool.get(model) + if err != nil { + return nil, nil, err + } + + task.CopyParams() + + release := func() { pool.put(task) } + return task, release, nil +} + // WithModel gets a task for the specified model and executes the function. // The task is automatically returned to the pool when done. func (m *Manager) WithModel(model *schema.Model, fn func(task *Task) error, opts ...Opt) error {