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

fix: now the input field can be array type now (close #149)

Browse Source
This commit is contained in:
JustSong 2023-06-12 16:11:57 +08:00
parent 8b2ef666ef
commit 7c7eb6b7ec
2 changed files with 16 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
controller/relay-utils.go
View File

@ -58,6 +58,20 @@ func countTokenMessages(messages []Message, model string) int {
return tokenNum
}
func countTokenInput(input any, model string) int {
switch input.(type) {
case string:
return countTokenText(input.(string), model)
case []string:
text := ""
for _, s := range input.([]string) {
text += s
}
return countTokenText(text, model)
}
return 0
}
func countTokenText(text string, model string) int {
tokenEncoder := getTokenEncoder(model)
token := tokenEncoder.Encode(text, nil, nil)

4
controller/relay.go
View File

@ -38,7 +38,7 @@ type GeneralOpenAIRequest struct {
Temperature float64 `json:"temperature"`
TopP float64 `json:"top_p"`
N int `json:"n"`
Input string `json:"input"`
Input any `json:"input"`
}
type ChatRequest struct {
@ -189,7 +189,7 @@ func relayHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
case RelayModeCompletions:
promptTokens = countTokenText(textRequest.Prompt, textRequest.Model)
case RelayModeModeration:
promptTokens = countTokenText(textRequest.Input, textRequest.Model)
promptTokens = countTokenInput(textRequest.Input, textRequest.Model)
}
preConsumedTokens := common.PreConsumedQuota
if textRequest.MaxTokens != 0 {