模块介绍 / 今日规划 / 下班日报 / 团队摘要 生成成功后没有上报,也几乎不写「AI」日志。
流结束时的 usage 用了非阻塞发送,通道一满就被丢掉;很多服务商默认还不在流里带 usage。 现在每次真正打到模型的调用都会: 在本地日志(分类 AI,消息 AI 调用完成)写下 scene / prompt / completion / cached 登录后同步上报到后台日汇总
This commit is contained in:
@@ -56,6 +56,51 @@ func TestChatStreamSSE(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
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)
|
||||
|
||||
@@ -26,9 +26,14 @@ type openAICompatible struct {
|
||||
func (p *openAICompatible) Name() string { return p.name }
|
||||
|
||||
type chatRequest struct {
|
||||
Model string `json:"model"`
|
||||
Messages []Message `json:"messages"`
|
||||
Stream bool `json:"stream"`
|
||||
Model string `json:"model"`
|
||||
Messages []Message `json:"messages"`
|
||||
Stream bool `json:"stream"`
|
||||
StreamOptions *streamOptions `json:"stream_options,omitempty"`
|
||||
}
|
||||
|
||||
type streamOptions struct {
|
||||
IncludeUsage bool `json:"include_usage"`
|
||||
}
|
||||
|
||||
type chatDelta struct {
|
||||
@@ -37,17 +42,42 @@ type chatDelta struct {
|
||||
Content string `json:"content"`
|
||||
} `json:"delta"`
|
||||
} `json:"choices"`
|
||||
Usage *struct {
|
||||
PromptTokens int64 `json:"prompt_tokens"`
|
||||
CompletionTokens int64 `json:"completion_tokens"`
|
||||
} `json:"usage"`
|
||||
Usage *chatUsage `json:"usage"`
|
||||
Error *struct {
|
||||
Message string `json:"message"`
|
||||
} `json:"error"`
|
||||
}
|
||||
|
||||
// chatUsage 兼容 OpenAI / DeepSeek / 星火:缓存命中字段名不统一。
|
||||
type chatUsage struct {
|
||||
PromptTokens int64 `json:"prompt_tokens"`
|
||||
CompletionTokens int64 `json:"completion_tokens"`
|
||||
PromptCacheHitTokens int64 `json:"prompt_cache_hit_tokens"`
|
||||
PromptTokensDetails *struct {
|
||||
CachedTokens int64 `json:"cached_tokens"`
|
||||
} `json:"prompt_tokens_details"`
|
||||
}
|
||||
|
||||
func parseUsage(u *chatUsage) *Usage {
|
||||
if u == nil {
|
||||
return nil
|
||||
}
|
||||
cached := u.PromptCacheHitTokens
|
||||
if u.PromptTokensDetails != nil && u.PromptTokensDetails.CachedTokens > cached {
|
||||
cached = u.PromptTokensDetails.CachedTokens
|
||||
}
|
||||
return &Usage{
|
||||
PromptTokens: u.PromptTokens,
|
||||
CompletionTokens: u.CompletionTokens,
|
||||
CachedTokens: cached,
|
||||
}
|
||||
}
|
||||
|
||||
func (p *openAICompatible) ChatStream(ctx context.Context, messages []Message) (<-chan Chunk, error) {
|
||||
body, e := json.Marshal(chatRequest{Model: p.model, Messages: messages, Stream: true})
|
||||
body, e := json.Marshal(chatRequest{
|
||||
Model: p.model, Messages: messages, Stream: true,
|
||||
StreamOptions: &streamOptions{IncludeUsage: true},
|
||||
})
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
@@ -100,10 +130,7 @@ func (p *openAICompatible) ChatStream(ctx context.Context, messages []Message) (
|
||||
return
|
||||
}
|
||||
if d.Usage != nil {
|
||||
lastUsage = &Usage{
|
||||
PromptTokens: d.Usage.PromptTokens,
|
||||
CompletionTokens: d.Usage.CompletionTokens,
|
||||
}
|
||||
lastUsage = parseUsage(d.Usage)
|
||||
}
|
||||
if len(d.Choices) > 0 && d.Choices[0].Delta.Content != "" {
|
||||
completionChars += len(d.Choices[0].Delta.Content)
|
||||
@@ -119,7 +146,6 @@ func (p *openAICompatible) ChatStream(ctx context.Context, messages []Message) (
|
||||
return
|
||||
}
|
||||
if lastUsage == nil {
|
||||
// 粗估:约 4 字符 ≈ 1 token
|
||||
lastUsage = &Usage{
|
||||
PromptTokens: int64((promptChars + 3) / 4),
|
||||
CompletionTokens: int64((completionChars + 3) / 4),
|
||||
@@ -128,7 +154,7 @@ func (p *openAICompatible) ChatStream(ctx context.Context, messages []Message) (
|
||||
}
|
||||
select {
|
||||
case out <- Chunk{Usage: lastUsage}:
|
||||
default:
|
||||
case <-ctx.Done():
|
||||
}
|
||||
}()
|
||||
return out, nil
|
||||
|
||||
@@ -22,9 +22,23 @@ type Chunk struct {
|
||||
type Usage struct {
|
||||
PromptTokens int64
|
||||
CompletionTokens int64
|
||||
CachedTokens int64 // 提示缓存命中(服务商返回时才有)
|
||||
Estimated bool
|
||||
}
|
||||
|
||||
// EstimateUsage 在服务商未回 usage 时按字符粗估(约 4 字符 ≈ 1 token)。
|
||||
func EstimateUsage(messages []Message, completion string) *Usage {
|
||||
n := 0
|
||||
for _, m := range messages {
|
||||
n += len(m.Content)
|
||||
}
|
||||
return &Usage{
|
||||
PromptTokens: int64((n + 3) / 4),
|
||||
CompletionTokens: int64((len(completion) + 3) / 4),
|
||||
Estimated: true,
|
||||
}
|
||||
}
|
||||
|
||||
// Provider 是 AI 服务商的统一抽象。
|
||||
type Provider interface {
|
||||
Name() string
|
||||
|
||||
Reference in New Issue
Block a user