count function call tokens

common/constants.go

@@ -97,6 +97,8 @@ var BatchUpdateInterval = GetOrDefault("BATCH_UPDATE_INTERVAL", 5)
 
 var RelayTimeout = GetOrDefault("RELAY_TIMEOUT", 0) // unit is second
 
+var LogPrompt = os.Getenv("LOG_PROMPT") == "true"
+
 const (
 	RequestIdKey = "X-Oneapi-Request-Id"
 )

controller/relay-openai.go

@@ -11,9 +11,11 @@ import (
 	"strings"
 )
 
-func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*OpenAIErrorWithStatusCode, string) {
+func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*OpenAIErrorWithStatusCode, string, string, string) {
-	// 1. 因为这个是空的
 	responseText := ""
+	responseFunctionCallName := ""
+	responseFunctionCallArguments := ""
+
 	scanner := bufio.NewScanner(resp.Body)
 	scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) {
 		if atEOF && len(data) == 0 {
@@ -32,7 +34,6 @@ func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*O
 	go func() {
 		for scanner.Scan() {
 			data := scanner.Text()
-			common.LogInfo(c, "stream received: "+data)
 			if len(data) < 6 { // ignore blank line or wrong format
 				continue
 			}
@@ -52,6 +53,10 @@ func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*O
 					}
 					for _, choice := range streamResponse.Choices {
 						responseText += choice.Delta.Content
+						if choice.Delta.FunctionCall != nil {
+							responseFunctionCallName += choice.Delta.FunctionCall.Name
+							responseFunctionCallArguments += choice.Delta.FunctionCall.Arguments
+						}
 					}
 				case RelayModeCompletions:
 					var streamResponse CompletionsStreamResponse
@@ -85,10 +90,9 @@ func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*O
 	})
 	err := resp.Body.Close()
 	if err != nil {
-		return errorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), ""
+		return errorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), "", "", ""
 	}
