From 04ce0405bf9f5bd8d808d5bc84628185ac0e2e7a Mon Sep 17 00:00:00 2001
From: bourne7 <lawrencepbr@gmail.com>
Date: Sat, 6 May 2023 18:26:58 +0800
Subject: [PATCH] feat: add http client timeout config

---
 code/config.example.yaml       | 2 ++
 code/initialization/config.go  | 2 ++
 code/main.go                   | 6 ++++++
 code/services/openai/common.go | 5 +++--
 docker-compose.yaml            | 2 ++
 5 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/code/config.example.yaml b/code/config.example.yaml
index 211c6ae..c249d65 100644
--- a/code/config.example.yaml
+++ b/code/config.example.yaml
@@ -19,6 +19,8 @@ KEY_FILE: key.pem
 API_URL: https://api.openai.com
 # 代理设置, 例如 "http://127.0.0.1:7890", ""代表不使用代理
 HTTP_PROXY: ""
+# 访问OpenAi的 普通 Http请求的超时时间，单位秒，不配置的话默认为 550 秒
+OPENAI_HTTP_CLIENT_TIMEOUT:
 
 # AZURE OPENAI
 AZURE_ON: true # set to true to use Azure rather than OpenAI
diff --git a/code/initialization/config.go b/code/initialization/config.go
index 6897c41..ce8fd12 100644
--- a/code/initialization/config.go
+++ b/code/initialization/config.go
@@ -35,6 +35,7 @@ type Config struct {
 	AzureOpenaiToken                   string
 	AccessControlEnable                bool
 	AccessControlMaxCountPerUserPerDay int
+	OpenAIHttpClientTimeOut            int
 }
 
 var (
@@ -92,6 +93,7 @@ func LoadConfig(cfg string) *Config {
 		AzureOpenaiToken:                   getViperStringValue("AZURE_OPENAI_TOKEN", ""),
 		AccessControlEnable:                getViperBoolValue("ACCESS_CONTROL_ENABLE", false),
 		AccessControlMaxCountPerUserPerDay: getViperIntValue("ACCESS_CONTROL_MAX_COUNT_PER_USER_PER_DAY", 0),
+		OpenAIHttpClientTimeOut:            getViperIntValue("OPENAI_HTTP_CLIENT_TIMEOUT", 550),
 	}
 
 	return config
diff --git a/code/main.go b/code/main.go
index 52a0130..527dbf1 100644
--- a/code/main.go
+++ b/code/main.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	larkcard "github.com/larksuite/oapi-sdk-go/v3/card"
 	larkim "github.com/larksuite/oapi-sdk-go/v3/service/im/v1"
@@ -26,6 +27,11 @@ func main() {
 	initialization.InitRoleList()
 	pflag.Parse()
 	globalConfig := initialization.GetConfig()
+
+	// 打印一下实际读取到的配置
+	globalConfigPrettyString, _ := json.MarshalIndent(globalConfig, "", "    ")
+	log.Println(string(globalConfigPrettyString))
+
 	initialization.LoadLarkClient(*globalConfig)
 	gpt := openai.NewChatGPT(*globalConfig)
 	handlers.InitHandlers(gpt, *globalConfig)
diff --git a/code/services/openai/common.go b/code/services/openai/common.go
index ea517ff..b8609d5 100644
--- a/code/services/openai/common.go
+++ b/code/services/openai/common.go
@@ -185,8 +185,9 @@ func (gpt *ChatGPT) sendRequestWithBodyType(link, method string,
 
 func GetProxyClient(proxyString string) (*http.Client, error) {
 	var client *http.Client
+	timeOutDuration := time.Duration(initialization.GetConfig().OpenAIHttpClientTimeOut) * time.Second
 	if proxyString == "" {
-		client = &http.Client{Timeout: 110 * time.Second}
+		client = &http.Client{Timeout: timeOutDuration}
 	} else {
 		proxyUrl, err := url.Parse(proxyString)
 		if err != nil {
@@ -197,7 +198,7 @@ func GetProxyClient(proxyString string) (*http.Client, error) {
 		}
 		client = &http.Client{
 			Transport: transport,
-			Timeout:   110 * time.Second,
+			Timeout:   timeOutDuration,
 		}
 	}
 	return client, nil
diff --git a/docker-compose.yaml b/docker-compose.yaml
index 6e94dee..79ccbb5 100644
--- a/docker-compose.yaml
+++ b/docker-compose.yaml
@@ -38,3 +38,5 @@ services:
       - ACCESS_CONTROL_ENABLE=false
       # 每个用户每天最多问多少个问题。默认为0. 配置成为小于等于0表示不限制。
       - ACCESS_CONTROL_MAX_COUNT_PER_USER_PER_DAY=0
+      # 访问OpenAi的 普通 Http请求的超时时间，单位秒，不配置的话默认为 550 秒
+      - OPENAI_HTTP_CLIENT_TIMEOUT