refactor: route audio models through wrappers (#20387)

internal/entity/models/types.go

@@ -417,6 +417,50 @@ func (r *RerankModel) Rerank(ctx context.Context, request RerankRequest, apiConf
 	return r.ModelDriver.Rerank(ctx, r.ModelName, request, apiConfig, rerankConfig, modelUsage)
 }
 
+// ASRModel wraps a ModelDriver with speech-to-text configuration.
+type ASRModel struct {
+	ModelDriver ModelDriver
+	ModelName   *string
+	APIConfig   *APIConfig
+}
+
+// NewASRModel creates a new ASRModel.
+func NewASRModel(driver ModelDriver, modelName *string, apiConfig *APIConfig) *ASRModel {
+	return &ASRModel{ModelDriver: driver, ModelName: modelName, APIConfig: apiConfig}
+}
+
+// Transcribe converts audio to text.
+func (m *ASRModel) Transcribe(ctx context.Context, audioFile *string, config *ASRConfig, usage *common.ModelUsage) (*ASRResponse, error) {
+	return m.ModelDriver.TranscribeAudio(ctx, m.ModelName, audioFile, m.APIConfig, config, usage)
+}
+
+// TranscribeWithSender streams transcription results through sender.
+func (m *ASRModel) TranscribeWithSender(ctx context.Context, audioFile *string, config *ASRConfig, usage *common.ModelUsage, sender func(*string, *string) error) error {
+	return m.ModelDriver.TranscribeAudioWithSender(ctx, m.ModelName, audioFile, m.APIConfig, config, usage, sender)
+}
+
+// TTSModel wraps a ModelDriver with text-to-speech configuration.
+type TTSModel struct {
+	ModelDriver ModelDriver
+	ModelName   *string
+	APIConfig   *APIConfig
+}
+
+// NewTTSModel creates a new TTSModel.
+func NewTTSModel(driver ModelDriver, modelName *string, apiConfig *APIConfig) *TTSModel {
+	return &TTSModel{ModelDriver: driver, ModelName: modelName, APIConfig: apiConfig}
+}
+
+// Speech converts text to audio.
+func (m *TTSModel) Speech(ctx context.Context, audioContent *string, config *TTSConfig, usage *common.ModelUsage) (*TTSResponse, error) {
+	return m.ModelDriver.AudioSpeech(ctx, m.ModelName, audioContent, m.APIConfig, config, usage)
+}
+
+// SpeechWithSender streams synthesized audio through sender.
+func (m *TTSModel) SpeechWithSender(ctx context.Context, audioContent *string, config *TTSConfig, usage *common.ModelUsage, sender func(*string, *string) error) error {
+	return m.ModelDriver.AudioSpeechWithSender(ctx, m.ModelName, audioContent, m.APIConfig, config, usage, sender)
+}
+
 // ToolConfig bundles tool-calling configuration for a ChatModel.
 type ToolConfig struct {
 	Tools           string          // JSON-encoded tools list

internal/handler/chat_audio.go

@@ -82,6 +82,7 @@ func (h *ChatHandler) ChatAudioSpeech(c *gin.Context) {
 	driver := target.Driver
 	modelName := target.ModelName
 	apiConfig := target.APIConfig
+	ttsModel := modelModule.NewTTSModel(driver, &modelName, apiConfig)
 
 	writeAudioHeaders := func(mediaType string) {
 		c.Header("Content-Type", mediaType)
@@ -108,7 +109,7 @@ func (h *ChatHandler) ChatAudioSpeech(c *gin.Context) {
 		if seg == "" {
 			continue
 		}
-		resp, err := driver.AudioSpeech(ctx, &modelName, &seg, apiConfig, &modelModule.TTSConfig{Format: "mp3"}, nil)
+		resp, err := ttsModel.Speech(ctx, &seg, &modelModule.TTSConfig{Format: "mp3"}, nil)
 		if err != nil {
 			if firstSynthErr == nil {
 				firstSynthErr = err
@@ -276,6 +277,7 @@ func (h *ChatHandler) ChatAudioTranscription(c *gin.Context) {
 	driver := target.Driver
 	modelName := target.ModelName
 	apiConfig := target.APIConfig
+	asrModel := modelModule.NewASRModel(driver, &modelName, apiConfig)
 
 	streamMode := strings.ToLower(c.PostForm("stream")) == "true"
 	if streamMode {
@@ -306,7 +308,7 @@ func (h *ChatHandler) ChatAudioTranscription(c *gin.Context) {
 			return nil
 		}
 
-		if err = driver.TranscribeAudioWithSender(ctx, &modelName, &tmpPath, apiConfig, &modelModule.ASRConfig{}, nil, sender); err != nil {
+		if err = asrModel.TranscribeWithSender(ctx, &tmpPath, &modelModule.ASRConfig{}, nil, sender); err != nil {
 			errEvent := map[string]interface{}{"event": "error", "text": err.Error()}
 			data, _ := json.Marshal(errEvent)
 			_, _ = c.Writer.WriteString(fmt.Sprintf("data: %s\n\n", data))
@@ -324,7 +326,7 @@ func (h *ChatHandler) ChatAudioTranscription(c *gin.Context) {
 		return
 	}
 
-	resp, err := driver.TranscribeAudio(ctx, &modelName, &tmpPath, apiConfig, &modelModule.ASRConfig{}, nil)
+	resp, err := asrModel.Transcribe(ctx, &tmpPath, &modelModule.ASRConfig{}, nil)
 	if err != nil {
 		common.ErrorWithCode(c, common.CodeServerError, err.Error())
 		return

internal/ingestion/component/media_dispatch.go

@@ -260,7 +260,8 @@ func maybeDispatchAudio(
 	}
 	defer os.Remove(tmpFile)
 
-	resp, err := driver.TranscribeAudio(ctx, &modelName, &tmpFile, apiConfig, nil, nil)
+	asrModel := modelModule.NewASRModel(driver, &modelName, apiConfig)
+	resp, err := asrModel.Transcribe(ctx, &tmpFile, nil, nil)
 	if err != nil {
 		return parser.ParseResult{}, true,
 			fmt.Errorf("Parser: audio transcription: %w", err)

internal/service/chat_pipeline.go

@@ -1662,7 +1662,7 @@ func (s *ChatPipelineService) AsyncChatSolo(
 		chatModel := modelModule.NewChatModel(target.Driver, &target.ModelName, target.APIConfig)
 
 		// 5. Resolve TTS model. Best-effort: warn and proceed without TTS on lookup failure.
-		var ttsModel *modelModule.ChatModel
+		var ttsModel *modelModule.TTSModel
 		if promptConfig != nil {
 			if useTTS, _ := promptConfig["tts"].(bool); useTTS {
 				target, ttsErr := s.ModelProviderSvc.modelSolver().ResolveDefaultModelConfig(ctx, chat.TenantID, entity.ModelTypeTTS)
@@ -1671,7 +1671,7 @@ func (s *ChatPipelineService) AsyncChatSolo(
 						zap.String("tenant_id", chat.TenantID),
 						zap.Error(ttsErr))
 				} else {
-					ttsModel = modelModule.NewChatModel(target.Driver, &target.ModelName, target.APIConfig)
+					ttsModel = modelModule.NewTTSModel(target.Driver, &target.ModelName, target.APIConfig)
 				}
 			}
 		}
@@ -2298,7 +2298,7 @@ func (s *ChatPipelineService) getModels(ctx context.Context, chat *entity.Chat)
 	*modelModule.EmbeddingModel,
 	*modelModule.RerankModel,
 	*modelModule.ChatModel,
-	*modelModule.ChatModel, // TTS model
+	*modelModule.TTSModel, // TTS model
 	error,
 ) {
 	kbDAO := dao.NewKnowledgebaseDAO()
@@ -2368,12 +2368,12 @@ func (s *ChatPipelineService) getModels(ctx context.Context, chat *entity.Chat)
 	}
 
 	// TTS model.
-	var ttsModel *modelModule.ChatModel
+	var ttsModel *modelModule.TTSModel
 	if chat.PromptConfig != nil {
 		if useTTS, _ := chat.PromptConfig["tts"].(bool); useTTS {
 			target, err := s.ModelProviderSvc.modelSolver().ResolveDefaultModelConfig(ctx, chat.TenantID, entity.ModelTypeTTS)
 			if err == nil {
-				ttsModel = modelModule.NewChatModel(target.Driver, &target.ModelName, target.APIConfig)
+				ttsModel = modelModule.NewTTSModel(target.Driver, &target.ModelName, target.APIConfig)
 			}
 		}
 	}
@@ -2596,7 +2596,7 @@ func cleanTTSText(text string) string {
 
 // synthesizeTTS calls the TTS model to convert text to audio.
 // Mirrors dialog_service.py:1426-1432.
-func (s *ChatPipelineService) synthesizeTTS(ctx context.Context, ttsModel *modelModule.ChatModel, text string) interface{} {
+func (s *ChatPipelineService) synthesizeTTS(ctx context.Context, ttsModel *modelModule.TTSModel, text string) interface{} {
 	if ttsModel == nil || text == "" {
 		return nil
 	}
@@ -2604,9 +2604,7 @@ func (s *ChatPipelineService) synthesizeTTS(ctx context.Context, ttsModel *model
 	if text == "" {
 		return nil
 	}
-	ttsResp, err := ttsModel.ModelDriver.AudioSpeech(
-		ctx, ttsModel.ModelName, &text, ttsModel.APIConfig, &modelModule.TTSConfig{Format: "mp3"}, nil,
-	)
+	ttsResp, err := ttsModel.Speech(ctx, &text, &modelModule.TTSConfig{Format: "mp3"}, nil)
 	if err != nil {
 		common.Warn("TTS synthesis failed", zap.Error(err))
 		return nil
@@ -3117,7 +3115,7 @@ func (s *ChatPipelineService) decorateAnswer(
 	embModel *modelModule.EmbeddingModel,
 	vectorSimilarityWeight float64,
 	quote bool,
-	ttsModel *modelModule.ChatModel,
+	ttsModel *modelModule.TTSModel,
 	langfuseTraceID string,
 	llmModelConfig map[string]interface{},
 	tenantID string,

internal/service/model_call.go

@@ -173,7 +173,8 @@ func (s *ModelCallService) TranscribeAudio(ctx context.Context, modelRef, userID
 	}
 
 	modelName := target.ModelName
-	response, err := target.Driver.TranscribeAudio(ctx, &modelName, audioFile, target.APIConfig, config, nil)
+	asrModel := modelModule.NewASRModel(target.Driver, &modelName, target.APIConfig)
+	response, err := asrModel.Transcribe(ctx, audioFile, config, nil)
 	if err != nil {
 		return nil, common.CodeServerError, err
 	}
@@ -195,7 +196,8 @@ func (s *ModelCallService) TranscribeAudioStream(ctx context.Context, modelRef,
 	}
 
 	modelName := target.ModelName
-	if err := target.Driver.TranscribeAudioWithSender(ctx, &modelName, audioFile, target.APIConfig, config, nil, sender); err != nil {
+	asrModel := modelModule.NewASRModel(target.Driver, &modelName, target.APIConfig)
+	if err := asrModel.TranscribeWithSender(ctx, audioFile, config, nil, sender); err != nil {
 		return common.CodeServerError, err
 	}
 	return common.CodeSuccess, nil
@@ -231,7 +233,8 @@ func (s *ModelCallService) AudioSpeech(ctx context.Context, modelRef, userID str
 	}
 
 	modelName := target.ModelName
-	response, err := target.Driver.AudioSpeech(ctx, &modelName, audioContent, target.APIConfig, config, nil)
+	ttsModel := modelModule.NewTTSModel(target.Driver, &modelName, target.APIConfig)
+	response, err := ttsModel.Speech(ctx, audioContent, config, nil)
 	if err != nil {
 		return nil, common.CodeServerError, err
 	}
@@ -263,7 +266,8 @@ func (s *ModelCallService) AudioSpeechForTenant(ctx context.Context, modelRef, t
 	}
 
 	modelName := target.ModelName
-	response, err := target.Driver.AudioSpeech(ctx, &modelName, audioContent, target.APIConfig, config, nil)
+	ttsModel := modelModule.NewTTSModel(target.Driver, &modelName, target.APIConfig)
+	response, err := ttsModel.Speech(ctx, audioContent, config, nil)
 	if err != nil {
 		return nil, common.CodeServerError, err
 	}
@@ -285,7 +289,8 @@ func (s *ModelCallService) AudioSpeechStream(ctx context.Context, modelRef, user
 	}
 
 	modelName := target.ModelName
-	if err := target.Driver.AudioSpeechWithSender(ctx, &modelName, audioContent, target.APIConfig, config, nil, sender); err != nil {
+	ttsModel := modelModule.NewTTSModel(target.Driver, &modelName, target.APIConfig)
+	if err := ttsModel.SpeechWithSender(ctx, audioContent, config, nil, sender); err != nil {
 		return common.CodeServerError, err
 	}
 	return common.CodeSuccess, nil

internal/service/model_service.go

@@ -977,7 +977,8 @@ func verifyProviderModel(ctx context.Context, driver modelModule.ModelDriver, pr
 				_, err = driver.Rerank(ctx, &modelName, rerankRequest, apiConfig, &modelModule.RerankConfig{}, nil)
 			case "tts":
 				content := "hello"
-				_, err = driver.AudioSpeech(ctx, &modelName, &content, apiConfig, nil, nil)
+				ttsModel := modelModule.NewTTSModel(driver, &modelName, apiConfig)
+				_, err = ttsModel.Speech(ctx, &content, nil, nil)
 			case "asr":
 				err = verifyASRModel(ctx, driver, modelName, apiConfig)
 			case "ocr":
@@ -1099,7 +1100,8 @@ func verifyASRModel(ctx context.Context, driver modelModule.ModelDriver, modelNa
 	}
 	tmpFile.Close()
 
-	resp, err := driver.TranscribeAudio(ctx, &modelName, &tmpPath, apiConfig, nil, nil)
+	asrModel := modelModule.NewASRModel(driver, &modelName, apiConfig)
+	resp, err := asrModel.Transcribe(ctx, &tmpPath, nil, nil)
 	if err != nil {
 		return err
 	}