Merge 47918f3143 into ea0721d525

common/config/config.go

@@ -161,4 +161,5 @@ var RelayProxy = env.String("RELAY_PROXY", "")
 var UserContentRequestProxy = env.String("USER_CONTENT_REQUEST_PROXY", "")
 var UserContentRequestTimeout = env.Int("USER_CONTENT_REQUEST_TIMEOUT", 30)
 
+// EnforceIncludeUsage is used to determine whether to include usage in the response
 var EnforceIncludeUsage = env.Bool("ENFORCE_INCLUDE_USAGE", false)

common/ctxkey/key.go

@@ -21,4 +21,5 @@ const (
 	AvailableModels   = "available_models"
 	KeyRequestBody    = "key_request_body"
 	SystemPrompt      = "system_prompt"
+	Meta              = "meta"
 )

common/helper/audio.go

@@ -0,0 +1,62 @@
+package helper
+
+import (
+	"bytes"
+	"context"
+	"io"
+	"os"
+	"os/exec"
+	"strconv"
+
+	"github.com/pkg/errors"
+)
+
+// SaveTmpFile saves data to a temporary file. The filename would be apppended with a random string.
+func SaveTmpFile(filename string, data io.Reader) (string, error) {
+	if data == nil {
+		return "", errors.New("data is nil")
+	}
+
+	f, err := os.CreateTemp("", "*-"+filename)
+	if err != nil {
+		return "", errors.Wrapf(err, "failed to create temporary file %s", filename)
+	}
+	defer f.Close()
+
+	_, err = io.Copy(f, data)
+	if err != nil {
+		return "", errors.Wrapf(err, "failed to copy data to temporary file %s", filename)
+	}
+
+	return f.Name(), nil
+}
+
+// GetAudioTokens returns the number of tokens in an audio file.
+func GetAudioTokens(ctx context.Context, audio io.Reader, tokensPerSecond float64) (float64, error) {
+	filename, err := SaveTmpFile("audio", audio)
+	if err != nil {
+		return 0, errors.Wrap(err, "failed to save audio to temporary file")
+	}
+	defer os.Remove(filename)
+
+	duration, err := GetAudioDuration(ctx, filename)
+	if err != nil {
+		return 0, errors.Wrap(err, "failed to get audio tokens")
+	}
+
+	return duration * tokensPerSecond, nil
+}
+
+// GetAudioDuration returns the duration of an audio file in seconds.
+func GetAudioDuration(ctx context.Context, filename string) (float64, error) {
+	// ffprobe -v error -show_entries format=duration -of default=noprint_wrappers=1:nokey=1 {{input}}
+	c := exec.CommandContext(ctx, "ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "default=noprint_wrappers=1:nokey=1", filename)
+	output, err := c.Output()
+	if err != nil {
+		return 0, errors.Wrap(err, "failed to get audio duration")
+	}
+
+	// Actually gpt-4-audio calculates tokens with 0.1s precision,
+	// while whisper calculates tokens with 1s precision
+	return strconv.ParseFloat(string(bytes.TrimSpace(output)), 64)
+}

common/helper/audio_test.go

@@ -0,0 +1,55 @@
+package helper
+
+import (
+	"context"
+	"io"
+	"net/http"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestGetAudioDuration(t *testing.T) {
+	t.Run("should return correct duration for a valid audio file", func(t *testing.T) {
+		tmpFile, err := os.CreateTemp("", "test_audio*.mp3")
+		require.NoError(t, err)
+		defer os.Remove(tmpFile.Name())
+
+		// download test audio file
+		resp, err := http.Get("https://s3.laisky.com/uploads/2025/01/audio-sample.m4a")
+		require.NoError(t, err)
+		defer resp.Body.Close()
+
+		_, err = io.Copy(tmpFile, resp.Body)
+		require.NoError(t, err)
+		require.NoError(t, tmpFile.Close())
+
+		duration, err := GetAudioDuration(context.Background(), tmpFile.Name())
+		require.NoError(t, err)
+		require.Equal(t, duration, 3.904)
+	})
+
+	t.Run("should return an error for a non-existent file", func(t *testing.T) {
+		_, err := GetAudioDuration(context.Background(), "non_existent_file.mp3")
+		require.Error(t, err)
+	})
+}
+
+func TestGetAudioTokens(t *testing.T) {
+	t.Run("should return correct tokens for a valid audio file", func(t *testing.T) {
+		// download test audio file
+		resp, err := http.Get("https://s3.laisky.com/uploads/2025/01/audio-sample.m4a")
+		require.NoError(t, err)
+		defer resp.Body.Close()
+
+		tokens, err := GetAudioTokens(context.Background(), resp.Body, 50)
+		require.NoError(t, err)
+		require.Equal(t, tokens, 200)
+	})
+
+	t.Run("should return an error for a non-existent file", func(t *testing.T) {
+		_, err := GetAudioTokens(context.Background(), nil, 1)
+		require.Error(t, err)
+	})
+}