majordomo/provider/llamaswap/audio.go

package llamaswap

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"mime"
	"mime/multipart"
	"net/http"
	"net/url"
	"strings"

	"gitea.stevedudenhoeffer.com/steve/majordomo/audio"
	"gitea.stevedudenhoeffer.com/steve/majordomo/llm"
)

// SpeechModel implements audio.SpeechProvider, binding a text-to-speech model
// served by llama-swap (routed to a kokoro/chatterbox-style OpenAI-compatible
// upstream). The id is passed through verbatim and selects which upstream
// llama-swap loads.
func (p *Provider) SpeechModel(id string, opts ...audio.SpeechModelOption) (audio.SpeechModel, error) {
	if err := p.requireBaseURL(); err != nil {
		return nil, err
	}
	_ = audio.ApplySpeechModelOptions(opts)
	return &speechModel{p: p, id: id}, nil
}

type speechModel struct {
	p  *Provider
	id string
}

// speechRequest is the OpenAI /v1/audio/speech request shape. llama-swap
// routes by the `model` field in the body.
type speechRequest struct {
	Model          string  `json:"model"`
	Input          string  `json:"input"`
	Voice          string  `json:"voice,omitempty"`
	ResponseFormat string  `json:"response_format,omitempty"`
	Speed          float64 `json:"speed,omitempty"`
}

