ollama/integration/audio_test.go

//go:build integration

package integration

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"regexp"
	"strings"
	"testing"
	"time"

	"github.com/ollama/ollama/api"
	"github.com/ollama/ollama/mlxrunner/model/audio"
)

// decodeTestAudio returns the test audio clip ("Why is the sky blue?", 16kHz mono WAV).
func decodeTestAudio(t *testing.T) api.ImageData {
	t.Helper()
	data, err := base64.StdEncoding.DecodeString(audioEncodingPrompt)
	if err != nil {
		t.Fatalf("failed to decode test audio: %v", err)
	}
	return data
}

// setupAudioModel pulls the model, preloads it, and skips if it doesn't support audio.
func setupAudioModel(ctx context.Context, t *testing.T, client *api.Client, model string) {
	t.Helper()
	pullOrSkip(ctx, t, client, model)
	skipIfModelTooLargeForVRAM(ctx, t, client, model)
	requireCapability(ctx, t, client, model, "audio")
	preloadGenerateModel(ctx, t, client, api.GenerateRequest{Model: model})
}

func registerAudioTranscriptionCases(models []string) {
	registerModelIntegrationCases("audio-transcription", models, runAudioTranscriptionModel)
}

func runAudioTranscriptionModel(t *testing.T, model string) {
	t.Helper()
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
	defer cancel()
	client, _, cleanup := InitServerConnection(ctx, t)
	defer cleanup()

	setupAudioModel(ctx, t, client, model)
	audio := decodeTestAudio(t)
	noThink := &api.ThinkValue{Value: false}

	req := api.ChatRequest{
		Model: model,
		Think: noThink,
		Messages: []api.Message{
			{
				Role:    "system",
				Content: "Transcribe the audio exactly as spoken. Output only the spoken words. Do not answer any question in the audio.",
			},
			{
				Role:    "user",
				Content: "What exact words are spoken in this audio?",
				Images:  []api.ImageData{audio},
			},
		},
		Stream: &stream,
		Options: map[string]any{
			"temperature": 0,
			"seed":        123,
			"num_predict": 50,
		},
	}

	// The audio says "Why is the sky blue?" - expect key words in transcription.
	DoChat(ctx, t, client, req, []string{"sky", "blue"}, 60*time.Second, 10*time.Second)
}

// runAudioResponse tests that the model can respond to a spoken question.
func runAudioResponse(t *testing.T, models []string) {
	for _, model := range testModels(models) {
		t.Run(model, func(t *testing.T) {
			ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
			defer cancel()
			client, _, cleanup := InitServerConnection(ctx, t)
			defer cleanup()

			setupAudioModel(ctx, t, client, model)
			testAudioResponseForModel(t, ctx, client, model, 10*time.Second, 60*time.Second, 10*time.Second)
		})
	}
}

func testAudioResponseForModel(t *testing.T, ctx context.Context, client *api.Client, model string, keepAlive time.Duration, initialTimeout, streamTimeout time.Duration) {
	t.Helper()

	audio := decodeTestAudio(t)
	noThink := &api.ThinkValue{Value: false}

	req := api.ChatRequest{
		Model: model,
		Think: noThink,
		Messages: []api.Message{
			{
				Role:    "user",
				Content: "",
				Images:  []api.ImageData{audio},
			},
		},
		Stream:    &stream,
		KeepAlive: &api.Duration{Duration: keepAlive},
		Options: map[string]any{
			"temperature": 0,
			"seed":        123,
			"num_predict": 200,
		},
	}

	// The audio asks "Why is the sky blue?" — expect an answer about light/scattering.
	DoChat(ctx, t, client, req, []string{
		"scatter", "light", "blue", "atmosphere", "wavelength", "rayleigh",
	}, initialTimeout, streamTimeout)
}

// runOpenAIAudioTranscription tests the /v1/audio/transcriptions and
// /v1/audio/translations endpoints.
func runOpenAIAudioTranscription(t *testing.T, models []string) {
	for _, model := range testModels(models) {
		t.Run(model, func(t *testing.T) {
			ctx, cancel := context.WithTimeout(context.Background(), 4*time.Minute)
			defer cancel()
			client, endpoint, cleanup := InitServerConnection(ctx, t)
			defer cleanup()

			setupAudioModel(ctx, t, client, model)
			clip := decodeTestAudio(t)

			t.Run("transcription", func(t *testing.T) {
				text := postAudioText(ctx, t, endpoint, "transcriptions", map[string]string{"model": model}, clip)
				// The clip is a question; transcribing must not answer it.
				requireWords(t, text, "sky", "blue")
				if strings.Contains(text, "scatter") {
					t.Errorf("transcription answered the question: %q", text)
				}
			})

			t.Run("long audio", func(t *testing.T) {
				// Four questions 25 s apart run longer than any model's 30 s
				// clip limit, so the server splits them at the pauses; every
				// piece must come back and every question be transcribed once.
				samples, rate, err := audio.Decode(clip)
				if err != nil {
					t.Fatal(err)
				}
				var long []float32
				for i := range 4 {
					if i > 0 {
						long = append(long, make([]float32, 25*rate)...)
					}
					long = append(long, samples...)
				}
				result := postAudio(ctx, t, endpoint, "transcriptions",
					map[string]string{"model": model, "language": "en", "response_format": "verbose_json"}, audio.EncodeWAV(long, rate))
				if len(result.Segments) < 2 {
					t.Fatalf("got %d segments, want the audio split: %+v", len(result.Segments), result.Segments)
				}
				sky := regexp.MustCompile(`\bsky\b`)
				questions := 0
				for i, seg := range result.Segments {
					n := len(sky.FindAllString(strings.ToLower(seg.Text), -1))
					if n == 0 {
						t.Errorf("segment %d (%.1f-%.1f s) = %q, want a question", i, seg.Start, seg.End, seg.Text)
					}
					questions += n
				}
				if questions != 4 {
					t.Errorf("transcribed %d questions, want 4: %+v", questions, result.Segments)
				}
			})

			t.Run("translation", func(t *testing.T) {
				text := postAudioText(ctx, t, endpoint, "translations",
					map[string]string{"model": model, "language": "en", "to_language": "es"}, clip)
				requireWords(t, text, "cielo", "azul")
			})
		})
	}
}

