Sakurasan
b2e500ff51
- 计费：19 个金额字段 float64→decimal.Decimal，DB 列保持 numeric(20,8) 零迁移； API 金额仍为 JSON number（MarshalJSONWithoutQuotes），前端零改动 - failover：保留时间预算，新增 OT_PROXY_FAILOVER_CAPS_NET/SERVER/RATE/QUOTA 按失败类型封顶换渠道次数（429/配额类瞬时失败时间预算形同虚设） - 预聚合：新增 usage_channel_hourlies（schema v10），记账事务内同步 upsert， 长窗口渠道统计走桶不再扫明细，spent_last_30d 改扫 usage_daily，v10 一次性回填幂等 - 守护：internal/arch_test.go 用 go/parser 强制分层依赖（fail-closed）
1149 lines
40 KiB
Go
1149 lines
40 KiB
Go
package proxy
|
||
|
||
import (
|
||
"bufio"
|
||
"bytes"
|
||
"context"
|
||
"crypto/rand"
|
||
"encoding/hex"
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"io"
|
||
"net/http"
|
||
"strconv"
|
||
"strings"
|
||
"time"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
"github.com/openteam/server/internal/channel"
|
||
"github.com/openteam/server/internal/pkg/tokenizer"
|
||
"github.com/openteam/server/internal/proxy/convert"
|
||
"github.com/openteam/server/internal/store"
|
||
"github.com/shopspring/decimal"
|
||
)
|
||
|
||
// newTraceID 生成请求 trace（用于日志与记账幂等 ref）。
|
||
func newTraceID() string {
|
||
b := make([]byte, 8)
|
||
_, _ = rand.Read(b)
|
||
return hex.EncodeToString(b)
|
||
}
|
||
|
||
// bodyReq 统一取出请求体并解析 model / stream 字段。
|
||
type bodyReq struct {
|
||
Model string `json:"model"`
|
||
Stream bool `json:"stream"`
|
||
}
|
||
|
||
// parseBody 读取并回填请求体，解析 model/stream。
|
||
func parseBody(c *gin.Context) (*bodyReq, []byte, error) {
|
||
body, err := io.ReadAll(c.Request.Body)
|
||
if err != nil {
|
||
return nil, nil, err
|
||
}
|
||
c.Request.Body = io.NopCloser(bytes.NewReader(body))
|
||
// 只留存原始请求体：输入文本提取（requestText）要整包 JSON 解析，
|
||
// 而它只在「上游没报用量」时才用得到（实测 100KB 请求体 ~506µs）。
|
||
// 亲和键与估算都改成用到时再提取。
|
||
c.Set("req_body", body)
|
||
br := &bodyReq{}
|
||
_ = json.Unmarshal(body, br) // 解析失败按空处理，直通仍可转发
|
||
return br, body, nil
|
||
}
|
||
|
||
// requestText 提取请求体中的用户输入文本（chat/messages 的 content、responses 的 input/instructions），
|
||
// 用于中断时估算输入 token。
|
||
func requestText(body []byte) string {
|
||
var m map[string]any
|
||
if json.Unmarshal(body, &m) != nil {
|
||
return ""
|
||
}
|
||
var parts []string
|
||
add := func(s string) {
|
||
if s = strings.TrimSpace(s); s != "" {
|
||
parts = append(parts, s)
|
||
}
|
||
}
|
||
if s, ok := m["instructions"].(string); ok {
|
||
add(s)
|
||
}
|
||
if s, ok := m["system"].(string); ok {
|
||
add(s)
|
||
}
|
||
switch input := m["input"].(type) {
|
||
case string:
|
||
add(input)
|
||
case []any:
|
||
for _, it := range input {
|
||
if im, ok := it.(map[string]any); ok {
|
||
if s, ok := im["content"].(string); ok {
|
||
add(s)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
if msgs, ok := m["messages"].([]any); ok {
|
||
for _, msg := range msgs {
|
||
mm, ok := msg.(map[string]any)
|
||
if !ok {
|
||
continue
|
||
}
|
||
switch c := mm["content"].(type) {
|
||
case string:
|
||
add(c)
|
||
case []any:
|
||
for _, b := range c {
|
||
if bm, ok := b.(map[string]any); ok {
|
||
if s, ok := bm["text"].(string); ok {
|
||
add(s)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
return strings.Join(parts, "\n")
|
||
}
|
||
|
||
// sseDataPayload 提取一条 SSE 的 JSON 载荷（去掉 data: 前缀与空白）。
|
||
// 兼容三种写法：
|
||
// - 单 data: 行：data: {...} 或 data:{...}（上游如火山方舟会省略 data: 后的空格）
|
||
// - event:+data: 多行块：转换器 eventLine 产出的块（event: xxx\ndata: {...} 拼在一个 []byte）
|
||
func sseDataPayload(line []byte) (string, bool) {
|
||
s := string(line)
|
||
idx := strings.LastIndex(s, "\ndata:")
|
||
if idx >= 0 {
|
||
s = s[idx+len("\ndata:"):] // 跳过 event: 头，落在 data: 之后
|
||
} else if strings.HasPrefix(s, "data:") {
|
||
s = strings.TrimPrefix(s, "data:")
|
||
} else {
|
||
return "", false
|
||
}
|
||
s = strings.TrimLeft(s, " ") // data: 后的可选空格
|
||
s = strings.TrimSpace(s)
|
||
if s == "" || s == "[DONE]" {
|
||
return "", false
|
||
}
|
||
return s, true
|
||
}
|
||
|
||
// sseContentText 提取一条 SSE 中的内容文本（chat delta.content / responses delta / messages delta.text）。
|
||
func sseContentText(line []byte) string {
|
||
s, ok := sseDataPayload(line)
|
||
if !ok {
|
||
return ""
|
||
}
|
||
var m map[string]any
|
||
if json.Unmarshal([]byte(s), &m) != nil {
|
||
return ""
|
||
}
|
||
// responses output_text.delta: {"delta":"..."}
|
||
if d, ok := m["delta"].(string); ok {
|
||
return d
|
||
}
|
||
// messages content_block_delta: {"delta":{"text":"..."}}
|
||
if dm, ok := m["delta"].(map[string]any); ok {
|
||
if t, ok := dm["text"].(string); ok {
|
||
return t
|
||
}
|
||
}
|
||
// chat chunk: {"choices":[{"delta":{"content":"..."}}]}（GLM 思考模型另有 reasoning_content）
|
||
if choices, ok := m["choices"].([]any); ok && len(choices) > 0 {
|
||
if c0, ok := choices[0].(map[string]any); ok {
|
||
if delta, ok := c0["delta"].(map[string]any); ok {
|
||
if t, ok := delta["content"].(string); ok {
|
||
return t
|
||
}
|
||
if t, ok := delta["reasoning_content"].(string); ok {
|
||
return t
|
||
}
|
||
}
|
||
}
|
||
}
|
||
return ""
|
||
}
|
||
|
||
// fastContentText 从一条 SSE 行里取出增量正文，走字节扫描而不是整行 JSON 解析。
|
||
// 逐条 JSON 反序列化到 map 大约 1–2µs、二十多次分配，一份长回答几千条事件就是
|
||
// 毫秒级纯开销；这里只在三种固定键位上定位字符串值，遇到转义再回退到 JSON 解析。
|
||
//
|
||
// 覆盖三协议的增量形态：
|
||
// - chat：{"choices":[{"delta":{"content":"…"}}]}
|
||
// - messages：{"type":"content_block_delta","delta":{"text":"…"}}
|
||
// - responses：{"type":"response.output_text.delta","delta":"…"}
|
||
func fastContentText(line []byte) string {
|
||
for _, key := range []string{`"content":"`, `"text":"`, `"delta":"`} {
|
||
idx := bytes.Index(line, []byte(key))
|
||
if idx < 0 {
|
||
continue
|
||
}
|
||
start := idx + len(key) // 指向值的第一个字符（引号已含在 key 里）
|
||
end := start
|
||
escaped := false
|
||
for end < len(line) {
|
||
switch line[end] {
|
||
case '\\':
|
||
escaped = true
|
||
end += 2
|
||
continue
|
||
case '"':
|
||
if escaped {
|
||
// 有转义：这一段交给 JSON 解析，保证还原正确
|
||
var out string
|
||
if json.Unmarshal(line[idx+len(key)-1:end+1], &out) == nil {
|
||
return out
|
||
}
|
||
return ""
|
||
}
|
||
return string(line[start:end])
|
||
}
|
||
end++
|
||
}
|
||
return "" // 行被截断：留给下一行
|
||
}
|
||
return ""
|
||
}
|
||
|
||
// retryAfter 解析上游 Retry-After（delta-秒或 HTTP 日期），用于按对端意愿延长渠道冷却。
|
||
// 解析不出或为非正返回 0，交给 channel.Service 用默认时长（并夹到上限）。
|
||
func retryAfter(h http.Header) time.Duration {
|
||
v := strings.TrimSpace(h.Get("Retry-After"))
|
||
if v == "" {
|
||
return 0
|
||
}
|
||
if secs, err := strconv.Atoi(v); err == nil {
|
||
if secs <= 0 {
|
||
return 0
|
||
}
|
||
return time.Duration(secs) * time.Second
|
||
}
|
||
if t, err := http.ParseTime(v); err == nil {
|
||
if d := time.Until(t); d > 0 {
|
||
return d
|
||
}
|
||
}
|
||
return 0
|
||
}
|
||
|
||
// upstreamURL 组装上游地址：按协议选 base_url 再拼资源路径（见 store.Channel.UpstreamURL）。
|
||
func upstreamURL(ch *store.Channel, proto, path string) string {
|
||
return ch.UpstreamURL(proto, path)
|
||
}
|
||
|
||
// doProxy 通用代理（M5）：加权随机选一个候选渠道；失败自动故障转移。
|
||
// 流式请求挂一条 streamHold：首个内容事件到达之前，上游的错误/空流/断连都还能
|
||
// 换渠道消化，客户端无感知（借鉴 magpie 的 holdWriter）。
|
||
func (g *Gateway) doProxy(c *gin.Context, cands []channel.Candidate, clientProto string, body []byte, stream bool, sink *usageSink) {
|
||
if len(cands) == 0 {
|
||
apiError(c, http.StatusBadGateway, "upstream_error", "no available channel")
|
||
g.recordError(c, nil, nil, now(), "no_available_channel")
|
||
return
|
||
}
|
||
// 会话亲和：上次应答过该会话的渠道优先（候选集已过滤健康/冷却/绑定），
|
||
// 让上游提示词缓存被再次命中；命中不了（首次/冷却/下线）走加权随机。
|
||
affKey, hasAff := g.affinityKeyFor(c)
|
||
startIdx := -1
|
||
if hasAff {
|
||
if chID, ok := g.aff.get(affKey); ok {
|
||
for i, cand := range cands {
|
||
if cand.Channel.ID == chID {
|
||
startIdx = i
|
||
break
|
||
}
|
||
}
|
||
}
|
||
}
|
||
if startIdx < 0 {
|
||
// 加权随机选择起始渠道
|
||
picked := g.ch.Pick(cands)
|
||
for i, cand := range cands {
|
||
if cand.Channel.ID == picked.ID {
|
||
startIdx = i
|
||
break
|
||
}
|
||
}
|
||
}
|
||
|
||
// 流式容错缓冲 + 心跳，整条请求一条，跨渠道重试共用
|
||
var sh *streamHold
|
||
stopHB := func() {}
|
||
if stream {
|
||
sh = newStreamHold(c)
|
||
stopHB = sh.heartbeat(g.streamHeartbeat)
|
||
defer stopHB()
|
||
}
|
||
|
||
var lastStatus = http.StatusBadGateway
|
||
var lastBody = []byte("all upstream channels failed")
|
||
// 最后真正试过的渠道：全部失败时把错误归到它头上，否则 usage_logs 里
|
||
// channel_id=0，渠道维度的统计与排障就看不到这类失败。
|
||
var lastCh *store.Channel
|
||
// saturated：有候选因并发满载被跳过。一次都没打到上游就全被跳过时，
|
||
// 那不是「上游坏了」（502），而是「我方满了」（503 + Retry-After）。
|
||
saturated := false
|
||
startAll := time.Now()
|
||
// 换渠道次数按失败类型计数：cap 到了就不再为同类失败换渠道（配置见 failover_caps）。
|
||
switches := map[channel.FailKind]int{}
|
||
var lastKind channel.FailKind
|
||
// 从选中的渠道开始遍历，到末尾后再从头遍历到选中渠道之前
|
||
for offset := 0; offset < len(cands); offset++ {
|
||
// 重试预算只约束第二个及以后的渠道：首个渠道不受限（长回答本就要跑几分钟），
|
||
// 但坏渠道一个接一个地等下去会让客户端挂到 N×timeout。
|
||
if offset > 0 {
|
||
if g.failoverBudget > 0 && time.Since(startAll) > g.failoverBudget {
|
||
lastBody = []byte("failover budget exhausted: " + string(lastBody))
|
||
break
|
||
}
|
||
// 429/配额类失败瞬时返回，时间预算对它们形同虚设；按同类失败次数封顶。
|
||
// switches 记该类型已发生的失败次数（= 已换渠道次数），本次是第
|
||
// switches+1 次换渠道，超过上限即停。
|
||
if cap := g.failoverCaps[lastKind]; cap > 0 && switches[lastKind] > cap {
|
||
lastBody = []byte("failover switch cap reached (" + lastKind.String() + "): " + string(lastBody))
|
||
break
|
||
}
|
||
}
|
||
idx := (startIdx + offset) % len(cands)
|
||
cand := cands[idx]
|
||
ch := cand.Channel
|
||
// 客户端已断开：上游请求会随 context 立刻失败，换渠道也没有意义
|
||
if c.Request.Context().Err() != nil {
|
||
g.finishUsage(c, ch, now(), store.UsageStatusCanceled, "client_disconnect")
|
||
return
|
||
}
|
||
if g.ch.InCooldown(ch.ID) {
|
||
continue // 本次请求期间已被并发请求判入运行期冷却
|
||
}
|
||
plan, err := prepareUpstream(ch, clientProto, body, cand.UpstreamModel)
|
||
if err != nil {
|
||
// 请求体转换失败与渠道健康无关，不换渠道计数
|
||
lastStatus, lastBody = http.StatusInternalServerError, []byte("conversion error: "+err.Error())
|
||
continue
|
||
}
|
||
release, ok := g.ch.Acquire(c.Request.Context(), ch, g.channelAcquireWait)
|
||
if !ok {
|
||
saturated = true // 满载：稍等仍拿不到槽，溢出到下一个
|
||
continue
|
||
}
|
||
// 失败/成功计数由 proxyOne 就地报：只有它读得到上游响应头（429 的 Retry-After）
|
||
lastCh = ch
|
||
done, served, retry, st, b, fkind := g.proxyOne(c, ch, plan, stream, sink, sh, clientProto)
|
||
release()
|
||
if done {
|
||
if served && hasAff {
|
||
g.aff.put(affKey, ch.ID, g.affinityTTL)
|
||
}
|
||
return
|
||
}
|
||
if !retry {
|
||
// 非重试性失败（如 400）：透传上游错误体。
|
||
// 响应头若已被心跳先行写出，状态码不可再改，只能以 SSE 错误事件收尾。
|
||
if sh != nil && sh.headSent() {
|
||
_ = sh.writeDirect(streamErrorEvent(clientProto, st, string(b)))
|
||
g.recordError(c, ch, nil, now(), "upstream_http_"+strconv.Itoa(st))
|
||
return
|
||
}
|
||
c.Header("Content-Type", "application/json")
|
||
c.Data(st, "application/json", b)
|
||
g.recordError(c, ch, nil, now(), "upstream_http_"+strconv.Itoa(st))
|
||
return
|
||
}
|
||
lastStatus, lastBody = st, b
|
||
switches[fkind]++
|
||
lastKind = fkind
|
||
}
|
||
// 全部候选都没打出去、且都是被并发满载挡住的：503 让客户端稍后重试，
|
||
// 而不是 502（那意味着上游故障，客户端通常会直接放弃）。
|
||
if saturated && lastCh == nil {
|
||
c.Header("Retry-After", "1")
|
||
msg := "all upstream channels are at capacity, retry shortly"
|
||
if sh != nil && sh.headSent() {
|
||
_ = sh.writeDirect(streamErrorEvent(clientProto, http.StatusServiceUnavailable, msg))
|
||
} else {
|
||
apiError(c, http.StatusServiceUnavailable, "channel_saturated", msg)
|
||
}
|
||
g.recordError(c, nil, nil, now(), "channel_saturated")
|
||
return
|
||
}
|
||
// 全部候选重试性失败
|
||
if sh != nil && sh.headSent() {
|
||
_ = sh.writeDirect(streamErrorEvent(clientProto, lastStatus, string(lastBody)))
|
||
g.recordError(c, lastCh, nil, now(), "all_channels_failed")
|
||
return
|
||
}
|
||
apiError(c, lastStatus, "upstream_error", string(lastBody))
|
||
g.recordError(c, lastCh, nil, now(), "all_channels_failed")
|
||
}
|
||
|
||
// affinityKeyFor 计算本请求的会话亲和键；亲和关闭或缺少身份要素时 has=false。
|
||
func (g *Gateway) affinityKeyFor(c *gin.Context) (key uint64, has bool) {
|
||
if g.affinityTTL <= 0 {
|
||
return 0, false
|
||
}
|
||
kid, ok := c.Get(CtxKeyID)
|
||
kID, ok2 := kid.(uint64)
|
||
if !ok || !ok2 {
|
||
return 0, false
|
||
}
|
||
mn, _ := c.Get("model_name")
|
||
model, _ := mn.(string)
|
||
// 亲和键取原始请求体的前缀：每轮重发全量上下文时开头是稳定的（模型、系统提示、
|
||
// 最早的几条消息），足够区分会话；不必为了它整包解析 JSON。
|
||
var prefix []byte
|
||
if v, ok := c.Get("req_body"); ok {
|
||
if raw, ok2 := v.([]byte); ok2 {
|
||
prefix = raw // affinityKey 内部只取前 affinityPrefix 字节
|
||
}
|
||
}
|
||
return affinityKey(kID, model, prefix), true
|
||
}
|
||
|
||
// streamErrorEvent 把错误构造成客户端协议的 SSE 错误事件（响应头已写出时的唯一交付方式）。
|
||
func streamErrorEvent(proto string, status int, msg string) []byte {
|
||
errObj := map[string]any{"message": msg, "type": "api_error", "code": "upstream_error", "status": status}
|
||
switch proto {
|
||
case convert.ProtoMessages:
|
||
body, _ := json.Marshal(map[string]any{"type": "error", "error": errObj})
|
||
return append([]byte("event: error\ndata: "), append(body, '\n', '\n')...)
|
||
case convert.ProtoResponses:
|
||
body, _ := json.Marshal(map[string]any{"type": "error", "code": "upstream_error", "message": msg})
|
||
return append([]byte("data: "), append(body, '\n', '\n')...)
|
||
default:
|
||
body, _ := json.Marshal(map[string]any{"error": errObj})
|
||
return append(append([]byte("data: "), append(body, '\n', '\n')...), []byte("data: [DONE]\n\n")...)
|
||
}
|
||
}
|
||
|
||
// proxyOne 对单个渠道执行一次代理。
|
||
// 返回：done=本次请求已终结（客户端响应已写出或已记账，停止故障转移）；
|
||
// served=该渠道成功应答（会话亲和据此绑定）；retry=是否可安全换渠道重试；
|
||
// status+respBody=失败信息。
|
||
// sh 为流式容错缓冲（非流式为 nil）；sink 为请求级记账 sink。
|
||
// 最后一个返回值 fkind 是本次尝试的失败类型（仅 retry=true 时被 doProxy 消费，
|
||
// 用于按失败类型统计换渠道次数）。
|
||
func (g *Gateway) proxyOne(c *gin.Context, ch *store.Channel, plan *upstreamPlan, stream bool, sink *usageSink, sh *streamHold, clientProto string) (done bool, served bool, retry bool, status int, respBody []byte, fkind channel.FailKind) {
|
||
upKey, err := g.ch.UpstreamKey(ch)
|
||
if err != nil {
|
||
g.ch.ReportFailure(ch, channel.FailNet, 0)
|
||
return false, false, true, http.StatusInternalServerError, []byte("failed to decrypt channel key"), channel.FailNet
|
||
}
|
||
|
||
upBody := plan.body
|
||
// 直通 chat 流式：注入 stream_options.include_usage，保证末块带 usage（OpenAI 行为）
|
||
if stream && plan.path == "/chat/completions" && plan.lineConv == nil && !bytes.Contains(upBody, []byte(`"include_usage"`)) {
|
||
// 定点追加，避免整包 unmarshal/marshal（大请求体上是几百微秒 + 一倍分配）
|
||
if b, ok := injectStreamOptions(upBody, `{"include_usage":true}`); ok {
|
||
upBody = b
|
||
}
|
||
}
|
||
|
||
// 上游请求绑定客户端 context：客户端断开即取消上游，不再空转生成、继续付费。
|
||
// 非流式再叠加渠道超时作为整段硬上限；流式只靠 Transport.ResponseHeaderTimeout
|
||
// 约束首包，正文时长不设限（长回答不能被超时砍断）。
|
||
limit := g.channelTimeout(ch)
|
||
ctx := c.Request.Context()
|
||
if !stream && limit > 0 {
|
||
var cancel context.CancelFunc
|
||
ctx, cancel = context.WithTimeout(ctx, limit)
|
||
defer cancel()
|
||
}
|
||
// 首 token 计时（可选）：流式在放行前迟迟不出内容就取消本次尝试、换下一个渠道。
|
||
// 计时器与放行互斥（cancelIfHeld），已放行的健康流不会被掐死。
|
||
tryCtx := ctx
|
||
var tryCancel context.CancelFunc
|
||
var ftTimer *time.Timer
|
||
if stream && g.firstTokenWait > 0 {
|
||
tryCtx, tryCancel = context.WithCancel(ctx)
|
||
defer tryCancel()
|
||
ftTimer = time.AfterFunc(g.firstTokenWait, func() { sh.cancelIfHeld(tryCancel) })
|
||
defer ftTimer.Stop()
|
||
}
|
||
onRelease := func() {
|
||
if ftTimer != nil {
|
||
ftTimer.Stop()
|
||
}
|
||
}
|
||
|
||
req, err := http.NewRequestWithContext(tryCtx, http.MethodPost, upstreamURL(ch, plan.proto, plan.path), bytes.NewReader(upBody))
|
||
if err != nil {
|
||
return false, false, false, http.StatusInternalServerError, []byte("failed to build upstream request"), channel.FailNet
|
||
}
|
||
req.Header.Set("Content-Type", "application/json")
|
||
req.Header.Set("Authorization", "Bearer "+upKey)
|
||
req.Header.Set("Accept", c.GetHeader("Accept"))
|
||
if ua := c.GetHeader("User-Agent"); ua != "" {
|
||
req.Header.Set("User-Agent", ua)
|
||
}
|
||
if plan.path == "/messages" {
|
||
// Anthropic 官方 API 只认 x-api-key（Bearer 是 OAuth token 的用法），
|
||
// 兼容网关两种都收；带上 anthropic-version 与客户端声明的 beta 头。
|
||
req.Header.Set("x-api-key", upKey)
|
||
req.Header.Set("anthropic-version", "2023-06-01")
|
||
if beta := c.GetHeader("anthropic-beta"); beta != "" {
|
||
req.Header.Set("anthropic-beta", beta)
|
||
}
|
||
}
|
||
for _, h := range []string{"OpenAI-Organization", "OpenAI-Project", "OpenAI-Beta"} {
|
||
if v := c.GetHeader(h); v != "" {
|
||
req.Header.Set(h, v)
|
||
}
|
||
}
|
||
|
||
start := time.Now()
|
||
resp, err := g.clientFor(limit).Do(req)
|
||
if err != nil {
|
||
// 客户端断开导致的取消：就地记账并终结本次请求（不计入渠道失败，也不换渠道）
|
||
if c.Request.Context().Err() != nil {
|
||
g.finishUsage(c, ch, start, store.UsageStatusCanceled, "client_disconnect")
|
||
return true, false, false, 0, nil, channel.FailNet
|
||
}
|
||
g.ch.ReportFailure(ch, channel.FailNet, 0) // 网络错误/超时
|
||
return false, false, true, http.StatusBadGateway, []byte("upstream request failed: " + err.Error()), channel.FailNet
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
// 非 2xx：按原因分类记账——429 不一定是限流（可能是余额/配额），402 也值得
|
||
// 换渠道（渠道 A 没余额不代表渠道 B 没有）；其余透传错误体。
|
||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||
errBody, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
|
||
kind, retryable := classifyUpstreamError(resp.StatusCode, errBody)
|
||
if retryable {
|
||
g.ch.ReportFailure(ch, kind, retryAfter(resp.Header))
|
||
return false, false, true, resp.StatusCode, errBody, kind
|
||
}
|
||
return false, false, false, resp.StatusCode, errBody, kind
|
||
}
|
||
g.ch.ReportSuccess(ch)
|
||
|
||
// 协议形态校验：要流给了整包（或反过来）是劣质中转的常见毛病，按失败换渠道；
|
||
// 反复如此的渠道会被健康计数剔除（didNotStream 不会因重试而变好）。
|
||
ct := resp.Header.Get("Content-Type")
|
||
isSSE := strings.HasPrefix(ct, "text/event-stream")
|
||
if stream != isSSE {
|
||
_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 4<<20)) // 排干以复用连接
|
||
g.ch.ReportFailure(ch, channel.FailServer, 0)
|
||
asked, got := "整包", "流式响应"
|
||
if stream {
|
||
asked, got = "流式", "整包响应"
|
||
}
|
||
return false, false, true, http.StatusBadGateway,
|
||
[]byte(ch.Name + " answered a " + asked + " request with " + got + " (didNotStream)"), channel.FailServer
|
||
}
|
||
|
||
if stream {
|
||
trySink := &usageSink{}
|
||
return g.streamHoldCopy(c, sh, ch, resp.Body, start, plan.lineConv, clientProto, trySink, sink, onRelease, tryCtx)
|
||
}
|
||
return g.copyAndCapture(c, ch, resp, start, plan.bodyConv, sink)
|
||
}
|
||
|
||
// copyAndCapture 非流式：先读完整上游响应再写客户端——响应头与正文都在"确定成功"后
|
||
// 才写出，读失败（网络错误/超时）发生在写出之前，可换渠道重试。
|
||
// 此前实现先写 200 头再读，超时只能给客户端留一个 200 空响应。
|
||
func (g *Gateway) copyAndCapture(c *gin.Context, ch *store.Channel, resp *http.Response, start time.Time, bodyConv func([]byte) ([]byte, error), sink *usageSink) (done bool, served bool, retry bool, status int, respBody []byte, fkind channel.FailKind) {
|
||
data, err := io.ReadAll(resp.Body)
|
||
if err != nil {
|
||
switch {
|
||
case c.Request.Context().Err() != nil:
|
||
// 客户端断开：就地记账终结
|
||
g.finishUsage(c, ch, start, store.UsageStatusCanceled, "client_disconnect")
|
||
return true, false, false, 0, nil, channel.FailNet
|
||
case errors.Is(err, context.DeadlineExceeded):
|
||
g.ch.ReportFailure(ch, channel.FailNet, 0)
|
||
return false, false, true, http.StatusBadGateway, []byte(ch.Name + " upstream timeout"), channel.FailNet
|
||
default:
|
||
g.ch.ReportFailure(ch, channel.FailNet, 0)
|
||
return false, false, true, http.StatusBadGateway, []byte(ch.Name + " read error: " + err.Error()), channel.FailNet
|
||
}
|
||
}
|
||
if usageRaw := extractUsage(data); usageRaw != nil && sink != nil {
|
||
sink.push(usageRaw)
|
||
}
|
||
if sink != nil && sink.servedModel == "" {
|
||
sink.servedModel = extractServedModel(data)
|
||
}
|
||
out := data
|
||
if bodyConv != nil {
|
||
if converted, cerr := bodyConv(data); cerr == nil {
|
||
out = converted
|
||
}
|
||
}
|
||
c.Header("Content-Type", resp.Header.Get("Content-Type"))
|
||
c.Status(http.StatusOK)
|
||
_, _ = c.Writer.Write(out)
|
||
if _, ok := c.Get("raw_request"); ok {
|
||
c.Set("raw_response", capRaw(data, rawResponseCap)) // 上游原始响应（未转换）
|
||
}
|
||
g.finishUsage(c, ch, start, store.UsageStatusSuccess, "")
|
||
return true, true, false, 0, nil, channel.FailNet
|
||
}
|
||
|
||
// streamHoldCopy 流式转发：边读上游 SSE 边经 hold 写客户端（放行前缓冲）。
|
||
// 放行前的一切失败（上游错误事件、空流、断连、首 token 超时）都可换渠道重试，
|
||
// 健康计数就地上报；放行后的失败与旧语义一致，就地记账终结。
|
||
// trySink 收集本次尝试扫到的 usage，只有成功走到 EOF 才并入请求级 reqSink——
|
||
// 失败渠道的 message_start 用量不能计入最终账单。
|
||
func (g *Gateway) streamHoldCopy(c *gin.Context, sh *streamHold, ch *store.Channel, r io.Reader, start time.Time,
|
||
lineConv func([]byte) []byte, clientProto string, trySink, reqSink *usageSink, onRelease func(), tryCtx context.Context,
|
||
) (done bool, served bool, retry bool, status int, respBody []byte, fkind channel.FailKind) {
|
||
// 原始响应捕获：仅管理员且开关开启（raw_request 已 set）时累积上游原始行
|
||
_, capture := c.Get("raw_request")
|
||
var rawResp strings.Builder
|
||
commitRaw := func() {
|
||
if capture {
|
||
c.Set("raw_response", capRaw([]byte(rawResp.String()), rawResponseCap))
|
||
}
|
||
}
|
||
|
||
scanner := newSSEScanner(r)
|
||
releaseSeen := false
|
||
fail := func(msg string, kind channel.FailKind) (bool, bool, bool, int, []byte, channel.FailKind) {
|
||
g.ch.ReportFailure(ch, kind, 0)
|
||
return false, false, true, http.StatusBadGateway, []byte(ch.Name + ": " + msg), kind
|
||
}
|
||
for {
|
||
line, err := scanner.Next()
|
||
if line != nil {
|
||
if capture {
|
||
rawResp.Write(line)
|
||
}
|
||
// 上游错误事件（200+流式形式的错误）：放行前换渠道消化
|
||
if !releaseSeen {
|
||
if msg, isErr := sseUpstreamError(line); isErr {
|
||
return fail("stream error: "+msg, channel.FailServer)
|
||
}
|
||
}
|
||
out := line
|
||
if lineConv != nil {
|
||
out = lineConv(line)
|
||
}
|
||
if out != nil {
|
||
if werr := sh.writeEvent(out); werr != nil {
|
||
// 客户端意外断开：按已生成部分收费（canceled）
|
||
commitRaw()
|
||
g.finishUsage(c, ch, start, store.UsageStatusCanceled, "client_disconnect")
|
||
return true, false, false, 0, nil, channel.FailNet
|
||
}
|
||
if trySink != nil {
|
||
trySink.addContent(out)
|
||
}
|
||
}
|
||
if !releaseSeen && sseEventHasContent(out, clientProto) {
|
||
releaseSeen = true
|
||
sh.release()
|
||
if onRelease != nil {
|
||
onRelease()
|
||
}
|
||
}
|
||
if usageRaw := scanUsage(line); usageRaw != nil && trySink != nil {
|
||
trySink.push(usageRaw)
|
||
}
|
||
if trySink != nil && trySink.servedModel == "" && bytes.Contains(line, []byte(`"model"`)) {
|
||
if payload, ok := sseDataPayload(line); ok {
|
||
trySink.servedModel = extractServedModel([]byte(payload))
|
||
}
|
||
}
|
||
}
|
||
if err != nil {
|
||
if releaseSeen {
|
||
// 已放行：客户端正在看内容，失败就地记账终结，与旧语义一致
|
||
commitRaw()
|
||
if err == io.EOF {
|
||
reqSink.merge(trySink)
|
||
g.finishUsage(c, ch, start, store.UsageStatusSuccess, "")
|
||
} else if c.Request.Context().Err() != nil {
|
||
g.finishUsage(c, ch, start, store.UsageStatusCanceled, "client_disconnect")
|
||
} else {
|
||
g.recordError(c, ch, nil, start, "stream_read_error")
|
||
}
|
||
return true, err == io.EOF, false, 0, nil, channel.FailNet
|
||
}
|
||
// 放行前失败：换渠道重试
|
||
if c.Request.Context().Err() != nil {
|
||
g.finishUsage(c, ch, start, store.UsageStatusCanceled, "client_disconnect")
|
||
return true, false, false, 0, nil, channel.FailNet
|
||
}
|
||
if tryCtx != nil && tryCtx.Err() != nil {
|
||
return fail(fmt.Sprintf("first token timeout (no content in %s)", g.firstTokenWait), channel.FailNet)
|
||
}
|
||
if err == io.EOF {
|
||
return fail("ended the stream before any content", channel.FailServer)
|
||
}
|
||
return fail("stream failed before any content: "+err.Error(), channel.FailNet)
|
||
}
|
||
}
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// usage 提取
|
||
|
||
// usageShape 兼容 chat (prompt/completion)、responses (input/output)、messages (input/output) 命名。
|
||
type usageShape struct {
|
||
PromptTokens int64 `json:"prompt_tokens"`
|
||
CompletionTokens int64 `json:"completion_tokens"`
|
||
InputTokens int64 `json:"input_tokens"`
|
||
OutputTokens int64 `json:"output_tokens"`
|
||
TotalTokens int64 `json:"total_tokens"`
|
||
CacheReadInputTokens int64 `json:"cache_read_input_tokens"`
|
||
CacheCreationInputTokens int64 `json:"cache_creation_input_tokens"`
|
||
|
||
// Reasoning：思考 token（chat 的 completion_tokens_details / responses 的
|
||
// output_tokens_details）。上游已把它计入 completion/output，只做展示口径，不重复计费。
|
||
CompletionDetails *usageTokenDetails `json:"completion_tokens_details,omitempty"`
|
||
OutputDetails *usageTokenDetails `json:"output_tokens_details,omitempty"`
|
||
// Anthropic 的 cache_creation 细分：1h 写价是输入价 2×、5m 是 1.25×，必须拆开算钱
|
||
CacheCreation *usageCacheCreation `json:"cache_creation,omitempty"`
|
||
}
|
||
|
||
type usageTokenDetails struct {
|
||
ReasoningTokens int64 `json:"reasoning_tokens"`
|
||
}
|
||
|
||
type usageCacheCreation struct {
|
||
Ephemeral5m int64 `json:"ephemeral_5m_input_tokens"`
|
||
Ephemeral1h int64 `json:"ephemeral_1h_input_tokens"`
|
||
}
|
||
|
||
// extractUsage 从完整响应体提取 usage 子对象（chat / responses / messages）。
|
||
func extractUsage(data []byte) json.RawMessage {
|
||
var m map[string]json.RawMessage
|
||
if json.Unmarshal(data, &m) != nil {
|
||
return nil
|
||
}
|
||
if u := usageFromMap(m); u != nil {
|
||
return u
|
||
}
|
||
if respRaw, ok := m["response"]; ok {
|
||
var resp map[string]json.RawMessage
|
||
if json.Unmarshal(respRaw, &resp) == nil {
|
||
if u := usageFromMap(resp); u != nil {
|
||
return u
|
||
}
|
||
}
|
||
}
|
||
if choices, ok := m["choices"]; ok {
|
||
var cs []map[string]json.RawMessage
|
||
if json.Unmarshal(choices, &cs) == nil {
|
||
for _, ch := range cs {
|
||
if msgRaw, ok := ch["message"]; ok {
|
||
var msg map[string]json.RawMessage
|
||
if json.Unmarshal(msgRaw, &msg) == nil {
|
||
if u := usageFromMap(msg); u != nil {
|
||
return u
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// scanUsage 从 SSE 一行中提取 usage（OpenAI 末块 / responses completed / messages message_delta 等）。
|
||
func scanUsage(line []byte) json.RawMessage {
|
||
if !bytes.Contains(line, []byte(`"usage"`)) {
|
||
return nil
|
||
}
|
||
s, ok := sseDataPayload(line)
|
||
if !ok {
|
||
return nil
|
||
}
|
||
var m map[string]json.RawMessage
|
||
if json.Unmarshal([]byte(s), &m) != nil {
|
||
return nil
|
||
}
|
||
if u := usageFromMap(m); u != nil {
|
||
return u
|
||
}
|
||
if respRaw, ok := m["response"]; ok {
|
||
var resp map[string]json.RawMessage
|
||
if json.Unmarshal(respRaw, &resp) == nil {
|
||
if u := usageFromMap(resp); u != nil {
|
||
return u
|
||
}
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// usageFromMap 从 map 顶层、message 或 delta 子对象中取 usage。
|
||
// 覆盖三种形态：
|
||
// - OpenAI 流式末块顶层 usage
|
||
// - Anthropic 非流式 / message_start 的 message.usage
|
||
// - Anthropic 流式 message_delta 的 delta.usage（真实 token 计数所在）
|
||
func usageFromMap(m map[string]json.RawMessage) json.RawMessage {
|
||
if u, ok := m["usage"]; ok && string(u) != "null" {
|
||
return u
|
||
}
|
||
if msgRaw, ok := m["message"]; ok {
|
||
var msg map[string]json.RawMessage
|
||
if json.Unmarshal(msgRaw, &msg) == nil {
|
||
if u, ok := msg["usage"]; ok && string(u) != "null" {
|
||
return u
|
||
}
|
||
}
|
||
}
|
||
if deltaRaw, ok := m["delta"]; ok {
|
||
var delta map[string]json.RawMessage
|
||
if json.Unmarshal(deltaRaw, &delta) == nil {
|
||
if u, ok := delta["usage"]; ok && string(u) != "null" {
|
||
return u
|
||
}
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// sseScanner 按 SSE 行边界读取（兼容 \n 与 \r\n），保留原始行内容。
|
||
type sseScanner struct {
|
||
r *bufio.Reader
|
||
}
|
||
|
||
func newSSEScanner(r io.Reader) *sseScanner { return &sseScanner{r: bufio.NewReaderSize(r, 32*1024)} }
|
||
|
||
func (s *sseScanner) Next() ([]byte, error) {
|
||
line, err := s.r.ReadBytes('\n')
|
||
if len(line) > 0 {
|
||
return line, nil
|
||
}
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
return nil, io.EOF
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// 记账
|
||
|
||
// extractServedModel 从响应体（或流式事件载荷）提取上游实际应答的模型名：
|
||
// chat/responses/messages 三协议都在顶层带 model；流式事件在 message.model /
|
||
// response.model。模型映射、聚合商场景下它与请求的模型名不同，落库便于对账。
|
||
func extractServedModel(data []byte) string {
|
||
var m map[string]json.RawMessage
|
||
if json.Unmarshal(data, &m) != nil {
|
||
return ""
|
||
}
|
||
if raw, ok := m["model"]; ok {
|
||
var s string
|
||
if json.Unmarshal(raw, &s) == nil && s != "" {
|
||
return s
|
||
}
|
||
}
|
||
for _, key := range []string{"response", "message"} {
|
||
raw, ok := m[key]
|
||
if !ok {
|
||
continue
|
||
}
|
||
var sub map[string]json.RawMessage
|
||
if json.Unmarshal(raw, &sub) != nil {
|
||
continue
|
||
}
|
||
if mraw, ok := sub["model"]; ok {
|
||
var s string
|
||
if json.Unmarshal(mraw, &s) == nil && s != "" {
|
||
return s
|
||
}
|
||
}
|
||
}
|
||
return ""
|
||
}
|
||
|
||
// usageSink 累积多次 usage：合并各事件字段（message_start 给 input，message_delta 给 output）。
|
||
// outputText 累积已转发内容文本，用于流式中断时按 tiktoken 估算输出 token。
|
||
type usageSink struct {
|
||
us usageShape
|
||
servedModel string // 上游实际应答的模型名（模型映射/聚合商场景下与请求的不同）
|
||
|
||
// text 累积已转发的正文，仅用于「上游没报输出用量」时按 tiktoken 估算。
|
||
// 用 Builder 而非字符串拼接：后者每块整体拷贝，长回答是 O(n²)
|
||
// （实测 10KB/500 块要 949µs、7.9MB 分配）。
|
||
text strings.Builder
|
||
// sawOutput：上游已经报过输出 token，估算不再需要，后面不再累积正文。
|
||
sawOutput bool
|
||
}
|
||
|
||
// textCap 正文累积上限：超过后不再累积（估算本来只是近似，不值得为它吃满内存）。
|
||
const textCap = 512 << 10
|
||
|
||
// addContent 从一条已转发的 SSE 行里取出正文并累积。
|
||
// 先做廉价的子串预过滤，只有可能携带正文的行才走 JSON 解析——一份流里
|
||
// role/usage/finish 之类的元事件占比不低，逐行 JSON 解析是纯浪费。
|
||
func (u *usageSink) addContent(line []byte) {
|
||
if u.sawOutput || u.text.Len() >= textCap {
|
||
return
|
||
}
|
||
if !bytes.Contains(line, []byte(`"content"`)) && !bytes.Contains(line, []byte(`"delta"`)) && !bytes.Contains(line, []byte(`"text"`)) {
|
||
return
|
||
}
|
||
if s := fastContentText(line); s != "" {
|
||
u.text.WriteString(s)
|
||
}
|
||
}
|
||
|
||
// textOf 已累积的正文。
|
||
func (u *usageSink) textOf() string { return u.text.String() }
|
||
|
||
func (u *usageSink) push(raw json.RawMessage) {
|
||
if len(raw) == 0 {
|
||
return
|
||
}
|
||
var t usageShape
|
||
if json.Unmarshal(raw, &t) != nil {
|
||
return
|
||
}
|
||
// messages 流式最终事件（message_delta 的 usage）带 cache_* 字段，是上游的最终计费口径，
|
||
// 其中 input_tokens 仅指"非缓存输入"（与 message_start 的"总输入"语义不同）。
|
||
// 整体替换而非字段合并，避免 delta 的非缓存 input 覆盖 start 的总 input 后语义错乱
|
||
// （实际消耗由 finishUsage 按 input + cache_read + cache_creation 汇总）。
|
||
if t.CacheReadInputTokens > 0 || t.CacheCreationInputTokens > 0 {
|
||
u.us = t
|
||
return
|
||
}
|
||
// 一旦上游报过输出 token，正文估算就没有意义了（也就不用再累积）
|
||
if t.CompletionTokens > 0 || t.OutputTokens > 0 {
|
||
u.sawOutput = true
|
||
u.text.Reset()
|
||
}
|
||
// 零值不覆盖：不同事件携带不同字段
|
||
if t.PromptTokens > 0 {
|
||
u.us.PromptTokens = t.PromptTokens
|
||
}
|
||
if t.CompletionTokens > 0 {
|
||
u.us.CompletionTokens = t.CompletionTokens
|
||
}
|
||
if t.InputTokens > 0 {
|
||
u.us.InputTokens = t.InputTokens
|
||
}
|
||
if t.OutputTokens > 0 {
|
||
u.us.OutputTokens = t.OutputTokens
|
||
}
|
||
if t.CacheReadInputTokens > 0 {
|
||
u.us.CacheReadInputTokens = t.CacheReadInputTokens
|
||
}
|
||
if t.CacheCreationInputTokens > 0 {
|
||
u.us.CacheCreationInputTokens = t.CacheCreationInputTokens
|
||
}
|
||
if t.CompletionDetails != nil {
|
||
u.us.CompletionDetails = t.CompletionDetails
|
||
}
|
||
if t.OutputDetails != nil {
|
||
u.us.OutputDetails = t.OutputDetails
|
||
}
|
||
if t.CacheCreation != nil {
|
||
u.us.CacheCreation = t.CacheCreation
|
||
}
|
||
}
|
||
|
||
// Shape 返回合并后的用量。
|
||
func (u *usageSink) Shape() usageShape { return u.us }
|
||
|
||
// merge 并入一次成功尝试的用量：失败渠道的用量整体丢弃（换渠道重试时），
|
||
// 成功渠道的 usage 事件（含放行前缓冲期扫到的 message_start）在这里落账。
|
||
func (u *usageSink) merge(o *usageSink) {
|
||
if o == nil {
|
||
return
|
||
}
|
||
if o.us.PromptTokens > 0 {
|
||
u.us.PromptTokens = o.us.PromptTokens
|
||
}
|
||
if o.us.CompletionTokens > 0 {
|
||
u.us.CompletionTokens = o.us.CompletionTokens
|
||
}
|
||
if o.us.InputTokens > 0 {
|
||
u.us.InputTokens = o.us.InputTokens
|
||
}
|
||
if o.us.OutputTokens > 0 {
|
||
u.us.OutputTokens = o.us.OutputTokens
|
||
}
|
||
if o.us.CacheReadInputTokens > 0 {
|
||
u.us.CacheReadInputTokens = o.us.CacheReadInputTokens
|
||
}
|
||
if o.us.CacheCreationInputTokens > 0 {
|
||
u.us.CacheCreationInputTokens = o.us.CacheCreationInputTokens
|
||
}
|
||
if !u.sawOutput && !o.sawOutput {
|
||
u.text.WriteString(o.text.String())
|
||
}
|
||
u.sawOutput = u.sawOutput || o.sawOutput
|
||
if o.servedModel != "" && u.servedModel == "" {
|
||
u.servedModel = o.servedModel
|
||
}
|
||
}
|
||
|
||
// finishUsage 落账：计算成本并异步写入。
|
||
func (g *Gateway) finishUsage(c *gin.Context, ch *store.Channel, start time.Time, status, errCode string) {
|
||
uid, _ := c.Get(CtxUserID)
|
||
kid, _ := c.Get(CtxKeyID)
|
||
trace, _ := c.Get(CtxTrace)
|
||
|
||
var us usageShape
|
||
var sink *usageSink
|
||
if h, ok := c.Get("usage_raw"); ok {
|
||
if holder, ok := h.(*sinkHolder); ok && holder.sink != nil {
|
||
us = holder.sink.Shape()
|
||
sink = holder.sink
|
||
}
|
||
}
|
||
|
||
in := us.PromptTokens + us.InputTokens
|
||
out := us.CompletionTokens + us.OutputTokens
|
||
cacheRead := us.CacheReadInputTokens
|
||
cacheCreate := us.CacheCreationInputTokens
|
||
reasoning := int64(0)
|
||
if us.CompletionDetails != nil {
|
||
reasoning += us.CompletionDetails.ReasoningTokens
|
||
}
|
||
if us.OutputDetails != nil {
|
||
reasoning += us.OutputDetails.ReasoningTokens
|
||
}
|
||
// Anthropic 缓存写细分：1h 是输入价 2×、5m 是 1.25×，倍率 1.6。
|
||
// CacheWritePriceOr() 是 5m 口径（缺省输入价 ×1.25），1h 按同比例放大。
|
||
cache1h := int64(0)
|
||
if us.CacheCreation != nil && us.CacheCreation.Ephemeral1h > 0 {
|
||
cache1h = us.CacheCreation.Ephemeral1h
|
||
if cache1h > cacheCreate {
|
||
cache1h = cacheCreate
|
||
}
|
||
}
|
||
cache5m := cacheCreate - cache1h
|
||
|
||
modelName, _ := c.Get("model_name")
|
||
mn, _ := modelName.(string)
|
||
servedModel := ""
|
||
if sink != nil {
|
||
servedModel = sink.servedModel
|
||
}
|
||
|
||
// 上游未返回真实 usage 时估算（tiktoken）：
|
||
// - 流式中断（canceled）：上游最终 usage 可能未返回
|
||
// - 流式正常结束（success）：部分上游（如火山方舟 Anthropic 流式）message_start.usage 恒为 0、
|
||
// message_delta 不带 usage，只能按已收发内容估算，否则记账为 0 消耗
|
||
// 非流式上游必返回 usage，此处 in/out 非 0 不受影响。
|
||
if status == store.UsageStatusCanceled || status == store.UsageStatusSuccess {
|
||
if in == 0 && cacheRead == 0 && cacheCreate == 0 {
|
||
if v, ok := c.Get("req_body"); ok {
|
||
if raw, ok2 := v.([]byte); ok2 && len(raw) > 0 {
|
||
if text := requestText(raw); text != "" {
|
||
in = int64(tokenizer.Count(text, mn))
|
||
}
|
||
}
|
||
}
|
||
}
|
||
if out == 0 && sink != nil && sink.textOf() != "" {
|
||
out = int64(tokenizer.Count(sink.textOf(), mn))
|
||
}
|
||
}
|
||
|
||
var model store.Model
|
||
var cost decimal.Decimal
|
||
var modelID uint64
|
||
model, _ = g.modelByName(mn)
|
||
if model.ID > 0 {
|
||
modelID = model.ID
|
||
// 计价口径：in=非缓存输入、cacheRead=缓存读、cacheCreate=缓存写（Anthropic 语义，
|
||
// messages 流式 message_delta 的 input_tokens 即非缓存部分）。
|
||
// 缓存写价未单独配置时按输入价 1.25×（Anthropic 5m 口径）。
|
||
// 金额用 decimal 精确运算：token×价/1e6，结果对齐 numeric(20,8) 列精度。
|
||
cwp := model.CacheWritePriceOr()
|
||
million := decimal.NewFromInt(1e6)
|
||
perMillion := func(tokens int64, price decimal.Decimal) decimal.Decimal {
|
||
return decimal.NewFromInt(tokens).Mul(price).Div(million)
|
||
}
|
||
cost = perMillion(in, model.InputPrice).
|
||
Add(perMillion(out, model.OutputPrice)).
|
||
Add(perMillion(cacheRead, model.CacheReadPrice)).
|
||
Add(perMillion(cache5m, cwp)).
|
||
Add(perMillion(cache1h, cwp.Mul(decimal.NewFromFloat(1.6)))).
|
||
Round(8)
|
||
} else {
|
||
// 模型在记账时已被删除（请求进行中被下架）：无法定价，按 0 记账留痕。
|
||
// 虚构单价扣用户的钱既不可审计也无法对账。
|
||
accessf("proxy: trace=%s model %q not found at billing time, cost recorded as 0", trace, mn)
|
||
}
|
||
// 落库的 input_tokens 存输入总量（含缓存读/写），与上游 message_start 口径一致，便于对账展示。
|
||
in += cacheRead + cacheCreate
|
||
|
||
proto, _ := c.Get("protocol")
|
||
p, _ := proto.(string)
|
||
if p == "" {
|
||
p = "chat"
|
||
}
|
||
traceStr, _ := trace.(string)
|
||
latency := int(time.Since(start).Milliseconds())
|
||
|
||
if status == store.UsageStatusSuccess && c.Writer.Status() >= 400 {
|
||
status = store.UsageStatusError
|
||
}
|
||
|
||
var errCodePtr *string
|
||
if errCode != "" {
|
||
errCodePtr = &errCode
|
||
}
|
||
|
||
var uidVal, kidVal uint64
|
||
if u, ok := uid.(uint64); ok {
|
||
uidVal = u
|
||
}
|
||
if k, ok := kid.(uint64); ok {
|
||
kidVal = k
|
||
}
|
||
var chID uint64
|
||
if ch != nil {
|
||
chID = ch.ID
|
||
}
|
||
|
||
var rawReq, rawResp string
|
||
if v, ok := c.Get("raw_request"); ok {
|
||
rawReq, _ = v.(string)
|
||
}
|
||
if v, ok := c.Get("raw_response"); ok {
|
||
rawResp, _ = v.(string)
|
||
}
|
||
|
||
// 密钥今日 token 用量累计（配额检查用）
|
||
if g.lim != nil && kidVal > 0 {
|
||
g.lim.AddTokens(kidVal, in+out)
|
||
}
|
||
|
||
g.rec.Record(&store.UsageLog{
|
||
RequestID: fmt.Sprintf("trace-%s", traceStr),
|
||
TraceID: traceStr,
|
||
UserID: uidVal,
|
||
KeyID: kidVal,
|
||
ChannelID: chID,
|
||
ModelID: modelID,
|
||
ModelName: mn,
|
||
Protocol: p,
|
||
InputTokens: in,
|
||
OutputTokens: out,
|
||
CacheReadTokens: cacheRead,
|
||
CacheCreationTokens: cacheCreate,
|
||
InputPrice: model.InputPrice,
|
||
OutputPrice: model.OutputPrice,
|
||
CacheReadPrice: model.CacheReadPrice,
|
||
Cost: cost,
|
||
ReasoningTokens: reasoning,
|
||
CacheWrite1hTokens: cache1h,
|
||
ServedModel: servedModel,
|
||
LatencyMS: latency,
|
||
Status: status,
|
||
ErrorCode: errCodePtr,
|
||
RawRequest: rawReq,
|
||
RawResponse: rawResp,
|
||
CreatedAt: time.Now().UTC(),
|
||
})
|
||
|
||
// 请求级一行日志：usage_logs 落库是异步的且可能丢，排障时日志是唯一即时依据
|
||
accessf("proxy: trace=%s user=%d key=%d model=%q served=%q channel=%d proto=%s in=%d out=%d r=%d cache_r=%d cache_w=%d(1h=%d) cost=%s status=%s http=%d latency=%dms err=%s",
|
||
traceStr, uidVal, kidVal, mn, servedModel, chID, p, in, out, reasoning, cacheRead, cacheCreate, cache1h, cost, status,
|
||
c.Writer.Status(), latency, errCode)
|
||
}
|
||
|
||
// recordError 失败请求的记账（不产生扣费，status=error）。
|
||
func (g *Gateway) recordError(c *gin.Context, ch *store.Channel, resp *http.Response, start time.Time, code string) {
|
||
_ = resp
|
||
g.finishUsage(c, ch, start, store.UsageStatusError, code)
|
||
}
|
||
|
||
func now() time.Time { return time.Now() }
|