// Speak implements audio.SpeechModel via POST {base}/v1/audio/speech.
func (m *speechModel) Speak(ctx context.Context, req audio.SpeechRequest, opts ...audio.SpeechOption) (*audio.SpeechResult, error) {
	req = req.Apply(opts...)
	if strings.TrimSpace(req.Input) == "" {
		return nil, fmt.Errorf("%w: speech synthesis requires input text", llm.ErrUnsupported)
	}
	if req.Speed < 0 {
		return nil, fmt.Errorf("%w: speech speed must be >= 0, got %g", llm.ErrUnsupported, req.Speed)
	}
	wire := speechRequest{
		Model:          m.id,
		Input:          req.Input,
		Voice:          req.Voice,
		ResponseFormat: req.Format,
		Speed:          req.Speed,
	}
	body, err := json.Marshal(wire)
	if err != nil {
		return nil, fmt.Errorf("llama-swap: encode speech request: %w", err)
	}
	audioBytes, contentType, err := m.p.doRaw(ctx, http.MethodPost, "/v1/audio/speech", m.id, "application/json", bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	if len(audioBytes) == 0 {
		return nil, &llm.APIError{Provider: m.p.name, Model: m.id, Message: "speech response contained no audio"}
	}
	return &audio.SpeechResult{Audio: audioBytes, MIME: speechMIME(contentType, req.Format)}, nil
}

// speechMIME resolves the result MIME type: the response Content-Type when it
// is a concrete audio type, else a mapping from the requested format, else
// audio/mpeg (the OpenAI endpoint's default container is mp3).
func speechMIME(contentType, format string) string {
	if mt, _, err := mime.ParseMediaType(contentType); err == nil && strings.HasPrefix(mt, "audio/") {
		return mt
	}
	switch strings.ToLower(strings.TrimSpace(format)) {
	case "wav":
		return "audio/wav"
	case "opus":
		return "audio/ogg"
	case "aac":
		return "audio/aac"
	case "flac":
		return "audio/flac"
	default: // "", "mp3", and anything unrecognized
		return "audio/mpeg"
	}
}

// TranscriptionModel implements audio.TranscriptionProvider, binding a
// speech-to-text model served by llama-swap (routed to a whisper.cpp-style
// upstream).
func (p *Provider) TranscriptionModel(id string, opts ...audio.TranscriptionModelOption) (audio.TranscriptionModel, error) {
	if err := p.requireBaseURL(); err != nil {
		return nil, err
	}
	_ = audio.ApplyTranscriptionModelOptions(opts)
	return &transcriptionModel{p: p, id: id}, nil
}

type transcriptionModel struct {
	p  *Provider
	id string
}

// Transcribe implements audio.TranscriptionModel via POST
// {base}/v1/audio/transcriptions (multipart/form-data — llama-swap routes by
// the `model` form field).
func (m *transcriptionModel) Transcribe(ctx context.Context, req audio.TranscriptionRequest, opts ...audio.TranscriptionOption) (*audio.TranscriptionResult, error) {
	req = req.Apply(opts...)
	if len(req.Audio) == 0 {
		return nil, fmt.Errorf("%w: transcription requires audio bytes", llm.ErrUnsupported)
	}

	var buf bytes.Buffer
	w := multipart.NewWriter(&buf)
	fw, err := w.CreateFormFile("file", transcriptionFilename(req))
	if err != nil {
		return nil, fmt.Errorf("llama-swap: build transcription form: %w", err)
	}
	if _, err := fw.Write(req.Audio); err != nil {
		return nil, fmt.Errorf("llama-swap: build transcription form: %w", err)
	}
	// Required fields always go on the wire (an empty model id should fail
	// loudly upstream, not silently vanish); optional ones only when set.
	fields := []struct {
		key, value string
		required   bool
	}{
		{"model", m.id, true},
		{"response_format", "json", true},
		{"language", req.Language, false},
		{"prompt", req.Prompt, false},
	}
	for _, f := range fields {
		if !f.required && f.value == "" {
			continue
		}
		if err := w.WriteField(f.key, f.value); err != nil {
			return nil, fmt.Errorf("llama-swap: build transcription form: %w", err)
		}
	}
	if err := w.Close(); err != nil {
		return nil, fmt.Errorf("llama-swap: build transcription form: %w", err)
	}

	raw, _, err := m.p.doRaw(ctx, http.MethodPost, "/v1/audio/transcriptions", m.id, w.FormDataContentType(), &buf)
	if err != nil {
		return nil, err
	}
	var out struct {
		Text string `json:"text"`
	}
	if err := json.Unmarshal(raw, &out); err != nil {
		return nil, fmt.Errorf("llama-swap: decode transcription response: %w", err)
	}
	return &audio.TranscriptionResult{Text: out.Text, Raw: json.RawMessage(raw)}, nil
}

// transcriptionFilename picks the multipart filename hint: the caller's
// (sanitized — upload metadata is untrusted and CR/LF would inject multipart
// headers), else one derived from the MIME subtype ("audio.mp3"), else
// "audio". MIME parameters ("audio/ogg; codecs=opus") are stripped before
// matching.
func transcriptionFilename(req audio.TranscriptionRequest) string {
	if name := sanitizeFilename(req.Filename); name != "" {
		return name
	}
	mt := strings.ToLower(strings.TrimSpace(req.MIME))
	if parsed, _, err := mime.ParseMediaType(mt); err == nil {
		mt = parsed
	}
	switch mt {
	case "audio/mpeg", "audio/mp3":
		return "audio.mp3"
	case "audio/wav", "audio/x-wav", "audio/wave":
		return "audio.wav"
	case "audio/ogg":
		return "audio.ogg"
	case "audio/opus":
		return "audio.opus"
	case "audio/flac", "audio/x-flac":
		return "audio.flac"
	case "audio/mp4", "audio/m4a", "audio/x-m4a":
		return "audio.m4a"
	case "audio/webm", "video/webm":
		return "audio.webm"
	default:
		return "audio"
	}
}

// sanitizeFilename strips characters that would corrupt or inject into the
// multipart Content-Disposition header. Quotes and backslashes are escaped
// by mime/multipart itself; CR/LF are not — they must go.
func sanitizeFilename(name string) string {
	name = strings.NewReplacer("\r", "", "\n", "").Replace(name)
	return strings.TrimSpace(name)
}

// ListVoices returns the voices a TTS model offers (GET
// /v1/audio/voices?model=...). The decode is tolerant: upstreams answer with
// either a bare string list or a list of {id|name} objects.
func (p *Provider) ListVoices(ctx context.Context, model string) ([]string, error) {
	if model == "" {
		return nil, fmt.Errorf("llama-swap: ListVoices requires a model id")
	}
	raw, _, err := p.doRaw(ctx, http.MethodGet, "/v1/audio/voices?model="+url.QueryEscape(model), model, "", nil)
	if err != nil {
		return nil, err
	}
	return parseVoices(raw)
}

// parseVoices extracts voice names from the various shapes upstreams use:
// {"voices":[...]} or {"data":[...]} envelopes (or a bare array), holding
// either strings or objects keyed by id/name/voice_id.
func parseVoices(raw []byte) ([]string, error) {
	var env struct {
		Voices json.RawMessage `json:"voices"`
		Data   json.RawMessage `json:"data"`
	}
	list := json.RawMessage(raw)
	if err := json.Unmarshal(raw, &env); err == nil {
		if len(env.Voices) > 0 {
			list = env.Voices
		} else if len(env.Data) > 0 {
			list = env.Data
		}
	}

	var names []string
	if err := json.Unmarshal(list, &names); err == nil {
		return names, nil
	}
	// Not a plain string list. The failed decode above may have partially
	// populated names (Unmarshal appends zero values before erroring on an
	// element type mismatch) — start fresh for the object shape.
	names = nil
	var objs []struct {
		ID      string `json:"id"`
		Name    string `json:"name"`
		VoiceID string `json:"voice_id"`
	}
	if err := json.Unmarshal(list, &objs); err == nil {
		for _, o := range objs {
			switch {
			case o.ID != "":
				names = append(names, o.ID)
			case o.Name != "":
				names = append(names, o.Name)
			case o.VoiceID != "":
				names = append(names, o.VoiceID)
			}
		}
		return names, nil
	}
	return nil, fmt.Errorf("llama-swap: unrecognized voices payload shape")
}

// doRaw performs a request to a llama-swap endpoint and returns the raw
// response body and its Content-Type — the sibling of doJSON for endpoints
// whose success payload is not JSON (audio bytes) or whose shape varies.
// contentType sets the request Content-Type when body is non-nil. A response
// larger than maxResponseBytes is an error, never a silent truncation.
func (p *Provider) doRaw(ctx context.Context, method, path, model, contentType string, body io.Reader) ([]byte, string, error) {
	if err := p.requireBaseURL(); err != nil {
		return nil, "", err
	}
	req, err := p.newRequest(ctx, method, path, contentType, body)
	if err != nil {
		return nil, "", err
	}
	resp, err := p.client.Do(req)
	if err != nil {
		return nil, "", fmt.Errorf("llama-swap: do request: %w", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode/100 != 2 {
		return nil, "", p.apiError(resp, model)
	}
	data, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes+1))
	if err != nil {
		return nil, "", fmt.Errorf("llama-swap: read response: %w", err)
	}
	if len(data) > maxResponseBytes {
		return nil, "", fmt.Errorf("llama-swap: response exceeds %d bytes", maxResponseBytes)
	}
	return data, resp.Header.Get("Content-Type"), nil
}