Merge remote-tracking branch 'origin/upstream/main'

common/constants.go

@@ -98,6 +98,10 @@ var BatchUpdateInterval = GetOrDefault("BATCH_UPDATE_INTERVAL", 5)
 
 var RelayTimeout = GetOrDefault("RELAY_TIMEOUT", 0) // unit is second
 
+var GeminiSafetySetting = GetOrDefaultString("GEMINI_SAFETY_SETTING", "BLOCK_NONE")
+
+var Theme = GetOrDefaultString("THEME", "default")
+
 const (
 	RequestIdKey = "X-Oneapi-Request-Id"
 )

common/image/image.go

@@ -15,6 +15,9 @@ import (
 	_ "golang.org/x/image/webp"
 )
 
+// Regex to match data URL pattern
+var dataURLPattern = regexp.MustCompile(`data:image/([^;]+);base64,(.*)`)
+
 func IsImageUrl(url string) (bool, error) {
 	resp, err := http.Head(url)
 	if err != nil {
@@ -44,9 +47,13 @@ func GetImageSizeFromUrl(url string) (width int, height int, err error) {
 }
 
 func GetImageFromUrl(url string) (mimeType string, data string, err error) {
-	// openai's image_url support base64 encoded image
-	if strings.HasPrefix(url, "data:image/jpeg;base64,") {
-		return "image/jpeg", strings.TrimPrefix(url, "data:image/jpeg;base64,"), nil
+	// Check if the URL is a data URL
+	matches := dataURLPattern.FindStringSubmatch(url)
+	if len(matches) == 3 {
+		// URL is a data URL
+		mimeType = "image/" + matches[1]
+		data = matches[2]
+		return
 	}
 
 	isImage, err := IsImageUrl(url)

common/model-ratio.go

@@ -52,6 +52,8 @@ var ModelRatio = map[string]float64{
 	"gpt-3.5-turbo-16k-0613":    1.5,
 	"gpt-3.5-turbo-instruct":    0.75, // $0.0015 / 1K tokens
 	"gpt-3.5-turbo-1106":        0.5,  // $0.001 / 1K tokens
+	"davinci-002":               1,    // $0.002 / 1K tokens
+	"babbage-002":               0.2,  // $0.0004 / 1K tokens
 	"text-ada-001":              0.2,
 	"text-babbage-001":          0.25,
 	"text-curie-001":            1,

common/utils.go

@@ -196,6 +196,13 @@ func GetOrDefault(env string, defaultValue int) int {
 	return num
 }
 
+func GetOrDefaultString(env string, defaultValue string) string {
+	if env == "" || os.Getenv(env) == "" {
+		return defaultValue
+	}
+	return os.Getenv(env)
+}
+
 func MessageWithRequestId(message string, id string) string {
 	return fmt.Sprintf("%s (request id: %s)", message, id)
 }