mirror of
https://github.com/BlueSkyXN/AI2API.git
synced 2026-04-22 15:47:19 +08:00
0.4.2
This commit is contained in:
BlueSkyXN
@@ -148,3 +148,48 @@ jobs:
|
||||
path: hunyuan2api-${{ matrix.os }}-${{ matrix.arch }}${{ matrix.os == 'windows' && '.exe' || '' }}
|
||||
retention-days: 90
|
||||
|
||||
build-api-pool-lite:
|
||||
name: Build api-pool-lite
|
||||
runs-on: ubuntu-latest
|
||||
strategy:
|
||||
matrix:
|
||||
os: [windows, linux]
|
||||
arch: [amd64]
|
||||
steps:
|
||||
- name: Check out code
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Go
|
||||
uses: actions/setup-go@v4
|
||||
with:
|
||||
go-version: '1.21'
|
||||
check-latest: true
|
||||
cache: true
|
||||
|
||||
- name: Set version
|
||||
id: set-version
|
||||
run: |
|
||||
echo "version=$(date +'%Y%m%d%H%M')" >> $GITHUB_OUTPUT
|
||||
|
||||
- name: Build for ${{ matrix.os }}-${{ matrix.arch }}
|
||||
env:
|
||||
GOOS: ${{ matrix.os }}
|
||||
GOARCH: ${{ matrix.arch }}
|
||||
VERSION: ${{ steps.set-version.outputs.version }}
|
||||
run: |
|
||||
# 设置文件扩展名（Windows为.exe，其他无扩展名）
|
||||
if [ "${{ matrix.os }}" == "windows" ]; then
|
||||
EXT=".exe"
|
||||
else
|
||||
EXT=""
|
||||
fi
|
||||
|
||||
BINARY_NAME="api-pool-lite-${{ matrix.os }}-${{ matrix.arch }}${EXT}"
|
||||
go build -ldflags "-X main.AppVersion=$VERSION" -o "$BINARY_NAME" ./api-pool.go
|
||||
|
||||
- name: Upload Artifact
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: api-pool-lite-${{ matrix.os }}-${{ matrix.arch }}-${{ steps.set-version.outputs.version }}
|
||||
path: api-pool-lite-${{ matrix.os }}-${{ matrix.arch }}${{ matrix.os == 'windows' && '.exe' || '' }}
|
||||
retention-days: 90
|
||||
@@ -0,0 +1,520 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"flag"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"path"
|
||||
"strings"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// Config 结构体用于存储命令行参数配置
|
||||
type Config struct {
|
||||
KeyFile string // API 密钥文件路径
|
||||
TargetURL string // 目标 API 基础 URL
|
||||
Port string // 代理服务器监听端口
|
||||
Address string // 代理服务器监听地址
|
||||
Password string // 客户端身份验证密码
|
||||
MaxWorkers int // 最大工作协程数
|
||||
MaxQueue int // 最大请求队列长度
|
||||
}
|
||||
|
||||
// parseFlags 解析命令行参数并返回 Config 实例
|
||||
func parseFlags() *Config {
|
||||
cfg := &Config{}
|
||||
flag.StringVar(&cfg.KeyFile, "key-file", "", "Path to the API key file")
|
||||
flag.StringVar(&cfg.TargetURL, "target-url", "", "Target API base URL")
|
||||
flag.StringVar(&cfg.Port, "port", "8080", "Port to listen on")
|
||||
flag.StringVar(&cfg.Address, "address", "localhost", "Address to listen on")
|
||||
flag.StringVar(&cfg.Password, "password", "", "Password for client authentication")
|
||||
|
||||
// 添加WorkerPool相关配置
|
||||
maxWorkers := flag.Int("max-workers", 50, "Maximum number of worker goroutines")
|
||||
maxQueue := flag.Int("max-queue", 500, "Maximum size of request queue")
|
||||
|
||||
flag.Parse()
|
||||
|
||||
// 将WorkerPool配置添加到Config结构体
|
||||
cfg.MaxWorkers = *maxWorkers
|
||||
cfg.MaxQueue = *maxQueue
|
||||
|
||||
return cfg
|
||||
}
|
||||
|
||||
// KeyPool 管理 API 密钥池
|
||||
type KeyPool struct {
|
||||
keys []string // 密钥列表
|
||||
mu sync.Mutex // 互斥锁，确保线程安全
|
||||
currentIndex int // 当前密钥索引，用于循环抽取
|
||||
}
|
||||
|
||||
// NewKeyPool 从文件中加载密钥并创建 KeyPool 实例
|
||||
func NewKeyPool(filePath string) (*KeyPool, error) {
|
||||
file, err := os.Open(filePath)
|
||||
if err != nil {
|
||||
log.Printf("[ERROR] Failed to open key file %s: %v", filePath, err)
|
||||
return nil, err
|
||||
}
|
||||
defer file.Close()
|
||||
|
||||
var keys []string
|
||||
scanner := bufio.NewScanner(file)
|
||||
for scanner.Scan() {
|
||||
key := strings.TrimSpace(scanner.Text())
|
||||
if key != "" {
|
||||
keys = append(keys, key)
|
||||
}
|
||||
}
|
||||
if err := scanner.Err(); err != nil {
|
||||
log.Printf("[ERROR] Failed to read key file %s: %v", filePath, err)
|
||||
return nil, err
|
||||
}
|
||||
log.Printf("[INFO] Loaded %d keys from file %s", len(keys), filePath)
|
||||
return &KeyPool{keys: keys, currentIndex: 0}, nil
|
||||
}
|
||||
|
||||
// GetRandomKey 按顺序循环返回一个密钥
|
||||
func (kp *KeyPool) GetRandomKey() string {
|
||||
kp.mu.Lock()
|
||||
defer kp.mu.Unlock()
|
||||
if len(kp.keys) == 0 {
|
||||
return ""
|
||||
}
|
||||
key := kp.keys[kp.currentIndex]
|
||||
kp.currentIndex = (kp.currentIndex + 1) % len(kp.keys) // 循环到下一个索引
|
||||
return key
|
||||
}
|
||||
|
||||
// 定义请求结构体
|
||||
type ProxyRequest struct {
|
||||
Request *http.Request
|
||||
Response http.ResponseWriter
|
||||
Done chan bool // 用于通知请求处理完成
|
||||
}
|
||||
|
||||
// Worker结构体，表示一个工作协程
|
||||
type Worker struct {
|
||||
ID int
|
||||
TaskQueue chan *ProxyRequest // 任务队列
|
||||
Quit chan bool // 退出信号
|
||||
WorkerPool *WorkerPool // 所属工作池
|
||||
}
|
||||
|
||||
// 创建新的Worker
|
||||
func NewWorker(id int, workerPool *WorkerPool) *Worker {
|
||||
return &Worker{
|
||||
ID: id,
|
||||
TaskQueue: make(chan *ProxyRequest),
|
||||
Quit: make(chan bool),
|
||||
WorkerPool: workerPool,
|
||||
}
|
||||
}
|
||||
|
||||
// Worker开始工作
|
||||
func (w *Worker) Start() {
|
||||
go func() {
|
||||
for {
|
||||
// 将worker注册到工作池的空闲队列
|
||||
w.WorkerPool.WorkerQueue <- w.TaskQueue
|
||||
|
||||
select {
|
||||
case task := <-w.TaskQueue:
|
||||
// 处理请求
|
||||
w.WorkerPool.HandleFunc(task.Response, task.Request)
|
||||
task.Done <- true
|
||||
case <-w.Quit:
|
||||
// 收到退出信号
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// Worker停止工作
|
||||
func (w *Worker) Stop() {
|
||||
go func() {
|
||||
w.Quit <- true
|
||||
}()
|
||||
}
|
||||
|
||||
// WorkerPool结构体，管理工作协程池
|
||||
type WorkerPool struct {
|
||||
WorkerQueue chan chan *ProxyRequest // 空闲Worker队列
|
||||
TaskQueue chan *ProxyRequest // 任务队列
|
||||
MaxWorkers int // 最大Worker数量
|
||||
MaxQueue int // 最大队列长度
|
||||
HandleFunc func(http.ResponseWriter, *http.Request) // 请求处理函数
|
||||
}
|
||||
|
||||
// 创建新的WorkerPool
|
||||
func NewWorkerPool(maxWorkers int, maxQueue int, handleFunc func(http.ResponseWriter, *http.Request)) *WorkerPool {
|
||||
pool := &WorkerPool{
|
||||
WorkerQueue: make(chan chan *ProxyRequest, maxWorkers),
|
||||
TaskQueue: make(chan *ProxyRequest, maxQueue),
|
||||
MaxWorkers: maxWorkers,
|
||||
MaxQueue: maxQueue,
|
||||
HandleFunc: handleFunc,
|
||||
}
|
||||
return pool
|
||||
}
|
||||
|
||||
// 启动WorkerPool
|
||||
func (wp *WorkerPool) Start() {
|
||||
// 创建并启动workers
|
||||
for i := 0; i < wp.MaxWorkers; i++ {
|
||||
worker := NewWorker(i, wp)
|
||||
worker.Start()
|
||||
log.Printf("[INFO] Started worker %d", i)
|
||||
}
|
||||
|
||||
// 启动任务分发协程
|
||||
go wp.dispatch()
|
||||
}
|
||||
|
||||
// 停止WorkerPool
|
||||
func (wp *WorkerPool) Stop() {
|
||||
// TODO: 实现停止逻辑
|
||||
}
|
||||
|
||||
// 将任务分发给空闲worker
|
||||
func (wp *WorkerPool) dispatch() {
|
||||
for {
|
||||
select {
|
||||
case task := <-wp.TaskQueue:
|
||||
// 等待空闲worker
|
||||
workerTaskQueue := <-wp.WorkerQueue
|
||||
// 将任务发送给worker
|
||||
workerTaskQueue <- task
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 将请求提交到WorkerPool
|
||||
func (wp *WorkerPool) Submit(response http.ResponseWriter, request *http.Request) bool {
|
||||
task := &ProxyRequest{
|
||||
Request: request,
|
||||
Response: response,
|
||||
Done: make(chan bool, 1),
|
||||
}
|
||||
|
||||
select {
|
||||
case wp.TaskQueue <- task:
|
||||
// 请求成功加入队列
|
||||
<-task.Done // 等待任务完成
|
||||
return true
|
||||
default:
|
||||
// 队列已满，实现背压
|
||||
log.Println("[WARN] Task queue is full, rejecting request")
|
||||
http.Error(response, "Server is busy, please try again later", http.StatusServiceUnavailable)
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// ProxyHandler 处理 HTTP 代理请求
|
||||
type ProxyHandler struct {
|
||||
cfg *Config // 配置信息
|
||||
keyPool *KeyPool // 密钥池
|
||||
client *http.Client // HTTP 客户端
|
||||
workerPool *WorkerPool // 工作协程池
|
||||
}
|
||||
|
||||
// NewProxyHandler 创建 ProxyHandler 实例
|
||||
func NewProxyHandler(cfg *Config, keyPool *KeyPool) *ProxyHandler {
|
||||
handler := &ProxyHandler{
|
||||
cfg: cfg,
|
||||
keyPool: keyPool,
|
||||
client: &http.Client{},
|
||||
}
|
||||
return handler
|
||||
}
|
||||
|
||||
// InitWorkerPool 初始化工作协程池
|
||||
func (ph *ProxyHandler) InitWorkerPool(maxWorkers int, maxQueue int) {
|
||||
ph.workerPool = NewWorkerPool(maxWorkers, maxQueue, ph.HandleRequest)
|
||||
ph.workerPool.Start()
|
||||
log.Printf("[INFO] Started worker pool with %d workers and queue size %d", maxWorkers, maxQueue)
|
||||
}
|
||||
|
||||
// ServeHTTP 实现 HTTP 处理逻辑
|
||||
func (ph *ProxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
// 记录接收到的请求
|
||||
log.Printf("[INFO] Received request: %s %s", r.Method, r.URL.String())
|
||||
|
||||
// 将请求提交到工作池处理
|
||||
ph.workerPool.Submit(w, r)
|
||||
}
|
||||
|
||||
// 简化的模型提取函数
|
||||
func extractModelInfo(bodyData []byte) (string, error) {
|
||||
var data map[string]interface{}
|
||||
if err := json.Unmarshal(bodyData, &data); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if model, ok := data["model"].(string); ok {
|
||||
return model, nil
|
||||
}
|
||||
return "", nil
|
||||
}
|
||||
|
||||
// 检测是否为流式请求
|
||||
func isStreamRequest(r *http.Request, bodyData []byte) bool {
|
||||
// 检查URL参数
|
||||
if strings.Contains(r.URL.RawQuery, "stream=true") {
|
||||
return true
|
||||
}
|
||||
|
||||
// 检查请求体
|
||||
var data map[string]interface{}
|
||||
if err := json.Unmarshal(bodyData, &data); err == nil {
|
||||
if stream, ok := data["stream"].(bool); ok && stream {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// HandleRequest 处理请求的方法，由Worker调用
|
||||
func (ph *ProxyHandler) HandleRequest(w http.ResponseWriter, r *http.Request) {
|
||||
// 验证客户端身份
|
||||
if !ph.authenticate(r) {
|
||||
log.Println("[WARN] Unauthorized access attempt")
|
||||
http.Error(w, "Unauthorized", http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
log.Println("[INFO] Authentication successful")
|
||||
|
||||
// 读取请求体 - 只读取一次
|
||||
var bodyBytes []byte
|
||||
var err error
|
||||
if r.Body != nil {
|
||||
bodyBytes, err = io.ReadAll(r.Body)
|
||||
if err != nil {
|
||||
log.Printf("[ERROR] Failed to read request body: %v", err)
|
||||
http.Error(w, "Internal Server Error", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
// 异步提取模型信息
|
||||
go func() {
|
||||
if model, err := extractModelInfo(bodyBytes); err == nil && model != "" {
|
||||
log.Printf("[INFO] Model specified in request: %s", model)
|
||||
}
|
||||
}()
|
||||
|
||||
// 检测是否为流式请求
|
||||
if isStreamRequest(r, bodyBytes) {
|
||||
log.Println("[INFO] Streaming request detected")
|
||||
}
|
||||
}
|
||||
|
||||
// 构建目标 URL
|
||||
targetURL, err := ph.buildTargetURL(r)
|
||||
if err != nil {
|
||||
log.Printf("[ERROR] Failed to build target URL: %v", err)
|
||||
http.Error(w, "Bad Request", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
log.Printf("[INFO] Target URL: %s", targetURL)
|
||||
|
||||
// 重试逻辑
|
||||
maxRetries := len(ph.keyPool.keys)
|
||||
attemptedKeys := make(map[string]bool)
|
||||
log.Printf("[INFO] Starting key selection process, total keys available: %d", maxRetries)
|
||||
|
||||
for i := 0; i < maxRetries; i++ {
|
||||
key := ph.getUnusedKey(attemptedKeys)
|
||||
if key == "" {
|
||||
log.Printf("[ERROR] No unused keys remaining after %d attempts", i)
|
||||
break
|
||||
}
|
||||
attemptedKeys[key] = true
|
||||
maskedKey := maskKey(key)
|
||||
log.Printf("[INFO] Attempt %d/%d: Selecting key %s", i+1, maxRetries, maskedKey)
|
||||
|
||||
// 创建请求 - 为每次尝试创建新的请求体副本
|
||||
req, err := ph.createRequest(r, targetURL, key, bodyBytes)
|
||||
if err != nil {
|
||||
log.Printf("[ERROR] Failed to create request with key %s: %v", maskedKey, err)
|
||||
log.Printf("[INFO] Switching to another key due to request creation failure")
|
||||
continue
|
||||
}
|
||||
|
||||
// 发送请求
|
||||
log.Printf("[INFO] Sending request to target API with key %s", maskedKey)
|
||||
resp, err := ph.client.Do(req)
|
||||
if err != nil {
|
||||
log.Printf("[ERROR] Failed to send request with key %s: %v", maskedKey, err)
|
||||
log.Printf("[INFO] Switching to another key due to network error")
|
||||
continue
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
// 处理响应
|
||||
log.Printf("[INFO] Received response with status code %d", resp.StatusCode)
|
||||
if resp.StatusCode >= 200 && resp.StatusCode < 300 {
|
||||
log.Println("[INFO] Request successful, forwarding response")
|
||||
ph.forwardResponse(w, resp)
|
||||
return
|
||||
} else if resp.StatusCode == 403 || resp.StatusCode == 429 {
|
||||
log.Printf("[WARN] Received %d status code with key %s", resp.StatusCode, maskedKey)
|
||||
log.Printf("[INFO] Switching to another key due to status code %d", resp.StatusCode)
|
||||
continue
|
||||
} else {
|
||||
log.Printf("[INFO] Forwarding response with status code %d", resp.StatusCode)
|
||||
ph.forwardResponse(w, resp)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// 所有密钥尝试后仍失败
|
||||
log.Printf("[ERROR] All %d keys failed after retries", maxRetries)
|
||||
http.Error(w, "Failed to get response from API after all retries", http.StatusBadGateway)
|
||||
}
|
||||
|
||||
// getUnusedKey 获取一个未使用过的密钥
|
||||
func (ph *ProxyHandler) getUnusedKey(attempted map[string]bool) string {
|
||||
key := ph.keyPool.GetRandomKey()
|
||||
// 如果获取到的密钥已使用过，则尝试其他密钥
|
||||
for attempted[key] && len(attempted) < len(ph.keyPool.keys) {
|
||||
key = ph.keyPool.GetRandomKey()
|
||||
}
|
||||
// 如果所有密钥都已尝试过，返回空字符串
|
||||
if attempted[key] {
|
||||
return ""
|
||||
}
|
||||
return key
|
||||
}
|
||||
|
||||
// authenticate 验证客户端身份
|
||||
func (ph *ProxyHandler) authenticate(r *http.Request) bool {
|
||||
authHeader := r.Header.Get("Authorization")
|
||||
if authHeader == "" {
|
||||
return false
|
||||
}
|
||||
parts := strings.Split(authHeader, " ")
|
||||
if len(parts) != 2 || parts[0] != "Bearer" {
|
||||
return false
|
||||
}
|
||||
return parts[1] == ph.cfg.Password
|
||||
}
|
||||
|
||||
// buildTargetURL 构建目标 API 的完整 URL
|
||||
func (ph *ProxyHandler) buildTargetURL(r *http.Request) (string, error) {
|
||||
u, err := url.Parse(ph.cfg.TargetURL)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
u.Path = path.Join(u.Path, r.URL.Path)
|
||||
u.RawQuery = r.URL.RawQuery
|
||||
return u.String(), nil
|
||||
}
|
||||
|
||||
// createRequest 创建转发请求 - 只修改Authorization头，为每次请求创建新的请求体
|
||||
func (ph *ProxyHandler) createRequest(r *http.Request, targetURL, key string, bodyBytes []byte) (*http.Request, error) {
|
||||
// 为每次请求创建新的请求体
|
||||
var bodyReader io.Reader
|
||||
if len(bodyBytes) > 0 {
|
||||
bodyReader = bytes.NewReader(bodyBytes)
|
||||
}
|
||||
|
||||
// 创建请求
|
||||
req, err := http.NewRequest(r.Method, targetURL, bodyReader)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 复制所有原始请求头
|
||||
for k, v := range r.Header {
|
||||
req.Header[k] = v
|
||||
}
|
||||
|
||||
// 仅替换Authorization头
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
return req, nil
|
||||
}
|
||||
|
||||
// forwardResponse 将响应转发给客户端，支持流式和非流式
|
||||
func (ph *ProxyHandler) forwardResponse(w http.ResponseWriter, resp *http.Response) {
|
||||
// 设置响应头
|
||||
for k, v := range resp.Header {
|
||||
w.Header()[k] = v
|
||||
}
|
||||
w.WriteHeader(resp.StatusCode)
|
||||
|
||||
// 处理流式响应
|
||||
if strings.Contains(resp.Header.Get("Content-Type"), "text/event-stream") || resp.Header.Get("Transfer-Encoding") == "chunked" {
|
||||
log.Println("[INFO] Handling streaming response")
|
||||
flusher, ok := w.(http.Flusher)
|
||||
if !ok {
|
||||
log.Println("[ERROR] Streaming unsupported by server")
|
||||
http.Error(w, "Streaming unsupported", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
reader := bufio.NewReader(resp.Body)
|
||||
for {
|
||||
line, err := reader.ReadBytes('\n')
|
||||
if err != nil {
|
||||
if err == io.EOF {
|
||||
log.Println("[INFO] Stream ended")
|
||||
break
|
||||
}
|
||||
log.Printf("[ERROR] Error reading stream: %v", err)
|
||||
http.Error(w, "Error reading stream", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
w.Write(line)
|
||||
flusher.Flush()
|
||||
}
|
||||
} else {
|
||||
// 非流式响应，直接复制
|
||||
_, err := io.Copy(w, resp.Body)
|
||||
if err != nil {
|
||||
log.Printf("[ERROR] Failed to forward response: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// maskKey 直接返回原始密钥，不再进行掩码处理
|
||||
func maskKey(key string) string {
|
||||
return key
|
||||
}
|
||||
|
||||
// main 函数，启动代理服务器
|
||||
func main() {
|
||||
// 解析配置
|
||||
cfg := parseFlags()
|
||||
if cfg.KeyFile == "" || cfg.TargetURL == "" || cfg.Password == "" {
|
||||
log.Println("[ERROR] Missing required flags: --key-file, --target-url, --password")
|
||||
os.Exit(1)
|
||||
}
|
||||
log.Printf("[INFO] Configuration loaded: KeyFile=%s, TargetURL=%s, Address=%s, Port=%s, MaxWorkers=%d, MaxQueue=%d",
|
||||
cfg.KeyFile, cfg.TargetURL, cfg.Address, cfg.Port, cfg.MaxWorkers, cfg.MaxQueue)
|
||||
|
||||
// 初始化密钥池
|
||||
keyPool, err := NewKeyPool(cfg.KeyFile)
|
||||
if err != nil {
|
||||
log.Printf("[ERROR] Failed to initialize key pool: %v", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
// 创建代理处理器
|
||||
proxyHandler := NewProxyHandler(cfg, keyPool)
|
||||
|
||||
// 初始化并启动工作池
|
||||
proxyHandler.InitWorkerPool(cfg.MaxWorkers, cfg.MaxQueue)
|
||||
|
||||
// 启动服务器
|
||||
addr := cfg.Address + ":" + cfg.Port
|
||||
log.Printf("[INFO] Starting proxy server on %s", addr)
|
||||
if err := http.ListenAndServe(addr, proxyHandler); err != nil {
|
||||
log.Printf("[ERROR] Failed to start server: %v", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user