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

package stt

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/gorilla/websocket"
	"go.uber.org/zap"
)

// DeepgramService 基于 Deepgram WebSocket API 的语音识别实现。
type DeepgramService struct {
	apiKey   string
	model    string
	endpoint string
	timeout  time.Duration
	logger   *zap.SugaredLogger
}

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

// deepgramResponse Deepgram WebSocket 响应。
type deepgramResponse struct {
	Channel struct {
		Alternatives []struct {
			Transcript string  `json:"transcript"`
			Confidence float64 `json:"confidence"`
		} `json:"alternatives"`
	} `json:"channel"`
	IsFinal bool `json:"is_final"`
}

// Recognize 实现 stt.Service。通过 WebSocket 发送音频到 Deepgram，返回最终识别文本。
func (d *DeepgramService) Recognize(ctx context.Context, audio []byte, opts Options) (string, error) {
	if len(audio) == 0 {
		return "", fmt.Errorf("stt: empty audio")
	}

	// 构建 WebSocket URL，附带查询参数
	wsURL := d.buildURL(opts)

	// 总超时
	ctx, cancel := context.WithTimeout(ctx, d.timeout)
	defer cancel()

	// 建立 WebSocket 连接
	conn, _, err := websocket.DefaultDialer.DialContext(ctx, wsURL, http.Header{
		"Authorization": []string{"Token " + d.apiKey},
	})
	if err != nil {
		return "", fmt.Errorf("stt: connect deepgram: %w", err)
	}
	defer conn.Close()

	// 发送音频数据（一次性）
	if err := conn.WriteMessage(websocket.BinaryMessage, audio); err != nil {
		return "", fmt.Errorf("stt: send audio: %w", err)
	}

	// 发送 Close 消息通知服务端音频已发送完毕
	closeMsg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")
	_ = conn.WriteMessage(websocket.CloseMessage, closeMsg)

	// 读取识别结果
	var transcript strings.Builder
	for {
		_, message, err := conn.ReadMessage()
		if err != nil {
			// Close 帧是正常的结束信号
			if websocket.IsCloseError(err, websocket.CloseNormalClosure) {
				break
			}
			if websocket.IsUnexpectedCloseError(err, websocket.CloseNormalClosure) {
				break
			}
			return "", fmt.Errorf("stt: read response: %w", err)
		}

		var resp deepgramResponse
		if err := json.Unmarshal(message, &resp); err != nil {
			d.logger.Warnw("stt: unmarshal response failed", "error", err)
			continue
		}

		// 只累积 final 结果，跳过中间结果
		if resp.IsFinal && len(resp.Channel.Alternatives) > 0 {
			text := strings.TrimSpace(resp.Channel.Alternatives[0].Transcript)
			if text != "" {
				transcript.WriteString(text)
			}
		}
	}

	return strings.TrimSpace(transcript.String()), nil
}

// buildURL 构建 Deepgram WebSocket URL，包含音频格式参数。
func (d *DeepgramService) buildURL(opts Options) string {
	u, _ := url.Parse(d.endpoint)

	encoding := opts.Encoding
	if encoding == "" {
		encoding = "pcm_s16le"
	}
	sampleRate := opts.SampleRate
	if sampleRate == 0 {
		sampleRate = 16000
	}
	language := opts.Language
	if language == "" {
		language = "zh-CN"
	}

	q := u.Query()
	q.Set("encoding", encoding)
	q.Set("sample_rate", fmt.Sprintf("%d", sampleRate))
	q.Set("language", language)
	q.Set("model", d.model)
	q.Set("punctuate", "true")
	u.RawQuery = q.Encode()

	return u.String()
}