Skip to content

Commit

Permalink
Merge remote-tracking branch 'aaakoako/https' into https
Browse files Browse the repository at this point in the history
  • Loading branch information
@Leizhenpeng
Leizhenpeng committed Mar 9, 2023
2 parents aeca69b + b317a5c commit 0d98cec
Showing 1 changed file with 101 additions and 112 deletions.
213 changes: 101 additions & 112 deletions code/services/gpt3.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -129,7 +129,10 @@ func (gpt *ChatGPT) Completions(msg []Messages) (resp Messages, err error) {
return resp, err
}
req.Header.Set("Content-Type", "application/json")
apiKey := gpt.getApiKey()
apiKey, err := gpt.getApiKey()
if err != nil {
return resp, err
}
req.Header.Set("Authorization", "Bearer "+apiKey)
client := &http.Client{Timeout: 110 * time.Second}
response, err := client.Do(req)
Expand Down Expand Up @@ -158,33 +161,33 @@ func (gpt *ChatGPT) Completions(msg []Messages) (resp Messages, err error) {
resp = gptResponseBody.Choices[0].Message
gpt.apiKeyUsage[currentApiKey]++
gpt.apiKeyUsageMutex.Unlock()
err = gpt.updateApiKeyUsage()
if err != nil {
log.Printf("update api key usage failed: %v", err)
}
gpt.updateApiKeyUsage(apiKey)
return resp, nil
}
}

