Go: add LLM usage (#17049)

### Summary

```
RAGFlow(api/default)> CHAT WITH 'glm-4-flash@new_test@zhipu-ai' MESSAGE '30 words describes LLM';
Answer: Hello! I'm ChatGLM, an AI assistant. Feel free to ask me any questions or request help with any tasks.
Input tokens: 5
Output tokens: 28
Time: 12.748241
```

Signed-off-by: Jin Hai <haijin.chn@gmail.com>

internal/cli/response.go

@@ -19,6 +19,7 @@ package cli
 import (
 	"encoding/json"
 	"fmt"
+	"ragflow/internal/entity/models"
 	"strings"
 )
 
@@ -529,10 +530,11 @@ func (r *MessageResponse) PrintOut() {
 }
 
 type NonStreamResponse struct {
-	Code             int    `json:"code"`
-	ReasoningContent string `json:"reasoning_content"`
-	Answer           string `json:"answer"`
-	Message          string `json:"message"`
+	Code             int               `json:"code"`
+	ReasoningContent string            `json:"reasoning_content"`
+	Answer           string            `json:"answer"`
+	Message          string            `json:"message"`
+	Usage            *models.ChatUsage `json:"usage"`
 	Duration         float64
 	OutputFormat     OutputFormat
 }
@@ -555,6 +557,10 @@ func (r *NonStreamResponse) PrintOut() {
 			fmt.Printf("Thinking: %s\n", r.ReasoningContent)
 		}
 		fmt.Printf("Answer: %s\n", r.Answer)
+		if r.Usage != nil {
+			fmt.Printf("Input tokens: %v\n", r.Usage.PromptTokens)
+			fmt.Printf("Output tokens: %v\n", r.Usage.CompletionTokens)
+		}
 		fmt.Printf("Time: %f\n", r.Duration)
 	} else {
 		fmt.Println("ERROR")

internal/engine/clickhouse/clickhouse.go

@@ -17,14 +17,8 @@
 package clickhouse
 
 import (
-	//"context"
-	//"crypto/tls"
-	//"fmt"
-	//"log"
-
 	"context"
 	"fmt"
-	"ragflow/internal/server"
 	"sync"
 	"time"
 
@@ -47,18 +41,18 @@ type Driver struct {
 	config *clickhouse.Options
 }
 
-func Init(config *server.ClickhouseConfig, ctx context.Context) error {
+func Init(ctx context.Context, host string, port int, user, password, database string) error {
 	var err error
 	once.Do(func() {
-		address := fmt.Sprintf("%s:%d", config.Host, config.Port)
+		address := fmt.Sprintf("%s:%d", host, port)
 
 		globalDriver = &Driver{}
 		globalDriver.config = &clickhouse.Options{
 			Addr: []string{address},
 			Auth: clickhouse.Auth{
-				Database: config.Database,
-				Username: config.User,
-				Password: config.Password,
+				Database: database,
+				Username: user,
+				Password: password,
 			},
 			DialTimeout:     5 * time.Second,
 			MaxOpenConns:    10,

internal/entity/models/base_model.go

@@ -193,6 +193,7 @@ func NewDriverHTTPClient() *http.Client {
 	t.IdleConnTimeout = 90 * time.Second
 	t.DisableCompression = false
 	t.ResponseHeaderTimeout = 60 * time.Second
+	t.TLSHandshakeTimeout = 30 * time.Second
 	return &http.Client{Transport: t}
 }
 

internal/entity/models/types.go

@@ -177,6 +177,28 @@ type ChatConfig struct {
 	StreamCallback func(contentDelta, reasoningDelta string) `json:"-"`
 }
 
+type ChatAppConfig struct {
+	UserID         string  `json:"user_id"`
+	UserEmail      string  `json:"user_email"`
+	TenantID       string  `json:"tenant_id"`
+	TenantEmail    string  `json:"tenant_email"`
+	GroupID        *string `json:"group_id"`
+	GroupName      *string `json:"group_name"`
+	AppID          string  `json:"app_id"`
+	AppName        string  `json:"app_name"`
+	SessionID      *string `json:"session_id"`
+	ProviderName   string  `json:"provider_name"`
+	InstanceID     string  `json:"instance_id"`
+	ModelName      string  `json:"model_name"`
+	Type           string  `json:"type"`
+	InputTokens    uint64  `json:"input_tokens"`
+	OutputTokens   uint64  `json:"output_tokens"`
+	TotalTokens    uint64  `json:"total_tokens"`
+	RequestID      *string `json:"request_id"`
+	ResponseTimeMS uint64  `json:"response_time_ms"`
+	ErrorMessage   *string `json:"error_message"`
+}
+
 type APIConfig struct {
 	ApiKey  *string
 	Region  *string

internal/entity/models/zhipu-ai.go

@@ -55,6 +55,31 @@ func (z *ZhipuAIModel) Name() string {
 	return "zhipu"
 }
 
+type ZhipuChatResponse struct {
+	Choices []struct {
+		FinishReason string `json:"finish_reason"`
+		Index        int    `json:"index"`
+		Message      struct {
+			Content          string `json:"content"`
+			ReasoningContent string `json:"reasoning_content"`
+			Role             string `json:"role"`
+		} `json:"message"`
+	} `json:"choices"`
+	Created   int    `json:"created"`
+	Id        string `json:"id"`
+	Model     string `json:"model"`
+	Object    string `json:"object"`
+	RequestId string `json:"request_id"`
+	Usage     struct {
+		CompletionTokens    int `json:"completion_tokens"`
+		PromptTokens        int `json:"prompt_tokens"`
+		PromptTokensDetails struct {
+			CachedTokens int `json:"cached_tokens"`
+		} `json:"prompt_tokens_details"`
+		TotalTokens int `json:"total_tokens"`
+	} `json:"usage"`
+}
+
 // ChatWithMessages sends multiple messages with roles and returns response
 func (z *ZhipuAIModel) ChatWithMessages(modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig) (*ChatResponse, error) {
 	if err := z.baseModel.APIConfigCheck(apiConfig); err != nil {
@@ -155,46 +180,28 @@ func (z *ZhipuAIModel) ChatWithMessages(modelName string, messages []Message, ap
 	}
 
 	// Parse response
-	var result map[string]interface{}
+	var result ZhipuChatResponse
 	if err = json.Unmarshal(body, &result); err != nil {
 		return nil, fmt.Errorf("failed to parse response: %w", err)
 	}
 
-	choices, ok := result["choices"].([]interface{})
-	if !ok || len(choices) == 0 {
-		return nil, fmt.Errorf("no choices in response")
-	}
+	content := &result.Choices[0].Message.Content
 
-	firstChoice, ok := choices[0].(map[string]interface{})
-	if !ok {
-		return nil, fmt.Errorf("invalid choice format")
-	}
-
-	messageMap, ok := firstChoice["message"].(map[string]interface{})
-	if !ok {
-		return nil, fmt.Errorf("invalid message format")
-	}
-
-	content, ok := messageMap["content"].(string)
-	if !ok {
-		return nil, fmt.Errorf("invalid content format")
-	}
-
-	var reasonContent string
+	var reasonContent *string
 	if chatModelConfig != nil && chatModelConfig.Thinking != nil && *chatModelConfig.Thinking {
-		reasonContent, ok = messageMap["reasoning_content"].(string)
-		if !ok {
-			return nil, fmt.Errorf("invalid content format")
-		}
-		// if first char of reasonContent is \n remove the '\n'
-		if reasonContent != "" && reasonContent[0] == '\n' {
-			reasonContent = reasonContent[1:]
-		}
+		reasonContent = &result.Choices[0].Message.ReasoningContent
+	}
+
+	usage := &ChatUsage{
+		PromptTokens:     result.Usage.PromptTokens,
+		CompletionTokens: result.Usage.CompletionTokens,
+		TotalTokens:      result.Usage.TotalTokens,
 	}
 
 	chatResponse := &ChatResponse{
-		Answer:        &content,
-		ReasonContent: &reasonContent,
+		Answer:        content,
+		ReasonContent: reasonContent,
+		Usage:         usage,
 	}
 
 	return chatResponse, nil

internal/handler/providers.go

@@ -820,6 +820,7 @@ func (h *ProviderHandler) ChatToModel(c *gin.Context) {
 		"code":              0,
 		"reasoning_content": response.ReasonContent,
 		"answer":            response.Answer,
+		"usage":             response.Usage,
 	})
 }
 

internal/service/toc_enhancer.go

@@ -495,7 +495,7 @@ func (e *TOCEnhancer) scoreEntries(ctx context.Context, entries []tocEntry, limi
 		if rerr != nil {
 			repaired = raw
 		}
-		if err := json.Unmarshal([]byte(repaired), &scores); err != nil {
+		if err = json.Unmarshal([]byte(repaired), &scores); err != nil {
 			lastErr = err.Error()
 			common.Warn("TOC enhancer: JSON parse failed, retrying",
 				zap.Error(err), zap.Int("attempt", attempt))