1.0.0

Author: helloxz

.github/workflows/build.yml

@@ -0,0 +1,42 @@
+name: 构建并发布Docker镜像
+
+on:
+  workflow_dispatch:  # 允许手动触发工作流
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    
+    steps:
+    - name: 检出代码
+      uses: actions/checkout@v3
+
+    - name: 设置Go环境
+      uses: actions/setup-go@v4
+      with:
+        go-version: '1.23'
+
+    - name: 安装UPX
+      run: sudo apt-get update && sudo apt-get install -y upx
+
+    - name: 构建Go程序
+      run: go build -o aiload -ldflags -w main.go
+
+    - name: 使用UPX压缩可执行文件
+      run: upx -9 aiload
+
+    - name: 登录到Docker Hub
+      uses: docker/login-action@v3
+      with:
+        username: ${{ secrets.DOCKERHUB_USERNAME }}
+        password: ${{ secrets.DOCKERHUB_TOKEN }}
+
+    - name: 设置Docker Buildx
+      uses: docker/setup-buildx-action@v3
+
+    - name: 构建并推送Docker镜像
+      uses: docker/build-push-action@v5
+      with:
+        context: .
+        push: true
+        tags: ${{ secrets.DOCKERHUB_USERNAME }}/aiload:latest

.gitignore

