faiss server

ai-chat-service/chat-server/server/app.go

@@ -149,6 +149,30 @@ func (a *app) buildChatCompletionResponse(msg string) *proto.ChatCompletionRespo
 	}
 }
 
+func (a *app) countMessageTokens(role, content string) (int, error) {
+	message := openai.ChatCompletionMessage{
+		Role:    role,
+		Content: content,
+	}
+	return tokenizer.GetTokens(&message, a.openaiConf.Model)
+}
+
+func (a *app) buildUsage(promptTokens int, answer string) (*proto.Usage, error) {
+	completionTokens := 0
+	if answer != "" {
+		tokens, err := a.countMessageTokens(openai.ChatMessageRoleAssistant, answer)
+		if err != nil {
+			return nil, err
+		}
+		completionTokens = tokens
+	}
+	return &proto.Usage{
+		PromptTokens:     int32(promptTokens),
+		CompletionTokens: int32(completionTokens),
+		TotalTokens:      int32(promptTokens + completionTokens),
+	}, nil
+}
+
 func (a *app) buildChatCompletionStreamResponse(id, delta, finishReason string) *proto.ChatCompletionStreamResponse {
 	return &proto.ChatCompletionStreamResponse{
 		Id:      id,