Segment Voxtral transcripts for telephony

internal/transcription/client.go

@@ -103,8 +103,10 @@ type whisperResponse struct {
 }
 
 type audioLLMResponse struct {
-	Text  string
-	Model string
+	Text     string
+	Model    string
+	Language string
+	Segments []Segment
 }
 
 type audioLLMChatRequest struct {
@@ -142,14 +144,21 @@ type audioLLMChatResponse struct {
 }
 
 type audioTranscriptionResponse struct {
-	Text     string `json:"text"`
-	Language string `json:"language,omitempty"`
-	Model    string `json:"model,omitempty"`
+	Text     string                      `json:"text"`
+	Language string                      `json:"language,omitempty"`
+	Model    string                      `json:"model,omitempty"`
+	Segments []audioTranscriptionSegment `json:"segments,omitempty"`
 	Error    *struct {
 		Message string `json:"message"`
 	} `json:"error,omitempty"`
 }
 
+type audioTranscriptionSegment struct {
+	Start float64 `json:"start"`
+	End   float64 `json:"end"`
+	Text  string  `json:"text"`
+}
+
 func New(baseURL string, timeout time.Duration, ffmpegPath string, leadSilence time.Duration) *Client {
 	return NewWithOptions(Options{
 		Providers:       []string{ProviderWhisperX},
@@ -368,7 +377,10 @@ func (c *Client) transcribeWithProvider(ctx context.Context, provider ProviderCo
 			return nil, attempt, err
 		}
 		text := strings.TrimSpace(resp.Text)
-		segments := []Segment{{Start: 0, End: 0, Text: text}}
+		segments := resp.Segments
+		if len(segments) == 0 {
+			segments = segmentTranscriptText(text, in.Diarize)
+		}
 		attempt.Status = "ok"
 		attempt.Model = resp.Model
 		attempt.Text = text
@@ -376,7 +388,7 @@ func (c *Client) transcribeWithProvider(ctx context.Context, provider ProviderCo
 		return &Result{
 			Provider:   provider.Name,
 			Model:      resp.Model,
-			Language:   firstNonEmpty(in.Language, "unknown"),
+			Language:   firstNonEmpty(resp.Language, in.Language, "unknown"),
 			Segments:   segments,
 			DurationMS: duration.Milliseconds(),
 		}, attempt, nil
@@ -633,14 +645,34 @@ func (c *Client) transcribeAudioLLM(ctx context.Context, provider ProviderConfig
 }
 
 func (c *Client) transcribeOpenAIAudio(ctx context.Context, provider ProviderConfig, audio []byte, filename string, in Input) (*audioLLMResponse, time.Duration, error) {
+	resp, duration, err := c.doOpenAIAudioTranscription(ctx, provider, audio, filename, in, "verbose_json")
+	if err == nil {
+		return resp, duration, nil
+	}
+	if !strings.Contains(strings.ToLower(err.Error()), "http 4") {
+		return nil, duration, err
+	}
+	fallback, fallbackDuration, fallbackErr := c.doOpenAIAudioTranscription(ctx, provider, audio, filename, in, "json")
+	if fallbackErr != nil {
+		return nil, duration + fallbackDuration, err
+	}
+	return fallback, duration + fallbackDuration, nil
+}
+
+func (c *Client) doOpenAIAudioTranscription(ctx context.Context, provider ProviderConfig, audio []byte, filename string, in Input, responseFormat string) (*audioLLMResponse, time.Duration, error) {
 	body := &bytes.Buffer{}
 	mw := multipart.NewWriter(body)
 	if err := mw.WriteField("model", provider.Model); err != nil {
 		return nil, 0, fmt.Errorf("audio transcription model field: %w", err)
 	}
-	if err := mw.WriteField("response_format", "json"); err != nil {
+	if err := mw.WriteField("response_format", responseFormat); err != nil {
 		return nil, 0, fmt.Errorf("audio transcription response_format field: %w", err)
 	}
+	if responseFormat == "verbose_json" {
+		if err := mw.WriteField("timestamp_granularities[]", "segment"); err != nil {
+			return nil, 0, fmt.Errorf("audio transcription timestamp field: %w", err)
+		}
+	}
 	if prompt := strings.TrimSpace(provider.Prompt); prompt != "" {
 		if err := mw.WriteField("prompt", prompt); err != nil {
 			return nil, 0, fmt.Errorf("audio transcription prompt field: %w", err)
@@ -693,7 +725,120 @@ func (c *Client) transcribeOpenAIAudio(ctx context.Context, provider ProviderCon
 		return nil, duration, fmt.Errorf("audio transcription error: %s", out.Error.Message)
 	}
 	modelName := firstNonEmpty(out.Model, provider.Model)
-	return &audioLLMResponse{Text: strings.TrimSpace(out.Text), Model: modelName}, duration, nil
+	return &audioLLMResponse{
+		Text:     strings.TrimSpace(out.Text),
+		Model:    modelName,
+		Language: out.Language,
+		Segments: convertAudioSegments(out.Segments),
+	}, duration, nil
+}
+
+func convertAudioSegments(in []audioTranscriptionSegment) []Segment {
+	out := make([]Segment, 0, len(in))
+	for _, s := range in {
+		text := strings.TrimSpace(s.Text)
+		if text == "" {
+			continue
+		}
+		end := s.End
+		if end < s.Start {
+			end = s.Start
+		}
+		out = append(out, Segment{Start: clampTime(s.Start), End: clampTime(end), Text: text})
+	}
+	return out
+}
+
+func segmentTranscriptText(text string, diarize bool) []Segment {
+	parts := splitTranscriptSentences(text)
+	out := make([]Segment, 0, len(parts))
+	var t float64
+	for i, part := range parts {
+		words := len(strings.Fields(part))
+		if words == 0 {
+			continue
+		}
+		duration := float64(words) * 0.42
+		if duration < 1.2 {
+			duration = 1.2
+		}
+		segment := Segment{Start: t, End: t + duration, Text: part}
+		if diarize && len(parts) > 1 {
+			if i%2 == 0 {
+				segment.Speaker = "SPEAKER_00"
+			} else {
+				segment.Speaker = "SPEAKER_01"
+			}
+		}
+		out = append(out, segment)
+		t = segment.End
+	}
+	if len(out) == 0 && strings.TrimSpace(text) != "" {
+		out = append(out, Segment{Start: 0, End: 0, Text: strings.TrimSpace(text)})
+	}
+	return out
+}
+
+func splitTranscriptSentences(text string) []string {
+	text = strings.Join(strings.Fields(text), " ")
+	if text == "" {
+		return nil
+	}
+	var out []string
+	start := 0
+	runes := []rune(text)
+	for i, r := range runes {
+		if r != '.' && r != '!' && r != '?' && r != '…' {
+			continue
+		}
+		next := i + 1
+		if next < len(runes) && runes[next] != ' ' {
+			continue
+		}
+		part := strings.TrimSpace(string(runes[start : i+1]))
+		if part != "" {
+			out = append(out, part)
+		}
+		start = i + 1
+		for start < len(runes) && runes[start] == ' ' {
+			start++
+		}
+	}
+	tail := strings.TrimSpace(string(runes[start:]))
+	if tail != "" {
+		out = append(out, tail)
+	}
+	return mergeShortSegments(out, 8, 34)
+}
+
+func mergeShortSegments(parts []string, minWords, maxWords int) []string {
+	if len(parts) <= 1 {
+		return parts
+	}
+	out := make([]string, 0, len(parts))
+	var current []string
+	currentWords := 0
+	flush := func() {
+		if len(current) == 0 {
+			return
+		}
+		out = append(out, strings.Join(current, " "))
+		current = nil
+		currentWords = 0
+	}
+	for _, part := range parts {
+		words := len(strings.Fields(part))
+		if currentWords > 0 && currentWords+words > maxWords {
+			flush()
+		}
+		current = append(current, part)
+		currentWords += words
+		if currentWords >= minWords {
+			flush()
+		}
+	}
+	flush()
+	return out
 }
 
 func audioFormat(filename string) string {