Merge branch 'main' into patch/images-edits

common/client/init.go

@@ -0,0 +1,60 @@
+package client
+
+import (
+	"fmt"
+	"github.com/songquanpeng/one-api/common/config"
+	"github.com/songquanpeng/one-api/common/logger"
+	"net/http"
+	"net/url"
+	"time"
+)
+
+var HTTPClient *http.Client
+var ImpatientHTTPClient *http.Client
+var UserContentRequestHTTPClient *http.Client
+
+func Init() {
+	if config.UserContentRequestProxy != "" {
+		logger.SysLog(fmt.Sprintf("using %s as proxy to fetch user content", config.UserContentRequestProxy))
+		proxyURL, err := url.Parse(config.UserContentRequestProxy)
+		if err != nil {
+			logger.FatalLog(fmt.Sprintf("USER_CONTENT_REQUEST_PROXY set but invalid: %s", config.UserContentRequestProxy))
+		}
+		transport := &http.Transport{
+			Proxy: http.ProxyURL(proxyURL),
+		}
+		UserContentRequestHTTPClient = &http.Client{
+			Transport: transport,
+			Timeout:   time.Second * time.Duration(config.UserContentRequestTimeout),
+		}
+	} else {
+		UserContentRequestHTTPClient = &http.Client{}
+	}
+	var transport http.RoundTripper
+	if config.RelayProxy != "" {
+		logger.SysLog(fmt.Sprintf("using %s as api relay proxy", config.RelayProxy))
+		proxyURL, err := url.Parse(config.RelayProxy)
+		if err != nil {
+			logger.FatalLog(fmt.Sprintf("USER_CONTENT_REQUEST_PROXY set but invalid: %s", config.UserContentRequestProxy))
+		}
+		transport = &http.Transport{
+			Proxy: http.ProxyURL(proxyURL),
+		}
+	}
+
+	if config.RelayTimeout == 0 {
+		HTTPClient = &http.Client{
+			Transport: transport,
+		}
+	} else {
+		HTTPClient = &http.Client{
+			Timeout:   time.Duration(config.RelayTimeout) * time.Second,
+			Transport: transport,
+		}
+	}
+
+	ImpatientHTTPClient = &http.Client{
+		Timeout:   5 * time.Second,
+		Transport: transport,
+	}
+}

common/config/config.go

@@ -117,10 +117,10 @@ var ValidThemes = map[string]bool{
 // All duration's unit is seconds
 // Shouldn't larger then RateLimitKeyExpirationDuration
 var (
-	GlobalApiRateLimitNum            = env.Int("GLOBAL_API_RATE_LIMIT", 180)
+	GlobalApiRateLimitNum            = env.Int("GLOBAL_API_RATE_LIMIT", 240)
 	GlobalApiRateLimitDuration int64 = 3 * 60
 
-	GlobalWebRateLimitNum            = env.Int("GLOBAL_WEB_RATE_LIMIT", 60)
+	GlobalWebRateLimitNum            = env.Int("GLOBAL_WEB_RATE_LIMIT", 120)
 	GlobalWebRateLimitDuration int64 = 3 * 60
 
 	UploadRateLimitNum            = 10
@@ -144,3 +144,7 @@ var MetricFailChanSize = env.Int("METRIC_FAIL_CHAN_SIZE", 128)
 var InitialRootToken = os.Getenv("INITIAL_ROOT_TOKEN")
 
 var GeminiVersion = env.String("GEMINI_VERSION", "v1")
+
+var RelayProxy = env.String("RELAY_PROXY", "")
+var UserContentRequestProxy = env.String("USER_CONTENT_REQUEST_PROXY", "")
+var UserContentRequestTimeout = env.Int("USER_CONTENT_REQUEST_TIMEOUT", 30)

common/image/image.go

@@ -3,6 +3,7 @@ package image
 import (
 	"bytes"
 	"encoding/base64"
+	"github.com/songquanpeng/one-api/common/client"
 	"image"
 	_ "image/gif"
 	_ "image/jpeg"
@@ -19,7 +20,7 @@ import (
 var dataURLPattern = regexp.MustCompile(`data:image/([^;]+);base64,(.*)`)
 
 func IsImageUrl(url string) (bool, error) {
-	resp, err := http.Head(url)
+	resp, err := client.UserContentRequestHTTPClient.Head(url)
 	if err != nil {
 		return false, err
 	}
@@ -34,7 +35,7 @@ func GetImageSizeFromUrl(url string) (width int, height int, err error) {
 	if !isImage {
 		return
 	}
-	resp, err := http.Get(url)
+	resp, err := client.UserContentRequestHTTPClient.Get(url)
 	if err != nil {
 		return
 	}

common/image/image_test.go

@@ -2,6 +2,7 @@ package image_test
 
 import (
 	"encoding/base64"
+	"github.com/songquanpeng/one-api/common/client"
 	"image"
 	_ "image/gif"
 	_ "image/jpeg"
@@ -44,6 +45,11 @@ var (
 	}
 )
 
+func TestMain(m *testing.M) {
+	client.Init()
+	m.Run()
+}
+
 func TestDecode(t *testing.T) {
 	// Bytes read: varies sometimes
 	// jpeg: 1063892

common/init.go

@@ -24,7 +24,7 @@ func printHelp() {
 	fmt.Println("Usage: one-api [--port <port>] [--log-dir <log directory>] [--version] [--help]")
 }
 
-func init() {
+func Init() {
 	flag.Parse()
 
 	if *PrintVersion {

common/logger/logger.go

@@ -44,11 +44,19 @@ func SysLog(s string) {
 	_, _ = fmt.Fprintf(gin.DefaultWriter, "[SYS] %v | %s \n", t.Format("2006/01/02 - 15:04:05"), s)
 }
 
+func SysLogf(format string, a ...any) {
+	SysLog(fmt.Sprintf(format, a...))
+}
+
 func SysError(s string) {
 	t := time.Now()
 	_, _ = fmt.Fprintf(gin.DefaultErrorWriter, "[SYS] %v | %s \n", t.Format("2006/01/02 - 15:04:05"), s)
 }
 
+func SysErrorf(format string, a ...any) {
+	SysError(fmt.Sprintf(format, a...))
+}
+
 func Debug(ctx context.Context, msg string) {
 	if config.DebugEnabled {
 		logHelper(ctx, loggerDEBUG, msg)