Sakurasan
ca4dc4b3b7
- 新增 ComputeCost：OpenAI 系 prompt 含缓存读需扣减（异常数据钳制为 0）； Anthropic input_tokens 不含缓存，按原值计费，缓存写按输入价 ×1.25（原实现误用输出价） - recordUsage 传入上游协议（用量语义跟随解析它的上游响应，而非客户端协议） - 补充单元测试覆盖两种协议口径与边界情况
667 lines
20 KiB
Go
667 lines
20 KiB
Go
package proxy
|
||
|
||
import (
|
||
"bufio"
|
||
"bytes"
|
||
"context"
|
||
"crypto/rand"
|
||
"encoding/hex"
|
||
"encoding/json"
|
||
"fmt"
|
||
"io"
|
||
"log"
|
||
"net/http"
|
||
"net/url"
|
||
"opencatd-open/internal/channel"
|
||
"opencatd-open/internal/dao"
|
||
"opencatd-open/internal/proxy/convert"
|
||
"opencatd-open/internal/store"
|
||
"opencatd-open/internal/usage"
|
||
"opencatd-open/pkg/config"
|
||
"os"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
"gorm.io/gorm"
|
||
)
|
||
|
||
type Gateway struct {
|
||
ctx context.Context
|
||
cfg *config.Config
|
||
db *gorm.DB
|
||
wg *sync.WaitGroup
|
||
httpClient *http.Client
|
||
|
||
userDAO *dao.UserDAO
|
||
apiKeyDAO *dao.ApiKeyDAO
|
||
usageDAO *dao.UsageDAO
|
||
dailyDAO *dao.DailyUsageDAO
|
||
modelDAO *dao.ModelDAO
|
||
channelSvc *channel.Service
|
||
usageRec *usage.Recorder
|
||
|
||
// 原始请求/响应记录开关（系统配置 log_raw_requests，带 TTL 缓存避免每次查库）。
|
||
rawLogMu sync.Mutex
|
||
rawLogVal bool
|
||
rawLogSet time.Time
|
||
}
|
||
|
||
func NewGateway(ctx context.Context, cfg *config.Config, db *gorm.DB, wg *sync.WaitGroup, userDAO *dao.UserDAO, apiKeyDAO *dao.ApiKeyDAO, usageDAO *dao.UsageDAO, dailyDAO *dao.DailyUsageDAO) *Gateway {
|
||
client := &http.Client{Timeout: 120 * time.Second}
|
||
if os.Getenv("LOCAL_PROXY") != "" {
|
||
proxyUrl, err := url.Parse(os.Getenv("LOCAL_PROXY"))
|
||
if err == nil {
|
||
tr := &http.Transport{
|
||
Proxy: http.ProxyURL(proxyUrl),
|
||
}
|
||
client.Transport = tr
|
||
}
|
||
}
|
||
|
||
return &Gateway{
|
||
ctx: ctx,
|
||
cfg: cfg,
|
||
db: db,
|
||
wg: wg,
|
||
httpClient: client,
|
||
userDAO: userDAO,
|
||
apiKeyDAO: apiKeyDAO,
|
||
usageDAO: usageDAO,
|
||
dailyDAO: dailyDAO,
|
||
modelDAO: dao.NewModelDAO(db),
|
||
channelSvc: nil,
|
||
}
|
||
}
|
||
|
||
func (g *Gateway) SetChannelService(svc *channel.Service) {
|
||
g.channelSvc = svc
|
||
}
|
||
|
||
// SetUsageRecorder 注入异步用量记录器；nil 时网关跳过用量上报。
|
||
func (g *Gateway) SetUsageRecorder(r *usage.Recorder) {
|
||
g.usageRec = r
|
||
}
|
||
|
||
// rawLogEnabled 读取系统配置 log_raw_requests（10s TTL 缓存），决定是否记录原始请求/响应。
|
||
func (g *Gateway) rawLogEnabled() bool {
|
||
g.rawLogMu.Lock()
|
||
defer g.rawLogMu.Unlock()
|
||
if time.Since(g.rawLogSet) < 10*time.Second {
|
||
return g.rawLogVal
|
||
}
|
||
var sc store.SystemConfig
|
||
g.rawLogVal = false
|
||
if err := g.db.Where("key = ?", "log_raw_requests").First(&sc).Error; err == nil {
|
||
g.rawLogVal = strings.TrimSpace(sc.Value) == "true"
|
||
}
|
||
g.rawLogSet = time.Now()
|
||
return g.rawLogVal
|
||
}
|
||
|
||
// generateRequestID 生成请求级唯一 ID，用于用量明细关联与排障。
|
||
func generateRequestID() string {
|
||
b := make([]byte, 12)
|
||
if _, err := rand.Read(b); err != nil {
|
||
return fmt.Sprintf("req-%d", time.Now().UnixNano())
|
||
}
|
||
return "req-" + hex.EncodeToString(b)
|
||
}
|
||
|
||
// Request represents a parsed incoming request
|
||
type Request struct {
|
||
Model string
|
||
Stream bool
|
||
Protocol string // "chat", "messages", "responses"
|
||
Body []byte
|
||
APIKey *store.APIKey
|
||
UserID uint64
|
||
KeyID uint64
|
||
RequestID string
|
||
|
||
CaptureRaw bool // 原始请求/响应记录（管理员 + 系统开关开启）
|
||
|
||
rawBuf *strings.Builder // 上游原始响应累积器（仅 CaptureRaw 时非 nil）
|
||
}
|
||
|
||
// ParseRequest parses the incoming request and extracts key fields
|
||
func (g *Gateway) ParseRequest(c *gin.Context, protocol string) (*Request, error) {
|
||
body, err := io.ReadAll(c.Request.Body)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("failed to read body: %w", err)
|
||
}
|
||
|
||
apiKey, _ := c.Get("api_key")
|
||
userID, _ := c.Get("user_id")
|
||
userRole, _ := c.Get("user_role")
|
||
|
||
req := &Request{
|
||
Protocol: protocol,
|
||
Body: body,
|
||
UserID: userID.(uint64),
|
||
RequestID: c.GetHeader("X-Request-Id"),
|
||
}
|
||
if req.RequestID == "" {
|
||
req.RequestID = generateRequestID()
|
||
}
|
||
|
||
if ak, ok := apiKey.(*store.APIKey); ok {
|
||
req.APIKey = ak
|
||
}
|
||
|
||
// 原始请求/响应记录：仅管理员 且 系统开关 log_raw_requests 开启。
|
||
if role, _ := userRole.(string); role == store.RoleAdmin && g.rawLogEnabled() {
|
||
req.CaptureRaw = true
|
||
}
|
||
|
||
// Parse model and stream based on protocol
|
||
switch protocol {
|
||
case "chat":
|
||
var parsed convert.ChatCompletionRequest
|
||
if err := json.Unmarshal(body, &parsed); err != nil {
|
||
return nil, fmt.Errorf("invalid chat request: %w", err)
|
||
}
|
||
req.Model = parsed.Model
|
||
req.Stream = parsed.Stream
|
||
case "messages":
|
||
var parsed convert.MessagesRequest
|
||
if err := json.Unmarshal(body, &parsed); err != nil {
|
||
return nil, fmt.Errorf("invalid messages request: %w", err)
|
||
}
|
||
req.Model = parsed.Model
|
||
req.Stream = parsed.Stream
|
||
case "responses":
|
||
var parsed convert.ResponsesRequest
|
||
if err := json.Unmarshal(body, &parsed); err != nil {
|
||
return nil, fmt.Errorf("invalid responses request: %w", err)
|
||
}
|
||
req.Model = parsed.Model
|
||
req.Stream = parsed.Stream
|
||
}
|
||
|
||
return req, nil
|
||
}
|
||
|
||
// Dispatch routes the request to the appropriate upstream.
|
||
// 遍历候选渠道（绑定优先，全局回退；按优先级/权重排序），可重试性失败自动故障转移。
|
||
func (g *Gateway) Dispatch(c *gin.Context, req *Request) {
|
||
if g.channelSvc == nil {
|
||
g.writeError(c, http.StatusBadGateway, "channel service not available")
|
||
return
|
||
}
|
||
|
||
// 原始请求/响应捕获：仅管理员 + 系统开关开启（req.CaptureRaw 已在 ParseRequest 判定）。
|
||
// 客户端原始请求体即 req.Body；上游原始响应由 stream/bufferResponse 累积进 rawBuf。
|
||
if req.CaptureRaw {
|
||
req.rawBuf = &strings.Builder{}
|
||
}
|
||
|
||
cands := g.channelSvc.Candidates(req.Model)
|
||
// 内存健康过滤：连续失败进入 cooldown 的渠道不再尝试（渠道级健康自愈靠冷却过期）。
|
||
cands = g.channelSvc.FilterHealthy(cands)
|
||
if len(cands) == 0 {
|
||
g.writeError(c, http.StatusServiceUnavailable, "no enabled channels for model: "+req.Model)
|
||
g.recordUsage(req, nil, nil, usage.Event{
|
||
IsError: true, ErrorCode: "no_channel",
|
||
}, convert.TokenUsage{}, "")
|
||
return
|
||
}
|
||
|
||
var lastCh *store.Channel
|
||
_ = lastCh // 保留变量名便于断点排查；失败渠道已在循环内各自 RecordFailure
|
||
lastErrStatus := http.StatusBadGateway
|
||
lastErrBody := "all upstream channels failed"
|
||
|
||
for i := range cands {
|
||
cand := &cands[i]
|
||
ch := cand.Channel
|
||
lastCh = ch
|
||
|
||
apiKey, err := g.channelSvc.GetAPIKey(ch)
|
||
if err != nil {
|
||
lastErrStatus, lastErrBody = http.StatusBadGateway, "failed to decrypt API key"
|
||
continue
|
||
}
|
||
|
||
// Determine target format: channel declares support for the client protocol
|
||
// then passthrough, otherwise convert to its first supported protocol
|
||
// (chat > messages > responses).
|
||
targetFormat := g.conversionTarget(ch, req.Protocol)
|
||
if targetFormat == "" {
|
||
continue // 渠道不支持该协议，换下一个
|
||
}
|
||
|
||
// Build upstream URL
|
||
upstreamURL := ch.UpstreamURL(targetFormat, g.getUpstreamPath(targetFormat))
|
||
|
||
// Convert request if needed
|
||
var requestBody []byte
|
||
if targetFormat != req.Protocol {
|
||
var err error
|
||
requestBody, err = convert.ConvertRequest(req.Body, req.Protocol, targetFormat)
|
||
if err != nil {
|
||
lastErrStatus, lastErrBody = http.StatusBadRequest, "conversion failed: "+err.Error()
|
||
continue
|
||
}
|
||
} else {
|
||
requestBody = req.Body
|
||
}
|
||
|
||
// 绑定了 upstream_model 时把请求体里的 model 重写为上游模型名（别名映射）。
|
||
if cand.Binding != nil && cand.Binding.UpstreamModel != "" &&
|
||
cand.Binding.UpstreamModel != req.Model {
|
||
requestBody = rewriteModel(requestBody, cand.Binding.UpstreamModel)
|
||
}
|
||
|
||
// Create upstream request
|
||
httpReq, err := http.NewRequestWithContext(g.ctx, "POST", upstreamURL, bytes.NewReader(requestBody))
|
||
if err != nil {
|
||
lastErrStatus, lastErrBody = http.StatusBadGateway, "failed to create request"
|
||
continue
|
||
}
|
||
g.setHeaders(httpReq, ch, apiKey, targetFormat)
|
||
|
||
// Execute request
|
||
start := time.Now()
|
||
resp, err := g.httpClient.Do(httpReq)
|
||
if err != nil {
|
||
g.channelSvc.RecordFailure(ch.ID)
|
||
lastErrStatus = http.StatusBadGateway
|
||
lastErrBody = fmt.Sprintf("upstream error: %v", err)
|
||
g.recordUsage(req, cand, ch, usage.Event{
|
||
IsError: true,
|
||
ErrorCode: "upstream_error",
|
||
LatencyMS: int(time.Since(start).Milliseconds()),
|
||
}, convert.TokenUsage{}, targetFormat)
|
||
continue // 可重试：换下一个渠道
|
||
}
|
||
|
||
// Handle upstream error responses
|
||
if resp.StatusCode >= 400 {
|
||
body, _ := io.ReadAll(resp.Body)
|
||
resp.Body.Close()
|
||
log.Printf("Upstream error: status=%d body=%s", resp.StatusCode, string(body))
|
||
if req.rawBuf != nil {
|
||
req.rawBuf.Write(body)
|
||
}
|
||
g.recordUsage(req, cand, ch, usage.Event{
|
||
IsError: true,
|
||
ErrorCode: fmt.Sprintf("upstream_%d", resp.StatusCode),
|
||
LatencyMS: int(time.Since(start).Milliseconds()),
|
||
}, convert.TokenUsage{}, targetFormat)
|
||
// 429/5xx 可换渠道重试；4xx 直接透传
|
||
if resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500 {
|
||
lastErrStatus, lastErrBody = resp.StatusCode, string(body)
|
||
continue
|
||
}
|
||
c.Data(resp.StatusCode, "application/json", body)
|
||
return
|
||
}
|
||
|
||
g.channelSvc.RecordSuccess(ch.ID)
|
||
|
||
// Stream or buffer response；tok 从上游响应（SSE usage 块或非流式 JSON）提取。
|
||
// 上游可能返回 HTTP 200 但 body/SSE 内带 error（OpenRouter 超时等），
|
||
// 此时按失败记账（errCode 非空），非流式错误体以 502 返回给客户端。
|
||
var tok convert.TokenUsage
|
||
var errCode string
|
||
if req.Stream {
|
||
tok, errCode = g.streamResponse(c, resp, req.Protocol, targetFormat, req.rawBuf)
|
||
} else {
|
||
tok, errCode, _ = g.bufferResponse(c, resp, req.Protocol, targetFormat, req.rawBuf)
|
||
}
|
||
resp.Body.Close()
|
||
if errCode != "" {
|
||
// 记账为失败（错误码），不产生费用；响应内容已由 buffer/stream 写出
|
||
g.recordUsage(req, cand, ch, usage.Event{
|
||
IsError: true,
|
||
ErrorCode: errCode,
|
||
LatencyMS: int(time.Since(start).Milliseconds()),
|
||
}, tok, targetFormat)
|
||
return
|
||
}
|
||
// 成功记录：用量 + 定价计费。
|
||
g.recordUsage(req, cand, ch, usage.Event{
|
||
LatencyMS: int(time.Since(start).Milliseconds()),
|
||
}, tok, targetFormat)
|
||
return
|
||
}
|
||
|
||
// 全部候选失败（每个候选失败时已各自 RecordFailure，不再重复计数）
|
||
g.writeError(c, lastErrStatus, lastErrBody)
|
||
}
|
||
|
||
// rewriteModel 把 JSON 请求体顶层的 model 字段替换为 upstreamModel。
|
||
func rewriteModel(body []byte, upstreamModel string) []byte {
|
||
var m map[string]json.RawMessage
|
||
if json.Unmarshal(body, &m) != nil {
|
||
return body
|
||
}
|
||
if _, ok := m["model"]; !ok {
|
||
return body
|
||
}
|
||
m["model"], _ = json.Marshal(upstreamModel)
|
||
out, err := json.Marshal(m)
|
||
if err != nil {
|
||
return body
|
||
}
|
||
return out
|
||
}
|
||
|
||
// recordUsage 汇总一次请求的用量事件并异步落库。tok 为从上游响应提取的用量，
|
||
// 其 token 语义由 upstreamProto（渠道实际使用的上游协议）决定。
|
||
// cand/ch 可为 nil（无可用渠道的失败场景）。
|
||
func (g *Gateway) recordUsage(req *Request, cand *channel.Candidate, ch *store.Channel, ev usage.Event, tok convert.TokenUsage, upstreamProto string) {
|
||
if g.usageRec == nil {
|
||
return
|
||
}
|
||
ev.UserID = req.UserID
|
||
ev.ModelName = req.Model
|
||
ev.Protocol = req.Protocol
|
||
ev.RequestID = req.RequestID
|
||
if req.APIKey != nil {
|
||
ev.KeyID = req.APIKey.ID
|
||
}
|
||
if ch != nil {
|
||
ev.ChannelID = ch.ID
|
||
}
|
||
if cand != nil && cand.Binding != nil {
|
||
ev.ModelID = cand.Binding.ModelID
|
||
}
|
||
ev.PromptTokens = tok.InputTokens
|
||
ev.CompletionTokens = tok.OutputTokens
|
||
ev.CacheReadTokens = tok.CacheReadTokens
|
||
ev.CacheCreationTokens = tok.CacheCreationTokens
|
||
// 原始请求/响应（仅管理员+开关开启时捕获）。
|
||
if req.CaptureRaw {
|
||
ev.RawRequest = string(req.Body)
|
||
if req.rawBuf != nil {
|
||
ev.RawResponse = req.rawBuf.String()
|
||
}
|
||
}
|
||
// 定价与成本（价格按每百万 token 的 USD 单价）。
|
||
// 成本口径按上游协议区分（详见 ComputeCost）：OpenAI 系 prompt 含缓存读需扣减；
|
||
// Anthropic 的 input_tokens 不含缓存，缓存写按输入价 ×1.25。
|
||
if ev.ModelID != 0 {
|
||
if m, err := g.modelDAO.GetByID(ev.ModelID); err == nil {
|
||
ev.InputPrice = m.InputPrice
|
||
ev.OutputPrice = m.OutputPrice
|
||
ev.CacheReadPrice = m.CacheReadPrice
|
||
}
|
||
}
|
||
if !ev.IsError {
|
||
ev.Cost = ComputeCost(upstreamProto, tok.InputTokens, tok.OutputTokens, tok.CacheReadTokens, tok.CacheCreationTokens,
|
||
ev.InputPrice, ev.OutputPrice, ev.CacheReadPrice)
|
||
}
|
||
g.usageRec.Record(ev)
|
||
}
|
||
|
||
// conversionTarget 决定客户端协议在渠道上的处理方式：
|
||
// 渠道声明支持该协议则直通；否则转为其首选支持协议（chat > messages > responses）。
|
||
func (g *Gateway) conversionTarget(ch *store.Channel, clientProto string) string {
|
||
formats := ch.FormatsEffective()
|
||
for _, f := range formats {
|
||
if f == clientProto {
|
||
return clientProto
|
||
}
|
||
}
|
||
for _, p := range []string{convert.ProtoChat, convert.ProtoMessages, convert.ProtoResponses} {
|
||
for _, f := range formats {
|
||
if f == p {
|
||
return p
|
||
}
|
||
}
|
||
}
|
||
return ""
|
||
}
|
||
|
||
func (g *Gateway) getUpstreamPath(protocol string) string {
|
||
switch protocol {
|
||
case "chat":
|
||
return "/chat/completions"
|
||
case "messages":
|
||
return "/messages"
|
||
case "responses":
|
||
return "/responses"
|
||
default:
|
||
return "/chat/completions"
|
||
}
|
||
}
|
||
|
||
func (g *Gateway) setHeaders(req *http.Request, ch *store.Channel, apiKey string, format string) {
|
||
req.Header.Set("Content-Type", "application/json")
|
||
|
||
switch ch.Provider {
|
||
case store.ChannelProviderOpenAI, store.ChannelProviderCompatible:
|
||
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||
case store.ChannelProviderAnthropic:
|
||
req.Header.Set("x-api-key", apiKey)
|
||
req.Header.Set("anthropic-version", "2023-06-01")
|
||
}
|
||
}
|
||
|
||
|
||
// streamResponse 流式响应：按 \n\n 分块零缓冲转发；跨协议时逐行转换。
|
||
// 返回从上游 SSE usage 块累计的 token 用量（按上游协议解析）。
|
||
// capture 非 nil 时把上游原始行累积进去（原始响应记录）。
|
||
// 上游部分实现（如 OpenRouter）在超时时返回 HTTP 200 但 SSE data 内带
|
||
// error 字段；检测到则返回错误码，供 Dispatch 按失败记账。
|
||
func (g *Gateway) streamResponse(c *gin.Context, resp *http.Response, clientProto, upstreamProto string, capture *strings.Builder) (convert.TokenUsage, string) {
|
||
w := c.Writer
|
||
c.Header("Content-Type", "text/event-stream")
|
||
c.Header("Cache-Control", "no-cache")
|
||
c.Header("Connection", "keep-alive")
|
||
c.Status(http.StatusOK)
|
||
|
||
flusher, _ := w.(http.Flusher)
|
||
|
||
// 跨协议时按行转换；同协议直通（lineConv 为 nil）。
|
||
var lineConv func([]byte) []byte
|
||
if upstreamProto != clientProto {
|
||
lineConv = convert.NewStreamTransformer(upstreamProto, clientProto)
|
||
}
|
||
|
||
// 上游原始行按 \n\n 分块，避免把 data 行内的转义换行当成事件边界。
|
||
// 同时喂入用量累计器（usage 块可能出现在任一事件）。
|
||
r := bufio.NewReaderSize(resp.Body, 32*1024)
|
||
accum := convert.NewStreamUsageAccum()
|
||
errCode := ""
|
||
for {
|
||
buf := []byte{}
|
||
for {
|
||
line, err := r.ReadSlice('\n')
|
||
if err == bufio.ErrBufferFull {
|
||
buf = append(buf, line...)
|
||
continue
|
||
}
|
||
buf = append(buf, line...)
|
||
if err == io.EOF {
|
||
if len(buf) == 0 {
|
||
return accum.Usage(), errCode
|
||
}
|
||
if !bytes.HasSuffix(buf, []byte("\n")) {
|
||
buf = append(buf, '\n')
|
||
}
|
||
} else if err != nil {
|
||
log.Printf("stream read error: %v", err)
|
||
return accum.Usage(), errCode
|
||
}
|
||
if len(buf) >= 2 && bytes.HasSuffix(buf, []byte("\n\n")) {
|
||
break
|
||
}
|
||
}
|
||
|
||
// 原始响应捕获（仅管理员+开关开启时启用）。
|
||
if capture != nil {
|
||
capture.Write(buf)
|
||
}
|
||
|
||
// 先解析用量（data: {...} 行），再决定转发内容。
|
||
for _, data := range sseDataPayloads(buf) {
|
||
accum.Feed(data, upstreamProto)
|
||
if errCode == "" && streamChunkHasError(data) {
|
||
errCode = "upstream_stream_error"
|
||
}
|
||
}
|
||
|
||
out := buf
|
||
if lineConv != nil {
|
||
out = lineConv(buf)
|
||
}
|
||
if len(out) == 0 {
|
||
continue
|
||
}
|
||
if _, err := w.Write(out); err != nil {
|
||
return accum.Usage(), errCode // 客户端已断开
|
||
}
|
||
if flusher != nil {
|
||
flusher.Flush()
|
||
}
|
||
|
||
// 流结束标记：chat/messages 上游以 data: [DONE] 收尾。部分上游（keep-alive）
|
||
// 发完 [DONE] 后不关连接，继续读会阻塞到超时；据此主动收尾。
|
||
// responses 协议没有 [DONE]，以 response.completed 事件收尾。
|
||
if streamTerminated(buf, upstreamProto) {
|
||
return accum.Usage(), errCode
|
||
}
|
||
}
|
||
}
|
||
|
||
// streamChunkHasError 判断一块 SSE data 载荷是否带 error 字段（OpenRouter 超时等）。
|
||
func streamChunkHasError(data []byte) bool {
|
||
var m map[string]any
|
||
if json.Unmarshal(data, &m) != nil {
|
||
return false
|
||
}
|
||
if _, ok := m["error"]; ok {
|
||
return true
|
||
}
|
||
// responses 协议错误事件可能形如 {"type":"error",...}
|
||
return m["type"] == "error"
|
||
}
|
||
|
||
// streamTerminated 判断一块 SSE 是否为上游流的结束事件。
|
||
func streamTerminated(chunk []byte, proto string) bool {
|
||
switch proto {
|
||
case convert.ProtoChat:
|
||
// chat 上游以 data: [DONE] 收尾；keep-alive 上游发完不关连接。
|
||
return bytes.Contains(chunk, []byte("data: [DONE]"))
|
||
case convert.ProtoMessages:
|
||
// messages 上游以 message_stop 事件结束（无 [DONE]）。
|
||
return bytes.Contains(chunk, []byte(`"type":"message_stop"`)) ||
|
||
bytes.Contains(chunk, []byte(`"type": "message_stop"`)) ||
|
||
bytes.Contains(chunk, []byte("data: [DONE]"))
|
||
case convert.ProtoResponses:
|
||
return bytes.Contains(chunk, []byte(`"response.completed"`)) ||
|
||
bytes.Contains(chunk, []byte(`"type":"response.completed"`))
|
||
}
|
||
return false
|
||
}
|
||
|
||
// sseDataPayloads 从一块 SSE（一个完整事件，\n\n 结尾）中取出所有 data 行的原始载荷。
|
||
func sseDataPayloads(chunk []byte) [][]byte {
|
||
var out [][]byte
|
||
for _, line := range bytes.Split(chunk, []byte("\n")) {
|
||
line = bytes.TrimSuffix(line, []byte("\r"))
|
||
if !bytes.HasPrefix(line, []byte("data:")) {
|
||
continue
|
||
}
|
||
payload := bytes.TrimPrefix(line, []byte("data:"))
|
||
payload = bytes.TrimPrefix(payload, []byte(" "))
|
||
if len(payload) == 0 || bytes.Equal(payload, []byte("[DONE]")) {
|
||
continue
|
||
}
|
||
out = append(out, payload)
|
||
}
|
||
return out
|
||
}
|
||
|
||
// bufferResponse 非流式响应：整体读取、可选转换后写回。
|
||
// 返回 (用量, 错误码, 是否错误)。部分上游（如 OpenRouter）在超时时返回
|
||
// HTTP 200 但 JSON 内含 error 字段，需要识别并让调用方按失败处理。
|
||
func (g *Gateway) bufferResponse(c *gin.Context, resp *http.Response, clientProto, upstreamProto string, capture *strings.Builder) (convert.TokenUsage, string, bool) {
|
||
body, err := io.ReadAll(resp.Body)
|
||
if err != nil {
|
||
g.writeError(c, http.StatusBadGateway, "failed to read response")
|
||
return convert.TokenUsage{}, "", false
|
||
}
|
||
|
||
// 原始响应捕获（仅管理员+开关开启时启用）。
|
||
if capture != nil {
|
||
capture.Write(body)
|
||
}
|
||
|
||
// 用量从上游原始响应体提取（先于转换，转换会改字段名）。
|
||
tok, _ := convert.ExtractUsageJSON(body, upstreamProto)
|
||
|
||
// HTTP 200 但带 error 字段（OpenRouter 超时 504 等）：识别并转失败。
|
||
errCode, isErr := bodyHasError(body)
|
||
|
||
out := body
|
||
if upstreamProto != clientProto {
|
||
if converted, cerr := convert.ConvertResponse(body, upstreamProto, clientProto); cerr == nil {
|
||
out = converted
|
||
} else {
|
||
// 转换失败时至少剥掉非 JSON 前缀，让客户端能解析出正文
|
||
out = convert.CleanJSON(body)
|
||
}
|
||
} else {
|
||
// 直通：部分上游（如 OpenRouter）的 non-stream 响应在 JSON 前夹带空白/注释
|
||
out = convert.CleanJSON(body)
|
||
}
|
||
// 上游错误体：用 502 返回，让客户端感知失败（不伪装成 200）。
|
||
if isErr {
|
||
c.Data(http.StatusBadGateway, "application/json", out)
|
||
return tok, errCode, true
|
||
}
|
||
c.Data(resp.StatusCode, "application/json", out)
|
||
return tok, errCode, false
|
||
}
|
||
|
||
// bodyHasError 判断 JSON 响应体是否带 error 字段（openai 风格 {"error":{...}} 或
|
||
// anthropic 风格 {"type":"error",...}）。返回 (错误码, 是否错误)。找不到 JSON 返回 ("", false)。
|
||
func bodyHasError(body []byte) (string, bool) {
|
||
var m map[string]any
|
||
if json.Unmarshal(bytes.TrimSpace(body), &m) != nil {
|
||
return "", false
|
||
}
|
||
if _, ok := m["error"]; ok {
|
||
return "upstream_error", true
|
||
}
|
||
if m["type"] == "error" {
|
||
return "upstream_error", true
|
||
}
|
||
return "", false
|
||
}
|
||
|
||
func (g *Gateway) writeError(c *gin.Context, status int, message string) {
|
||
protocol := c.GetHeader("X-Protocol")
|
||
if protocol == "" {
|
||
protocol = "chat"
|
||
}
|
||
|
||
switch {
|
||
case strings.Contains(c.GetHeader("Accept"), "text/event-stream"):
|
||
c.Header("Content-Type", "text/event-stream")
|
||
c.Status(status)
|
||
fmt.Fprintf(c.Writer, "data: {\"error\":{\"message\":\"%s\"}}\n\n", message)
|
||
fmt.Fprintf(c.Writer, "data: [DONE]\n\n")
|
||
case protocol == "messages":
|
||
c.JSON(status, gin.H{
|
||
"type": "error",
|
||
"error": gin.H{
|
||
"type": "api_error",
|
||
"message": message,
|
||
},
|
||
})
|
||
default:
|
||
c.JSON(status, gin.H{
|
||
"error": gin.H{
|
||
"message": message,
|
||
"type": "invalid_request_error",
|
||
},
|
||
})
|
||
}
|
||
}
|