e71ce451ff
流结束时的 usage 用了非阻塞发送，通道一满就被丢掉；很多服务商默认还不在流里带 usage。 现在每次真正打到模型的调用都会： 在本地日志（分类 AI，消息 AI 调用完成）写下 scene / prompt / completion / cached 登录后同步上报到后台日汇总
48 lines
1.4 KiB
Go
48 lines
1.4 KiB
Go
// Package ai 定义 AI 聊天 Provider 抽象与工厂：
|
||
// 上层只依赖 Provider 接口，通过 New 按设置创建讯飞星火 Lite 或 DeepSeek 实例。
|
||
package ai
|
||
|
||
import "context"
|
||
|
||
// Message 是一条对话消息（role: system | user | assistant）。
|
||
type Message struct {
|
||
Role string `json:"role"`
|
||
Content string `json:"content"`
|
||
}
|
||
|
||
// Chunk 是流式返回的一段增量内容；Err 非空表示流异常中止。
|
||
// Usage 在流结束时可选附带 token 用量（若服务商提供）。
|
||
type Chunk struct {
|
||
Content string
|
||
Err error
|
||
Usage *Usage
|
||
}
|
||
|
||
// Usage 一次调用的 token 用量。
|
||
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
|
||
// ChatStream 发起流式对话，返回增量内容通道；通道关闭即流结束。
|
||
ChatStream(ctx context.Context, messages []Message) (<-chan Chunk, error)
|
||
}
|