libnovel/backend/internal/kokoro/client.go

// Package kokoro provides a client for the Kokoro-FastAPI TTS service.
//
// The Kokoro API is an OpenAI-compatible audio speech API that returns a
// download link (X-Download-Path header) instead of streaming audio directly.
// GenerateAudio handles the two-step flow: POST /v1/audio/speech → GET /v1/download/{file}.
package kokoro

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"
)

// Client is the interface for interacting with the Kokoro TTS service.
type Client interface {
	// GenerateAudio synthesises text using voice and returns raw MP3 bytes.
	GenerateAudio(ctx context.Context, text, voice string) ([]byte, error)

	// StreamAudioMP3 synthesises text and returns an io.ReadCloser that streams
	// MP3-encoded audio incrementally. Uses the kokoro-fastapi streaming mode
	// (stream:true), which delivers MP3 frames as they are generated without
	// waiting for the full output. The caller must always close the ReadCloser.
	StreamAudioMP3(ctx context.Context, text, voice string) (io.ReadCloser, error)

	// StreamAudioWAV synthesises text and returns an io.ReadCloser that streams
	// WAV-encoded audio incrementally using kokoro-fastapi's streaming mode with
	// response_format:"wav". The caller must always close the ReadCloser.
	StreamAudioWAV(ctx context.Context, text, voice string) (io.ReadCloser, error)

	// ListVoices returns the available voice IDs. Falls back to an empty slice
	// on error — callers should treat an empty list as "service unavailable".
	ListVoices(ctx context.Context) ([]string, error)
}

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

// New returns a Kokoro Client targeting baseURL (e.g. "https://kokoro.example.com").
func New(baseURL string) Client {
	return &httpClient{
		baseURL: strings.TrimRight(baseURL, "/"),
		http:    &http.Client{Timeout: 10 * time.Minute},
	}
}

// GenerateAudio calls POST /v1/audio/speech (return_download_link=true) and then
// downloads the resulting MP3 from GET /v1/download/{filename}.
func (c *httpClient) GenerateAudio(ctx context.Context, text, voice string) ([]byte, error) {
	if text == "" {
		return nil, fmt.Errorf("kokoro: empty text")
	}
	if voice == "" {
		voice = "af_bella"
	}

	// ── Step 1: request generation ────────────────────────────────────────────
	reqBody, err := json.Marshal(map[string]any{
		"model":                "kokoro",
		"input":                text,
		"voice":                voice,
		"response_format":      "mp3",
		"speed":                1.0,
		"stream":               false,
		"return_download_link": true,
	})
	if err != nil {
		return nil, fmt.Errorf("kokoro: marshal request: %w", err)
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost,
		c.baseURL+"/v1/audio/speech", bytes.NewReader(reqBody))
	if err != nil {
		return nil, fmt.Errorf("kokoro: build speech request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := c.http.Do(req)
	if err != nil {
		return nil, fmt.Errorf("kokoro: speech request: %w", err)
	}
	defer resp.Body.Close()
	_, _ = io.Copy(io.Discard, resp.Body)

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("kokoro: speech returned %d", resp.StatusCode)
	}

	dlPath := resp.Header.Get("X-Download-Path")
	if dlPath == "" {
		return nil, fmt.Errorf("kokoro: no X-Download-Path header in response")
	}
	filename := dlPath
	if idx := strings.LastIndex(dlPath, "/"); idx >= 0 {
		filename = dlPath[idx+1:]
	}
	if filename == "" {
		return nil, fmt.Errorf("kokoro: empty filename in X-Download-Path: %q", dlPath)
	}

	// ── Step 2: download the generated file ───────────────────────────────────
	dlURL := c.baseURL + "/v1/download/" + filename
	dlReq, err := http.NewRequestWithContext(ctx, http.MethodGet, dlURL, nil)
	if err != nil {
		return nil, fmt.Errorf("kokoro: build download request: %w", err)
	}

	dlResp, err := c.http.Do(dlReq)
	if err != nil {
		return nil, fmt.Errorf("kokoro: download request: %w", err)
	}
	defer dlResp.Body.Close()

	if dlResp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("kokoro: download returned %d", dlResp.StatusCode)
	}

	data, err := io.ReadAll(dlResp.Body)
	if err != nil {
		return nil, fmt.Errorf("kokoro: read download body: %w", err)
	}
	return data, nil
}

// StreamAudioMP3 calls POST /v1/audio/speech with stream:true and returns an
// io.ReadCloser that delivers MP3 frames as kokoro generates them.
// kokoro-fastapi emits raw MP3 bytes when stream mode is enabled — no download
// redirect; the response body IS the audio stream.
func (c *httpClient) StreamAudioMP3(ctx context.Context, text, voice string) (io.ReadCloser, error) {
	if text == "" {
		return nil, fmt.Errorf("kokoro: empty text")
	}
	if voice == "" {
		voice = "af_bella"
	}

	reqBody, err := json.Marshal(map[string]any{
		"model":           "kokoro",
		"input":           text,
		"voice":           voice,
		"response_format": "mp3",
		"speed":           1.0,
		"stream":          true,
	})
	if err != nil {
		return nil, fmt.Errorf("kokoro: marshal stream request: %w", err)
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost,
		c.baseURL+"/v1/audio/speech", bytes.NewReader(reqBody))
	if err != nil {
		return nil, fmt.Errorf("kokoro: build stream request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")

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

// StreamAudioWAV calls POST /v1/audio/speech with stream:true and response_format:wav,
// returning an io.ReadCloser that delivers WAV bytes as kokoro generates them.
func (c *httpClient) StreamAudioWAV(ctx context.Context, text, voice string) (io.ReadCloser, error) {
	if text == "" {
		return nil, fmt.Errorf("kokoro: empty text")
	}
	if voice == "" {
		voice = "af_bella"
	}

	reqBody, err := json.Marshal(map[string]any{
		"model":           "kokoro",
		"input":           text,
		"voice":           voice,
		"response_format": "wav",
		"speed":           1.0,
		"stream":          true,
	})
	if err != nil {
		return nil, fmt.Errorf("kokoro: marshal wav stream request: %w", err)
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost,
		c.baseURL+"/v1/audio/speech", bytes.NewReader(reqBody))
	if err != nil {
		return nil, fmt.Errorf("kokoro: build wav stream request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")

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

// ListVoices calls GET /v1/audio/voices and returns the list of voice IDs.
func (c *httpClient) ListVoices(ctx context.Context) ([]string, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet,
		c.baseURL+"/v1/audio/voices", nil)
	if err != nil {
		return nil, fmt.Errorf("kokoro: build voices request: %w", err)
	}

	resp, err := c.http.Do(req)
	if err != nil {
		return nil, fmt.Errorf("kokoro: voices request: %w", err)
	}
	defer resp.Body.Close()

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

	var result struct {
		Voices []string `json:"voices"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return nil, fmt.Errorf("kokoro: decode voices response: %w", err)
	}
	return result.Voices, nil
}

// VoiceSampleKey returns the MinIO object key for a voice sample MP3.
// Key: _voice-samples/{voice}.mp3 (sanitised).
func VoiceSampleKey(voice string) string {
	safe := strings.Map(func(r rune) rune {
		if (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') ||
			(r >= '0' && r <= '9') || r == '_' || r == '-' {
			return r
		}
		return '_'
	}, voice)
	return fmt.Sprintf("_voice-samples/%s.mp3", safe)
}