-	common.LogInfo(c, "stream ended, responseText: "+responseText)
+	return nil, responseText, responseFunctionCallName, responseFunctionCallArguments
-	return nil, responseText
 }
 
 func openaiHandler(c *gin.Context, resp *http.Response, consumeQuota bool, promptTokens int, model string) (*OpenAIErrorWithStatusCode, *Usage) {

controller/relay-text.go

@@ -202,6 +202,8 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 	switch relayMode {
 	case RelayModeChatCompletions:
 		promptTokens = countTokenMessages(textRequest.Messages, textRequest.Model)
+		promptTokens += countTokenFunctions(textRequest.Functions, textRequest.Model)
+		promptTokens += countTokenFunctionCall(textRequest.FunctionCall, textRequest.Model)
 	case RelayModeCompletions:
 		promptTokens = countTokenInput(textRequest.Prompt, textRequest.Model)
 	case RelayModeModerations:
@@ -332,6 +334,16 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 	isStream := textRequest.Stream
 
 	if apiType != APITypeXunfei { // cause xunfei use websocket
+		if common.LogPrompt {
+			requestRaw, err := io.ReadAll(requestBody)
+			var logContent string
+			if err != nil {
+				logContent = fmt.Sprintf("failed to read request body, err: %s", err)
+			} else {
+				logContent = "request content: " + string(requestRaw)
+			}
+			common.LogInfo(c, logContent)
+		}
 		req, err = http.NewRequest(c.Request.Method, fullRequestURL, requestBody)
 		if err != nil {
 			return errorWrapper(err, "new_request_failed", http.StatusInternalServerError)
@@ -447,12 +459,19 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 	switch apiType {
 	case APITypeOpenAI:
 		if isStream {
-			err, responseText := openaiStreamHandler(c, resp, relayMode)
+			err, responseText, responseFunctionCallName, responseFunctionCallArguments := openaiStreamHandler(c, resp, relayMode)
 			if err != nil {
 				return err
 			}
 			textResponse.Usage.PromptTokens = promptTokens
 			textResponse.Usage.CompletionTokens = countTokenText(responseText, textRequest.Model)
+			if responseFunctionCallName != "" {
+				textResponse.Usage.CompletionTokens += countTokenFunctionCall(responseFunctionCallName, textRequest.Model)
+			}
+			if responseFunctionCallArguments != "" {
+				responseFunctionCallArguments = strings.Replace(responseFunctionCallArguments, "\\\"", "\"", -1)
+				textResponse.Usage.CompletionTokens += countTokenFunctionCall(responseFunctionCallArguments, textRequest.Model)
+			}
 			return nil
 		} else {
 			err, usage := openaiHandler(c, resp, consumeQuota, promptTokens, textRequest.Model)

controller/relay-utils.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"github.com/gin-gonic/gin"
 	"github.com/pkoukk/tiktoken-go"
+	"gorm.io/gorm/utils"
 	"io"
 	"net/http"
 	"one-api/common"
@@ -114,6 +115,71 @@ func countTokenText(text string, model string) int {
 	return getTokenNum(tokenEncoder, text)
 }
 
+func countTokenFunctionCall(functionCall any, model string) int {
+	tokenEncoder := getTokenEncoder(model)
+	jsonBytes, err := json.Marshal(functionCall)
+	if err != nil {
+		return 0
+	}
+	return getTokenNum(tokenEncoder, string(jsonBytes))
+}
+
+func countTokenFunctions(functions []Function, model string) int {
+	// https://community.openai.com/t/how-to-know-of-tokens-beforehand-when-i-make-function-calling-chat-history-request-witn-nodejs/289060/6
+	if len(functions) == 0 {
+		return 0
+	}
+	tokenEncoder := getTokenEncoder(model)
+
+	paramSignature := func(name string, pSpec Property, pRequired []string) string {
+		var requiredString string
+		if utils.Contains(pRequired, name) == false {
+			requiredString = "?"
+		}
+		var enumString string
+		if len(pSpec.Enum) > 0 {
+			enumValues := make([]string, len(pSpec.Enum))
+			for i, v := range pSpec.Enum {
+				enumValues[i] = fmt.Sprintf("\"%s\"", v)
+			}
+			enumString = strings.Join(enumValues, " | ")
+		} else {
+			enumString = pSpec.Type
+		}
+		signature := fmt.Sprintf("%s%s: %s, ", name, requiredString, enumString)
+		if pSpec.Description != "" {
+			signature = fmt.Sprintf("// %s\n%s", pSpec.Description, signature)
+		}
+		return signature
+	}
+
+	functionSignature := func(fSpec Function) string {
+		var params []string
+		for name, p := range fSpec.Parameters.Properties {
+			params = append(params, paramSignature(name, p, fSpec.Parameters.Required))
+		}
+		var descriptionString string
+		if fSpec.Description != "" {
+			descriptionString = fmt.Sprintf("// %s\n", fSpec.Description)
+		}
+
+		var paramString string
+		if len(params) > 0 {
+			paramString = fmt.Sprintf("_: {\n%s\n}", strings.Join(params, "\n"))
+		}
+
+		return fmt.Sprintf("%stype %s = (%s) => any;", descriptionString, fSpec.Name, paramString)
+	}
+
+	var functionSignatures []string
+	for _, f := range functions {
+		functionSignatures = append(functionSignatures, functionSignature(f))
+	}
+	functionString := fmt.Sprintf("# Tools\n\n## functions\n\nnamespace functions {\n\n%s\n\n} // namespace functions", strings.Join(functionSignatures, "\n\n"))
+
+	return getTokenNum(tokenEncoder, functionString)
+}
+
 func errorWrapper(err error, code string, statusCode int) *OpenAIErrorWithStatusCode {
 	openAIError := OpenAIError{
 		Message: err.Error(),

controller/relay.go

@@ -29,6 +29,24 @@ const (
 
 // https://platform.openai.com/docs/api-reference/chat
 
+type Property struct {
+	Type        string   `json:"type"`
+	Description string   `json:"description"`
+	Enum        []string `json:"enum"`
+}
+
+type Parameter struct {
+	Type       string              `json:"type"`
+	Properties map[string]Property `json:"properties"`
+	Required   []string            `json:"required"`
+}
+
+type Function struct {
+	Name        string    `json:"name"`
+	Description string    `json:"description"`
+	Parameters  Parameter `json:"parameters"`
+}
+
 type GeneralOpenAIRequest struct {
 	Model        string     `json:"model,omitempty"`
 	Messages     []Message  `json:"messages,omitempty"`
@@ -41,7 +59,8 @@ type GeneralOpenAIRequest struct {
 	Input        any        `json:"input,omitempty"`
 	Instruction  string     `json:"instruction,omitempty"`
 	Size         string     `json:"size,omitempty"`
-	Functions   any       `json:"functions,omitempty"`
+	Functions    []Function `json:"functions,omitempty"`
+	FunctionCall any        `json:"functioncall,omitempty"`
 }
 
 func (r GeneralOpenAIRequest) ParseInput() []string {
@@ -111,6 +130,11 @@ type TextResponse struct {
 	Error   OpenAIError `json:"error"`
 }
 
+type FunctionCall struct {
+	Name      string `json:"name,omitempty"`
+	Arguments string `json:"arguments,omitempty"`
+}
+
 type OpenAITextResponseChoice struct {
 	Index        int `json:"index"`
 	Message      `json:"message"`
@@ -148,6 +172,7 @@ type ImageResponse struct {
 type ChatCompletionsStreamResponseChoice struct {
 	Delta struct {
 		Content      string        `json:"content"`
+		FunctionCall *FunctionCall `json:"function_call,omitempty"`
 	} `json:"delta"`
 	FinishReason *string `json:"finish_reason"`
 }