diff --git a/conf/models/aliyun.json b/conf/models/aliyun.json index d2a39c272f..37c5db20bb 100644 --- a/conf/models/aliyun.json +++ b/conf/models/aliyun.json @@ -11,7 +11,7 @@ "embedding": "compatible-mode/v1/embeddings", "rerank": "compatible-api/v1/reranks", "models": "compatible-mode/v1/models", - "tts": "compatible-mode/v1/audio/speech" + "tts": "api/v1/services/aigc/multimodal-generation/generation" }, "models": [ { diff --git a/internal/entity/models/aliyun.go b/internal/entity/models/aliyun.go index 7e5350c69a..ffffb019f1 100644 --- a/internal/entity/models/aliyun.go +++ b/internal/entity/models/aliyun.go @@ -395,9 +395,39 @@ func (a *AliyunModel) TranscribeAudioWithSender(ctx context.Context, modelName * // DashScope's Qwen TTS models require one. const aliyunTTSDefaultVoice = "Cherry" -// newAliyunTTSRequest builds the OpenAI-compatible audio/speech request -// against DashScope's compatible mode. -func (a *AliyunModel) newAliyunTTSRequest(ctx context.Context, modelName *string, audioContent *string, apiConfig *APIConfig, ttsConfig *TTSConfig) (*http.Request, error) { +// aliyunTTSRequest is the DashScope multimodal-generation request for Qwen +// TTS models (qwen-tts / qwen3-tts-flash family). +type aliyunTTSRequest struct { + Model string `json:"model"` + Input aliyunTTSInput `json:"input"` +} + +type aliyunTTSInput struct { + Text string `json:"text"` + // Voice is required by Qwen TTS models (e.g. "Cherry"). + Voice string `json:"voice"` + // LanguageType hints the text language (e.g. "Chinese", "English"); + // omitted to let the model auto-detect. + LanguageType string `json:"language_type,omitempty"` +} + +// aliyunTTSResponse is the non-streaming DashScope multimodal-generation +// response. The synthesized audio is not inlined; output.audio.url points +// to a downloadable file (valid for 24h). +type aliyunTTSResponse struct { + Output struct { + Audio struct { + URL string `json:"url"` + } `json:"audio"` + FinishReason string `json:"finish_reason"` + } `json:"output"` + Code string `json:"code"` + Message string `json:"message"` + RequestID string `json:"request_id"` +} + +// AudioSpeech convert text to audio +func (a *AliyunModel) AudioSpeech(ctx context.Context, modelName *string, audioContent *string, apiConfig *APIConfig, ttsConfig *TTSConfig, modelUsage *common.ModelUsage) (*TTSResponse, error) { if err := a.baseModel.APIConfigCheck(apiConfig); err != nil { return nil, err } @@ -417,43 +447,30 @@ func (a *AliyunModel) newAliyunTTSRequest(ctx context.Context, modelName *string } url := fmt.Sprintf("%s/%s", strings.TrimSuffix(resolvedBaseURL, "/"), strings.TrimPrefix(a.baseModel.URLSuffix.TTS, "/")) - reqBody := map[string]interface{}{ - "model": *modelName, - "input": *audioContent, - "voice": aliyunTTSDefaultVoice, - } + input := aliyunTTSInput{Text: *audioContent, Voice: aliyunTTSDefaultVoice} if ttsConfig != nil { - for key, value := range ttsConfig.Params { - reqBody[key] = value + if voice, ok := ttsConfig.Params["voice"].(string); ok && strings.TrimSpace(voice) != "" { + input.Voice = voice } - if ttsConfig.Format != "" { - reqBody["response_format"] = ttsConfig.Format + if lang, ok := ttsConfig.Params["language_type"].(string); ok && strings.TrimSpace(lang) != "" { + input.LanguageType = lang } } - jsonData, err := json.Marshal(reqBody) + jsonData, err := json.Marshal(aliyunTTSRequest{Model: *modelName, Input: input}) if err != nil { return nil, fmt.Errorf("failed to marshal request: %w", err) } + ctx, cancel := context.WithTimeout(ctx, longOpCallTimeout) + defer cancel() + req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData)) if err != nil { return nil, fmt.Errorf("failed to create request: %w", err) } req.Header.Set("Content-Type", "application/json") req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey)) - return req, nil -} - -// AudioSpeech convert text to audio -func (a *AliyunModel) AudioSpeech(ctx context.Context, modelName *string, audioContent *string, apiConfig *APIConfig, ttsConfig *TTSConfig, modelUsage *common.ModelUsage) (*TTSResponse, error) { - ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout) - defer cancel() - - req, err := a.newAliyunTTSRequest(ctx, modelName, audioContent, apiConfig, ttsConfig) - if err != nil { - return nil, err - } resp, err := a.baseModel.httpClient.Do(req) if err != nil { @@ -465,15 +482,60 @@ func (a *AliyunModel) AudioSpeech(ctx context.Context, modelName *string, audioC if err != nil { return nil, fmt.Errorf("failed to read response body: %w", err) } - if resp.StatusCode != http.StatusOK { return nil, fmt.Errorf("aliyun TTS API error: %s, body: %s", resp.Status, string(body)) } - if len(body) == 0 { - return nil, fmt.Errorf("aliyun TTS API returned empty audio") + + var parsed aliyunTTSResponse + if err := json.Unmarshal(body, &parsed); err != nil { + return nil, fmt.Errorf("failed to parse aliyun TTS response: %w, body: %s", err, string(body)) + } + if parsed.Code != "" { + return nil, fmt.Errorf("aliyun TTS API error: %s: %s", parsed.Code, parsed.Message) + } + if parsed.Output.Audio.URL == "" { + return nil, fmt.Errorf("aliyun TTS response has no audio url, body: %s", string(body)) } - return &TTSResponse{Audio: body}, nil + audio, err := a.downloadAliyunTTSAudio(ctx, parsed.Output.Audio.URL) + if err != nil { + return nil, err + } + // Qwen TTS audio files are WAV. + return &TTSResponse{Audio: audio, MediaType: "audio/wav"}, nil +} + +// aliyunTTSAudioMaxBytes caps a synthesized audio download. DashScope TTS +// audio is far smaller; this only guards against runaway responses. +const aliyunTTSAudioMaxBytes int64 = 64 << 20 // 64 MiB + +// downloadAliyunTTSAudio fetches the synthesized audio file referenced by a +// DashScope TTS response. +func (a *AliyunModel) downloadAliyunTTSAudio(ctx context.Context, audioURL string) ([]byte, error) { + req, err := http.NewRequestWithContext(ctx, "GET", audioURL, nil) + if err != nil { + return nil, fmt.Errorf("failed to create audio download request: %w", err) + } + resp, err := a.baseModel.httpClient.Do(req) + if err != nil { + return nil, fmt.Errorf("failed to download aliyun TTS audio: %w", err) + } + defer resp.Body.Close() + + audio, err := io.ReadAll(io.LimitReader(resp.Body, aliyunTTSAudioMaxBytes+1)) + if err != nil { + return nil, fmt.Errorf("failed to read aliyun TTS audio: %w", err) + } + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("failed to download aliyun TTS audio: %s", resp.Status) + } + if int64(len(audio)) > aliyunTTSAudioMaxBytes { + return nil, fmt.Errorf("aliyun TTS audio download exceeds %d bytes", aliyunTTSAudioMaxBytes) + } + if len(audio) == 0 { + return nil, fmt.Errorf("aliyun TTS audio download is empty") + } + return audio, nil } func (a *AliyunModel) AudioSpeechWithSender(ctx context.Context, modelName *string, audioContent *string, apiConfig *APIConfig, ttsConfig *TTSConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { @@ -481,8 +543,8 @@ func (a *AliyunModel) AudioSpeechWithSender(ctx context.Context, modelName *stri return fmt.Errorf("sender is required") } - // DashScope's compatible-mode audio/speech endpoint returns the whole - // audio in one response; forward it as a single chunk. + // The non-streaming DashScope TTS endpoint returns the whole audio via + // a downloadable URL; forward it as a single chunk. resp, err := a.AudioSpeech(ctx, modelName, audioContent, apiConfig, ttsConfig, modelUsage) if err != nil { return err diff --git a/internal/entity/models/aliyun_test.go b/internal/entity/models/aliyun_test.go index 4cd7074435..3959425beb 100644 --- a/internal/entity/models/aliyun_test.go +++ b/internal/entity/models/aliyun_test.go @@ -19,8 +19,10 @@ package models import ( "encoding/json" "errors" + "fmt" "net/http" "net/http/httptest" + "strings" "testing" ) @@ -307,27 +309,54 @@ func TestAliyunChatStreamlyWithSenderRejectsStreamFalse(t *testing.T) { } } -func TestAliyunAudioSpeechSynthesizesViaCompatibleEndpoint(t *testing.T) { - withSSRFBypass(t) - requestBody := make(chan map[string]interface{}, 1) - requestPath := make(chan string, 1) - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { +// newAliyunTTSTestServer stubs the DashScope multimodal-generation endpoint: +// POST returns a JSON body whose output.audio.url points back at the same +// server, GET returns the synthesized WAV bytes. +func newAliyunTTSTestServer(t *testing.T, requestBody chan<- map[string]interface{}, requestPath chan<- string) *httptest.Server { + t.Helper() + var server *httptest.Server + server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodGet { + w.Header().Set("Content-Type", "audio/wav") + _, _ = w.Write([]byte("fake-wav-bytes")) + return + } + if r.Method != http.MethodPost { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } var body map[string]interface{} if err := json.NewDecoder(r.Body).Decode(&body); err != nil { http.Error(w, err.Error(), http.StatusBadRequest) return } - requestPath <- r.URL.Path - requestBody <- body - w.Header().Set("Content-Type", "audio/mpeg") - _, _ = w.Write([]byte("fake-mp3-bytes")) + if requestPath != nil { + requestPath <- r.URL.Path + } + if requestBody != nil { + requestBody <- body + } + w.Header().Set("Content-Type", "application/json") + if _, err := fmt.Fprintf(w, `{"output":{"audio":{"url":%q},"finish_reason":"stop"},"request_id":"req-1"}`, server.URL+"/audio.wav"); err != nil { + t.Errorf("failed to write TTS response: %v", err) + } })) - defer server.Close() + t.Cleanup(server.Close) + return server +} + +const aliyunTTSTestSuffix = "api/v1/services/aigc/multimodal-generation/generation" + +func TestAliyunAudioSpeechSynthesizesViaNativeEndpoint(t *testing.T) { + withSSRFBypass(t) + requestBody := make(chan map[string]interface{}, 1) + requestPath := make(chan string, 1) + server := newAliyunTTSTestServer(t, requestBody, requestPath) ctx := t.Context() model := NewAliyunModel( map[string]string{"default": server.URL}, - URLSuffix{TTS: "compatible-mode/v1/audio/speech"}, + URLSuffix{TTS: aliyunTTSTestSuffix}, ) apiKey := "test-key" modelName := "qwen-tts-flash" @@ -344,46 +373,44 @@ func TestAliyunAudioSpeechSynthesizesViaCompatibleEndpoint(t *testing.T) { if err != nil { t.Fatalf("AudioSpeech: %v", err) } - if string(response.Audio) != "fake-mp3-bytes" { - t.Errorf("audio = %q, want fake-mp3-bytes", string(response.Audio)) + if string(response.Audio) != "fake-wav-bytes" { + t.Errorf("audio = %q, want fake-wav-bytes", string(response.Audio)) + } + if response.MediaType != "audio/wav" { + t.Errorf("media type = %q, want audio/wav", response.MediaType) } - if got := <-requestPath; got != "/compatible-mode/v1/audio/speech" { - t.Errorf("request path = %q, want /compatible-mode/v1/audio/speech", got) + if got := <-requestPath; got != "/"+aliyunTTSTestSuffix { + t.Errorf("request path = %q, want /%s", got, aliyunTTSTestSuffix) } body := <-requestBody if body["model"] != "qwen-tts-flash" { t.Errorf("model = %v, want qwen-tts-flash", body["model"]) } - if body["input"] != "你好，世界" { - t.Errorf("input = %v, want 你好，世界", body["input"]) + input, ok := body["input"].(map[string]interface{}) + if !ok { + t.Fatalf("input = %T, want JSON object", body["input"]) } - if body["voice"] != aliyunTTSDefaultVoice { - t.Errorf("voice = %v, want default %s", body["voice"], aliyunTTSDefaultVoice) + if input["text"] != "你好，世界" { + t.Errorf("input.text = %v, want 你好，世界", input["text"]) } - if body["response_format"] != "mp3" { - t.Errorf("response_format = %v, want mp3", body["response_format"]) + if input["voice"] != aliyunTTSDefaultVoice { + t.Errorf("input.voice = %v, want default %s", input["voice"], aliyunTTSDefaultVoice) + } + if _, ok := input["language_type"]; ok { + t.Errorf("language_type = %v, want omitted by default", input["language_type"]) } } -func TestAliyunAudioSpeechHonorsExplicitVoice(t *testing.T) { +func TestAliyunAudioSpeechHonorsExplicitVoiceAndLanguage(t *testing.T) { withSSRFBypass(t) requestBody := make(chan map[string]interface{}, 1) - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - var body map[string]interface{} - if err := json.NewDecoder(r.Body).Decode(&body); err != nil { - http.Error(w, err.Error(), http.StatusBadRequest) - return - } - requestBody <- body - _, _ = w.Write([]byte("audio")) - })) - defer server.Close() + server := newAliyunTTSTestServer(t, requestBody, nil) ctx := t.Context() model := NewAliyunModel( map[string]string{"default": server.URL}, - URLSuffix{TTS: "compatible-mode/v1/audio/speech"}, + URLSuffix{TTS: aliyunTTSTestSuffix}, ) apiKey := "test-key" modelName := "qwen-tts-flash" @@ -394,27 +421,70 @@ func TestAliyunAudioSpeechHonorsExplicitVoice(t *testing.T) { &modelName, &text, &APIConfig{ApiKey: &apiKey}, - &TTSConfig{Params: map[string]any{"voice": "Serena"}}, + &TTSConfig{Params: map[string]any{"voice": "Serena", "language_type": "English"}}, nil, ); err != nil { t.Fatalf("AudioSpeech: %v", err) } - if got := (<-requestBody)["voice"]; got != "Serena" { + input := (<-requestBody)["input"].(map[string]interface{}) + if got := input["voice"]; got != "Serena" { t.Errorf("voice = %v, want Serena", got) } + if got := input["language_type"]; got != "English" { + t.Errorf("language_type = %v, want English", got) + } +} + +func TestAliyunAudioSpeechRejectsOversizedAudio(t *testing.T) { + withSSRFBypass(t) + oversized := make([]byte, aliyunTTSAudioMaxBytes+1) + var server *httptest.Server + server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodGet { + w.Header().Set("Content-Type", "audio/wav") + _, _ = w.Write(oversized) + return + } + w.Header().Set("Content-Type", "application/json") + if _, err := fmt.Fprintf(w, `{"output":{"audio":{"url":%q},"finish_reason":"stop"},"request_id":"req-1"}`, server.URL+"/audio.wav"); err != nil { + t.Errorf("failed to write TTS response: %v", err) + } + })) + t.Cleanup(server.Close) + ctx := t.Context() + + model := NewAliyunModel( + map[string]string{"default": server.URL}, + URLSuffix{TTS: aliyunTTSTestSuffix}, + ) + apiKey := "test-key" + modelName := "qwen-tts-flash" + text := "hello" + + _, err := model.AudioSpeech( + ctx, + &modelName, + &text, + &APIConfig{ApiKey: &apiKey}, + &TTSConfig{Format: "mp3"}, + nil, + ) + if err == nil || !strings.Contains(err.Error(), "exceeds") { + t.Fatalf("error = %v, want oversized audio error", err) + } } func TestAliyunAudioSpeechSurfacesAPIError(t *testing.T) { withSSRFBypass(t) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - http.Error(w, `{"error":{"message":"Invalid api-key"}}`, http.StatusUnauthorized) + http.Error(w, `{"code":"InvalidApiKey","message":"Invalid API-key provided."}`, http.StatusUnauthorized) })) defer server.Close() ctx := t.Context() model := NewAliyunModel( map[string]string{"default": server.URL}, - URLSuffix{TTS: "compatible-mode/v1/audio/speech"}, + URLSuffix{TTS: aliyunTTSTestSuffix}, ) apiKey := "bad-key" modelName := "qwen-tts-flash" @@ -426,6 +496,29 @@ func TestAliyunAudioSpeechSurfacesAPIError(t *testing.T) { } } +func TestAliyunAudioSpeechRejectsMissingAudioURL(t *testing.T) { + withSSRFBypass(t) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"output":{"finish_reason":"stop"},"request_id":"req-1"}`)) + })) + defer server.Close() + ctx := t.Context() + + model := NewAliyunModel( + map[string]string{"default": server.URL}, + URLSuffix{TTS: aliyunTTSTestSuffix}, + ) + apiKey := "test-key" + modelName := "qwen-tts-flash" + text := "hello" + + _, err := model.AudioSpeech(ctx, &modelName, &text, &APIConfig{ApiKey: &apiKey}, nil, nil) + if err == nil { + t.Fatal("error = nil, want missing audio url error") + } +} + func TestAliyunAudioSpeechRequiresTTSSuffix(t *testing.T) { withSSRFBypass(t) ctx := t.Context() @@ -445,15 +538,12 @@ func TestAliyunAudioSpeechRequiresTTSSuffix(t *testing.T) { func TestAliyunAudioSpeechWithSenderSendsSingleChunk(t *testing.T) { withSSRFBypass(t) - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - _, _ = w.Write([]byte("whole-audio")) - })) - defer server.Close() + server := newAliyunTTSTestServer(t, nil, nil) ctx := t.Context() model := NewAliyunModel( map[string]string{"default": server.URL}, - URLSuffix{TTS: "compatible-mode/v1/audio/speech"}, + URLSuffix{TTS: aliyunTTSTestSuffix}, ) apiKey := "test-key" modelName := "qwen-tts-flash" @@ -477,7 +567,7 @@ func TestAliyunAudioSpeechWithSenderSendsSingleChunk(t *testing.T) { if err != nil { t.Fatalf("AudioSpeechWithSender: %v", err) } - if len(chunks) != 1 || chunks[0] != "whole-audio" { - t.Fatalf("chunks = %v, want [whole-audio]", chunks) + if len(chunks) != 1 || chunks[0] != "fake-wav-bytes" { + t.Fatalf("chunks = %v, want [fake-wav-bytes]", chunks) } } diff --git a/internal/entity/models/types.go b/internal/entity/models/types.go index 22ba50636f..9cb3977681 100644 --- a/internal/entity/models/types.go +++ b/internal/entity/models/types.go @@ -98,6 +98,9 @@ type ASRResponse struct { type TTSResponse struct { Audio []byte `json:"audio"` + // MediaType is the MIME type of Audio (e.g. "audio/mpeg", "audio/wav"). + // Empty means the caller's default (audio/mpeg). + MediaType string `json:"media_type,omitempty"` } type OCRFileResponse struct { diff --git a/internal/handler/chat_audio.go b/internal/handler/chat_audio.go index c2090734fd..e23c2433e5 100644 --- a/internal/handler/chat_audio.go +++ b/internal/handler/chat_audio.go @@ -16,6 +16,7 @@ package handler import ( + "bytes" "encoding/json" "fmt" "net/http" @@ -78,8 +79,25 @@ func (h *ChatHandler) ChatAudioSpeech(c *gin.Context) { return } + writeAudioHeaders := func(mediaType string) { + c.Header("Content-Type", mediaType) + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + c.Header("X-Accel-Buffering", "no") + c.Writer.WriteHeader(http.StatusOK) + } + segments := ttsSegmentSplitRegex.Split(req.Text, -1) headerWritten := false + // mediaType locks the media type of the first successfully synthesized + // segment; later segments of a different type are skipped so the + // response never mixes encodings under a single Content-Type. + mediaType := "" + // WAV segments cannot be concatenated byte-wise (each carries its own + // RIFF header), so they are accumulated as raw PCM and re-wrapped into + // a single WAV before being written. + var wavFmt *wavFormat + var wavPCM []byte for i, seg := range segments { seg = strings.TrimSpace(seg) if seg == "" { @@ -96,14 +114,41 @@ func (h *ChatHandler) ChatAudioSpeech(c *gin.Context) { if resp == nil || len(resp.Audio) == 0 { continue } + segMediaType := resp.MediaType + if segMediaType == "" { + segMediaType = "audio/mpeg" + } + if mediaType == "" { + mediaType = segMediaType + } else if segMediaType != mediaType { + common.Warn("chat TTS segment media type differs from the first segment, skipping", + zap.Int("segmentIndex", i), + zap.String("segmentMediaType", segMediaType), + zap.String("mediaType", mediaType)) + continue + } + if segMediaType == "audio/wav" { + format, pcm, werr := splitWAV(resp.Audio) + if werr != nil { + common.Warn("chat TTS wav segment unparsable", + zap.Int("segmentIndex", i), + zap.Error(werr)) + continue + } + if wavFmt == nil { + wavFmt = format + } else if !bytes.Equal(wavFmt.fmtChunk, format.fmtChunk) { + common.Warn("chat TTS wav segment format differs from the first segment, skipping", + zap.Int("segmentIndex", i)) + continue + } + wavPCM = append(wavPCM, pcm...) + continue + } if !headerWritten { // Commit the audio headers only once the first chunk is available, // so a fully failed synthesis can still return a JSON error status. - c.Header("Content-Type", "audio/mpeg") - c.Header("Cache-Control", "no-cache") - c.Header("Connection", "keep-alive") - c.Header("X-Accel-Buffering", "no") - c.Writer.WriteHeader(http.StatusOK) + writeAudioHeaders(mediaType) headerWritten = true } if _, werr := c.Writer.Write(resp.Audio); werr != nil { @@ -111,6 +156,14 @@ func (h *ChatHandler) ChatAudioSpeech(c *gin.Context) { } c.Writer.Flush() } + if wavFmt != nil && !headerWritten { + writeAudioHeaders("audio/wav") + headerWritten = true + if _, werr := c.Writer.Write(buildWAV(wavFmt, wavPCM)); werr != nil { + return + } + c.Writer.Flush() + } if !headerWritten { common.ErrorWithCode(c, common.CodeServerError, "TTS synthesis produced no audio") } diff --git a/internal/handler/chat_audio_wav.go b/internal/handler/chat_audio_wav.go new file mode 100644 index 0000000000..92b2e14828 --- /dev/null +++ b/internal/handler/chat_audio_wav.go @@ -0,0 +1,86 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package handler + +import ( + "encoding/binary" + "fmt" +) + +// wavFormat carries the fmt chunk payload of a WAV file so concatenated PCM +// can be re-wrapped with the original encoding parameters. +type wavFormat struct { + fmtChunk []byte +} + +// splitWAV parses a RIFF/WAVE file and returns its fmt chunk payload and raw +// PCM data chunk. Chunks are 2-byte aligned per the RIFF spec. +func splitWAV(b []byte) (*wavFormat, []byte, error) { + if len(b) < 12 || string(b[0:4]) != "RIFF" || string(b[8:12]) != "WAVE" { + return nil, nil, fmt.Errorf("not a RIFF/WAVE file (%d bytes)", len(b)) + } + + var fmtChunk, data []byte + for off := 12; off+8 <= len(b); { + chunkID := string(b[off : off+4]) + size := int(binary.LittleEndian.Uint32(b[off+4 : off+8])) + off += 8 + if size < 0 || off+size > len(b) { + return nil, nil, fmt.Errorf("malformed WAV chunk %q: size %d exceeds buffer", chunkID, size) + } + switch chunkID { + case "fmt ": + fmtChunk = b[off : off+size] + case "data": + data = b[off : off+size] + } + off += size + (size & 1) // chunks are padded to even sizes + } + + if fmtChunk == nil { + return nil, nil, fmt.Errorf("WAV file has no fmt chunk") + } + if data == nil { + return nil, nil, fmt.Errorf("WAV file has no data chunk") + } + return &wavFormat{fmtChunk: fmtChunk}, data, nil +} + +// buildWAV wraps raw PCM data into a RIFF/WAVE file reusing the given fmt +// chunk (encoding, sample rate, channel count). +func buildWAV(f *wavFormat, pcm []byte) []byte { + fmtSize := len(f.fmtChunk) + dataPad := len(pcm) & 1 + total := 4 + (8 + fmtSize + fmtSize&1) + (8 + len(pcm) + dataPad) + + out := make([]byte, 0, 8+total) + out = append(out, "RIFF"...) + out = binary.LittleEndian.AppendUint32(out, uint32(total)) + out = append(out, "WAVE"...) + out = append(out, "fmt "...) + out = binary.LittleEndian.AppendUint32(out, uint32(fmtSize)) + out = append(out, f.fmtChunk...) + if fmtSize&1 == 1 { + out = append(out, 0) + } + out = append(out, "data"...) + out = binary.LittleEndian.AppendUint32(out, uint32(len(pcm))) + out = append(out, pcm...) + if dataPad == 1 { + out = append(out, 0) + } + return out +} diff --git a/internal/handler/chat_audio_wav_test.go b/internal/handler/chat_audio_wav_test.go new file mode 100644 index 0000000000..e3ff71bc2a --- /dev/null +++ b/internal/handler/chat_audio_wav_test.go @@ -0,0 +1,96 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package handler + +import ( + "bytes" + "encoding/binary" + "testing" +) + +// testWAV builds a minimal 24kHz/16-bit/mono PCM WAV. +func testWAV(pcm []byte) []byte { + fmtChunk := make([]byte, 16) + binary.LittleEndian.PutUint16(fmtChunk[0:], 1) // PCM + binary.LittleEndian.PutUint16(fmtChunk[2:], 1) // mono + binary.LittleEndian.PutUint32(fmtChunk[4:], 24000) // sample rate + binary.LittleEndian.PutUint32(fmtChunk[8:], 48000) // byte rate + binary.LittleEndian.PutUint16(fmtChunk[12:], 2) // block align + binary.LittleEndian.PutUint16(fmtChunk[14:], 16) // bits per sample + return buildWAV(&wavFormat{fmtChunk: fmtChunk}, pcm) +} + +func TestSplitWAVRoundTrip(t *testing.T) { + pcm := []byte{1, 2, 3, 4, 5} + format, gotPCM, err := splitWAV(testWAV(pcm)) + if err != nil { + t.Fatalf("splitWAV: %v", err) + } + if !bytes.Equal(gotPCM, pcm) { + t.Errorf("pcm = %v, want %v", gotPCM, pcm) + } + rebuilt := buildWAV(format, gotPCM) + if len(rebuilt)%2 != 0 { + t.Errorf("rebuilt WAV length = %d, want even (RIFF chunks are 2-byte aligned)", len(rebuilt)) + } + if got := binary.LittleEndian.Uint32(rebuilt[4:8]); got != uint32(len(rebuilt)-8) { + t.Errorf("RIFF size = %d, want %d", got, len(rebuilt)-8) + } + format2, pcm2, err := splitWAV(rebuilt) + if err != nil { + t.Fatalf("splitWAV(rebuilt): %v", err) + } + if !bytes.Equal(format.fmtChunk, format2.fmtChunk) { + t.Errorf("fmt chunk changed after round trip") + } + if !bytes.Equal(pcm2, pcm) { + t.Errorf("pcm changed after round trip: %v", pcm2) + } +} + +func TestBuildWAVConcatenatedPCM(t *testing.T) { + seg1 := []byte{1, 2} + seg2 := []byte{3, 4, 5} + fmt1, pcm1, err := splitWAV(testWAV(seg1)) + if err != nil { + t.Fatalf("splitWAV seg1: %v", err) + } + _, pcm2, err := splitWAV(testWAV(seg2)) + if err != nil { + t.Fatalf("splitWAV seg2: %v", err) + } + + combined := buildWAV(fmt1, append(pcm1, pcm2...)) + _, pcm, err := splitWAV(combined) + if err != nil { + t.Fatalf("splitWAV combined: %v", err) + } + if !bytes.Equal(pcm, []byte{1, 2, 3, 4, 5}) { + t.Errorf("combined pcm = %v, want [1 2 3 4 5]", pcm) + } + if got := binary.LittleEndian.Uint32(combined[4:8]); got != uint32(len(combined)-8) { + t.Errorf("RIFF size = %d, want %d", got, len(combined)-8) + } +} + +func TestSplitWAVRejectsNonWAV(t *testing.T) { + if _, _, err := splitWAV([]byte("not a wav")); err == nil { + t.Fatal("error = nil, want parse error") + } + if _, _, err := splitWAV(nil); err == nil { + t.Fatal("error = nil, want parse error for empty input") + } +}