From a661b6c33c5bb6489aba3f72e675e3bc166fcd6f Mon Sep 17 00:00:00 2001
From: Zhang Minghan <zmh13054618081@dingtalk.com>
Date: Sat, 2 Sep 2023 15:30:13 +0800
Subject: [PATCH] update apikey filter

---
 api/anonymous.go | 16 ++++++++++++++++
 api/selector.go  | 24 ++++++++++++++++++++++++
 2 files changed, 40 insertions(+)
 create mode 100644 api/selector.go

diff --git a/api/anonymous.go b/api/anonymous.go
index 133052a..622faab 100644
--- a/api/anonymous.go
+++ b/api/anonymous.go
@@ -41,6 +41,22 @@ func GetChatGPTResponse(message []types.ChatGPTMessage, token int) (string, erro
 	return data.(string), nil
 }
 
+func TestKey(key string) bool {
+	res, err := utils.Post(viper.GetString("openai.anonymous_endpoint")+"/chat/completions", map[string]string{
+		"Content-Type":  "application/json",
+		"Authorization": "Bearer " + key,
+	}, types.ChatGPTRequest{
+		Model:    "gpt-3.5-turbo",
+		Messages: []types.ChatGPTMessage{{Role: "user", Content: "hi"}},
+		MaxToken: 2,
+	})
+	if err != nil || res == nil {
+		panic(err)
+	}
+
+	return res.(map[string]interface{})["choices"] != nil
+}
+
 func GetAnonymousResponse(message string) (string, string, error) {
 	keyword, source := ChatWithWeb([]types.ChatGPTMessage{{Role: "user", Content: message}}, false)
 	resp, err := GetChatGPTResponse(source, 1000)
diff --git a/api/selector.go b/api/selector.go
new file mode 100644
index 0000000..a112048
--- /dev/null
+++ b/api/selector.go
@@ -0,0 +1,24 @@
+package api
+
+import "strings"
+
+func FilterKeys(keys string) string {
+	stack := make(chan string, len(strings.Split(keys, "|")))
+	for _, key := range strings.Split(keys, "|") {
+		go func(key string) {
+			if TestKey(key) {
+				stack <- key
+			} else {
+				stack <- ""
+			}
+		}(key)
+	}
+
+	var result string
+	for i := 0; i < len(strings.Split(keys, "|")); i++ {
+		if res := <-stack; res != "" {
+			result += res + "|"
+		}
+	}
+	return strings.Trim(result, "|")
+}