faiss server

ai-chat-backend/pkg/controllers/chat.go

@@ -45,6 +45,18 @@ type ChatMessageRequestOptions struct {
 	ParentMessageId string `json:"parentMessageId"`
 }
 
+type TokenUsage struct {
+	PromptTokens     int32 `json:"prompt_tokens"`
+	CompletionTokens int32 `json:"completion_tokens"`
+	TotalTokens      int32 `json:"total_tokens"`
+}
+
+type MessageMeta struct {
+	Source    string      `json:"source,omitempty"`
+	TokenUsed *bool       `json:"tokenUsed,omitempty"`
+	Usage     *TokenUsage `json:"usage,omitempty"`
+}
+
 type ChatMessage struct {
 	ID              string                                              `json:"id"`
 	Text            string                                              `json:"text"`
@@ -52,6 +64,10 @@ type ChatMessage struct {
 	Name            string                                              `json:"name"`
 	Delta           string                                              `json:"delta"`
 	Detail          *ai_chat_service_proto.ChatCompletionStreamResponse `json:"detail"`
+	Usage           *TokenUsage                                         `json:"usage,omitempty"`
+	Source          string                                              `json:"source,omitempty"`
+	TokenUsed       *bool                                               `json:"tokenUsed,omitempty"`
+	Meta            *MessageMeta                                        `json:"meta,omitempty"`
 	TokenCount      int                                                 `json:"tokenCount"`
 	ParentMessageId string                                              `json:"parentMessageId"`
 }
@@ -153,6 +169,23 @@ func (chat *ChatService) ChatProcess(ctx *gin.Context) {
 			}
 			result.Detail = rsp
 		}
+		if usage := toTokenUsage(rsp.GetUsage()); usage != nil {
+			tokenUsed := usage.TotalTokens > 0
+			result.Usage = usage
+			result.TokenUsed = &tokenUsed
+			if result.Meta == nil {
+				result.Meta = &MessageMeta{}
+			}
+			result.Meta.Usage = usage
+			result.Meta.TokenUsed = &tokenUsed
+		}
+		if rsp.GetSource() != "" {
+			result.Source = rsp.GetSource()
+			if result.Meta == nil {
+				result.Meta = &MessageMeta{}
+			}
+			result.Meta.Source = rsp.GetSource()
+		}
 
 		bts, err := json.Marshal(result)
 		if err != nil {
@@ -180,6 +213,17 @@ func (chat *ChatService) ChatProcess(ctx *gin.Context) {
 	}
 }
 
+func toTokenUsage(usage *ai_chat_service_proto.Usage) *TokenUsage {
+	if usage == nil {
+		return nil
+	}
+	return &TokenUsage{
+		PromptTokens:     usage.GetPromptTokens(),
+		CompletionTokens: usage.GetCompletionTokens(),
+		TotalTokens:      usage.GetTotalTokens(),
+	}
+}
+
 func (chat *ChatService) topP() float32 {
 	model := strings.ToLower(chat.config.Chat.Model)
 	if strings.HasPrefix(model, "kimi-") || strings.HasPrefix(model, "moonshot-") {