SakurasanandClaude
db83972b6f
- 渠道支持分协议 base_url(chat/responses/messages 各一), 网关按协议选 base 直通 (如智谱三种格式不同 base, 一个渠道即可), UpstreamURL 按 proto 拼接 - 渠道模型改为下方抽屉: 当前绑定列表(内联改上游/解除)、从接口拉取(remote 预览+勾选添加)、手动添加 - 新增 /channels/:id/models/remote 预览接口; 操作按钮加 Phosphor 图标 - 修复 formats jsonb 更新未序列化问题; 手机端渠道卡片化 Co-Authored-By: Claude <noreply@anthropic.com>
474 lines
14 KiB
Go
474 lines
14 KiB
Go
package proxy
|
||
|
||
import (
|
||
"bufio"
|
||
"bytes"
|
||
"context"
|
||
"crypto/rand"
|
||
"encoding/hex"
|
||
"encoding/json"
|
||
"fmt"
|
||
"io"
|
||
"net/http"
|
||
"strconv"
|
||
"strings"
|
||
"time"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
"github.com/openteam/server/internal/channel"
|
||
"github.com/openteam/server/internal/store"
|
||
)
|
||
|
||
// 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))
|
||
br := &bodyReq{}
|
||
_ = json.Unmarshal(body, br) // 解析失败按空处理，直通仍可转发
|
||
return br, body, nil
|
||
}
|
||
|
||
// upstreamURL 组装上游地址：按协议选 base_url 再拼资源路径（见 store.Channel.UpstreamURL）。
|
||
func upstreamURL(ch *store.Channel, proto, path string) string {
|
||
return ch.UpstreamURL(proto, path)
|
||
}
|
||
|
||
// doProxy 通用代理（M5）：遍历候选渠道，按需转换；可安全重试的失败自动故障转移。
|
||
func (g *Gateway) doProxy(c *gin.Context, cands []channel.Candidate, clientProto string, body []byte, stream bool, sink *usageSink) {
|
||
var lastStatus = http.StatusBadGateway
|
||
var lastBody = []byte("all upstream channels failed")
|
||
for _, cand := range cands {
|
||
ch := cand.Channel
|
||
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.TryAcquire(ch)
|
||
if !ok {
|
||
continue // 渠道满载，溢出到下一个
|
||
}
|
||
written, retry, st, b := g.proxyOne(c, ch, plan, stream, sink)
|
||
release()
|
||
if written {
|
||
return
|
||
}
|
||
if !retry {
|
||
// 非重试性失败（如 400）：透传上游错误体
|
||
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
|
||
}
|
||
// 全部候选重试性失败
|
||
apiError(c, lastStatus, "upstream_error", string(lastBody))
|
||
g.recordError(c, nil, nil, now(), "all_channels_failed")
|
||
}
|
||
|
||
// proxyOne 对单个渠道执行一次代理。
|
||
// 返回：written=是否已写客户端响应；retry=是否可安全换渠道重试；status+respBody=失败信息。
|
||
func (g *Gateway) proxyOne(c *gin.Context, ch *store.Channel, plan *upstreamPlan, stream bool, sink *usageSink) (written bool, retry bool, status int, respBody []byte) {
|
||
upKey, err := g.ch.UpstreamKey(ch)
|
||
if err != nil {
|
||
return false, true, http.StatusInternalServerError, []byte("failed to decrypt channel key")
|
||
}
|
||
|
||
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"`)) {
|
||
var m map[string]any
|
||
if json.Unmarshal(upBody, &m) == nil {
|
||
m["stream_options"] = map[string]any{"include_usage": true}
|
||
if b, err := json.Marshal(m); err == nil {
|
||
upBody = b
|
||
}
|
||
}
|
||
}
|
||
|
||
ctx, cancel := context.WithTimeout(c.Request.Context(), time.Duration(ch.TimeoutMS)*time.Millisecond)
|
||
defer cancel()
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, upstreamURL(ch, plan.proto, plan.path), bytes.NewReader(upBody))
|
||
if err != nil {
|
||
return false, false, http.StatusInternalServerError, []byte("failed to build upstream request")
|
||
}
|
||
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" {
|
||
req.Header.Set("anthropic-version", "2023-06-01")
|
||
}
|
||
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.hc.Do(req)
|
||
if err != nil {
|
||
if ctx.Err() == context.DeadlineExceeded {
|
||
return false, true, http.StatusGatewayTimeout, []byte("upstream request timed out")
|
||
}
|
||
return false, true, http.StatusBadGateway, []byte("upstream request failed: " + err.Error())
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
// 非 2xx：429/5xx 可重试；其余透传错误体
|
||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||
errBody, _ := io.ReadAll(resp.Body)
|
||
if resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500 {
|
||
return false, true, resp.StatusCode, errBody
|
||
}
|
||
return false, false, resp.StatusCode, errBody
|
||
}
|
||
|
||
c.Header("Content-Type", resp.Header.Get("Content-Type"))
|
||
c.Status(http.StatusOK)
|
||
if stream {
|
||
g.streamCopy(c, ch, resp.Body, start, plan.lineConv, sink)
|
||
} else {
|
||
g.copyAndCapture(c, ch, resp.Body, start, plan.bodyConv, sink)
|
||
}
|
||
return true, false, http.StatusOK, nil
|
||
}
|
||
|
||
// copyAndCapture 非流式：整体转发（可转换）+ 解析 usage + 记账。
|
||
func (g *Gateway) copyAndCapture(c *gin.Context, ch *store.Channel, r io.Reader, start time.Time, bodyConv func([]byte) ([]byte, error), sink *usageSink) {
|
||
data, err := io.ReadAll(r)
|
||
if err != nil {
|
||
apiError(c, http.StatusBadGateway, "upstream_error", "failed reading upstream response")
|
||
g.recordError(c, ch, nil, start, "read_error")
|
||
return
|
||
}
|
||
if usageRaw := extractUsage(data); usageRaw != nil && sink != nil {
|
||
sink.push(usageRaw)
|
||
}
|
||
out := data
|
||
if bodyConv != nil {
|
||
if converted, cerr := bodyConv(data); cerr == nil {
|
||
out = converted
|
||
}
|
||
}
|
||
_, _ = c.Writer.Write(out)
|
||
g.finishUsage(c, ch, start, store.UsageStatusSuccess, "")
|
||
}
|
||
|
||
// streamCopy 流式：边读上游 SSE 边写客户端，零缓冲转发；按 lineConv 转换；扫描 usage 记账。
|
||
func (g *Gateway) streamCopy(c *gin.Context, ch *store.Channel, r io.Reader, start time.Time, lineConv func([]byte) []byte, sink *usageSink) {
|
||
w := c.Writer
|
||
flusher, ok := w.(http.Flusher)
|
||
if !ok {
|
||
flusher = nopFlusher{}
|
||
}
|
||
|
||
scanner := newSSEScanner(r)
|
||
for {
|
||
line, err := scanner.Next()
|
||
if line != nil {
|
||
out := line
|
||
if lineConv != nil {
|
||
out = lineConv(line)
|
||
}
|
||
if out != nil {
|
||
if _, werr := w.Write(out); werr != nil {
|
||
g.recordError(c, ch, nil, start, "client_disconnect")
|
||
return
|
||
}
|
||
flusher.Flush()
|
||
}
|
||
if usageRaw := scanUsage(line); usageRaw != nil && sink != nil {
|
||
sink.push(usageRaw)
|
||
}
|
||
}
|
||
if err != nil {
|
||
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")
|
||
} else {
|
||
g.recordError(c, ch, nil, start, "stream_read_error")
|
||
}
|
||
return
|
||
}
|
||
}
|
||
}
|
||
|
||
type nopFlusher struct{}
|
||
|
||
func (nopFlusher) Flush() {}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// 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"`
|
||
}
|
||
|
||
// 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 {
|
||
s := string(line)
|
||
if !strings.Contains(s, `"usage"`) {
|
||
return nil
|
||
}
|
||
if strings.HasPrefix(s, "data: ") {
|
||
s = strings.TrimPrefix(s, "data: ")
|
||
}
|
||
s = strings.TrimSpace(s)
|
||
if s == "[DONE]" || s == "" {
|
||
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 子对象中取 usage。
|
||
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
|
||
}
|
||
}
|
||
}
|
||
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
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// 记账
|
||
|
||
// usageSink 累积多次 usage：合并各事件字段（message_start 给 input，message_delta 给 output）。
|
||
type usageSink struct {
|
||
us usageShape
|
||
}
|
||
|
||
func (u *usageSink) push(raw json.RawMessage) {
|
||
if len(raw) == 0 {
|
||
return
|
||
}
|
||
var t usageShape
|
||
if json.Unmarshal(raw, &t) != nil {
|
||
return
|
||
}
|
||
// 零值不覆盖：不同事件携带不同字段
|
||
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
|
||
}
|
||
}
|
||
|
||
// Shape 返回合并后的用量。
|
||
func (u *usageSink) Shape() usageShape { return u.us }
|
||
|
||
// 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
|
||
if h, ok := c.Get("usage_raw"); ok {
|
||
if holder, ok := h.(*sinkHolder); ok && holder.sink != nil {
|
||
us = holder.sink.Shape()
|
||
}
|
||
}
|
||
|
||
in := us.PromptTokens + us.InputTokens
|
||
out := us.CompletionTokens + us.OutputTokens
|
||
cacheRead := us.CacheReadInputTokens
|
||
cacheCreate := us.CacheCreationInputTokens
|
||
|
||
modelName, _ := c.Get("model_name")
|
||
mn, _ := modelName.(string)
|
||
|
||
var model store.Model
|
||
var cost float64
|
||
var modelID uint64
|
||
_ = g.db.Where("name = ?", mn).First(&model).Error
|
||
if model.ID > 0 {
|
||
modelID = model.ID
|
||
cost = float64(in)/1e6*model.InputPrice +
|
||
float64(out)/1e6*model.OutputPrice +
|
||
float64(cacheRead)/1e6*model.CacheReadPrice
|
||
} else {
|
||
cost = float64(in)/1e6*0.15 + float64(out)/1e6*0.60 // 无定价模型时按示例价
|
||
}
|
||
|
||
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
|
||
}
|
||
|
||
// 密钥今日 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,
|
||
LatencyMS: latency,
|
||
Status: status,
|
||
ErrorCode: errCodePtr,
|
||
CreatedAt: time.Now().UTC(),
|
||
})
|
||
}
|
||
|
||
// 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() }
|