diff --git a/server/go.mod b/server/go.mod index 5357a04..b9e45c5 100644 --- a/server/go.mod +++ b/server/go.mod @@ -1,6 +1,6 @@ module github.com/openteam/server -go 1.25.0 +go 1.26 require ( github.com/gin-gonic/gin v1.12.0 @@ -8,6 +8,7 @@ require ( github.com/go-webauthn/webauthn v0.17.4 github.com/golang-jwt/jwt/v5 v5.3.1 github.com/spf13/viper v1.21.0 + github.com/tiktoken-go/tokenizer v0.8.1 golang.org/x/crypto v0.55.0 gorm.io/driver/postgres v1.6.2 gorm.io/gorm v1.31.2 @@ -18,6 +19,7 @@ require ( github.com/bytedance/sonic v1.15.0 // indirect github.com/bytedance/sonic/loader v0.5.0 // indirect github.com/cloudwego/base64x v0.1.6 // indirect + github.com/dlclark/regexp2/v2 v2.5.1 // indirect github.com/dustin/go-humanize v1.0.1 // indirect github.com/fsnotify/fsnotify v1.9.0 // indirect github.com/fxamacker/cbor/v2 v2.9.2 // indirect diff --git a/server/go.sum b/server/go.sum index a1bef40..e338004 100644 --- a/server/go.sum +++ b/server/go.sum @@ -9,6 +9,8 @@ github.com/cloudwego/base64x v0.1.6/go.mod h1:OFcloc187FXDaYHvrNIjxSe8ncn0OOM8gE github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/dlclark/regexp2/v2 v2.5.1 h1:E5Ug7Dh264W1ymdySmiHNcDG7fmsR307APCE5R07a20= +github.com/dlclark/regexp2/v2 v2.5.1/go.mod h1:avUrQvPaLz2DrFNHJF0taWAFFX2C1GMSSoeiqFjcBmU= github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8= @@ -130,6 +132,8 @@ github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= github.com/subosito/gotenv v1.6.0 h1:9NlTDc1FTs4qu0DDq7AEtTPNw6SVm7uBMsUCUjABIf8= github.com/subosito/gotenv v1.6.0/go.mod h1:Dk4QP5c2W3ibzajGcXpNraDfq2IrhjMIvMSWPKKo0FU= +github.com/tiktoken-go/tokenizer v0.8.1 h1:4obDoB6/dhdBt9xMweX4nww5cjdOq/nYF4ecwPq2+mg= +github.com/tiktoken-go/tokenizer v0.8.1/go.mod h1:eLA0t6nGvn9mDc7gt90qt7pMat+gE9ViqwQ6l9B+tA4= github.com/tinylib/msgp v1.6.4 h1:mOwYbyYDLPj35mkA2BjjYejgJk9BuHxDdvRnb6v2ZcQ= github.com/tinylib/msgp v1.6.4/go.mod h1:RSp0LW9oSxFut3KzESt5Voq4GVWyS+PSulT77roAqEA= github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS4MhqMhdFk5YI= diff --git a/server/internal/pkg/tokenizer/tokenizer.go b/server/internal/pkg/tokenizer/tokenizer.go new file mode 100644 index 0000000..4946489 --- /dev/null +++ b/server/internal/pkg/tokenizer/tokenizer.go @@ -0,0 +1,24 @@ +// Package tokenizer 按模型估算 token 数(tiktoken-go)。 +// 用于流式中断时对已生成内容做近似计费;未知模型回退 cl100k_base。 +package tokenizer + +import "github.com/tiktoken-go/tokenizer" + +// Count 估算文本 token 数;空文本返回 0。 +func Count(text, model string) int { + if text == "" { + return 0 + } + enc, err := tokenizer.ForModel(tokenizer.Model(model)) + if err != nil { + enc, err = tokenizer.Get(tokenizer.Cl100kBase) + if err != nil { + return 0 + } + } + toks, _, err := enc.Encode(text) + if err != nil { + return 0 + } + return len(toks) +} diff --git a/server/internal/proxy/passthrough.go b/server/internal/proxy/passthrough.go index 32dffd7..c6af94d 100644 --- a/server/internal/proxy/passthrough.go +++ b/server/internal/proxy/passthrough.go @@ -16,6 +16,7 @@ import ( "github.com/gin-gonic/gin" "github.com/openteam/server/internal/channel" + "github.com/openteam/server/internal/pkg/tokenizer" "github.com/openteam/server/internal/store" ) @@ -39,11 +40,107 @@ func parseBody(c *gin.Context) (*bodyReq, []byte, error) { return nil, nil, err } c.Request.Body = io.NopCloser(bytes.NewReader(body)) + // 中断请求输入文本(token 估算用,仅存文本不立即分词) + c.Set("est_input_text", requestText(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") +} + +// sseContentText 提取一条 SSE data 行中的内容文本(chat delta.content / responses delta / messages delta.text)。 +func sseContentText(line []byte) string { + s := string(line) + if strings.HasPrefix(s, "data: ") { + s = strings.TrimPrefix(s, "data: ") + } + s = strings.TrimSpace(s) + if s == "" || s == "[DONE]" { + 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 "" +} + // upstreamURL 组装上游地址:按协议选 base_url 再拼资源路径(见 store.Channel.UpstreamURL)。 func upstreamURL(ch *store.Channel, proto, path string) string { return ch.UpstreamURL(proto, path) @@ -192,10 +289,14 @@ func (g *Gateway) streamCopy(c *gin.Context, ch *store.Channel, r io.Reader, sta } if out != nil { if _, werr := w.Write(out); werr != nil { - g.recordError(c, ch, nil, start, "client_disconnect") + // 客户端意外断开:按已生成部分收费(canceled) + g.finishUsage(c, ch, start, store.UsageStatusCanceled, "client_disconnect") return } flusher.Flush() + if sink != nil { + sink.outputText += sseContentText(out) + } } if usageRaw := scanUsage(line); usageRaw != nil && sink != nil { sink.push(usageRaw) @@ -205,7 +306,8 @@ func (g *Gateway) streamCopy(c *gin.Context, ch *store.Channel, r io.Reader, sta if err == io.EOF { g.finishUsage(c, ch, start, store.UsageStatusSuccess, "") } else if c.Request.Context().Err() != nil { - g.recordError(c, ch, nil, start, "client_disconnect") + // 客户端意外断开:按已生成部分收费(canceled) + g.finishUsage(c, ch, start, store.UsageStatusCanceled, "client_disconnect") } else { g.recordError(c, ch, nil, start, "stream_read_error") } @@ -336,8 +438,10 @@ func (s *sseScanner) Next() ([]byte, error) { // 记账 // usageSink 累积多次 usage:合并各事件字段(message_start 给 input,message_delta 给 output)。 +// outputText 累积已转发内容文本,用于流式中断时按 tiktoken 估算输出 token。 type usageSink struct { - us usageShape + us usageShape + outputText string } func (u *usageSink) push(raw json.RawMessage) { @@ -379,9 +483,11 @@ func (g *Gateway) finishUsage(c *gin.Context, ch *store.Channel, start time.Time 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 } } @@ -393,6 +499,20 @@ func (g *Gateway) finishUsage(c *gin.Context, ch *store.Channel, start time.Time modelName, _ := c.Get("model_name") mn, _ := modelName.(string) + // 流式中断(canceled):上游最终 usage 可能未返回,按已生成内容用 tiktoken 估算 + if status == store.UsageStatusCanceled { + if in == 0 { + if est, ok := c.Get("est_input_text"); ok { + if v, ok2 := est.(string); ok2 { + in = int64(tokenizer.Count(v, mn)) + } + } + } + if out == 0 && sink != nil && sink.outputText != "" { + out = int64(tokenizer.Count(sink.outputText, mn)) + } + } + var model store.Model var cost float64 var modelID uint64 diff --git a/server/internal/usage/usage.go b/server/internal/usage/usage.go index c2cb62b..84f1a94 100644 --- a/server/internal/usage/usage.go +++ b/server/internal/usage/usage.go @@ -97,7 +97,8 @@ func (r *Recorder) flush(logs []*store.UsageLog) error { return err } for _, l := range logs { - if l.Status != store.UsageStatusSuccess || l.Cost <= 0 { + // 计费范围:success(正常完成)与 canceled(流式中断,按已生成部分收费) + if (l.Status != store.UsageStatusSuccess && l.Status != store.UsageStatusCanceled) || l.Cost <= 0 { continue } // 扣余额(余额可为负:流式请求不中断;后续请求被拒)