majordomo/provider/llamaswap/diarize.go

// diarize.go implements audio.DiarizationProvider against a
// whisper-asr-webservice upstream (WhisperX engine) reached through
// llama-swap's /upstream passthrough (ADR-0020):
//
//	POST /upstream/<id>/asr?output=json&diarize=true[&language&min_speakers&max_speakers]
//
// output=json is load-bearing: the vtt/srt output formats DROP the speaker
// labels; only the json shape carries per-segment `speaker` fields.
package llamaswap

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strconv"

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

// DiarizationModel implements audio.DiarizationProvider. The id selects
// which upstream llama-swap loads (the pyannote-equipped WhisperX sidecar,
// not the plain whisper.cpp model).
func (p *Provider) DiarizationModel(id string, opts ...audio.DiarizationModelOption) (audio.DiarizationModel, error) {
	if err := p.requireBaseURL(); err != nil {
		return nil, err
	}
	_ = audio.ApplyDiarizationModelOptions(opts)
	return &diarizationModel{p: p, id: id}, nil
}

type diarizationModel struct {
	p  *Provider
	id string
}

// asrResponse is whisper-asr-webservice's output=json shape (the subset this
// package relies on; per-word entries are ignored — segment granularity is
// the contract).
type asrResponse struct {
	Text     string `json:"text"`
	Language string `json:"language"`
	Segments []struct {
		Start   float64 `json:"start"`
		End     float64 `json:"end"`
		Text    string  `json:"text"`
		Speaker string  `json:"speaker"`
	} `json:"segments"`
}

// Diarize implements audio.DiarizationModel.
func (m *diarizationModel) Diarize(ctx context.Context, req audio.DiarizationRequest, opts ...audio.DiarizationOption) (*audio.DiarizationResult, error) {
	req = req.Apply(opts...)
	if len(req.Audio) == 0 {
		return nil, fmt.Errorf("%w: diarization requires audio bytes", llm.ErrUnsupported)
	}
	if req.MinSpeakers < 0 || req.MaxSpeakers < 0 ||
		(req.MaxSpeakers > 0 && req.MinSpeakers > req.MaxSpeakers) {
		return nil, fmt.Errorf("%w: invalid speaker bounds [%d,%d]", llm.ErrUnsupported, req.MinSpeakers, req.MaxSpeakers)
	}

	q := url.Values{}
	q.Set("output", "json")
	q.Set("diarize", "true")
	if req.Language != "" {
		q.Set("language", req.Language)
	}
	if req.MinSpeakers > 0 {
		q.Set("min_speakers", strconv.Itoa(req.MinSpeakers))
	}
	if req.MaxSpeakers > 0 {
		q.Set("max_speakers", strconv.Itoa(req.MaxSpeakers))
	}
	path, err := upstreamPath(m.id, "/asr?"+q.Encode())
	if err != nil {
		return nil, err
	}

	body, contentType, err := buildMultipart("build diarization form",
		filePart{field: "audio_file", filename: diarizationFilename(req), data: req.Audio},
		nil)
	if err != nil {
		return nil, err
	}
	raw, _, err := m.p.doRaw(ctx, http.MethodPost, path, m.id, contentType, body, maxResponseBytes)
	if err != nil {
		return nil, err
	}

	var out asrResponse
	if err := json.Unmarshal(raw, &out); err != nil {
		return nil, fmt.Errorf("llama-swap: decode diarization response: %w", err)
	}
	res := &audio.DiarizationResult{
		Text:     out.Text,
		Language: out.Language,
		Raw:      json.RawMessage(raw),
	}
	for _, s := range out.Segments {
		res.Segments = append(res.Segments, audio.DiarizationSegment{
			Start:   s.Start,
			End:     s.End,
			Speaker: s.Speaker,
			Text:    s.Text,
		})
	}
	if len(res.Segments) == 0 && res.Text == "" {
		return nil, &llm.APIError{Provider: m.p.name, Model: m.id, Message: "diarization response contained no transcript"}
	}
	return res, nil
}

// diarizationFilename mirrors transcriptionFilename for the diarization
// request shape (same untrusted-metadata sanitization rules).
func diarizationFilename(req audio.DiarizationRequest) string {
	return transcriptionFilename(audio.TranscriptionRequest{
		Filename: req.Filename,
		MIME:     req.MIME,
	})
}