sundynix-agentix/sundynix-gateway/internal/voice/asr.go

package voice

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/gorilla/websocket"
)

// 火山引擎 V3 大模型流式语音识别（SeedASR / SAUC）客户端。协议见 voice/frame.go 与 voice-jarvis 记忆。
// 端点固定；鉴权走**新版 API Key**（Authorization: Bearer <APIKey>）。

const asrEndpoint = "wss://openspeech.bytedance.com/api/v3/sauc/bigmodel"

// ASRResult 是一次识别回传。Err 非空表示识别流出错/结束（此后 Results 关闭）。
type ASRResult struct {
	Text  string
	Final bool
	Err   error
}

// ASRSession 是一路流式识别会话：PushAudio 喂 PCM、Results 出转写、Finish 收尾、Close 关闭。
type ASRSession struct {
	conn    *websocket.Conn
	results chan ASRResult
}

// ---- 初始配置请求 JSON ----

type asrReq struct {
	User    asrUser    `json:"user"`
	Audio   asrAudio   `json:"audio"`
	Request asrOptions `json:"request"`
}
type asrUser struct {
	UID string `json:"uid"`
}
type asrAudio struct {
	Format  string `json:"format"`
	Rate    int    `json:"rate"`
	Bits    int    `json:"bits"`
	Channel int    `json:"channel"`
	Codec   string `json:"codec"`
}
type asrVAD struct {
	VadEnable     bool `json:"vad_enable"`
	EndWindowSize int  `json:"end_window_size"`
}
type asrOptions struct {
	ModelName  string `json:"model_name"`
	Language   string `json:"language"`
	EnableITN  bool   `json:"enable_itn"`
	EnablePunc bool   `json:"enable_punc"`
	ResultType string `json:"result_type"`
	VAD        asrVAD `json:"vad"`
}

// buildASRRequest 组初始配置 JSON。音频格式与客户端采集一致（PCM 16k 单声道 raw）。
// format="pcm" 是联调可调点（火山对 raw PCM 也接受 "raw"）。
func buildASRRequest(uid string) []byte {
	if uid == "" {
		uid = "sundynix"
	}
	r := asrReq{
		User:  asrUser{UID: uid},
		Audio: asrAudio{Format: "pcm", Rate: AudioSampleRate, Bits: AudioBits, Channel: AudioChannels, Codec: "raw"},
		Request: asrOptions{
			ModelName: "bigmodel", Language: "zh", EnableITN: true, EnablePunc: true,
			ResultType: "0", VAD: asrVAD{VadEnable: true, EndWindowSize: 800},
		},
	}
	b, _ := json.Marshal(r)
	return b
}

func newConnectID() string {
	b := make([]byte, 16)
	_, _ = rand.Read(b)
	return hex.EncodeToString(b)
}

// StartASR 连火山流式识别、发初始配置帧，返回会话。读 goroutine 持续把结果推入 Results。
func StartASR(ctx context.Context, cfg Config, uid string) (*ASRSession, error) {
	if !cfg.ASREnabled() {
		return nil, fmt.Errorf("ASR 未配置")
	}
	hdr := http.Header{}
	hdr.Set("Authorization", "Bearer "+cfg.APIKey) // 新版 API Key 鉴权
	hdr.Set("X-Api-Resource-Id", cfg.ASRResourceID)
	hdr.Set("X-Api-Connect-Id", newConnectID())

	dialer := websocket.Dialer{HandshakeTimeout: 10 * time.Second}
	conn, _, err := dialer.DialContext(ctx, asrEndpoint, hdr)
	if err != nil {
		return nil, fmt.Errorf("连接火山 ASR 失败: %w", err)
	}
	if err := conn.WriteMessage(websocket.BinaryMessage, jsonFrame(msgFullClientReq, flagNone, buildASRRequest(uid))); err != nil {
		_ = conn.Close()
		return nil, fmt.Errorf("发送 ASR 配置失败: %w", err)
	}
	s := &ASRSession{conn: conn, results: make(chan ASRResult, 16)}
	go s.readLoop()
	return s, nil
}

func (s *ASRSession) readLoop() {
	defer close(s.results)
	for {
		mt, data, err := s.conn.ReadMessage()
		if err != nil {
			s.results <- ASRResult{Err: err}
			return
		}
		if mt != websocket.BinaryMessage {
			continue
		}
		msgType, payload, ok := parseServerFrame(data)
		if !ok || len(payload) == 0 {
			continue
		}
		if msgType == msgServerError {
			s.results <- ASRResult{Err: fmt.Errorf("火山 ASR 错误: %s", strings.TrimSpace(string(payload)))}
			return
		}
		if text, final, ok := parseASRResult(payload); ok {
			s.results <- ASRResult{Text: text, Final: final}
		}
	}
}

// PushAudio 喂一帧 PCM。
func (s *ASRSession) PushAudio(pcm []byte) error {
	return s.conn.WriteMessage(websocket.BinaryMessage, audioFrame(pcm, false))
}

// Finish 发结束标记（空音频 + 最终帧），告知火山本轮说完。
func (s *ASRSession) Finish() error {
	return s.conn.WriteMessage(websocket.BinaryMessage, audioFrame(nil, true))
}

// Results 返回识别结果流（部分/最终；出错或结束时推一条 Err 后关闭）。
func (s *ASRSession) Results() <-chan ASRResult { return s.results }

// Close 关闭底层连接（读 goroutine 随之退出）。
func (s *ASRSession) Close() { _ = s.conn.Close() }

// parseASRResult 从响应 JSON 提取转写文本 + 是否最终。空文本且非最终时 ok=false。
func parseASRResult(payload []byte) (text string, final bool, ok bool) {
	var r struct {
		Type   string          `json:"type"`
		Result json.RawMessage `json:"result"`
	}
	if json.Unmarshal(payload, &r) != nil {
		return "", false, false
	}
	final = r.Type == "final"
	text = extractText(r.Result)
	return text, final, text != "" || final
}

// extractText 从 result 里取转写文本；result 可为 [{text}]/{text}/"string"。
func extractText(raw json.RawMessage) string {
	if len(raw) == 0 {
		return ""
	}
	var arr []struct {
		Text string `json:"text"`
	}
	if json.Unmarshal(raw, &arr) == nil && len(arr) > 0 {
		var sb strings.Builder
		for _, a := range arr {
			sb.WriteString(a.Text)
		}
		if sb.Len() > 0 {
			return sb.String()
		}
	}
	var obj struct {
		Text string `json:"text"`
	}
	if json.Unmarshal(raw, &obj) == nil && obj.Text != "" {
		return obj.Text
	}
	var s string
	if json.Unmarshal(raw, &s) == nil {
		return s
	}
	return ""
}