From 8ebdc02cf62de42cefe7d1ce6b4ba702e98911a3 Mon Sep 17 00:00:00 2001 From: Jin Hai Date: Fri, 17 Jul 2026 21:28:43 +0800 Subject: [PATCH] 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 --- internal/cli/response.go | 14 +++-- internal/engine/clickhouse/clickhouse.go | 16 ++---- internal/entity/models/base_model.go | 1 + internal/entity/models/types.go | 22 ++++++++ internal/entity/models/zhipu-ai.go | 69 +++++++++++++----------- internal/handler/providers.go | 1 + internal/service/toc_enhancer.go | 2 +- 7 files changed, 78 insertions(+), 47 deletions(-) diff --git a/internal/cli/response.go b/internal/cli/response.go index c90c030e18..7fff66b8fc 100644 --- a/internal/cli/response.go +++ b/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") diff --git a/internal/engine/clickhouse/clickhouse.go b/internal/engine/clickhouse/clickhouse.go index 3848598f33..5fe5029980 100644 --- a/internal/engine/clickhouse/clickhouse.go +++ b/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, diff --git a/internal/entity/models/base_model.go b/internal/entity/models/base_model.go index c3ab60c70b..5dcde22c4d 100644 --- a/internal/entity/models/base_model.go +++ b/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} } diff --git a/internal/entity/models/types.go b/internal/entity/models/types.go index 38b56a2301..9d9730073e 100644 --- a/internal/entity/models/types.go +++ b/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 diff --git a/internal/entity/models/zhipu-ai.go b/internal/entity/models/zhipu-ai.go index 70db58f7d0..e108e25a37 100644 --- a/internal/entity/models/zhipu-ai.go +++ b/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 diff --git a/internal/handler/providers.go b/internal/handler/providers.go index 6ff4c9592e..8f35a1c6b7 100644 --- a/internal/handler/providers.go +++ b/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, }) } diff --git a/internal/service/toc_enhancer.go b/internal/service/toc_enhancer.go index 84c7c2856a..ef0fd4c5c5 100644 --- a/internal/service/toc_enhancer.go +++ b/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))