Skip to content

Commit 3450a05

Browse files
committed
feat: Multi-Target Load Balancing
1 parent b97bd11 commit 3450a05

File tree

5 files changed

+63
-27
lines changed

5 files changed

+63
-27
lines changed

README.md

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ A high-performance proxy server for OpenAI-compatible APIs with multi-key rotati
1111
## Features
1212

1313
- **Multi-key Rotation**: Automatic API key rotation with load balancing
14+
- **Multi-Target Load Balancing**: Supports round-robin load balancing across multiple upstream API targets
1415
- **Intelligent Blacklisting**: Distinguishes between permanent and temporary errors for smart key management
1516
- **Real-time Monitoring**: Comprehensive statistics, health checks, and blacklist management
1617
- **Flexible Configuration**: Environment-based configuration with .env file support
@@ -98,7 +99,7 @@ cp .env.example .env
9899
| Keys File | `KEYS_FILE` | keys.txt | API keys file path |
99100
| Start Index | `START_INDEX` | 0 | Starting key index for rotation |
100101
| Blacklist Threshold | `BLACKLIST_THRESHOLD` | 1 | Error count before blacklisting |
101-
| Upstream URL | `OPENAI_BASE_URL` | `https://api.openai.com` | OpenAI-compatible API base URL |
102+
| Upstream URL | `OPENAI_BASE_URL` | `https://api.openai.com` | OpenAI-compatible API base URL. Supports multiple, comma-separated URLs for load balancing. |
102103
| Request Timeout | `REQUEST_TIMEOUT` | 30000 | Request timeout in milliseconds |
103104
| Auth Key | `AUTH_KEY` | - | Optional authentication key |
104105
| CORS | `ENABLE_CORS` | true | Enable CORS support |
@@ -125,9 +126,16 @@ OPENAI_BASE_URL=https://your-resource.openai.azure.com
125126
```bash
126127
OPENAI_BASE_URL=https://api.your-provider.com
127128
# Use provider-specific API keys
128-
```
129-
130-
## API Key Validation
129+
```
130+
131+
#### Multi-Target Load Balancing
132+
133+
```bash
134+
# Use a comma-separated list of target URLs
135+
OPENAI_BASE_URL=https://gateway.ai.cloudflare.com/v1/.../openai,https://api.openai.com/v1,https://api.another-provider.com/v1
136+
```
137+
138+
## API Key Validation
131139

132140
The project includes a high-performance API key validation tool:
133141

README_CN.md

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
## 功能特性
1212

1313
- **多密钥轮询**: 自动 API 密钥轮换和负载均衡
14+
- **多目标负载均衡**: 支持轮询多个上游 API 地址
1415
- **智能拉黑**: 区分永久性和临时性错误,智能密钥管理
1516
- **实时监控**: 全面的统计信息、健康检查和黑名单管理
1617
- **灵活配置**: 基于环境变量的配置,支持 .env 文件
@@ -98,7 +99,7 @@ cp .env.example .env
9899
| 密钥文件 | `KEYS_FILE` | keys.txt | API 密钥文件路径 |
99100
| 起始索引 | `START_INDEX` | 0 | 密钥轮换起始索引 |
100101
| 拉黑阈值 | `BLACKLIST_THRESHOLD` | 1 | 拉黑前的错误次数 |
101-
| 上游地址 | `OPENAI_BASE_URL` | `https://api.openai.com` | OpenAI 兼容 API 基础地址 |
102+
| 上游地址 | `OPENAI_BASE_URL` | `https://api.openai.com` | OpenAI 兼容 API 基础地址。支持多个地址,用逗号分隔 |
102103
| 请求超时 | `REQUEST_TIMEOUT` | 30000 | 请求超时时间(毫秒) |
103104
| 认证密钥 | `AUTH_KEY` | - | 可选的认证密钥 |
104105
| CORS | `ENABLE_CORS` | true | 启用 CORS 支持 |
@@ -125,9 +126,16 @@ OPENAI_BASE_URL=https://your-resource.openai.azure.com
125126
```bash
126127
OPENAI_BASE_URL=https://api.your-provider.com
127128
# 使用提供商特定的 API 密钥
128-
```
129-
130-
## API 密钥验证
129+
```
130+
131+
#### 多目标负载均衡
132+
133+
```bash
134+
# 使用逗号分隔多个目标地址
135+
OPENAI_BASE_URL=https://gateway.ai.cloudflare.com/v1/.../openai,https://api.openai.com/v1,https://api.another-provider.com/v1
136+
```
137+
138+
## API 密钥验证
131139

132140
项目包含高性能的 API 密钥验证工具:
133141

internal/config/manager.go

Lines changed: 22 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"os"
88
"strconv"
99
"strings"
10+
"sync/atomic"
1011

1112
"gpt-load/internal/errors"
1213
"gpt-load/pkg/types"
@@ -37,7 +38,8 @@ var DefaultConstants = Constants{
3738

3839
// Manager implements the ConfigManager interface
3940
type Manager struct {
40-
config *Config
41+
config *Config
42+
roundRobinCounter uint64
4143
}
4244

4345
// Config represents the application configuration
@@ -70,8 +72,8 @@ func NewManager() (types.ConfigManager, error) {
7072
MaxRetries: parseInteger(os.Getenv("MAX_RETRIES"), 3),
7173
},
7274
OpenAI: types.OpenAIConfig{
73-
BaseURL: getEnvOrDefault("OPENAI_BASE_URL", "https://api.openai.com"),
74-
Timeout: parseInteger(os.Getenv("REQUEST_TIMEOUT"), DefaultConstants.DefaultTimeout),
75+
BaseURLs: parseArray(os.Getenv("OPENAI_BASE_URL"), []string{"https://api.openai.com"}),
76+
Timeout: parseInteger(os.Getenv("REQUEST_TIMEOUT"), DefaultConstants.DefaultTimeout),
7577
},
7678
Auth: types.AuthConfig{
7779
Key: os.Getenv("AUTH_KEY"),
@@ -120,7 +122,15 @@ func (m *Manager) GetKeysConfig() types.KeysConfig {
120122

121123
// GetOpenAIConfig returns OpenAI configuration
122124
func (m *Manager) GetOpenAIConfig() types.OpenAIConfig {
123-
return m.config.OpenAI
125+
config := m.config.OpenAI
126+
if len(config.BaseURLs) > 1 {
127+
// Use atomic counter for thread-safe round-robin
128+
index := atomic.AddUint64(&m.roundRobinCounter, 1) - 1
129+
config.BaseURL = config.BaseURLs[index%uint64(len(config.BaseURLs))]
130+
} else if len(config.BaseURLs) == 1 {
131+
config.BaseURL = config.BaseURLs[0]
132+
}
133+
return config
124134
}
125135

126136
// GetAuthConfig returns authentication configuration
@@ -168,8 +178,13 @@ func (m *Manager) Validate() error {
168178
}
169179

170180
// Validate upstream URL format
171-
if _, err := url.Parse(m.config.OpenAI.BaseURL); err != nil {
172-
validationErrors = append(validationErrors, "invalid upstream API URL format")
181+
if len(m.config.OpenAI.BaseURLs) == 0 {
182+
validationErrors = append(validationErrors, "at least one upstream API URL is required")
183+
}
184+
for _, baseURL := range m.config.OpenAI.BaseURLs {
185+
if _, err := url.Parse(baseURL); err != nil {
186+
validationErrors = append(validationErrors, fmt.Sprintf("invalid upstream API URL format: %s", baseURL))
187+
}
173188
}
174189

175190
// Validate performance configuration
@@ -196,7 +211,7 @@ func (m *Manager) DisplayConfig() {
196211
logrus.Infof(" Start index: %d", m.config.Keys.StartIndex)
197212
logrus.Infof(" Blacklist threshold: %d errors", m.config.Keys.BlacklistThreshold)
198213
logrus.Infof(" Max retries: %d", m.config.Keys.MaxRetries)
199-
logrus.Infof(" Upstream URL: %s", m.config.OpenAI.BaseURL)
214+
logrus.Infof(" Upstream URLs: %s", strings.Join(m.config.OpenAI.BaseURLs, ", "))
200215
logrus.Infof(" Request timeout: %dms", m.config.OpenAI.Timeout)
201216

202217
authStatus := "disabled"

internal/proxy/server.go

Lines changed: 14 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,6 @@ type ProxyServer struct {
4444
configManager types.ConfigManager
4545
httpClient *http.Client
4646
streamClient *http.Client // Dedicated client for streaming
47-
upstreamURL *url.URL
4847
requestCount int64
4948
startTime time.Time
5049
}
@@ -54,11 +53,6 @@ func NewProxyServer(keyManager types.KeyManager, configManager types.ConfigManag
5453
openaiConfig := configManager.GetOpenAIConfig()
5554
perfConfig := configManager.GetPerformanceConfig()
5655

57-
// Parse upstream URL
58-
upstreamURL, err := url.Parse(openaiConfig.BaseURL)
59-
if err != nil {
60-
return nil, errors.NewAppErrorWithCause(errors.ErrConfigInvalid, "Failed to parse upstream URL", err)
61-
}
6256

6357
// Create high-performance HTTP client
6458
transport := &http.Transport{
@@ -104,7 +98,6 @@ func NewProxyServer(keyManager types.KeyManager, configManager types.ConfigManag
10498
configManager: configManager,
10599
httpClient: httpClient,
106100
streamClient: streamClient,
107-
upstreamURL: upstreamURL,
108101
startTime: time.Now(),
109102
}, nil
110103
}
@@ -205,8 +198,20 @@ func (ps *ProxyServer) executeRequestWithRetry(c *gin.Context, startTime time.Ti
205198
c.Set("retryCount", retryCount)
206199
}
207200

201+
// Get a base URL from the config manager (handles round-robin)
202+
openaiConfig := ps.configManager.GetOpenAIConfig()
203+
upstreamURL, err := url.Parse(openaiConfig.BaseURL)
204+
if err != nil {
205+
logrus.Errorf("Failed to parse upstream URL: %v", err)
206+
c.JSON(http.StatusInternalServerError, gin.H{
207+
"error": "Invalid upstream URL configured",
208+
"code": errors.ErrConfigInvalid,
209+
})
210+
return
211+
}
212+
208213
// Build upstream request URL
209-
targetURL := *ps.upstreamURL
214+
targetURL := *upstreamURL
210215
// Correctly append path instead of replacing it
211216
if strings.HasSuffix(targetURL.Path, "/") {
212217
targetURL.Path = targetURL.Path + strings.TrimPrefix(c.Request.URL.Path, "/")
@@ -223,8 +228,7 @@ func (ps *ProxyServer) executeRequestWithRetry(c *gin.Context, startTime time.Ti
223228
// Streaming requests only set response header timeout, no overall timeout
224229
ctx, cancel = context.WithCancel(c.Request.Context())
225230
} else {
226-
// Non-streaming requests use configured timeout
227-
openaiConfig := ps.configManager.GetOpenAIConfig()
231+
// Non-streaming requests use configured timeout from the already fetched config
228232
timeout := time.Duration(openaiConfig.Timeout) * time.Millisecond
229233
ctx, cancel = context.WithTimeout(c.Request.Context(), timeout)
230234
}

pkg/types/interfaces.go

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -54,8 +54,9 @@ type KeysConfig struct {
5454

5555
// OpenAIConfig represents OpenAI API configuration
5656
type OpenAIConfig struct {
57-
BaseURL string `json:"baseUrl"`
58-
Timeout int `json:"timeout"`
57+
BaseURL string `json:"baseUrl"`
58+
BaseURLs []string `json:"baseUrls"`
59+
Timeout int `json:"timeout"`
5960
}
6061

6162
// AuthConfig represents authentication configuration

0 commit comments

Comments
 (0)