e71ce451ff
流结束时的 usage 用了非阻塞发送，通道一满就被丢掉；很多服务商默认还不在流里带 usage。 现在每次真正打到模型的调用都会： 在本地日志（分类 AI，消息 AI 调用完成）写下 scene / prompt / completion / cached 登录后同步上报到后台日汇总
114 lines
3.5 KiB
Go
114 lines
3.5 KiB
Go
package ai
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
func TestFactory(t *testing.T) {
|
|
if _, e := New("bogus", "k"); e == nil || e.Error() != "AI_PROVIDER_INVALID" {
|
|
t.Fatalf("expected AI_PROVIDER_INVALID, got %v", e)
|
|
}
|
|
if _, e := New(ProviderSpark, ""); e == nil || e.Error() != "AI_KEY_MISSING" {
|
|
t.Fatalf("expected AI_KEY_MISSING, got %v", e)
|
|
}
|
|
p, e := New(ProviderDeepSeek, "key")
|
|
if e != nil || p.Name() != ProviderDeepSeek {
|
|
t.Fatalf("unexpected: %v %v", p, e)
|
|
}
|
|
p, e = New(ProviderSpark, "key")
|
|
if e != nil || p.Name() != ProviderSpark {
|
|
t.Fatalf("unexpected: %v %v", p, e)
|
|
}
|
|
}
|
|
|
|
func TestChatStreamSSE(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path != "/chat/completions" {
|
|
t.Errorf("unexpected path %s", r.URL.Path)
|
|
}
|
|
if got := r.Header.Get("Authorization"); got != "Bearer test-key" {
|
|
t.Errorf("unexpected auth header %q", got)
|
|
}
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"你好\"}}]}\n\n"))
|
|
w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"，世界\"}}]}\n\n"))
|
|
w.Write([]byte("data: [DONE]\n\n"))
|
|
}))
|
|
defer srv.Close()
|
|
p := &openAICompatible{name: "test", baseURL: srv.URL, model: "m", apiKey: "test-key", client: srv.Client()}
|
|
stream, e := p.ChatStream(context.Background(), []Message{{Role: "user", Content: "hi"}})
|
|
if e != nil {
|
|
t.Fatal(e)
|
|
}
|
|
var sb strings.Builder
|
|
for c := range stream {
|
|
if c.Err != nil {
|
|
t.Fatal(c.Err)
|
|
}
|
|
sb.WriteString(c.Content)
|
|
}
|
|
if sb.String() != "你好，世界" {
|
|
t.Fatalf("unexpected content %q", sb.String())
|
|
}
|
|
}
|
|
|
|
func TestParseUsageCachedTokens(t *testing.T) {
|
|
u := parseUsage(&chatUsage{
|
|
PromptTokens: 100,
|
|
CompletionTokens: 20,
|
|
PromptCacheHitTokens: 80,
|
|
})
|
|
if u == nil || u.CachedTokens != 80 || u.PromptTokens != 100 {
|
|
t.Fatalf("deepseek cache field: %+v", u)
|
|
}
|
|
u = parseUsage(&chatUsage{
|
|
PromptTokens: 90,
|
|
CompletionTokens: 10,
|
|
PromptTokensDetails: &struct {
|
|
CachedTokens int64 `json:"cached_tokens"`
|
|
}{CachedTokens: 64},
|
|
})
|
|
if u == nil || u.CachedTokens != 64 {
|
|
t.Fatalf("openai cached_tokens: %+v", u)
|
|
}
|
|
}
|
|
|
|
func TestChatStreamSendsUsage(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"ok\"}}]}\n\n"))
|
|
w.Write([]byte("data: {\"usage\":{\"prompt_tokens\":12,\"completion_tokens\":3,\"prompt_cache_hit_tokens\":8}}\n\n"))
|
|
w.Write([]byte("data: [DONE]\n\n"))
|
|
}))
|
|
defer srv.Close()
|
|
p := &openAICompatible{name: "test", baseURL: srv.URL, model: "m", apiKey: "k", client: srv.Client()}
|
|
stream, e := p.ChatStream(context.Background(), []Message{{Role: "user", Content: "hi"}})
|
|
if e != nil {
|
|
t.Fatal(e)
|
|
}
|
|
var usage *Usage
|
|
for c := range stream {
|
|
if c.Usage != nil {
|
|
usage = c.Usage
|
|
}
|
|
}
|
|
if usage == nil || usage.PromptTokens != 12 || usage.CachedTokens != 8 {
|
|
t.Fatalf("usage not delivered: %+v", usage)
|
|
}
|
|
}
|
|
|
|
func TestChatStreamAuthError(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.WriteHeader(http.StatusUnauthorized)
|
|
}))
|
|
defer srv.Close()
|
|
p := &openAICompatible{name: "test", baseURL: srv.URL, model: "m", apiKey: "bad", client: srv.Client()}
|
|
if _, e := p.ChatStream(context.Background(), nil); e == nil || e.Error() != "AI_KEY_INVALID" {
|
|
t.Fatalf("expected AI_KEY_INVALID, got %v", e)
|
|
}
|
|
}
|