@@ -0,0 +1 @@
+data/*

Dockerfile

@@ -0,0 +1,8 @@
+FROM alpine:3.21
+RUN mkdir -p /opt/aiload/data
+WORKDIR /opt/aiload
+COPY aiload config.json /opt/aiload/
+# 暴露文件夹和端口
+EXPOSE 2081
+# 启动程序
+CMD ["./aiload", "start"]

README.md

@@ -1 +1,22 @@
-# ai-relay
+# AILoad
+
+AILoad 是一个轻量级且高效的工具，通过负载均衡的方式将请求分发到多个 AI 模型接口，从而有效解决单一 AI 模型频率上限问题。目前支持轮询和基于 IP 的负载均衡策略，能够在短时间内确保同一 IP 的请求被分配到同一个 AI 模型，提升一致性和用户体验。
+
+## 主要功能
+
+* 负载均衡 ：自动将请求随机分发到多个 AI 模型接口，避免单个模型频率上限，提升整体吞吐量。
+* 统一API入口 ：提供单一入口调用
+* 兼容OpenAI：完全兼容OpenAI API接口格式
+* 流式返回：支持流式传输
+* IP 亲和性：同一个IP在一定时间内固定分配到一个AI模型，提升连续对话一致性
+
+
+## 使用场景
+
+目前市面部分厂商提供了免费的AI模型接口，但往往存在频率限制，通过AILoad可以将多个AI接口整合到一起，并通过负载均衡的方式来缓解单一AI模型频率上限的问题。
+
+## 部署
+
+## 联系我们
+
+如有任何问题或反馈，请通过 GitHub Issues提交。

api/home.go

@@ -0,0 +1,16 @@
+package api
+
+import (
+	"github.com/gin-gonic/gin"
+)
+
+func Home(c *gin.Context) {
+	// 准备 HTML 内容
+	htmlContent := `<p>Please visit <a href="https://github.com/helloxz/aiload" target="_blank">helloxz/aiload</a> for usage instructions.</p>`
+
+	// 设置 Content-Type 为 text/html
+	c.Header("Content-Type", "text/html; charset=utf-8")
+
+	// 返回 HTML 内容
+	c.String(200, htmlContent)
+}

api/relay.go

@@ -0,0 +1,160 @@
+package api
+
+import (
+	"aiload/utils"
+	"math/rand"
+	"os"
+	"strconv"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"github.com/spf13/viper"
+	"github.com/tidwall/gjson"
+)
+
+// 声明一个结构体
+type Model struct {
+	BaseURL string `json:"base_url"`
+	Model   string `json:"model"`
+	ApiKey  string `json:"api_key"`
+}
+
+func Relay(c *gin.Context) {
+	// 读取 JSON 文件内容
+	data, err := os.ReadFile("data/config/config.json")
+	if err != nil {
+		c.JSON(200, gin.H{
+			"code": 500,
+			"msg":  "Failed to read configuration file!",
+			"data": "",
+		})
+		c.Abort()
+		return
+	}
+
+	// 获取原始请求体数据
+	rawData, err := c.GetRawData()
+	if err != nil {
+		c.JSON(200, gin.H{
+			"code": 500,
+			"msg":  "Failed to read request body!",
+			"data": "",
+		})
+		c.Abort()
+		return
+	}
+
+	stream := gjson.GetBytes(rawData, "stream").Bool()
+
+	// 将文件内容转换为字符串
+	jsonData := string(data)
+
+	// 重要：将请求体放回，这样后续处理器仍然可以读取它
+	c.Request.Body = utils.CreateReadCloser(rawData)
+
+	// 使用 gjson 提取 models 列表
+	models := gjson.Get(jsonData, "models")
+	if !models.Exists() {
+		c.JSON(200, gin.H{
+			"code": 500,
+			"msg":  "Models list not found!",
+			"data": "",
+		})
+		c.Abort()
+		return
+	}
+
+	var modelList []Model
+
+	// 遍历 models 列表
+	models.ForEach(func(_, model gjson.Result) bool {
+		baseURL := model.Get("base_url").String()
+		modelName := model.Get("model").String()
+		api_key := model.Get("api_key").String()
+
+		// 将每个模型的信息添加到 modelList 中
+		modelList = append(modelList, Model{
+			BaseURL: baseURL,
+			Model:   modelName,
+			ApiKey:  api_key,
+		})
+		return true // 继续遍历下一个元素
+	})
+
+	// 随机从 modelList 中选择一个模型
+	randIndex := rand.Intn(len(modelList))
+	// 检查内存中是否存在index
+	index := getModelIndex(c)
+	// 如果不存在
+	if index == 0 {
+		// 将选择的模型的 index 存入缓存
+		setModelIndex(c, randIndex)
+	} else {
+		randIndex = index
+	}
+
+	selectedModel := modelList[randIndex]
+
+	// fmt.Println(selectedModel)
+	timeout := viper.GetInt("settins.timeout")
+	// 如果为0，则设置默认值
+	if timeout == 0 {
+		timeout = 120
+	}
+	config := utils.ProxyConfig{
+		TargetURL:    selectedModel.BaseURL,
+		Timeout:      time.Duration(timeout) * time.Second, // 设置超时时间
+		CustomHeader: map[string]string{"Authorization": "Bearer " + selectedModel.ApiKey},
+		RemoveHeader: []string{},
+		AddHeader:    map[string]string{"X-Model": selectedModel.Model}, // 添加自定义响应头
+		ModelName:    selectedModel.Model,
+	}
+
+	if stream {
+		// 请求转发服务
+		utils.ReverseStreamProxy(c, config)
+		return
+	} else {
+		// 请求转发服务
+		utils.ReverseProxy(c, config)
+	}
+
+}
+
+// 获取模型的index
+func getModelIndex(c *gin.Context) int {
+	// 获取用户IP
+	ip := utils.GetClientIP(c)
+	// 生成key
+	keyStr := "model_index_" + ip
+	// 转为[]byte
+	key := []byte(keyStr)
+	// 读取缓存
+	value, err := utils.Cache.Get(key)
+	if err != nil {
+		return 0
+	}
+	if value == nil {
+		return 0
+	}
+	// 转为int
+	index, err := strconv.Atoi(string(value))
+	if err != nil {
+		return 0
+	}
+	return index
+}
+
+// 设置模型的index
+func setModelIndex(c *gin.Context, index int) {
+	// 获取用户IP
+	ip := utils.GetClientIP(c)
+	// 生成key
+	keyStr := "model_index_" + ip
+	// 转为[]byte
+	key := []byte(keyStr)
+	// 转为[]byte
+	value := []byte(strconv.Itoa(index))
+	// 设置缓存，有效期5分钟
+	utils.Cache.Set(key, value, 60*5)
+}

config.json

@@ -0,0 +1,17 @@
+{
+    "auth_token": "",
+    "models": [
+      {
+        "api_key": "sk-xxx",
+        "base_url": "https://api.openai.com/v1/chat/completions",
+        "model": "gpt-4o"
+      }
+    ],
+    "server": {
+      "model": "release",
+      "port": 2081
+    },
+    "settings":{
+      "timeout":120
+    }
+  }

go.mod

@@ -0,0 +1,59 @@
+module aiload
+
+go 1.23.2
+
+require (
+	github.com/gin-gonic/gin v1.10.0
+	github.com/spf13/viper v1.19.0
+)
+
+require (
+	github.com/bytedance/sonic v1.11.6 // indirect
+	github.com/bytedance/sonic/loader v0.1.1 // indirect
+	github.com/cespare/xxhash/v2 v2.3.0 // indirect
+	github.com/cloudwego/base64x v0.1.4 // indirect
+	github.com/cloudwego/iasm v0.2.0 // indirect
+	github.com/coocood/freecache v1.2.4 // indirect
+	github.com/fsnotify/fsnotify v1.7.0 // indirect
+	github.com/gabriel-vasile/mimetype v1.4.3 // indirect
+	github.com/gin-contrib/sse v0.1.0 // indirect
+	github.com/go-playground/locales v0.14.1 // indirect
+	github.com/go-playground/universal-translator v0.18.1 // indirect
+	github.com/go-playground/validator/v10 v10.20.0 // indirect
+	github.com/go-resty/resty/v2 v2.16.5 // indirect
+	github.com/goccy/go-json v0.10.2 // indirect
+	github.com/hashicorp/hcl v1.0.0 // indirect
+	github.com/json-iterator/go v1.1.12 // indirect
+	github.com/klauspost/cpuid/v2 v2.2.7 // indirect
+	github.com/leodido/go-urn v1.4.0 // indirect
+	github.com/magiconair/properties v1.8.7 // indirect
+	github.com/mattn/go-isatty v0.0.20 // indirect
+	github.com/mitchellh/mapstructure v1.5.0 // indirect
+	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
+	github.com/modern-go/reflect2 v1.0.2 // indirect
+	github.com/pelletier/go-toml/v2 v2.2.2 // indirect
+	github.com/sagikazarmark/locafero v0.4.0 // indirect
+	github.com/sagikazarmark/slog-shim v0.1.0 // indirect
+	github.com/sourcegraph/conc v0.3.0 // indirect
+	github.com/spf13/afero v1.11.0 // indirect
+	github.com/spf13/cast v1.6.0 // indirect
+	github.com/spf13/pflag v1.0.5 // indirect
+	github.com/subosito/gotenv v1.6.0 // indirect
+	github.com/tidwall/gjson v1.18.0 // indirect
+	github.com/tidwall/match v1.1.1 // indirect
+	github.com/tidwall/pretty v1.2.1 // indirect
+	github.com/tidwall/sjson v1.2.5 // indirect
+	github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
+	github.com/ugorji/go/codec v1.2.12 // indirect
+	go.uber.org/atomic v1.9.0 // indirect
+	go.uber.org/multierr v1.9.0 // indirect
+	golang.org/x/arch v0.8.0 // indirect
+	golang.org/x/crypto v0.31.0 // indirect
+	golang.org/x/exp v0.0.0-20230905200255-921286631fa9 // indirect
+	golang.org/x/net v0.33.0 // indirect
+	golang.org/x/sys v0.28.0 // indirect
+	golang.org/x/text v0.21.0 // indirect
+	google.golang.org/protobuf v1.34.1 // indirect
+	gopkg.in/ini.v1 v1.67.0 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
+)