Merge pull request #73 from aaakoako/https

feat: support https and multi openai-keys to reduce access pressure

.gitignore

@@ -37,3 +37,4 @@ docker.md
 # Mac OS
 .DS_Store
 **/.DS_Store
+*.pem

code/.gitignore

@@ -0,0 +1,2 @@
+/apikey_usage.json
+*.pem

code/config.example.yaml

@@ -1,10 +1,16 @@
-#  飞书
+#  飞书  # 不要随意修改example中配置的顺序和空格，否则docker的sed脚本会执行出错
 APP_ID: cli_axxx
 APP_SECRET: xxx
 APP_ENCRYPT_KEY: xxx
 APP_VERIFICATION_TOKEN: xxx
-# 请确保和飞书应用管理平台中的设置一致 
+# 请确保和飞书应用管理平台中的设置一致
 BOT_NAME: chatGpt
-# openAI
-OPENAI_KEY: sk-xxx
+# openAI key 支持负载均衡 可以填写多个key 用逗号分隔
+OPENAI_KEY: sk-xxx,sk-xxx,sk-xxx
+# 服务器配置
+HTTP_PORT: 9000
+HTTPS_PORT: 9001
+USE_HTTPS: false
+CERT_FILE: cert.pem
+KEY_FILE: key.pem
 

code/handlers/common.go

@@ -8,7 +8,7 @@ import (
 	"strings"
 )
 
