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.
//
// 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)

	// 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"
	}

	// ── Build multipart form ──────────────────────────────────────────────────
	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)
	}
	// pocket-tts accepts a predefined voice name as voice_url.
	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)
	}
	defer resp.Body.Close()

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

	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
}

// 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
}

// 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
}