添加完整项目文件

包含Go API项目的所有源代码、配置文件、Docker配置、文档和前端资源 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
2025-08-21 23:33:10 +08:00
parent 69af9723d9
commit fb44c8bf03
513 changed files with 90387 additions and 0 deletions
--- a/relay/helper/common.go
+++ b/relay/helper/common.go
@@ -0,0 +1,167 @@
+package helper
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"net/http"
+	"one-api/common"
+	"one-api/dto"
+	"one-api/types"
+
+	"github.com/gin-gonic/gin"
+	"github.com/gorilla/websocket"
+)
+
+func SetEventStreamHeaders(c *gin.Context) {
+	// 检查是否已经设置过头部
+	if _, exists := c.Get("event_stream_headers_set"); exists {
+		return
+	}
+
+	c.Writer.Header().Set("Content-Type", "text/event-stream")
+	c.Writer.Header().Set("Cache-Control", "no-cache")
+	c.Writer.Header().Set("Connection", "keep-alive")
+	c.Writer.Header().Set("Transfer-Encoding", "chunked")
+	c.Writer.Header().Set("X-Accel-Buffering", "no")
+
+	// 设置标志，表示头部已经设置过
+	c.Set("event_stream_headers_set", true)
+}
+
+func ClaudeData(c *gin.Context, resp dto.ClaudeResponse) error {
+	jsonData, err := json.Marshal(resp)
+	if err != nil {
+		common.SysError("error marshalling stream response: " + err.Error())
+	} else {
+		c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", resp.Type)})
+		c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonData)})
+	}
+	if flusher, ok := c.Writer.(http.Flusher); ok {
+		flusher.Flush()
+	} else {
+		return errors.New("streaming error: flusher not found")
+	}
+	return nil
+}
+
+func ClaudeChunkData(c *gin.Context, resp dto.ClaudeResponse, data string) {
+	c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", resp.Type)})
+	c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("data: %s\n", data)})
+	if flusher, ok := c.Writer.(http.Flusher); ok {
+		flusher.Flush()
+	}
+}
+
+func ResponseChunkData(c *gin.Context, resp dto.ResponsesStreamResponse, data string) {
+	c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", resp.Type)})
+	c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("data: %s", data)})
+	if flusher, ok := c.Writer.(http.Flusher); ok {
+		flusher.Flush()
+	}
+}
+
+func StringData(c *gin.Context, str string) error {
+	//str = strings.TrimPrefix(str, "data: ")
+	//str = strings.TrimSuffix(str, "\r")
+	c.Render(-1, common.CustomEvent{Data: "data: " + str})
+	if flusher, ok := c.Writer.(http.Flusher); ok {
+		flusher.Flush()
+	} else {
+		return errors.New("streaming error: flusher not found")
+	}
+	return nil
+}
+
+func PingData(c *gin.Context) error {
+	c.Writer.Write([]byte(": PING\n\n"))
+	if flusher, ok := c.Writer.(http.Flusher); ok {
+		flusher.Flush()
+	} else {
+		return errors.New("streaming error: flusher not found")
+	}
+	return nil
+}
+
+func ObjectData(c *gin.Context, object interface{}) error {
+	if object == nil {
+		return errors.New("object is nil")
+	}
+	jsonData, err := common.Marshal(object)
+	if err != nil {
+		return fmt.Errorf("error marshalling object: %w", err)
+	}
+	return StringData(c, string(jsonData))
+}
+
+func Done(c *gin.Context) {
+	_ = StringData(c, "[DONE]")
+}
+
+func WssString(c *gin.Context, ws *websocket.Conn, str string) error {
+	if ws == nil {
+		common.LogError(c, "websocket connection is nil")
+		return errors.New("websocket connection is nil")
+	}
+	//common.LogInfo(c, fmt.Sprintf("sending message: %s", str))
+	return ws.WriteMessage(1, []byte(str))
+}
+
+func WssObject(c *gin.Context, ws *websocket.Conn, object interface{}) error {
+	jsonData, err := json.Marshal(object)
+	if err != nil {
+		return fmt.Errorf("error marshalling object: %w", err)
+	}
+	if ws == nil {
+		common.LogError(c, "websocket connection is nil")
+		return errors.New("websocket connection is nil")
+	}
+	//common.LogInfo(c, fmt.Sprintf("sending message: %s", jsonData))
+	return ws.WriteMessage(1, jsonData)
+}
+
+func WssError(c *gin.Context, ws *websocket.Conn, openaiError types.OpenAIError) {
+	errorObj := &dto.RealtimeEvent{
+		Type:    "error",
+		EventId: GetLocalRealtimeID(c),
+		Error:   &openaiError,
+	}
+	_ = WssObject(c, ws, errorObj)
+}
+
+func GetResponseID(c *gin.Context) string {
+	logID := c.GetString(common.RequestIdKey)
+	return fmt.Sprintf("chatcmpl-%s", logID)
+}
+
+func GetLocalRealtimeID(c *gin.Context) string {
+	logID := c.GetString(common.RequestIdKey)
+	return fmt.Sprintf("evt_%s", logID)
+}
+
+func GenerateStopResponse(id string, createAt int64, model string, finishReason string) *dto.ChatCompletionsStreamResponse {
+	return &dto.ChatCompletionsStreamResponse{
+		Id:                id,
+		Object:            "chat.completion.chunk",
+		Created:           createAt,
+		Model:             model,
+		SystemFingerprint: nil,
+		Choices: []dto.ChatCompletionsStreamResponseChoice{
+			{
+				FinishReason: &finishReason,
+			},
+		},
+	}
+}
+
+func GenerateFinalUsageResponse(id string, createAt int64, model string, usage dto.Usage) *dto.ChatCompletionsStreamResponse {
+	return &dto.ChatCompletionsStreamResponse{
+		Id:                id,
+		Object:            "chat.completion.chunk",
+		Created:           createAt,
+		Model:             model,
+		SystemFingerprint: nil,
+		Choices:           make([]dto.ChatCompletionsStreamResponseChoice, 0),
+		Usage:             &usage,
+	}
+}
--- a/relay/helper/model_mapped.go
+++ b/relay/helper/model_mapped.go
@@ -0,0 +1,92 @@
+package helper
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	common2 "one-api/common"
+	"one-api/dto"
+	"one-api/relay/common"
+
+	"github.com/gin-gonic/gin"
+)
+
+func ModelMappedHelper(c *gin.Context, info *common.RelayInfo, request any) error {
+	// map model name
+	modelMapping := c.GetString("model_mapping")
+	if modelMapping != "" && modelMapping != "{}" {
+		modelMap := make(map[string]string)
+		err := json.Unmarshal([]byte(modelMapping), &modelMap)
+		if err != nil {
+			return fmt.Errorf("unmarshal_model_mapping_failed")
+		}
+
+		// 支持链式模型重定向，最终使用链尾的模型
+		currentModel := info.OriginModelName
+		visitedModels := map[string]bool{
+			currentModel: true,
+		}
+		for {
+			if mappedModel, exists := modelMap[currentModel]; exists && mappedModel != "" {
+				// 模型重定向循环检测，避免无限循环
+				if visitedModels[mappedModel] {
+					if mappedModel == currentModel {
+						if currentModel == info.OriginModelName {
+							info.IsModelMapped = false
+							return nil
+						} else {
+							info.IsModelMapped = true
+							break
+						}
+					}
+					return errors.New("model_mapping_contains_cycle")
+				}
+				visitedModels[mappedModel] = true
+				currentModel = mappedModel
+				info.IsModelMapped = true
+			} else {
+				break
+			}
+		}
+		if info.IsModelMapped {
+			info.UpstreamModelName = currentModel
+		}
+	}
+	if request != nil {
+		switch info.RelayFormat {
+		case common.RelayFormatGemini:
+			// Gemini 模型映射
+		case common.RelayFormatClaude:
+			if claudeRequest, ok := request.(*dto.ClaudeRequest); ok {
+				claudeRequest.Model = info.UpstreamModelName
+			}
+		case common.RelayFormatOpenAIResponses:
+			if openAIResponsesRequest, ok := request.(*dto.OpenAIResponsesRequest); ok {
+				openAIResponsesRequest.Model = info.UpstreamModelName
+			}
+		case common.RelayFormatOpenAIAudio:
+			if openAIAudioRequest, ok := request.(*dto.AudioRequest); ok {
+				openAIAudioRequest.Model = info.UpstreamModelName
+			}
+		case common.RelayFormatOpenAIImage:
+			if imageRequest, ok := request.(*dto.ImageRequest); ok {
+				imageRequest.Model = info.UpstreamModelName
+			}
+		case common.RelayFormatRerank:
+			if rerankRequest, ok := request.(*dto.RerankRequest); ok {
+				rerankRequest.Model = info.UpstreamModelName
+			}
+		case common.RelayFormatEmbedding:
+			if embeddingRequest, ok := request.(*dto.EmbeddingRequest); ok {
+				embeddingRequest.Model = info.UpstreamModelName
+			}
+		default:
+			if openAIRequest, ok := request.(*dto.GeneralOpenAIRequest); ok {
+				openAIRequest.Model = info.UpstreamModelName
+			} else {
+				common2.LogWarn(c, fmt.Sprintf("model mapped but request type %T not supported", request))
+			}
+		}
+	}
+	return nil
+}
--- a/relay/helper/price.go
+++ b/relay/helper/price.go
@@ -0,0 +1,161 @@
+package helper
+
+import (
+	"fmt"
+	"one-api/common"
+	relaycommon "one-api/relay/common"
+	"one-api/setting/ratio_setting"
+
+	"github.com/gin-gonic/gin"
+)
+
+type GroupRatioInfo struct {
+	GroupRatio        float64
+	GroupSpecialRatio float64
+	HasSpecialRatio   bool
+}
+
+type PriceData struct {
+	ModelPrice             float64
+	ModelRatio             float64
+	CompletionRatio        float64
+	CacheRatio             float64
+	CacheCreationRatio     float64
+	ImageRatio             float64
+	UsePrice               bool
+	ShouldPreConsumedQuota int
+	GroupRatioInfo         GroupRatioInfo
+}
+
+func (p PriceData) ToSetting() string {
+	return fmt.Sprintf("ModelPrice: %f, ModelRatio: %f, CompletionRatio: %f, CacheRatio: %f, GroupRatio: %f, UsePrice: %t, CacheCreationRatio: %f, ShouldPreConsumedQuota: %d, ImageRatio: %f", p.ModelPrice, p.ModelRatio, p.CompletionRatio, p.CacheRatio, p.GroupRatioInfo.GroupRatio, p.UsePrice, p.CacheCreationRatio, p.ShouldPreConsumedQuota, p.ImageRatio)
+}
+
+// HandleGroupRatio checks for "auto_group" in the context and updates the group ratio and relayInfo.UsingGroup if present
+func HandleGroupRatio(ctx *gin.Context, relayInfo *relaycommon.RelayInfo) GroupRatioInfo {
+	groupRatioInfo := GroupRatioInfo{
+		GroupRatio:        1.0, // default ratio
+		GroupSpecialRatio: -1,
+	}
+
+	// check auto group
+	autoGroup, exists := ctx.Get("auto_group")
+	if exists {
+		if common.DebugEnabled {
+			println(fmt.Sprintf("final group: %s", autoGroup))
+		}
+		relayInfo.UsingGroup = autoGroup.(string)
+	}
+
+	// check user group special ratio
+	userGroupRatio, ok := ratio_setting.GetGroupGroupRatio(relayInfo.UserGroup, relayInfo.UsingGroup)
+	if ok {
+		// user group special ratio
+		groupRatioInfo.GroupSpecialRatio = userGroupRatio
+		groupRatioInfo.GroupRatio = userGroupRatio
+		groupRatioInfo.HasSpecialRatio = true
+	} else {
+		// normal group ratio
+		groupRatioInfo.GroupRatio = ratio_setting.GetGroupRatio(relayInfo.UsingGroup)
+	}
+
+	return groupRatioInfo
+}
+
+func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens int, maxTokens int) (PriceData, error) {
+	modelPrice, usePrice := ratio_setting.GetModelPrice(info.OriginModelName, false)
+
+	groupRatioInfo := HandleGroupRatio(c, info)
+
+	var preConsumedQuota int
+	var modelRatio float64
+	var completionRatio float64
+	var cacheRatio float64
+	var imageRatio float64
+	var cacheCreationRatio float64
+	if !usePrice {
+		preConsumedTokens := common.PreConsumedQuota
+		if maxTokens != 0 {
+			preConsumedTokens = promptTokens + maxTokens
+		}
+		var success bool
+		var matchName string
+		modelRatio, success, matchName = ratio_setting.GetModelRatio(info.OriginModelName)
+		if !success {
+			acceptUnsetRatio := false
+			if info.UserSetting.AcceptUnsetRatioModel {
+				acceptUnsetRatio = true
+			}
+			if !acceptUnsetRatio {
+				return PriceData{}, fmt.Errorf("模型 %s 倍率或价格未配置，请联系管理员设置或开始自用模式；Model %s ratio or price not set, please set or start self-use mode", matchName, matchName)
+			}
+		}
+		completionRatio = ratio_setting.GetCompletionRatio(info.OriginModelName)
+		cacheRatio, _ = ratio_setting.GetCacheRatio(info.OriginModelName)
+		cacheCreationRatio, _ = ratio_setting.GetCreateCacheRatio(info.OriginModelName)
+		imageRatio, _ = ratio_setting.GetImageRatio(info.OriginModelName)
+		ratio := modelRatio * groupRatioInfo.GroupRatio
+		preConsumedQuota = int(float64(preConsumedTokens) * ratio)
+	} else {
+		preConsumedQuota = int(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio)
+	}
+
+	priceData := PriceData{
+		ModelPrice:             modelPrice,
+		ModelRatio:             modelRatio,
+		CompletionRatio:        completionRatio,
+		GroupRatioInfo:         groupRatioInfo,
+		UsePrice:               usePrice,
+		CacheRatio:             cacheRatio,
+		ImageRatio:             imageRatio,
+		CacheCreationRatio:     cacheCreationRatio,
+		ShouldPreConsumedQuota: preConsumedQuota,
+	}
+
+	if common.DebugEnabled {
+		println(fmt.Sprintf("model_price_helper result: %s", priceData.ToSetting()))
+	}
+
+	return priceData, nil
+}
+
+type PerCallPriceData struct {
+	ModelPrice     float64
+	Quota          int
+	GroupRatioInfo GroupRatioInfo
+}
+
+// ModelPriceHelperPerCall 按次计费的 PriceHelper (MJ、Task)
+func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) PerCallPriceData {
+	groupRatioInfo := HandleGroupRatio(c, info)
+
+	modelPrice, success := ratio_setting.GetModelPrice(info.OriginModelName, true)
+	// 如果没有配置价格，则使用默认价格
+	if !success {
+		defaultPrice, ok := ratio_setting.GetDefaultModelRatioMap()[info.OriginModelName]
+		if !ok {
+			modelPrice = 0.1
+		} else {
+			modelPrice = defaultPrice
+		}
+	}
+	quota := int(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio)
+	priceData := PerCallPriceData{
+		ModelPrice:     modelPrice,
+		Quota:          quota,
+		GroupRatioInfo: groupRatioInfo,
+	}
+	return priceData
+}
+
+func ContainPriceOrRatio(modelName string) bool {
+	_, ok := ratio_setting.GetModelPrice(modelName, false)
+	if ok {
+		return true
+	}
+	_, ok, _ = ratio_setting.GetModelRatio(modelName)
+	if ok {
+		return true
+	}
+	return false
+}
--- a/relay/helper/stream_scanner.go
+++ b/relay/helper/stream_scanner.go
@@ -0,0 +1,259 @@
+package helper
+
+import (
+	"bufio"
+	"context"
+	"fmt"
+	"io"
+	"net/http"
+	"one-api/common"
+	"one-api/constant"
+	relaycommon "one-api/relay/common"
+	"one-api/setting/operation_setting"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/bytedance/gopkg/util/gopool"
+
+	"github.com/gin-gonic/gin"
+)
+
+const (
+	InitialScannerBufferSize = 64 << 10 // 64KB (64*1024)
+	MaxScannerBufferSize     = 10 << 20 // 10MB (10*1024*1024)
+	DefaultPingInterval      = 10 * time.Second
+)
+
+func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo, dataHandler func(data string) bool) {
+
+	if resp == nil || dataHandler == nil {
+		return
+	}
+
+	// 确保响应体总是被关闭
+	defer func() {
+		if resp.Body != nil {
+			resp.Body.Close()
+		}
+	}()
+
+	streamingTimeout := time.Duration(constant.StreamingTimeout) * time.Second
+	if strings.HasPrefix(info.UpstreamModelName, "o") {
+		// twice timeout for thinking model
+		streamingTimeout *= 2
+	}
+
+	var (
+		stopChan   = make(chan bool, 3) // 增加缓冲区避免阻塞
+		scanner    = bufio.NewScanner(resp.Body)
+		ticker     = time.NewTicker(streamingTimeout)
+		pingTicker *time.Ticker
+		writeMutex sync.Mutex     // Mutex to protect concurrent writes
+		wg         sync.WaitGroup // 用于等待所有 goroutine 退出
+	)
+
+	generalSettings := operation_setting.GetGeneralSetting()
+	pingEnabled := generalSettings.PingIntervalEnabled
+	pingInterval := time.Duration(generalSettings.PingIntervalSeconds) * time.Second
+	if pingInterval <= 0 {
+		pingInterval = DefaultPingInterval
+	}
+
+	if pingEnabled {
+		pingTicker = time.NewTicker(pingInterval)
+	}
+
+	if common.DebugEnabled {
+		// print timeout and ping interval for debugging
+		println("relay timeout seconds:", common.RelayTimeout)
+		println("streaming timeout seconds:", int64(streamingTimeout.Seconds()))
+		println("ping interval seconds:", int64(pingInterval.Seconds()))
+	}
+
+	// 改进资源清理，确保所有 goroutine 正确退出
+	defer func() {
+		// 通知所有 goroutine 停止
+		common.SafeSendBool(stopChan, true)
+
+		ticker.Stop()
+		if pingTicker != nil {
+			pingTicker.Stop()
+		}
+
+		// 等待所有 goroutine 退出，最多等待5秒
+		done := make(chan struct{})
+		go func() {
+			wg.Wait()
+			close(done)
+		}()
+
+		select {
+		case <-done:
+		case <-time.After(5 * time.Second):
+			common.LogError(c, "timeout waiting for goroutines to exit")
+		}
+
+		close(stopChan)
+	}()
+
+	scanner.Buffer(make([]byte, InitialScannerBufferSize), MaxScannerBufferSize)
+	scanner.Split(bufio.ScanLines)
+	SetEventStreamHeaders(c)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	ctx = context.WithValue(ctx, "stop_chan", stopChan)
+
+	// Handle ping data sending with improved error handling
+	if pingEnabled && pingTicker != nil {
+		wg.Add(1)
+		gopool.Go(func() {
+			defer func() {
+				wg.Done()
+				if r := recover(); r != nil {
+					common.LogError(c, fmt.Sprintf("ping goroutine panic: %v", r))
+					common.SafeSendBool(stopChan, true)
+				}
+				if common.DebugEnabled {
+					println("ping goroutine exited")
+				}
+			}()
+
+			// 添加超时保护，防止 goroutine 无限运行
+			maxPingDuration := 30 * time.Minute // 最大 ping 持续时间
+			pingTimeout := time.NewTimer(maxPingDuration)
+			defer pingTimeout.Stop()
+
+			for {
+				select {
+				case <-pingTicker.C:
+					// 使用超时机制防止写操作阻塞
+					done := make(chan error, 1)
+					go func() {
+						writeMutex.Lock()
+						defer writeMutex.Unlock()
+						done <- PingData(c)
+					}()
+
+					select {
+					case err := <-done:
+						if err != nil {
+							common.LogError(c, "ping data error: "+err.Error())
+							return
+						}
+						if common.DebugEnabled {
+							println("ping data sent")
+						}
+					case <-time.After(10 * time.Second):
+						common.LogError(c, "ping data send timeout")
+						return
+					case <-ctx.Done():
+						return
+					case <-stopChan:
+						return
+					}
+				case <-ctx.Done():
+					return
+				case <-stopChan:
+					return
+				case <-c.Request.Context().Done():
+					// 监听客户端断开连接
+					return
+				case <-pingTimeout.C:
+					common.LogError(c, "ping goroutine max duration reached")
+					return
+				}
+			}
+		})
+	}
+
+	// Scanner goroutine with improved error handling
+	wg.Add(1)
+	common.RelayCtxGo(ctx, func() {
+		defer func() {
+			wg.Done()
+			if r := recover(); r != nil {
+				common.LogError(c, fmt.Sprintf("scanner goroutine panic: %v", r))
+			}
+			common.SafeSendBool(stopChan, true)
+			if common.DebugEnabled {
+				println("scanner goroutine exited")
+			}
+		}()
+
+		for scanner.Scan() {
+			// 检查是否需要停止
+			select {
+			case <-stopChan:
+				return
+			case <-ctx.Done():
+				return
+			case <-c.Request.Context().Done():
+				return
+			default:
+			}
+
+			ticker.Reset(streamingTimeout)
+			data := scanner.Text()
+			if common.DebugEnabled {
+				println(data)
+			}
+
+			if len(data) < 6 {
+				continue
+			}
+			if data[:5] != "data:" && data[:6] != "[DONE]" {
+				continue
+			}
+			data = data[5:]
+			data = strings.TrimLeft(data, " ")
+			data = strings.TrimSuffix(data, "\r")
+			if !strings.HasPrefix(data, "[DONE]") {
+				info.SetFirstResponseTime()
+
+				// 使用超时机制防止写操作阻塞
+				done := make(chan bool, 1)
+				go func() {
+					writeMutex.Lock()
+					defer writeMutex.Unlock()
+					done <- dataHandler(data)
+				}()
+
+				select {
+				case success := <-done:
+					if !success {
+						return
+					}
+				case <-time.After(10 * time.Second):
+					common.LogError(c, "data handler timeout")
+					return
+				case <-ctx.Done():
+					return
+				case <-stopChan:
+					return
+				}
+			}
+		}
+
+		if err := scanner.Err(); err != nil {
+			if err != io.EOF {
+				common.LogError(c, "scanner error: "+err.Error())
+			}
+		}
+	})
+
+	// 主循环等待完成或超时
+	select {
+	case <-ticker.C:
+		// 超时处理逻辑
+		common.LogError(c, "streaming timeout")
+	case <-stopChan:
+		// 正常结束
+		common.LogInfo(c, "streaming finished")
+	case <-c.Request.Context().Done():
+		// 客户端断开连接
+		common.LogInfo(c, "client disconnected")
+	}
+}