zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-08-07 08:01:13 +08:00
Code Issues Packages Projects Releases Wiki Activity

Go: add stats (#17061)

Browse Source 
### 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
Download Patch File Download Diff File Expand all files Collapse all files

30
internal/entity/models/astraflow_test.go
View File

@@ -124,6 +124,7 @@ func TestAstraflowChatHappyPath(t *testing.T) {
[]Message{{Role: "user", Content: "ping"}},
&APIConfig{ApiKey: &apiKey},
&ChatConfig{MaxTokens: &mt, Temperature: &temp},
nil,
)
if err != nil {
t.Fatalf("ChatWithMessages: %v", err)
@@ -150,7 +151,7 @@ func TestAstraflowChatNoReasoning(t *testing.T) {
resp, err := newAstraflowForTest(srv.URL).ChatWithMessages(
"gpt-4o-mini",
[]Message{{Role: "user", Content: "hi"}},
&APIConfig{ApiKey: &apiKey}, nil)
&APIConfig{ApiKey: &apiKey}, nil, nil)
if err != nil {
t.Fatalf("Chat: %v", err)
}
@@ -166,7 +167,7 @@ func TestAstraflowChatRequiresAPIKey(t *testing.T) {
_, err := newAstraflowForTest("http://unused").ChatWithMessages(
"claude-opus-4-7",
[]Message{{Role: "user", Content: "x"}},
&APIConfig{}, nil)
&APIConfig{}, nil, nil)
if err == nil || !strings.Contains(err.Error(), "api key is required") {
t.Errorf("expected api-key error, got %v", err)
}
@@ -175,7 +176,7 @@ func TestAstraflowChatRequiresAPIKey(t *testing.T) {
func TestAstraflowChatRequiresMessages(t *testing.T) {
apiKey := "test-key"
_, err := newAstraflowForTest("http://unused").ChatWithMessages(
"claude-opus-4-7", nil, &APIConfig{ApiKey: &apiKey}, nil)
"claude-opus-4-7", nil, &APIConfig{ApiKey: &apiKey}, nil, nil)
if err == nil || !strings.Contains(err.Error(), "messages is empty") {
t.Errorf("expected messages-empty error, got %v", err)
}
@@ -192,7 +193,7 @@ func TestAstraflowChatPropagatesHTTPError(t *testing.T) {
_, err := newAstraflowForTest(srv.URL).ChatWithMessages(
"claude-opus-4-7",
[]Message{{Role: "user", Content: "x"}},
&APIConfig{ApiKey: &apiKey}, nil)
&APIConfig{ApiKey: &apiKey}, nil, nil)
if err == nil || !strings.Contains(err.Error(), "401") {
t.Errorf("expected 401 propagated, got %v", err)
}
@@ -213,7 +214,7 @@ func TestAstraflowStreamHappyPath(t *testing.T) {
err := newAstraflowForTest(srv.URL).ChatStreamlyWithSender(
"claude-opus-4-7",
[]Message{{Role: "user", Content: "hi"}},
&APIConfig{ApiKey: &apiKey}, nil,
&APIConfig{ApiKey: &apiKey}, nil, nil,
func(c *string, _ *string) error {
if c == nil {
return nil
@@ -251,7 +252,7 @@ func TestAstraflowStreamSplitsReasoning(t *testing.T) {
err := newAstraflowForTest(srv.URL).ChatStreamlyWithSender(
"kimi-k2.6",
[]Message{{Role: "user", Content: "x"}},
&APIConfig{ApiKey: &apiKey}, nil,
&APIConfig{ApiKey: &apiKey}, nil, nil,
func(c *string, r *string) error {
if c != nil && r != nil {
t.Errorf("sender called with both args non-nil")
@@ -283,6 +284,7 @@ func TestAstraflowStreamRejectsExplicitFalse(t *testing.T) {
[]Message{{Role: "user", Content: "x"}},
&APIConfig{ApiKey: &apiKey},
&ChatConfig{Stream: &stream},
nil,
func(*string, *string) error { return nil })
if err == nil || !strings.Contains(err.Error(), "stream must be true") {
t.Errorf("expected stream-true guard, got %v", err)
@@ -294,7 +296,7 @@ func TestAstraflowStreamRequiresSender(t *testing.T) {
err := newAstraflowForTest("http://unused").ChatStreamlyWithSender(
"claude-opus-4-7",
[]Message{{Role: "user", Content: "x"}},
&APIConfig{ApiKey: &apiKey}, nil, nil)
&APIConfig{ApiKey: &apiKey}, nil, nil, nil)
if err == nil || !strings.Contains(err.Error(), "sender is required") {
t.Errorf("expected sender-required error, got %v", err)
}
@@ -310,7 +312,7 @@ func TestAstraflowStreamFailsWithoutTerminal(t *testing.T) {
err := newAstraflowForTest(srv.URL).ChatStreamlyWithSender(
"claude-opus-4-7",
[]Message{{Role: "user", Content: "x"}},
&APIConfig{ApiKey: &apiKey}, nil,
&APIConfig{ApiKey: &apiKey}, nil, nil,
func(*string, *string) error { return nil })
if err == nil || !strings.Contains(err.Error(), "stream ended before") {
t.Errorf("expected truncation error, got %v", err)
@@ -328,7 +330,7 @@ func TestAstraflowStreamRejectsMalformedFrame(t *testing.T) {
err := newAstraflowForTest(srv.URL).ChatStreamlyWithSender(
"claude-opus-4-7",
[]Message{{Role: "user", Content: "x"}},
&APIConfig{ApiKey: &apiKey}, nil,
&APIConfig{ApiKey: &apiKey}, nil, nil,
func(*string, *string) error { return nil })
if err == nil || !strings.Contains(err.Error(), "invalid SSE event") {
t.Errorf("expected invalid-SSE error, got %v", err)
@@ -346,7 +348,7 @@ func TestAstraflowStreamSurfacesUpstreamError(t *testing.T) {
err := newAstraflowForTest(srv.URL).ChatStreamlyWithSender(
"claude-opus-4-7",
[]Message{{Role: "user", Content: "x"}},
&APIConfig{ApiKey: &apiKey}, nil,
&APIConfig{ApiKey: &apiKey}, nil, nil,
func(*string, *string) error { return nil })
if err == nil || !strings.Contains(err.Error(), "upstream stream error") {
t.Errorf("expected upstream-error surfacing, got %v", err)
@@ -428,7 +430,7 @@ func TestAstraflowEmbedReturnsNoSuchMethod(t *testing.T) {
// Embed IS implemented (not a stub). It should NOT be blocked by APIConfigCheck.
// With empty input texts it short-circuits to empty result (no error).
apiKey := "test-key"
embeddings, err := newAstraflowForTest("http://unused").Embed(nil, nil, &APIConfig{ApiKey: &apiKey}, nil)
embeddings, err := newAstraflowForTest("http://unused").Embed(nil, nil, &APIConfig{ApiKey: &apiKey}, nil, nil)
if err != nil || len(embeddings) != 0 {
t.Errorf("Embed: want empty result (no error), got embeddings=%v err=%v", embeddings, err)
}
@@ -439,16 +441,16 @@ func TestAstraflowAudioOCRReturnNoSuchMethod(t *testing.T) {
model := "x"
apiKey := "test-key"
// TranscribeAudio is a stub → "no such method"
if _, err := m.TranscribeAudio(&model, &model, &APIConfig{ApiKey: &apiKey}, nil); err == nil || !strings.Contains(err.Error(), "no such method") {
if _, err := m.TranscribeAudio(&model, &model, &APIConfig{ApiKey: &apiKey}, nil, nil); err == nil || !strings.Contains(err.Error(), "no such method") {
t.Errorf("TranscribeAudio: %v", err)
}
// AudioSpeech IS implemented; pass nil content to hit input validation,
// not api-key check (which would mean APIConfigCheck still blocks it).
if _, err := m.AudioSpeech(&model, nil, &APIConfig{ApiKey: &apiKey}, nil); err == nil || strings.Contains(err.Error(), "api key is required") {
if _, err := m.AudioSpeech(&model, nil, &APIConfig{ApiKey: &apiKey}, nil, nil); err == nil || strings.Contains(err.Error(), "api key is required") {
t.Errorf("AudioSpeech: expected non-api-key error, got %v", err)
}
// OCRFile is a stub → "no such method"
if _, err := m.OCRFile(&model, nil, &model, &APIConfig{ApiKey: &apiKey}, nil); err == nil || !strings.Contains(err.Error(), "no such method") {
if _, err := m.OCRFile(&model, nil, &model, &APIConfig{ApiKey: &apiKey}, nil, nil); err == nil || !strings.Contains(err.Error(), "no such method") {
t.Errorf("OCRFile: %v", err)
}
}