-//func sendCard
+// func sendCard
 func msgFilter(msg string) string {
 	//replace @到下一个非空的字段 为 ''
 	regex := regexp.MustCompile(`@[^ ]*`)
@@ -50,14 +50,14 @@ func processQuote(msg string) string {
 	return strings.Replace(msg, "\\\"", "\"", -1)
 }
 
-//将字符中 \u003c 替换为 <  等等
+// 将字符中 \u003c 替换为 <  等等
 func processUnicode(msg string) string {
 	regex := regexp.MustCompile(`\\u[0-9a-fA-F]{4}`)
 	return regex.ReplaceAllStringFunc(msg, func(s string) string {
 		r, _ := regexp.Compile(`\\u`)
 		s = r.ReplaceAllString(s, "")
 		i, _ := strconv.ParseInt(s, 16, 32)
-		return string(i)
+		return strconv.Itoa(int(i))
 	})
 }
 

code/handlers/handler.go

@@ -13,7 +13,7 @@ import (
 	larkim "github.com/larksuite/oapi-sdk-go/v3/service/im/v1"
 )
 
-//责任链
+// 责任链
 func chain(data *ActionInfo, actions ...Action) bool {
 	for _, v := range actions {
 		if !v.Execute(data) {
@@ -26,7 +26,7 @@ func chain(data *ActionInfo, actions ...Action) bool {
 type MessageHandler struct {
 	sessionCache services.SessionServiceCacheInterface
 	msgCache     services.MsgCacheInterface
-	gpt          services.ChatGPT
+	gpt          *services.ChatGPT
 	config       initialization.Config
 }
 
@@ -153,7 +153,7 @@ func (m MessageHandler) msgReceivedHandler(ctx context.Context, event *larkim.P2
 
 var _ MessageHandlerInterface = (*MessageHandler)(nil)
 
-func NewMessageHandler(gpt services.ChatGPT,
+func NewMessageHandler(gpt *services.ChatGPT,
 	config initialization.Config) MessageHandlerInterface {
 	return &MessageHandler{
 		sessionCache: services.GetSessionCache(),

code/handlers/init.go

@@ -25,7 +25,7 @@ const (
 // handlers 所有消息类型类型的处理器
 var handlers MessageHandlerInterface
 
-func InitHandlers(gpt services.ChatGPT, config initialization.Config) {
+func InitHandlers(gpt *services.ChatGPT, config initialization.Config) {
 	handlers = NewMessageHandler(gpt, config)
 }
 
@@ -50,7 +50,7 @@ func CardHandler() func(ctx context.Context,
 func judgeCardType(cardAction *larkcard.CardAction) HandlerType {
 	actionValue := cardAction.Action.Value
 	chatType := actionValue["chatType"]
-	fmt.Printf("chatType: %v", chatType)
+	//fmt.Printf("chatType: %v", chatType)
 	if chatType == "group" {
 		return GroupHandler
 	}

code/initialization/config.go

@@ -2,6 +2,9 @@ package initialization
 
 import (
 	"fmt"
+	"os"
+	"strconv"
+	"strings"
 
 	"github.com/spf13/viper"
 )
@@ -12,29 +15,103 @@ type Config struct {
 	FeishuAppEncryptKey        string
 	FeishuAppVerificationToken string
 	FeishuBotName              string
-	OpenaiApiKey               string
+	OpenaiApiKeys              []string
+	HttpPort                   int
+	HttpsPort                  int
+	UseHttps                   bool
+	CertFile                   string
+	KeyFile                    string
 }
 
 func LoadConfig(cfg string) *Config {
 	viper.SetConfigFile(cfg)
 	viper.ReadInConfig()
 	viper.AutomaticEnv()
+	//content, err := ioutil.ReadFile("config.yaml")
+	//if err != nil {
+	//	fmt.Println("Error reading file:", err)
+	//}
+	//fmt.Println(string(content))
 
-	return &Config{
-		FeishuAppId:                getViperStringValue("APP_ID"),
-		FeishuAppSecret:            getViperStringValue("APP_SECRET"),
-		FeishuAppEncryptKey:        getViperStringValue("APP_ENCRYPT_KEY"),
-		FeishuAppVerificationToken: getViperStringValue("APP_VERIFICATION_TOKEN"),
-		FeishuBotName:              getViperStringValue("BOT_NAME"),
-		OpenaiApiKey:               getViperStringValue("OPENAI_KEY"),
+	config := &Config{
+		FeishuAppId:                getViperStringValue("APP_ID", ""),
+		FeishuAppSecret:            getViperStringValue("APP_SECRET", ""),
+		FeishuAppEncryptKey:        getViperStringValue("APP_ENCRYPT_KEY", ""),
+		FeishuAppVerificationToken: getViperStringValue("APP_VERIFICATION_TOKEN", ""),
+		FeishuBotName:              getViperStringValue("BOT_NAME", ""),
+		OpenaiApiKeys:              getViperStringArray("OPENAI_KEY", nil),
+		HttpPort:                   getViperIntValue("HTTP_PORT", 9000),
+		HttpsPort:                  getViperIntValue("HTTPS_PORT", 9001),
+		UseHttps:                   getViperBoolValue("USE_HTTPS", false),
+		CertFile:                   getViperStringValue("CERT_FILE", "cert.pem"),
+		KeyFile:                    getViperStringValue("KEY_FILE", "key.pem"),
 	}
 
+	return config
 }
 
-func getViperStringValue(key string) string {
+func getViperStringValue(key string, defaultValue string) string {
 	value := viper.GetString(key)
 	if value == "" {
-		panic(fmt.Errorf("%s MUST be provided in environment or config.yaml file", key))
+		return defaultValue
 	}
 	return value
 }
+
+//OPENAI_KEY: sk-xxx,sk-xxx,sk-xxx
+//result:[sk-xxx sk-xxx sk-xxx]
+func getViperStringArray(key string, defaultValue []string) []string {
+	value := viper.GetString(key)
+	if value == "" {
+		return defaultValue
+	}
+	return strings.Split(value, ",")
+}
+
+func getViperIntValue(key string, defaultValue int) int {
+	value := viper.GetString(key)
+	if value == "" {
+		return defaultValue
+	}
+	intValue, err := strconv.Atoi(value)
+	if err != nil {
+		fmt.Printf("Invalid value for %s, using default value %d\n", key, defaultValue)
+		return defaultValue
+	}
+	return intValue
+}
+
+func getViperBoolValue(key string, defaultValue bool) bool {
+	value := viper.GetString(key)
+	if value == "" {
+		return defaultValue
+	}
+	boolValue, err := strconv.ParseBool(value)
+	if err != nil {
+		fmt.Printf("Invalid value for %s, using default value %v\n", key, defaultValue)
+		return defaultValue
+	}
+	return boolValue
+}
+
+func (config *Config) GetCertFile() string {
+	if config.CertFile == "" {
+		return "cert.pem"
+	}
+	if _, err := os.Stat(config.CertFile); err != nil {
+		fmt.Printf("Certificate file %s does not exist, using default file cert.pem\n", config.CertFile)
+		return "cert.pem"
+	}
+	return config.CertFile
+}
+
+func (config *Config) GetKeyFile() string {
+	if config.KeyFile == "" {
+		return "key.pem"
+	}
+	if _, err := os.Stat(config.KeyFile); err != nil {
+		fmt.Printf("Key file %s does not exist, using default file key.pem\n", config.KeyFile)
+		return "key.pem"
+	}
+	return config.KeyFile
+}

code/initialization/gin.go

@@ -0,0 +1,58 @@
+package initialization
+
+import (
+	"crypto/tls"
+	"fmt"
+	"github.com/gin-gonic/gin"
+	"log"
+	"net/http"
+	"time"
+)
+
+func loadCertificate(config Config) (cert tls.Certificate, err error) {
+	cert, err = tls.LoadX509KeyPair(config.CertFile, config.KeyFile)
+	if err != nil {
+		return cert, fmt.Errorf("failed to load certificate: %v", err)
+	}
+	// check certificate expiry
+	certExpiry := cert.Leaf.NotAfter
+	if certExpiry.Before(time.Now()) {
+		return cert, fmt.Errorf("certificate expired on %v", certExpiry)
+	}
+	return cert, nil
+}
+func startHTTPServer(config Config, r *gin.Engine) (err error) {
+	log.Printf("http server started: http://localhost:%d/webhook/event\n", config.HttpPort)
+	err = r.Run(fmt.Sprintf(":%d", config.HttpPort))
+	if err != nil {
+		return fmt.Errorf("failed to start http server: %v", err)
+	}
+	return nil
+}
+func startHTTPSServer(config Config, r *gin.Engine) (err error) {
+	cert, err := loadCertificate(config)
+	if err != nil {
+		return fmt.Errorf("failed to load certificate: %v", err)
+	}
+	server := &http.Server{
+		Addr:    fmt.Sprintf(":%d", config.HttpsPort),
+		Handler: r,
+		TLSConfig: &tls.Config{
+			Certificates: []tls.Certificate{cert},
+		},
+	}
+	fmt.Printf("https server started: https://localhost:%d/webhook/event\n", config.HttpsPort)
+	err = server.ListenAndServeTLS("", "")
+	if err != nil {
+		return fmt.Errorf("failed to start https server: %v", err)
+	}
+	return nil
+}
+func StartServer(config Config, r *gin.Engine) (err error) {
+	if config.UseHttps {
+		err = startHTTPSServer(config, r)
+	} else {
+		err = startHTTPServer(config, r)
+	}
+	return err
+}

code/main.go

@@ -2,8 +2,8 @@ package main
 
 import (
 	"context"
-	"fmt"
 	larkim "github.com/larksuite/oapi-sdk-go/v3/service/im/v1"
+	"log"
 	"start-feishubot/handlers"
 	"start-feishubot/initialization"
 	"start-feishubot/services"
@@ -26,9 +26,8 @@ func main() {
 	pflag.Parse()
 	config := initialization.LoadConfig(*cfg)
 	initialization.LoadLarkClient(*config)
-
-	gpt := services.NewChatGPT(config.OpenaiApiKey)
-	handlers.InitHandlers(*gpt, *config)
+	gpt := services.NewChatGPT(config.OpenaiApiKeys)
+	handlers.InitHandlers(gpt, *config)
 
 	eventHandler := dispatcher.NewEventDispatcher(
 		config.FeishuAppVerificationToken, config.FeishuAppEncryptKey).
@@ -53,8 +52,9 @@ func main() {
 		sdkginext.NewCardActionHandlerFunc(
 			cardHandler))
 
-	fmt.Println("http server started",
-		"http://localhost:9000/webhook/event")
-	r.Run(":9000")
+	err := initialization.StartServer(*config, r)
+	if err != nil {
+		log.Fatalf("failed to start server: %v", err)
+	}
 
 }