Ecosystem/ai-gateway
5
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: Ensure concurrency safety for tokenEncoder caching (close #459)

Browse Source
This commit is contained in:
shao0222 2023-08-22 19:09:55 +08:00
parent efeb9a16ce
commit b23870cb7b
1 changed files with 5 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
controller/relay-utils.go
View File

@ -2,6 +2,7 @@ package controller
import ( import (
"fmt" "fmt"
"sync"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/pkoukk/tiktoken-go" "github.com/pkoukk/tiktoken-go"
"one-api/common" "one-api/common"
@ -10,8 +11,12 @@ import (
var stopFinishReason = "stop" var stopFinishReason = "stop"
var tokenEncoderMap = map[string]*tiktoken.Tiktoken{} var tokenEncoderMap = map[string]*tiktoken.Tiktoken{}
var tokenEncoderMapMux sync.Mutex
func getTokenEncoder(model string) *tiktoken.Tiktoken { func getTokenEncoder(model string) *tiktoken.Tiktoken {
tokenEncoderMapMux.Lock()
defer tokenEncoderMapMux.Unlock()
if tokenEncoder, ok := tokenEncoderMap[model]; ok { if tokenEncoder, ok := tokenEncoderMap[model]; ok {
return tokenEncoder return tokenEncoder
} }