From db7f0eb316ef347ff05481c1d4703ac26f4da4e3 Mon Sep 17 00:00:00 2001
From: heimoshuiyu <heimoshuiyu@gmail.com>
Date: Sun, 18 Feb 2024 16:45:37 +0800
Subject: [PATCH] timeout

---
 process.go   | 11 ++---------
 structure.go | 17 +++++++++++++++++
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/process.go b/process.go
index 48d9dfa..4b3701c 100644
--- a/process.go
+++ b/process.go
@@ -91,16 +91,9 @@ func processRequest(c *gin.Context, upstream *OPENAI_UPSTREAM, record *Record, s
 		}
 
 		// set timeout, default is 60 second
-		timeout := 60 * time.Second
+		timeout := time.Duration(upstream.Timeout) * time.Second
 		if requestBodyOK == nil && requestBody.Stream {
-			timeout = 5 * time.Second
-		}
-		if len(inBody) > 1024*128 {
-			timeout = 20 * time.Second
-		}
-		if upstream.Timeout > 0 {
-			// convert upstream.Timeout(second) to nanosecond
-			timeout = time.Duration(upstream.Timeout) * time.Second
+			timeout = time.Duration(upstream.StreamTimeout) * time.Second
 		}
 
 		// timeout out request
diff --git a/structure.go b/structure.go
index 67dbeef..93845f3 100644
--- a/structure.go
+++ b/structure.go
@@ -14,12 +14,15 @@ type Config struct {
 	DBType        string            `yaml:"dbtype"`
 	DBAddr        string            `yaml:"dbaddr"`
 	Authorization string            `yaml:"authorization"`
+	Timeout       int64             `yaml:"timeout"`
+	StreamTimeout int64             `yaml:"stream_timeout"`
 	Upstreams     []OPENAI_UPSTREAM `yaml:"upstreams"`
 }
 type OPENAI_UPSTREAM struct {
 	SK            string   `yaml:"sk"`
 	Endpoint      string   `yaml:"endpoint"`
 	Timeout       int64    `yaml:"timeout"`
+	StreamTimeout int64    `yaml:"stream_timeout"`
 	Allow         []string `yaml:"allow"`
 	Deny          []string `yaml:"deny"`
 	Type          string   `yaml:"type"`
@@ -57,6 +60,14 @@ func readConfig(filepath string) Config {
 		log.Println("DBAddr not set, use default value: ./db.sqlite")
 		config.DBAddr = "./db.sqlite"
 	}
+	if config.Timeout == 0 {
+		log.Println("Timeout not set, use default value: 120")
+		config.Timeout = 120
+	}
+	if config.StreamTimeout == 0 {
+		log.Println("StreamTimeout not set, use default value: 10")
+		config.StreamTimeout = 10
+	}
 
 	for i, upstream := range config.Upstreams {
 		// parse upstream endpoint URL
@@ -75,6 +86,12 @@ func readConfig(filepath string) Config {
 		if config.Upstreams[i].Authorization == "" && !config.Upstreams[i].Noauth {
 			config.Upstreams[i].Authorization = config.Authorization
 		}
+		if config.Upstreams[i].Timeout == 0 {
+			config.Upstreams[i].Timeout = config.Timeout
+		}
+		if config.Upstreams[i].StreamTimeout == 0 {
+			config.Upstreams[i].StreamTimeout = config.StreamTimeout
+		}
 	}
 
 	return config