Go: add stats (#17061)

### Summary

Add LLM token stats framework

---------

Signed-off-by: Jin Hai <haijin.chn@gmail.com>
This commit is contained in:
Jin Hai
2026-07-18 21:02:07 +08:00
committed by GitHub
parent cafbbd467d
commit b8d06d02e6
144 changed files with 1578 additions and 1280 deletions

View File

@@ -23,6 +23,7 @@ import (
"fmt"
"io"
"net/http"
"ragflow/internal/common"
"strings"
)
@@ -67,7 +68,7 @@ func (a *AstraflowModel) Name() string {
}
// ChatWithMessages sends a non-streaming chat request
func (a *AstraflowModel) ChatWithMessages(modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig) (*ChatResponse, error) {
func (a *AstraflowModel) ChatWithMessages(modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) {
if err := a.baseModel.APIConfigCheck(apiConfig); err != nil {
return nil, err
}
@@ -177,7 +178,7 @@ func (a *AstraflowModel) ChatWithMessages(modelName string, messages []Message,
}
// ChatStreamlyWithSender opens the SSE chat-completions
func (a *AstraflowModel) ChatStreamlyWithSender(modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, sender func(*string, *string) error) error {
func (a *AstraflowModel) ChatStreamlyWithSender(modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error {
if err := a.baseModel.APIConfigCheck(apiConfig); err != nil {
return err
}
@@ -352,7 +353,7 @@ func (a *AstraflowModel) CheckConnection(apiConfig *APIConfig) error {
return err
}
func (a *AstraflowModel) Embed(modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig) ([]EmbeddingData, error) {
func (a *AstraflowModel) Embed(modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig, modelUsage *common.ModelUsage) ([]EmbeddingData, error) {
if err := a.baseModel.APIConfigCheck(apiConfig); err != nil {
return nil, err
}
@@ -429,7 +430,7 @@ func (a *AstraflowModel) Embed(modelName *string, texts []string, apiConfig *API
return embeddings, nil
}
func (a *AstraflowModel) Rerank(modelName *string, query string, documents []string, apiConfig *APIConfig, rerankConfig *RerankConfig) (*RerankResponse, error) {
func (a *AstraflowModel) Rerank(modelName *string, query string, documents []string, apiConfig *APIConfig, rerankConfig *RerankConfig, modelUsage *common.ModelUsage) (*RerankResponse, error) {
if err := a.baseModel.APIConfigCheck(apiConfig); err != nil {
return nil, err
}
@@ -511,15 +512,15 @@ func (a *AstraflowModel) Balance(apiConfig *APIConfig) (map[string]interface{},
return nil, fmt.Errorf("%s, no such method", a.Name())
}
func (a *AstraflowModel) TranscribeAudio(modelName *string, file *string, apiConfig *APIConfig, asrConfig *ASRConfig) (*ASRResponse, error) {
func (a *AstraflowModel) TranscribeAudio(modelName *string, file *string, apiConfig *APIConfig, asrConfig *ASRConfig, modelUsage *common.ModelUsage) (*ASRResponse, error) {
return nil, fmt.Errorf("%s, no such method", a.Name())
}
func (a *AstraflowModel) TranscribeAudioWithSender(modelName *string, file *string, apiConfig *APIConfig, asrConfig *ASRConfig, sender func(*string, *string) error) error {
func (a *AstraflowModel) TranscribeAudioWithSender(modelName *string, file *string, apiConfig *APIConfig, asrConfig *ASRConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error {
return fmt.Errorf("%s, no such method", a.Name())
}
func (a *AstraflowModel) AudioSpeech(modelName *string, audioContent *string, apiConfig *APIConfig, ttsConfig *TTSConfig) (*TTSResponse, error) {
func (a *AstraflowModel) AudioSpeech(modelName *string, audioContent *string, apiConfig *APIConfig, ttsConfig *TTSConfig, modelUsage *common.ModelUsage) (*TTSResponse, error) {
if err := a.baseModel.APIConfigCheck(apiConfig); err != nil {
return nil, err
}
@@ -579,15 +580,15 @@ func (a *AstraflowModel) AudioSpeech(modelName *string, audioContent *string, ap
return &TTSResponse{Audio: body}, nil
}
func (a *AstraflowModel) AudioSpeechWithSender(modelName *string, audioContent *string, apiConfig *APIConfig, ttsConfig *TTSConfig, sender func(*string, *string) error) error {
func (a *AstraflowModel) AudioSpeechWithSender(modelName *string, audioContent *string, apiConfig *APIConfig, ttsConfig *TTSConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error {
return fmt.Errorf("%s, no such method", a.Name())
}
func (a *AstraflowModel) OCRFile(modelName *string, content []byte, url *string, apiConfig *APIConfig, ocrConfig *OCRConfig) (*OCRFileResponse, error) {
func (a *AstraflowModel) OCRFile(modelName *string, content []byte, url *string, apiConfig *APIConfig, ocrConfig *OCRConfig, modelUsage *common.ModelUsage) (*OCRFileResponse, error) {
return nil, fmt.Errorf("%s, no such method", a.Name())
}
func (a *AstraflowModel) ParseFile(modelName *string, content []byte, url *string, apiConfig *APIConfig, parseFileConfig *ParseFileConfig) (*ParseFileResponse, error) {
func (a *AstraflowModel) ParseFile(modelName *string, content []byte, url *string, apiConfig *APIConfig, parseFileConfig *ParseFileConfig, modelUsage *common.ModelUsage) (*ParseFileResponse, error) {
return nil, fmt.Errorf("%s, no such method", a.Name())
}