func (gpt *ChatGPT) getApiKey() string {
func (gpt *ChatGPT) getApiKey() (string, error) {
// 获取可用的apikey
if len(gpt.ApiKeys) == 0 {
return "", fmt.Errorf("no api keys provided")
}
if len(gpt.ApiKeys) == 1 {
// 如果只有一个apikey，则直接返回该apikey
return gpt.ApiKeys[0]
} else {
totalWeight := 0
for _, weight := range gpt.apiKeyWeights {
totalWeight += weight
}
randNum := rand.Intn(totalWeight)
for i, weight := range gpt.apiKeyWeights {
if randNum < weight {
gpt.currentApiKeyIndex = i
break
}
randNum -= weight
return gpt.ApiKeys[0], nil
}
totalWeight := 0
for _, weight := range gpt.apiKeyWeights {
totalWeight += weight
}
randNum := rand.Intn(totalWeight)
for i, weight := range gpt.apiKeyWeights {
if randNum < weight {
gpt.currentApiKeyIndex = i
break
}
return gpt.ApiKeys[gpt.currentApiKeyIndex]
randNum -= weight
}
return gpt.ApiKeys[gpt.currentApiKeyIndex], nil
}

func (gpt *ChatGPT) calculateApiKeyWeights() {
Expand All @@ -195,65 +198,92 @@ func (gpt *ChatGPT) calculateApiKeyWeights() {
}
gpt.apiKeyWeights = make([]int, len(gpt.ApiKeys))
for i, apiKey := range gpt.ApiKeys {
usage, ok := gpt.apiKeyUsage[apiKey]
if ok {
if usage >= 0 {
gpt.apiKeyWeights[i] = totalUsage - usage + 1
} else {
gpt.apiKeyWeights[i] = 0
}
usage := gpt.apiKeyUsage[apiKey]
if usage >= 0 {
gpt.apiKeyWeights[i] = totalUsage - usage + 1
} else {
gpt.apiKeyWeights[i] = 0
}
}
}

func (gpt *ChatGPT) initApiKeyUsage() {
// 初始化apikey的调用次数和权重
gpt.apiKeyUsage = make(map[string]int)
gpt.apiKeyWeights = make([]int, len(gpt.ApiKeys))
for i := range gpt.ApiKeys {
gpt.apiKeyWeights[i] = 1
for _, apiKey := range gpt.ApiKeys {
gpt.apiKeyUsage[apiKey] = 0
}
}

func (gpt *ChatGPT) checkApiKeyAvailability() {
for _, apiKey := range gpt.ApiKeys {
// 检查apikey的可用性
req, err := http.NewRequest("GET", BASEURL+"engines", nil)
if err != nil {
log.Printf("check api key %s failed: %v", apiKey, err)
continue
func (gpt *ChatGPT) checkApiKeyAvailability(apiKey string) bool {
// 检查apikey的可用性
req, err := http.NewRequest("GET", BASEURL+"engines", nil)
if err != nil {
log.Printf("check api key %s failed: %v", apiKey, err)
return false
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+apiKey)
client := &http.Client{Timeout: 10 * time.Second}
response, err := client.Do(req)
if err != nil {
log.Printf("check api key %s failed: %v", apiKey, err)
return false
}
if response.StatusCode/2 != 100 {
log.Printf("api key %s is not available: %s", apiKey, response.Status)
if response.StatusCode == 429 {
gpt.apiKeyUsage[apiKey] = -1
log.Printf("api key %s has been permanently disabled due to too many requests", apiKey)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+apiKey)
client := &http.Client{Timeout: 10 * time.Second}
response, err := client.Do(req)
if err != nil {
log.Printf("check api key %s failed: %v", apiKey, err)
return false
}
return true
}

func (gpt *ChatGPT) setApiKeyAvailability() {
// 检查apikey的可用性
for apiKey := range gpt.apiKeyUsage {
if usage := gpt.apiKeyUsage[apiKey]; usage < 0 {
// 如果apikey的调用次数小于0，则禁用它
continue
}
// 如果apikey不可用，则将其调用次数设为-1
if response.StatusCode/2 != 100 {
log.Printf("api key %s is not available: %s", apiKey, response.Status)
gpt.apiKeyUsageMutex.Lock()
if !gpt.checkApiKeyAvailability(apiKey) {
// 如果apikey不可用，则将其调用次数设为-1
gpt.apiKeyUsage[apiKey] = -1
gpt.apiKeyUsageMutex.Unlock()
} else {
// 如果apikey可用，则将其调用次数设为0
gpt.apiKeyUsageMutex.Lock()
gpt.apiKeyUsage[apiKey] = 0
gpt.apiKeyUsageMutex.Unlock()
gpt.calculateApiKeyWeights()
}
}
}

func (gpt *ChatGPT) StartApiKeyAvailabilityCheck() {
// 初始化apikey的调用次数和权重，检查apikey的可用性，并定时检查apikey的可用性
if len(gpt.ApiKeys) == 0 {
log.Fatalf("no api keys provided")
}
gpt.initApiKeyUsage()
gpt.setApiKeyAvailability()
ticker := time.NewTicker(30 * time.Minute)
go func() {
for {
select {
case <-ticker.C:
gpt.setApiKeyAvailability()
}
}
}()
}

func (gpt *ChatGPT) updateApiKeyUsage(apiKey string) {
// 更新ApiKeyUsage并写入文件
gpt.apiKeyUsage[apiKey]++
gpt.calculateApiKeyWeights()
}

func (gpt *ChatGPT) loadApiKeyUsage() error {
// 从本地文件中读取apikey的调用次数和权重信息
file, err := os.Open("apikey_usage.json")
file, err := os.OpenFile("apikey_usage.json", os.O_RDWR|os.O_CREATE, 0644)
if err != nil {
if os.IsNotExist(err) {
gpt.initApiKeyUsage()
return nil
}
return err
}
defer file.Close()
Expand All @@ -263,6 +293,11 @@ func (gpt *ChatGPT) loadApiKeyUsage() error {
}
if fileInfo.Size() == 0 {
gpt.initApiKeyUsage()
gpt.calculateApiKeyWeights()
err = gpt.saveApiKeyUsage(file)
if err != nil {
log.Printf("save api key usage failed: %v", err)
}
return nil
}
decoder := json.NewDecoder(file)
Expand All @@ -274,67 +309,18 @@ func (gpt *ChatGPT) loadApiKeyUsage() error {
return nil
}

func (gpt *ChatGPT) saveApiKeyUsage() error {
func (gpt *ChatGPT) saveApiKeyUsage(file *os.File) error {
// 将apikey的调用次数和权重信息输出到本地文件
file, err := os.Create("apikey_usage.json")
if err != nil {
return err
}
defer file.Close()
file.Truncate(0)
file.Seek(0, 0)
encoder := json.NewEncoder(file)
err = encoder.Encode(gpt.apiKeyUsage)
if err != nil {
return err
}
return nil
}

func (gpt *ChatGPT) updateApiKeyUsage() error {
// 更新ApiKeyUsage并写入文件
gpt.apiKeyUsageMutex.Lock()
defer gpt.apiKeyUsageMutex.Unlock()
gpt.apiKeyUsage[gpt.ApiKeys[gpt.currentApiKeyIndex]]++
err := gpt.saveApiKeyUsage()
err := encoder.Encode(gpt.apiKeyUsage)
if err != nil {
return err
}
return nil
}

func (gpt *ChatGPT) setApiKeyAvailability() {
// 检查apikey的可用性
gpt.checkApiKeyAvailability()
for apiKey, usage := range gpt.apiKeyUsage {
// 如果apikey的调用次数小于0，则禁用它，并输出信息
if usage < 0 {
log.Printf("api key %s is disabled", apiKey)
gpt.currentApiKeyIndex = (gpt.currentApiKeyIndex + 1) % len(gpt.ApiKeys)
}
}
}

func (gpt *ChatGPT) StartApiKeyAvailabilityCheck() {
// 初始化apikey的调用次数和权重，检查apikey的可用性，并定时检查apikey的可用性
err := gpt.loadApiKeyUsage()
if err != nil {
log.Printf("load api key usage failed: %v", err)
}
gpt.setApiKeyAvailability()
ticker := time.NewTicker(30 * time.Minute)
go func() {
for {
select {
case <-ticker.C:
gpt.setApiKeyAvailability()
err := gpt.saveApiKeyUsage()
if err != nil {
log.Printf("save api key usage failed: %v", err)
}
}
}
}()
}

type ImageGenerationRequestBody struct {
Prompt string `json:"prompt"`
N int `json:"n"`
Expand Down Expand Up @@ -369,7 +355,10 @@ func (gpt *ChatGPT) GenerateImage(prompt string, size string,

req.Header.Set("Content-Type", "application/json")
// 获取apikey并设置到header中
apiKey := gpt.getApiKey()
apiKey, err := gpt.getApiKey()
if err != nil {
return nil, err
}
req.Header.Set("Authorization", "Bearer "+apiKey)
client := &http.Client{Timeout: 110 * time.Second}
response, err := client.Do(req)
Expand Down

0 comments on commit 0d98cec

Please sign in to comment.