feat: 算力引擎 Go 后端核心（new-api）
- 统一 OpenAI 兼容 /v1 中继 + 渠道/额度/令牌/流水 - 各领域包：relay 模型网关、model 数据层、controller 管理 API
This commit is contained in:
@@ -0,0 +1,368 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"embed"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log"
|
||||
"net/http"
|
||||
"os"
|
||||
"os/signal"
|
||||
"strconv"
|
||||
"strings"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/controller"
|
||||
"github.com/QuantumNous/new-api/i18n"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/QuantumNous/new-api/middleware"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/oauth"
|
||||
perfmetrics "github.com/QuantumNous/new-api/pkg/perf_metrics"
|
||||
"github.com/QuantumNous/new-api/relay"
|
||||
kitutil "github.com/QuantumNous/new-api/relaykit/relayconvert/kitutil"
|
||||
"github.com/QuantumNous/new-api/router"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/service/authz"
|
||||
_ "github.com/QuantumNous/new-api/setting/performance_setting"
|
||||
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
||||
|
||||
"github.com/bytedance/gopkg/util/gopool"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/joho/godotenv"
|
||||
|
||||
_ "net/http/pprof"
|
||||
)
|
||||
|
||||
//go:embed web/dist
|
||||
var buildFS embed.FS
|
||||
|
||||
//go:embed web/dist/index.html
|
||||
var indexPage []byte
|
||||
|
||||
func main() {
|
||||
startTime := time.Now()
|
||||
kitutil.SetLogging(common.SysLog, func(message string) {
|
||||
logger.LogError(nil, message)
|
||||
})
|
||||
kitutil.SetSystemErrorLogging(common.SysError)
|
||||
|
||||
err := InitResources()
|
||||
if err != nil {
|
||||
common.FatalLog("failed to initialize resources: " + err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
common.SysLog("New API " + common.Version + " started")
|
||||
if os.Getenv("GIN_MODE") != "debug" {
|
||||
gin.SetMode(gin.ReleaseMode)
|
||||
}
|
||||
if common.DebugEnabled {
|
||||
common.SysLog("running in debug mode")
|
||||
}
|
||||
|
||||
kitutil.Debug.Store(common.DebugEnabled)
|
||||
|
||||
defer func() {
|
||||
err := model.CloseDB()
|
||||
if err != nil {
|
||||
common.FatalLog("failed to close database: " + err.Error())
|
||||
}
|
||||
}()
|
||||
|
||||
if common.RedisEnabled {
|
||||
// for compatibility with old versions
|
||||
common.MemoryCacheEnabled = true
|
||||
}
|
||||
if common.MemoryCacheEnabled {
|
||||
common.SysLog("memory cache enabled")
|
||||
common.SysLog(fmt.Sprintf("sync frequency: %d seconds", common.SyncFrequency))
|
||||
|
||||
// Add panic recovery and retry for InitChannelCache
|
||||
func() {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
common.SysLog(fmt.Sprintf("InitChannelCache panic: %v, retrying once", r))
|
||||
// Retry once
|
||||
_, _, fixErr := model.FixAbility()
|
||||
if fixErr != nil {
|
||||
common.FatalLog(fmt.Sprintf("InitChannelCache failed: %s", fixErr.Error()))
|
||||
}
|
||||
}
|
||||
}()
|
||||
model.InitChannelCache()
|
||||
}()
|
||||
|
||||
go model.SyncChannelCache(common.SyncFrequency)
|
||||
}
|
||||
|
||||
// Warm pricing after channel cache initialization so Advanced Custom
|
||||
// endpoint inference can read cached route settings on first request.
|
||||
model.GetPricing()
|
||||
|
||||
// 热更新配置
|
||||
go model.SyncOptions(common.SyncFrequency)
|
||||
|
||||
// 周期性重载授权策略，保证多节点/多 master 部署下权限变更能传播到每个实例
|
||||
go authz.StartPolicySync(common.SyncFrequency)
|
||||
|
||||
// 数据看板
|
||||
go model.UpdateQuotaData()
|
||||
|
||||
if os.Getenv("CHANNEL_UPDATE_FREQUENCY") != "" {
|
||||
frequency, err := strconv.Atoi(os.Getenv("CHANNEL_UPDATE_FREQUENCY"))
|
||||
if err != nil {
|
||||
common.FatalLog("failed to parse CHANNEL_UPDATE_FREQUENCY: " + err.Error())
|
||||
}
|
||||
go controller.AutomaticallyUpdateChannels(frequency)
|
||||
}
|
||||
|
||||
// Codex credential auto-refresh check every 10 minutes, refresh when expires within 1 day
|
||||
service.StartCodexCredentialAutoRefreshTask()
|
||||
|
||||
// Subscription quota reset task (daily/weekly/monthly/custom)
|
||||
service.StartSubscriptionQuotaResetTask()
|
||||
|
||||
// Report this process as a system instance so the System Info page can show
|
||||
// all currently alive nodes in multi-instance deployments.
|
||||
service.StartSystemInstanceReporter()
|
||||
|
||||
// Wire task polling adaptor factory (breaks service -> relay import cycle).
|
||||
// Must run before the system task runner starts: the async_task_poll handler
|
||||
// calls service.RunTaskPollingOnce, which needs this factory set.
|
||||
service.GetTaskAdaptorFunc = func(platform constant.TaskPlatform) service.TaskPollingAdaptor {
|
||||
a := relay.GetTaskAdaptor(platform)
|
||||
if a == nil {
|
||||
return nil
|
||||
}
|
||||
return a
|
||||
}
|
||||
|
||||
// Register the periodic channel test, upstream model update, and async task
|
||||
// polling (Midjourney / Suno / video) jobs as scheduled system tasks
|
||||
// (DB-lease dedup across masters + run history), then start the runner that
|
||||
// schedules and executes them. Master-only execution and the UpdateTask
|
||||
// switch are enforced inside the runner and each handler's Enabled().
|
||||
controller.RegisterScheduledSystemTasks()
|
||||
service.StartSystemTaskRunner()
|
||||
|
||||
if os.Getenv("BATCH_UPDATE_ENABLED") == "true" {
|
||||
common.BatchUpdateEnabled = true
|
||||
common.SysLog("batch update enabled with interval " + strconv.Itoa(common.BatchUpdateInterval) + "s")
|
||||
model.InitBatchUpdater()
|
||||
}
|
||||
|
||||
if os.Getenv("ENABLE_PPROF") == "true" {
|
||||
gopool.Go(func() {
|
||||
log.Println(http.ListenAndServe("0.0.0.0:8005", nil))
|
||||
})
|
||||
go common.Monitor()
|
||||
common.SysLog("pprof enabled")
|
||||
}
|
||||
|
||||
err = common.StartPyroScope()
|
||||
if err != nil {
|
||||
common.SysError(fmt.Sprintf("start pyroscope error : %v", err))
|
||||
}
|
||||
|
||||
// Initialize HTTP server
|
||||
server := gin.New()
|
||||
if err := middleware.ConfigureTrustedProxies(server); err != nil {
|
||||
common.FatalLog("failed to configure trusted proxies: " + err.Error())
|
||||
return
|
||||
}
|
||||
server.Use(gin.CustomRecovery(func(c *gin.Context, err any) {
|
||||
common.SysLog(fmt.Sprintf("panic detected: %v", err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{
|
||||
"error": gin.H{
|
||||
"message": fmt.Sprintf("Panic detected, error: %v. Please submit a issue here: https://github.com/Calcium-Ion/new-api", err),
|
||||
"type": "new_api_panic",
|
||||
},
|
||||
})
|
||||
}))
|
||||
// This will cause SSE not to work!!!
|
||||
//server.Use(gzip.Gzip(gzip.DefaultCompression))
|
||||
server.Use(middleware.RequestId())
|
||||
server.Use(middleware.Version())
|
||||
server.Use(middleware.I18n())
|
||||
middleware.SetUpLogger(server)
|
||||
InjectUmamiAnalytics()
|
||||
InjectGoogleAnalytics()
|
||||
|
||||
// 设置路由
|
||||
router.SetRouter(server, router.WebAssets{
|
||||
BuildFS: buildFS,
|
||||
IndexPage: indexPage,
|
||||
})
|
||||
var port = os.Getenv("PORT")
|
||||
if port == "" {
|
||||
port = strconv.Itoa(*common.Port)
|
||||
}
|
||||
|
||||
srv := &http.Server{
|
||||
Addr: ":" + port,
|
||||
Handler: server,
|
||||
}
|
||||
|
||||
go func() {
|
||||
if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
|
||||
common.FatalLog("failed to start HTTP server: " + err.Error())
|
||||
}
|
||||
}()
|
||||
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
|
||||
common.LogStartupSuccess(startTime, port)
|
||||
|
||||
quit := make(chan os.Signal, 1)
|
||||
signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
|
||||
sig := <-quit
|
||||
common.SysLog(fmt.Sprintf("received signal: %v, shutting down...", sig))
|
||||
|
||||
// SSE streams may run for minutes; give them time to finish before forced exit
|
||||
shutdownTimeout := time.Duration(common.GetEnvOrDefault("SHUTDOWN_TIMEOUT_SECONDS", 120)) * time.Second
|
||||
ctx, cancel := context.WithTimeout(context.Background(), shutdownTimeout)
|
||||
defer cancel()
|
||||
if err := srv.Shutdown(ctx); err != nil {
|
||||
common.SysError(fmt.Sprintf("server forced to shutdown: %v", err))
|
||||
}
|
||||
// 内存中的看板数据保存入库，避免重启丢失未落库数据 (issue #5679)
|
||||
if common.DataExportEnabled {
|
||||
model.SaveQuotaDataCache()
|
||||
}
|
||||
common.SysLog("server exited")
|
||||
}
|
||||
|
||||
func InjectUmamiAnalytics() {
|
||||
analyticsInjectBuilder := &strings.Builder{}
|
||||
if os.Getenv("UMAMI_WEBSITE_ID") != "" {
|
||||
umamiSiteID := os.Getenv("UMAMI_WEBSITE_ID")
|
||||
umamiScriptURL := os.Getenv("UMAMI_SCRIPT_URL")
|
||||
if umamiScriptURL == "" {
|
||||
umamiScriptURL = "https://analytics.umami.is/script.js"
|
||||
}
|
||||
analyticsInjectBuilder.WriteString("<script defer src=\"")
|
||||
analyticsInjectBuilder.WriteString(umamiScriptURL)
|
||||
analyticsInjectBuilder.WriteString("\" data-website-id=\"")
|
||||
analyticsInjectBuilder.WriteString(umamiSiteID)
|
||||
analyticsInjectBuilder.WriteString("\"></script>")
|
||||
}
|
||||
analyticsInjectBuilder.WriteString("<!--Umami QuantumNous-->\n")
|
||||
analyticsInject := []byte(analyticsInjectBuilder.String())
|
||||
placeholder := []byte("<!--umami-->\n")
|
||||
indexPage = bytes.ReplaceAll(indexPage, placeholder, analyticsInject)
|
||||
}
|
||||
|
||||
func InjectGoogleAnalytics() {
|
||||
analyticsInjectBuilder := &strings.Builder{}
|
||||
if os.Getenv("GOOGLE_ANALYTICS_ID") != "" {
|
||||
gaID := os.Getenv("GOOGLE_ANALYTICS_ID")
|
||||
// Google Analytics 4 (gtag.js)
|
||||
analyticsInjectBuilder.WriteString("<script async src=\"https://www.googletagmanager.com/gtag/js?id=")
|
||||
analyticsInjectBuilder.WriteString(gaID)
|
||||
analyticsInjectBuilder.WriteString("\"></script>")
|
||||
analyticsInjectBuilder.WriteString("<script>")
|
||||
analyticsInjectBuilder.WriteString("window.dataLayer = window.dataLayer || [];")
|
||||
analyticsInjectBuilder.WriteString("function gtag(){dataLayer.push(arguments);}")
|
||||
analyticsInjectBuilder.WriteString("gtag('js', new Date());")
|
||||
analyticsInjectBuilder.WriteString("gtag('config', '")
|
||||
analyticsInjectBuilder.WriteString(gaID)
|
||||
analyticsInjectBuilder.WriteString("');")
|
||||
analyticsInjectBuilder.WriteString("</script>")
|
||||
}
|
||||
analyticsInjectBuilder.WriteString("<!--Google Analytics QuantumNous-->\n")
|
||||
analyticsInject := []byte(analyticsInjectBuilder.String())
|
||||
placeholder := []byte("<!--Google Analytics-->\n")
|
||||
indexPage = bytes.ReplaceAll(indexPage, placeholder, analyticsInject)
|
||||
}
|
||||
|
||||
func InitResources() error {
|
||||
// Initialize resources here if needed
|
||||
// This is a placeholder function for future resource initialization
|
||||
err := godotenv.Load(".env")
|
||||
if err != nil {
|
||||
if common.DebugEnabled {
|
||||
common.SysLog("No .env file found, using default environment variables. If needed, please create a .env file and set the relevant variables.")
|
||||
}
|
||||
}
|
||||
|
||||
// 加载环境变量
|
||||
common.InitEnv()
|
||||
|
||||
logger.SetupLogger()
|
||||
|
||||
// Initialize model settings
|
||||
ratio_setting.InitRatioSettings()
|
||||
|
||||
service.InitHttpClient()
|
||||
|
||||
service.InitTokenEncoders()
|
||||
|
||||
// Initialize SQL Database
|
||||
err = model.InitDB()
|
||||
if err != nil {
|
||||
common.FatalLog("failed to initialize database: " + err.Error())
|
||||
return err
|
||||
}
|
||||
if err = authz.Init(model.DB); err != nil {
|
||||
common.FatalLog("failed to initialize authorization: " + err.Error())
|
||||
return err
|
||||
}
|
||||
|
||||
model.CheckSetup()
|
||||
|
||||
// Initialize options, should after model.InitDB()
|
||||
if common.IsMasterNode {
|
||||
if err := model.MigrateRetiredFrontendOptions(); err != nil {
|
||||
common.SysError("failed to migrate retired frontend options: " + err.Error())
|
||||
}
|
||||
}
|
||||
model.InitOptionMap()
|
||||
|
||||
// 清理旧的磁盘缓存文件
|
||||
common.CleanupOldCacheFiles()
|
||||
|
||||
// Initialize SQL Database
|
||||
err = model.InitLogDB()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Initialize Redis
|
||||
err = common.InitRedisClient()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
perfmetrics.Init()
|
||||
|
||||
// 启动系统监控
|
||||
common.StartSystemMonitor()
|
||||
|
||||
// Initialize i18n
|
||||
err = i18n.Init()
|
||||
if err != nil {
|
||||
common.SysError("failed to initialize i18n: " + err.Error())
|
||||
// Don't return error, i18n is not critical
|
||||
} else {
|
||||
common.SysLog("i18n initialized with languages: " + strings.Join(i18n.SupportedLanguages(), ", "))
|
||||
}
|
||||
// Register user language loader for lazy loading
|
||||
i18n.SetUserLangLoader(model.GetUserLanguage)
|
||||
|
||||
// Load custom OAuth providers from database
|
||||
err = oauth.LoadCustomProviders()
|
||||
if err != nil {
|
||||
common.SysError("failed to load custom OAuth providers: " + err.Error())
|
||||
// Don't return error, custom OAuth is not critical
|
||||
}
|
||||
|
||||
service.StartAuthArtifactCleanup()
|
||||
|
||||
return nil
|
||||
}
|
||||
Reference in New Issue
Block a user