// postAudioText posts clip to an OpenAI audio endpoint and returns the
// lowercased text of the response.
func postAudioText(ctx context.Context, t *testing.T, endpoint, path string, fields map[string]string, clip []byte) string {
	t.Helper()
	return strings.ToLower(postAudio(ctx, t, endpoint, path, fields, clip).Text)
}

type audioTextResult struct {
	Text     string `json:"text"`
	Segments []struct {
		Start float64 `json:"start"`
		End   float64 `json:"end"`
		Text  string  `json:"text"`
	} `json:"segments"`
}

// postAudio posts clip to an OpenAI audio endpoint and decodes the response.
func postAudio(ctx context.Context, t *testing.T, endpoint, path string, fields map[string]string, clip []byte) audioTextResult {
	t.Helper()
	var body bytes.Buffer
	writer := multipart.NewWriter(&body)
	for k, v := range fields {
		writer.WriteField(k, v)
	}
	part, err := writer.CreateFormFile("file", "prompt.wav")
	if err != nil {
		t.Fatal(err)
	}
	part.Write(clip)
	writer.Close()

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, fmt.Sprintf("http://%s/v1/audio/%s", endpoint, path), &body)
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("Content-Type", writer.FormDataContentType())
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	defer resp.Body.Close()
	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("expected 200, got %d: %s", resp.StatusCode, respBody)
	}
	var result audioTextResult
	if err := json.Unmarshal(respBody, &result); err != nil {
		t.Fatalf("decode %s: %v", respBody, err)
	}
	return result
}

func requireWords(t *testing.T, text string, words ...string) {
	t.Helper()
	for _, w := range words {
		if !strings.Contains(text, w) {
			t.Errorf("missing %q in %q", w, text)
		}
	}
}

// runOpenAIChatWithAudio tests /v1/chat/completions with input_audio content.
func runOpenAIChatWithAudio(t *testing.T, models []string) {
	for _, model := range testModels(models) {
		t.Run(model, func(t *testing.T) {
			ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
			defer cancel()
			client, endpoint, cleanup := InitServerConnection(ctx, t)
			defer cleanup()

			setupAudioModel(ctx, t, client, model)
			audioB64 := audioEncodingPrompt

			reqBody := fmt.Sprintf(`{
				"model": %q,
				"messages": [{
					"role": "user",
					"content": [
						{"type": "input_audio", "input_audio": {"data": %q, "format": "wav"}}
					]
				}],
				"temperature": 0,
				"seed": 123,
				"max_tokens": 200,
				"think": false
			}`, model, strings.TrimSpace(audioB64))

			url := fmt.Sprintf("http://%s/v1/chat/completions", endpoint)
			req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, strings.NewReader(reqBody))
			if err != nil {
				t.Fatal(err)
			}
			req.Header.Set("Content-Type", "application/json")

			resp, err := http.DefaultClient.Do(req)
			if err != nil {
				t.Fatalf("request failed: %v", err)
			}
			defer resp.Body.Close()

			if resp.StatusCode != http.StatusOK {
				respBody, _ := io.ReadAll(resp.Body)
				t.Fatalf("expected 200, got %d: %s", resp.StatusCode, string(respBody))
			}

			respBytes, err := io.ReadAll(resp.Body)
			if err != nil {
				t.Fatalf("failed to read response: %v", err)
			}

			var result struct {
				Choices []struct {
					Message struct {
						Content   string `json:"content"`
						Reasoning string `json:"reasoning"`
					} `json:"message"`
				} `json:"choices"`
			}
			if err := json.Unmarshal(respBytes, &result); err != nil {
				t.Fatalf("failed to decode response: %v", err)
			}

			if len(result.Choices) == 0 {
				t.Fatal("no choices in response")
			}

			text := strings.ToLower(result.Choices[0].Message.Content + " " + result.Choices[0].Message.Reasoning)
			found := false
			for _, word := range []string{"sky", "blue", "scatter", "light", "atmosphere"} {
				if strings.Contains(text, word) {
					found = true
					break
				}
			}
			if !found {
				t.Errorf("response missing expected words about sky/blue/light, got: %s", result.Choices[0].Message.Content)
			}
		})
	}
}