ai-agent/workflow/service/flow/subtitle.go

package flow

import (
	flowDto "ai-agent/workflow/model/dto/flow"
	"regexp"
	"strings"
	"unicode/utf8"
)

// punctRe 切分/剥离用的中文标点（含顿号、）
var punctRe = regexp.MustCompile(`[，。；！？、]`)

// BuildSubtitles 核心工具：单个sentence生成多条subtitle
func BuildSubtitles(sents *[]flowDto.Sentence) ([]flowDto.Subtitle, error) {
	var subtitles []flowDto.Subtitle

	for _, sent := range *sents {
		// 1. 先按标点把文本拆成多个片段
		segList := splitTextByPunct(sent.Text)
		if len(segList) == 0 {
			continue
		}

		// 去标点后得到纯净片段（纯空白/纯标点片段跳过）
		var cleans []string
		for _, seg := range segList {
			c := strings.TrimSpace(cleanPunct(seg))
			if c != "" {
				cleans = append(cleans, c)
			}
		}
		if len(cleans) == 0 || len(sent.Words) == 0 {
			continue
		}

		// 2. 词级文本与句子文本一致时，按词精确对齐取首尾词时间（最准）
		if spans, ok := alignAllSegments(sent.Words, cleans); ok {
			for i, span := range spans {
				subtitles = append(subtitles, flowDto.Subtitle{
					Start: sent.Words[span[0]].StartTime,
					End:   sent.Words[span[1]].EndTime,
					Text:  cleans[i],
				})
			}
			continue
		}

		// 3. ASR 词级转写与句子文本不一致时（如 血→谑、数字写法不一），
		//    整句回退为按片段字符占比分配时间，避免整句被吞成一条字幕
		segWords := allocWordsByProportion(sent.Words, cleans)
		for i, ws := range segWords {
			if len(ws) == 0 {
				continue
			}
			subtitles = append(subtitles, flowDto.Subtitle{
				Start: ws[0].StartTime,
				End:   ws[len(ws)-1].EndTime,
				Text:  cleans[i],
			})
		}
	}

	return subtitles, nil
}

// splitTextByPunct 按中文标点分割句子，同时保留标点在分段内
// 例如："这个叫高血压调理方，注意是根源调理不是临时缓解，"
// 会变成：["这个叫高血压调理方，", "注意是根源调理不是临时缓解，"]
func splitTextByPunct(raw string) []string {
	// 匹配中文标点并保留在文本中，按标点位置切分
	indexes := punctRe.FindAllStringIndex(raw, -1)
	if len(indexes) == 0 {
		return []string{raw}
	}

	var res []string
	prev := 0
	for _, idx := range indexes {
		end := idx[1] // 标点的结束位置
		seg := raw[prev:end]
		res = append(res, seg)
		prev = end
	}
	// 处理最后一段没有标点的文本
	if prev < len(raw) {
		res = append(res, raw[prev:])
	}
	return res
}

// cleanPunct 去掉中文标点，得到纯净文本
func cleanPunct(raw string) string {
	return punctRe.ReplaceAllString(raw, "")
}

// alignAllSegments 按顺序把各纯净片段与词级文本逐字符对齐（允许个别字符不一致）。
// 全部片段对齐成功且词被完整覆盖时返回各片段对应的词区间，否则 ok=false，
// 由调用方回退到时间占比分配。
func alignAllSegments(words []flowDto.Word, cleans []string) ([][2]int, bool) {
	spans := make([][2]int, len(cleans))
	wordIdx := 0
	for i, seg := range cleans {
		start := wordIdx
		segRunes := []rune(seg)
		s := 0
		for wordIdx < len(words) && s < len(segRunes) {
			for _, r := range []rune(words[wordIdx].Word) {
				if s < len(segRunes) && r == segRunes[s] {
					s++
				}
			}
			wordIdx++
		}
		// 片段文本没被完整匹配，或该片段没吃到任何词 → 无法精确对齐
		if s < len(segRunes) || start == wordIdx {
			return nil, false
		}
		spans[i] = [2]int{start, wordIdx - 1}
	}
	// 有剩余词未被任何片段覆盖，说明对齐失败，避免吞掉剩余时间
	if wordIdx < len(words) {
		return nil, false
	}
	return spans, true
}

// allocWordsByProportion 按纯净片段字符占比把整句时间区间切成段，再按时间中点把
// 每个 word 归属到所属片段（对词级转写与句子文本不一致的情况兜底）。
func allocWordsByProportion(words []flowDto.Word, cleans []string) [][]flowDto.Word {
	runes := make([]int, len(cleans))
	totalChars := 0
	for i, c := range cleans {
		runes[i] = utf8.RuneCountInString(c)
		totalChars += runes[i]
	}

	sentStart := words[0].StartTime
	sentEnd := words[len(words)-1].EndTime
	duration := sentEnd - sentStart
	if duration < 0 {
		duration = 0
	}

	bounds := make([]float64, len(cleans)+1)
	bounds[0] = sentStart
	accum := 0.0
	for i := range cleans {
		if totalChars > 0 {
			accum += float64(runes[i]) / float64(totalChars)
		}
		bounds[i+1] = sentStart + accum*duration
	}

	segWords := make([][]flowDto.Word, len(cleans))
	for _, w := range words {
		mid := (w.StartTime + w.EndTime) / 2
		idx := 0
		for b := 0; b < len(bounds)-1; b++ {
			if mid >= bounds[b+1] {
				idx = b + 1
			}
		}
		segWords[idx] = append(segWords[idx], w)
	}
	return segWords
}