78 lines
2.7 KiB
Go
78 lines
2.7 KiB
Go
|
package main
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"encoding/json"
|
||
|
|
"errors"
|
||
|
|
"strings"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/models/chat"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
)
|
||
|
|
|
||
|
|
type fakeProbeChat struct {
|
||
|
|
chatDelegate
|
||
|
|
response *types.ChatResponse
|
||
|
|
err error
|
||
|
|
}
|
||
|
|
|
||
|
|
func (f *fakeProbeChat) Chat(context.Context, []chat.Message, *chat.ChatOptions) (*types.ChatResponse, error) {
|
||
|
|
return f.response, f.err
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestRecordingChatPreservesPerCallDiagnosticsWithoutContent(t *testing.T) {
|
||
|
|
fake := &fakeProbeChat{response: &types.ChatResponse{
|
||
|
|
Content: "answer", ReasoningContent: "private reasoning", FinishReason: "length",
|
||
|
|
Usage: types.TokenUsage{PromptTokens: 10, CompletionTokens: 4, TotalTokens: 14},
|
||
|
|
}}
|
||
|
|
recorder := &recordingChat{chatDelegate: fake}
|
||
|
|
thinking := false
|
||
|
|
options := &chat.ChatOptions{MaxTokens: 4096, Thinking: &thinking}
|
||
|
|
first, err := recorder.Chat(context.Background(), nil, options)
|
||
|
|
if err != nil || first != fake.response {
|
||
|
|
t.Fatal("recorder changed the response")
|
||
|
|
}
|
||
|
|
thinking = true
|
||
|
|
fake.response = &types.ChatResponse{Content: "\u4e2d\u6587", FinishReason: "stop"}
|
||
|
|
if _, err := recorder.Chat(context.Background(), nil, options); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if recorder.calls != 2 || len(recorder.records) != 2 {
|
||
|
|
t.Fatal("missing call diagnostics")
|
||
|
|
}
|
||
|
|
a, b := recorder.records[0], recorder.records[1]
|
||
|
|
if a.Number != 1 || a.MaxTokens != 4096 || a.Thinking == nil || *a.Thinking ||
|
||
|
|
!a.ResponseReceived || a.Failed || a.FinishReason != "length" ||
|
||
|
|
a.ContentRunes != 6 || a.ReasoningRunes != 17 || a.Usage.TotalTokens != 14 {
|
||
|
|
t.Fatalf("unexpected first call metadata: %+v", a)
|
||
|
|
}
|
||
|
|
if b.Number != 2 || b.Thinking == nil || !*b.Thinking || b.ContentRunes != 2 || b.ReasoningRunes != 0 {
|
||
|
|
t.Fatalf("unexpected second call metadata: %+v", b)
|
||
|
|
}
|
||
|
|
encoded, err := json.Marshal(recorder.records)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if strings.Contains(string(encoded), "private reasoning") || strings.Contains(string(encoded), "answer") {
|
||
|
|
t.Fatal("diagnostics exposed response content")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestRecordingChatRecordsNoResponseWithoutLeakingError(t *testing.T) {
|
||
|
|
want := errors.New("private upstream error")
|
||
|
|
recorder := &recordingChat{chatDelegate: &fakeProbeChat{err: want}}
|
||
|
|
response, err := recorder.Chat(context.Background(), nil, nil)
|
||
|
|
if response != nil || err != want || len(recorder.records) != 1 {
|
||
|
|
t.Fatal("recorder changed the failure")
|
||
|
|
}
|
||
|
|
record := recorder.records[0]
|
||
|
|
if record.ResponseReceived || !record.Failed || record.Thinking != nil || record.ContentRunes != 0 || record.ReasoningRunes != 0 {
|
||
|
|
t.Fatalf("unexpected failure metadata: %+v", record)
|
||
|
|
}
|
||
|
|
encoded, _ := json.Marshal(record)
|
||
|
|
if strings.Contains(string(encoded), want.Error()) {
|
||
|
|
t.Fatal("diagnostics exposed upstream error")
|
||
|
|
}
|
||
|
|
}
|