libnovel/backend/internal/pockettts/client.go

// Package pockettts provides a client for the kyutai-labs/pocket-tts TTS service.
//
// pocket-tts exposes a non-OpenAI API:
//
//	POST /tts  (multipart form: text, voice_url)  → streaming WAV
//	GET  /health                                   → {"status":"healthy"}
//
// GenerateAudio streams the WAV response and transcodes it to MP3 using ffmpeg,
// so callers receive MP3 bytes — the same format as the kokoro client — and the
// rest of the pipeline does not need to care which TTS engine was used.
//
// StreamAudioMP3 is the streaming variant: it returns an io.ReadCloser that
// yields MP3-encoded audio incrementally as pocket-tts generates it, without
// buffering the full output.
//
// Predefined voices (pass the bare name as the voice parameter):
//
//	alba, marius, javert, jean, fantine, cosette, eponine, azelma,
//	anna, vera, charles, paul, george, mary, jane, michael, eve,
//	bill_boerst, peter_yearsley, stuart_bell
package pockettts

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"os/exec"
	"strings"
	"time"
)

// PredefinedVoices is the set of voice names built into pocket-tts.
// The runner uses this to decide which TTS engine to route a task to.
var PredefinedVoices = map[string]struct{}{
	"alba": {}, "marius": {}, "javert": {}, "jean": {},
	"fantine": {}, "cosette": {}, "eponine": {}, "azelma": {},
	"anna": {}, "vera": {}, "charles": {}, "paul": {},
	"george": {}, "mary": {}, "jane": {}, "michael": {},
	"eve": {}, "bill_boerst": {}, "peter_yearsley": {}, "stuart_bell": {},
}

// IsPocketTTSVoice reports whether voice is served by pocket-tts.
func IsPocketTTSVoice(voice string) bool {
	_, ok := PredefinedVoices[voice]
	return ok
}

// Client is the interface for interacting with the pocket-tts service.
type Client interface {
	// GenerateAudio synthesises text using the given voice and returns MP3 bytes.
	// Voice must be one of the predefined pocket-tts voice names.
	GenerateAudio(ctx context.Context, text, voice string) ([]byte, error)

	// StreamAudioMP3 synthesises text and returns an io.ReadCloser that streams
	// MP3-encoded audio incrementally via a live ffmpeg transcode pipe.
	// The caller must always close the returned ReadCloser.
	StreamAudioMP3(ctx context.Context, text, voice string) (io.ReadCloser, error)

	// StreamAudioWAV synthesises text and returns an io.ReadCloser that streams
	// raw WAV audio directly from pocket-tts without any transcoding.
	// The stream begins with a WAV header followed by 16-bit PCM frames at 16 kHz.
	// The caller must always close the returned ReadCloser.
	StreamAudioWAV(ctx context.Context, text, voice string) (io.ReadCloser, error)

	// ListVoices returns the available predefined voice names.
	ListVoices(ctx context.Context) ([]string, error)
}

// httpClient is the concrete pocket-tts HTTP client.
type httpClient struct {
	baseURL string
	http    *http.Client
}

// New returns a Client targeting baseURL (e.g. "https://pocket-tts.libnovel.cc").
func New(baseURL string) Client {
	return &httpClient{
		baseURL: strings.TrimRight(baseURL, "/"),
		http:    &http.Client{Timeout: 10 * time.Minute},
	}
}

// GenerateAudio posts to POST /tts and transcodes the WAV response to MP3
// using the system ffmpeg binary. Requires ffmpeg to be on PATH (available in
// the runner Docker image via Alpine's ffmpeg package).
func (c *httpClient) GenerateAudio(ctx context.Context, text, voice string) ([]byte, error) {
	if text == "" {
		return nil, fmt.Errorf("pockettts: empty text")
	}
	if voice == "" {
		voice = "alba"
	}

	resp, err := c.postTTS(ctx, text, voice)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	wavData, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("pockettts: read response body: %w", err)
	}

	// ── Transcode WAV → MP3 via ffmpeg ────────────────────────────────────────
	mp3Data, err := wavToMP3(ctx, wavData)
	if err != nil {
		return nil, fmt.Errorf("pockettts: transcode to mp3: %w", err)
	}
	return mp3Data, nil
}

