CamTalk/backend/internal/ai/stt/mimo.go

package stt

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"go.uber.org/zap"
)

// MiMoService 基于 Xiaomi MiMo ASR HTTP API 的语音识别实现。
// 接口兼容 OpenAI chat/completions 格式，音频仅支持 mp3/wav。
type MiMoService struct {
	apiKey   string
	model    string
	endpoint string
	timeout  time.Duration
	logger   *zap.SugaredLogger
}

// NewMiMoService 创建 MiMo STT 服务。
// model、endpoint 由 config 层保证非空，timeoutSec 为 0 时默认 10 秒。
func NewMiMoService(apiKey, model, endpoint string, timeoutSec int, logger *zap.SugaredLogger) *MiMoService {
	timeout := time.Duration(timeoutSec) * time.Second
	if timeout <= 0 {
		timeout = 10 * time.Second
	}
	return &MiMoService{
		apiKey:   apiKey,
		model:    model,
		endpoint: endpoint,
		timeout:  timeout,
		logger:   logger,
	}
}

// mimoRequest MiMo ASR 请求体。
type mimoRequest struct {
	Model      string         `json:"model"`
	Messages   []mimoMessage  `json:"messages"`
	ASROptions *mimoASROptions `json:"asr_options,omitempty"`
}

type mimoMessage struct {
	Role    string        `json:"role"`
	Content []mimoContent `json:"content"`
}

type mimoContent struct {
	Type       string       `json:"type"`
	InputAudio *mimoAudioIn `json:"input_audio,omitempty"`
}

type mimoAudioIn struct {
	Data string `json:"data"` // data URL: data:{mime};base64,{data}
}

type mimoASROptions struct {
	Language string `json:"language"`
}

// mimoResponse MiMo ASR 非流式响应。
type mimoResponse struct {
	Choices []struct {
		Message struct {
			Content string `json:"content"`
		} `json:"message"`
	} `json:"choices"`
}

// Recognize 实现 stt.Service。将音频发送到 MiMo ASR API，返回识别文本。
func (m *MiMoService) Recognize(ctx context.Context, audio []byte, opts Options) (string, error) {
	if len(audio) == 0 {
		return "", fmt.Errorf("stt: empty audio")
	}

	// MiMo 仅支持 mp3/wav，若输入为原始 PCM 则封装为 WAV
	audioData := audio
	mimeType := "audio/wav"
	if !isWAV(audio) && !isMP3(audio) {
		wav, err := pcmToWAV(audio, opts.SampleRate, 1)
		if err != nil {
			return "", fmt.Errorf("stt: pcm to wav: %w", err)
		}
		audioData = wav
	} else if isMP3(audio) {
		mimeType = "audio/mpeg"
	}

	b64 := base64.StdEncoding.EncodeToString(audioData)
	dataURL := fmt.Sprintf("data:%s;base64,%s", mimeType, b64)

	// 映射语言代码
	language := mapLanguage(opts.Language)

	reqBody := mimoRequest{
		Model: m.model,
		Messages: []mimoMessage{
			{
				Role: "user",
				Content: []mimoContent{
					{
						Type: "input_audio",
						InputAudio: &mimoAudioIn{
							Data: dataURL,
						},
					},
				},
			},
		},
	}
	if language != "" {
		reqBody.ASROptions = &mimoASROptions{Language: language}
	}

	body, err := json.Marshal(reqBody)
	if err != nil {
		return "", fmt.Errorf("stt: marshal request: %w", err)
	}

	url := strings.TrimRight(m.endpoint, "/") + "/chat/completions"

	ctx, cancel := context.WithTimeout(ctx, m.timeout)
	defer cancel()

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
	if err != nil {
		return "", fmt.Errorf("stt: create request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+m.apiKey)

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return "", fmt.Errorf("stt: request mimo: %w", err)
	}
	defer resp.Body.Close()

	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("stt: read response: %w", err)
	}

	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("stt: mimo returned %d: %s", resp.StatusCode, string(respBody))
	}

	var mResp mimoResponse
	if err := json.Unmarshal(respBody, &mResp); err != nil {
		return "", fmt.Errorf("stt: unmarshal response: %w", err)
	}

	if len(mResp.Choices) == 0 {
		return "", fmt.Errorf("stt: mimo returned empty choices")
	}

	text := strings.TrimSpace(mResp.Choices[0].Message.Content)
	return text, nil
}

// mapLanguage 将标准语言代码映射为 MiMo 支持的值（auto/zh/en）。
func mapLanguage(lang string) string {
	switch {
	case lang == "":
		return "auto"
	case strings.HasPrefix(lang, "zh"):
		return "zh"
	case strings.HasPrefix(lang, "en"):
		return "en"
	default:
		return "auto"
	}
}

// isWAV 检查数据是否为 WAV 格式（RIFF 头）。
func isWAV(data []byte) bool {
	return len(data) > 4 && string(data[:4]) == "RIFF"
}

// isMP3 检查数据是否为 MP3 格式（ID3 标签或帧同步字）。
func isMP3(data []byte) bool {
	if len(data) > 3 && string(data[:3]) == "ID3" {
		return true
	}
	// 帧同步字：0xFF 0xFB/0xF3/0xF2
	return len(data) > 2 && data[0] == 0xFF && (data[1]&0xE0) == 0xE0
}

// pcmToWAV 将原始 PCM 数据封装为 WAV 文件。
func pcmToWAV(pcm []byte, sampleRate, channels int) ([]byte, error) {
	if sampleRate == 0 {
		sampleRate = 16000
	}
	if channels == 0 {
		channels = 1
	}

	bitsPerSample := 16
	byteRate := sampleRate * channels * bitsPerSample / 8
	blockAlign := channels * bitsPerSample / 8
	dataSize := len(pcm)

	var buf bytes.Buffer

	// RIFF header
	buf.WriteString("RIFF")
	binary.Write(&buf, binary.LittleEndian, uint32(36+dataSize))
	buf.WriteString("WAVE")

	// fmt 子块
	buf.WriteString("fmt ")
	binary.Write(&buf, binary.LittleEndian, uint32(16))                // 子块大小
	binary.Write(&buf, binary.LittleEndian, uint16(1))                 // PCM 格式
	binary.Write(&buf, binary.LittleEndian, uint16(channels))          // 通道数
	binary.Write(&buf, binary.LittleEndian, uint32(sampleRate))        // 采样率
	binary.Write(&buf, binary.LittleEndian, uint32(byteRate))          // 字节率
	binary.Write(&buf, binary.LittleEndian, uint16(blockAlign))        // 块对齐
	binary.Write(&buf, binary.LittleEndian, uint16(bitsPerSample))     // 每样本位数

	// data 子块
	buf.WriteString("data")
	binary.Write(&buf, binary.LittleEndian, uint32(dataSize))
	buf.Write(pcm)

	return buf.Bytes(), nil
}