Sambert WebSocket API 参考

package main

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"time"

	"github.com/google/uuid"
	"github.com/gorilla/websocket"
)

const (
	wsURL      = "wss://dashscope.aliyuncs.com/api-ws/v1/inference/" // WebSocket服务器地址
	outputFile = "output.mp3"                                        // 输出文件路径
)

func main() {
	// 若没有将API Key配置到环境变量，可将下行替换为：apiKey := "your_api_key"。不建议在生产环境中直接将API Key硬编码到代码中，以减少API Key泄露风险。
	apiKey := os.Getenv("DASHSCOPE_API_KEY")

	if err := clearOutputFile(outputFile); err != nil {
		fmt.Println("清空输出文件失败：", err)
		return
	}

	conn, err := connectWebSocket(apiKey)
	if err != nil {
		fmt.Println("连接WebSocket失败：", err)
		return
	}
	defer closeConnection(conn)

	done := make(chan struct{})
	go receiveMessage(conn, done)

	if err := sendRunTaskMsg(conn); err != nil {
		fmt.Println("发送run-task指令失败：", err)
		return
	}

	select {
	case <-done:
		fmt.Println("任务结束")
	case <-time.After(5 * time.Minute):
		fmt.Println("任务超时")
	}
}

type Message struct {
	Header  Header  `json:"header"`
	Payload Payload `json:"payload"`
}

type Header struct {
	Action       string                 `json:"action,omitempty"`
	TaskID       string                 `json:"task_id"`
	Streaming    string                 `json:"streaming,omitempty"`
	Event        string                 `json:"event,omitempty"`
	ErrorCode    string                 `json:"error_code,omitempty"`
	ErrorMessage string                 `json:"error_message,omitempty"`
	Attributes   map[string]interface{} `json:"attributes"`
}

type Payload struct {
	Model      string     `json:"model,omitempty"`
	TaskGroup  string     `json:"task_group,omitempty"`
	Task       string     `json:"task,omitempty"`
	Function   string     `json:"function,omitempty"`
	Input      Input      `json:"input,omitempty"`
	Parameters Parameters `json:"parameters,omitempty"`
	Output     Output     `json:"output,omitempty"`
	Usage      Usage      `json:"usage,omitempty"`
}

type Input struct {
	Text string `json:"text"`
}

type Parameters struct {
	TextType                string  `json:"text_type"`
	Format                  string  `json:"format"`
	SampleRate              int     `json:"sample_rate"`
	Volume                  int     `json:"volume"`
	Rate                    float64 `json:"rate"`
	Pitch                   float64 `json:"pitch"`
	WordTimestampEnabled    bool    `json:"word_timestamp_enabled"`
	PhonemeTimestampEnabled bool    `json:"phoneme_timestamp_enabled"`
}

type Output struct {
	Sentence Sentence `json:"sentence"`
}

type Sentence struct {
	BeginTime int    `json:"begin_time"`
	EndTime   int    `json:"end_time"`
	Words     []Word `json:"words"`
}

type Word struct {
	Text      string    `json:"text"`
	BeginTime int       `json:"begin_time"`
	EndTime   int       `json:"end_time"`
	Phonemes  []Phoneme `json:"phonemes"`
}

type Phoneme struct {
	BeginTime int    `json:"begin_time"`
	EndTime   int    `json:"end_time"`
	Text      string `json:"text"`
	Tone      int    `json:"tone"`
}

type Usage struct {
	Characters int `json:"characters"`
}

func receiveMessage(conn *websocket.Conn, done chan struct{}) {
	for {
		msgType, message, err := conn.ReadMessage()
		if err != nil {
			fmt.Println("解析服务器消息失败：", err)
			close(done)
			break
		}

		if msgType == websocket.BinaryMessage {
			if err := writeBinaryDataToFile(message, outputFile); err != nil {
				fmt.Println("写入二进制数据失败：", err)
				close(done)
				break
			}
			fmt.Println("音频片段已写入本地文件")
		} else {
			var msg Message
			if err := json.Unmarshal(message, &msg); err != nil {
				fmt.Println("解析事件失败：", err)
				continue
			}
			if handleMessage(conn, msg, done) {
				break
			}
		}
	}
}

func handleMessage(conn *websocket.Conn, msg Message, done chan struct{}) bool {
	switch msg.Header.Event {
	case "task-started":
		fmt.Println("任务已启动")
	case "result-generated":
		// 如需获取附加消息，可在此处添加相应代码
	case "task-finished":
		fmt.Println("任务已完成")
		close(done)
		return true
	case "task-failed":
		if msg.Header.ErrorMessage != "" {
			fmt.Printf("任务失败：%s\n", msg.Header.ErrorMessage)
		} else {
			fmt.Println("未知原因导致任务失败")
		}
		close(done)
		return true
	default:
		fmt.Printf("预料之外的事件：%v\n", msg)
		close(done)
	}
	return false
}

func sendRunTaskMsg(conn *websocket.Conn) error {
	runTaskMsg, err := generateRunTaskMsg()
	if err != nil {
		return err
	}
	return conn.WriteMessage(websocket.TextMessage, []byte(runTaskMsg))
}

func generateRunTaskMsg() (string, error) {
	runTaskMessage := Message{
		Header: Header{
			Action:    "run-task",
			TaskID:    uuid.New().String(),
			Streaming: "out",
		},
		Payload: Payload{
			Model:     "sambert-zhichu-v1",
			TaskGroup: "audio",
			Task:      "tts",
			Function:  "SpeechSynthesizer",
			Input: Input{
				Text: "白日依山尽，黄河入海流。欲穷千里目，更上一层楼。",
			},
			Parameters: Parameters{
				TextType:                "PlainText",
				Format:                  "mp3",
				SampleRate:              16000,
				Volume:                  50,
				Rate:                    1.0,
				Pitch:                   1.0,
				WordTimestampEnabled:    true,
				PhonemeTimestampEnabled: true,
			},
		},
	}
	runTaskMsgJSON, err := json.Marshal(runTaskMessage)
	return string(runTaskMsgJSON), err
}

func connectWebSocket(apiKey string) (*websocket.Conn, error) {
	header := make(http.Header)
	header.Add("X-DashScope-DataInspection", "enable")
	header.AddPOST