// StreamAudioMP3 posts to POST /tts and returns an io.ReadCloser that delivers
// MP3 bytes as pocket-tts generates WAV frames. ffmpeg runs as a subprocess
// with stdin connected to the live WAV stream and stdout piped to the caller.
// The caller must always close the returned ReadCloser.
func (c *httpClient) StreamAudioMP3(ctx context.Context, text, voice string) (io.ReadCloser, error) {
	if text == "" {
		return nil, fmt.Errorf("pockettts: empty text")
	}
	if voice == "" {
		voice = "alba"
	}

	resp, err := c.postTTS(ctx, text, voice)
	if err != nil {
		return nil, err
	}

	// Start ffmpeg: read WAV from stdin (the live HTTP body), write MP3 to stdout.
	cmd := exec.CommandContext(ctx,
		"ffmpeg",
		"-hide_banner", "-loglevel", "error",
		"-i", "pipe:0", // WAV from stdin
		"-f", "mp3", // output format
		"-q:a", "2", // VBR ~190 kbps
		"pipe:1", // MP3 to stdout
	)
	cmd.Stdin = resp.Body

	pr, pw := io.Pipe()
	cmd.Stdout = pw

	var stderrBuf bytes.Buffer
	cmd.Stderr = &stderrBuf

	if err := cmd.Start(); err != nil {
		resp.Body.Close()
		return nil, fmt.Errorf("pockettts: start ffmpeg: %w", err)
	}

	// Close the write end of the pipe when ffmpeg exits, propagating any error.
	go func() {
		waitErr := cmd.Wait()
		resp.Body.Close()
		if waitErr != nil {
			pw.CloseWithError(fmt.Errorf("ffmpeg: %w (stderr: %s)", waitErr, stderrBuf.String()))
		} else {
			pw.Close()
		}
	}()

	return pr, nil
}

// StreamAudioWAV posts to POST /tts and returns an io.ReadCloser that delivers
// raw WAV bytes directly from pocket-tts — no ffmpeg transcoding required.
// The first bytes will be a WAV header (RIFF/fmt chunk) followed by PCM frames.
// The caller must always close the returned ReadCloser.
func (c *httpClient) StreamAudioWAV(ctx context.Context, text, voice string) (io.ReadCloser, error) {
	if text == "" {
		return nil, fmt.Errorf("pockettts: empty text")
	}
	if voice == "" {
		voice = "alba"
	}

	resp, err := c.postTTS(ctx, text, voice)
	if err != nil {
		return nil, err
	}
	return resp.Body, nil
}

// ListVoices returns the statically known predefined voice names.
// pocket-tts has no REST endpoint for listing voices.
func (c *httpClient) ListVoices(_ context.Context) ([]string, error) {
	voices := make([]string, 0, len(PredefinedVoices))
	for v := range PredefinedVoices {
		voices = append(voices, v)
	}
	return voices, nil
}

// postTTS sends a multipart POST /tts request and returns the raw response.
// The caller is responsible for closing resp.Body.
func (c *httpClient) postTTS(ctx context.Context, text, voice string) (*http.Response, error) {
	var body bytes.Buffer
	mw := multipart.NewWriter(&body)

	if err := mw.WriteField("text", text); err != nil {
		return nil, fmt.Errorf("pockettts: write text field: %w", err)
	}
	if err := mw.WriteField("voice_url", voice); err != nil {
		return nil, fmt.Errorf("pockettts: write voice_url field: %w", err)
	}
	if err := mw.Close(); err != nil {
		return nil, fmt.Errorf("pockettts: close multipart writer: %w", err)
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost,
		c.baseURL+"/tts", &body)
	if err != nil {
		return nil, fmt.Errorf("pockettts: build request: %w", err)
	}
	req.Header.Set("Content-Type", mw.FormDataContentType())

	resp, err := c.http.Do(req)
	if err != nil {
		return nil, fmt.Errorf("pockettts: request: %w", err)
	}
	if resp.StatusCode != http.StatusOK {
		_, _ = io.Copy(io.Discard, resp.Body)
		resp.Body.Close()
		return nil, fmt.Errorf("pockettts: server returned %d", resp.StatusCode)
	}
	return resp, nil
}

// wavToMP3 converts raw WAV bytes to MP3 using ffmpeg.
// ffmpeg reads from stdin (pipe:0) and writes to stdout (pipe:1).
func wavToMP3(ctx context.Context, wav []byte) ([]byte, error) {
	cmd := exec.CommandContext(ctx,
		"ffmpeg",
		"-hide_banner", "-loglevel", "error",
		"-i", "pipe:0", // read WAV from stdin
		"-f", "mp3", // output format
		"-q:a", "2", // VBR quality ~190 kbps
		"pipe:1", // write MP3 to stdout
	)
	cmd.Stdin = bytes.NewReader(wav)

	var out, stderr bytes.Buffer
	cmd.Stdout = &out
	cmd.Stderr = &stderr

	if err := cmd.Run(); err != nil {
		return nil, fmt.Errorf("ffmpeg: %w (stderr: %s)", err, stderr.String())
	}
	return out.Bytes(), nil
}