File size: 5,289 Bytes
ca50bff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | package main
import (
"encoding/json"
"net/http"
"time"
"zai-proxy/internal"
)
func corsMiddleware(next http.HandlerFunc) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Access-Control-Allow-Origin", "*")
w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS")
w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization")
w.Header().Set("Access-Control-Allow-Credentials", "true")
if r.Method == "OPTIONS" {
w.WriteHeader(http.StatusOK)
return
}
next(w, r)
}
}
func loggingMiddleware(next http.HandlerFunc) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
start := time.Now()
wrapped := &responseWriter{ResponseWriter: w, statusCode: http.StatusOK}
// 获取客户端 IP
clientIP := internal.GetClientIP(r)
next(wrapped, r)
duration := time.Since(start)
internal.LogInfo("%s %s %d %v [%s]", r.Method, r.URL.Path, wrapped.statusCode, duration, clientIP)
}
}
type responseWriter struct {
http.ResponseWriter
statusCode int
}
func (rw *responseWriter) WriteHeader(code int) {
rw.statusCode = code
rw.ResponseWriter.WriteHeader(code)
}
func (rw *responseWriter) Flush() {
if flusher, ok := rw.ResponseWriter.(http.Flusher); ok {
flusher.Flush()
}
}
func handleRoot(w http.ResponseWriter, r *http.Request) {
telemetry := internal.GetTelemetryData()
response := map[string]interface{}{
"message": "Chat-Glm API",
"version": "2.0.0",
"telemetry": map[string]interface{}{
"uptime": telemetry.Uptime,
"total_requests": telemetry.TotalRequests,
"rpm": telemetry.RPM,
"total_input_tokens": telemetry.TotalInputTok,
"total_output_tokens": telemetry.TotalOutputTok,
"avg_input_tokens": telemetry.AvgInputTok,
"avg_output_tokens": telemetry.AvgOutputTok,
"valid_tokens": telemetry.ValidTokens,
"multimodal_calls": telemetry.MultimodalCalls,
"total_calls": telemetry.TotalCalls,
"success_calls": telemetry.SuccessCalls,
"success_rate": telemetry.SuccessRate,
"model_stats": telemetry.ModelStats,
},
}
if len(internal.Cfg.Note) > 0 {
response["note"] = internal.Cfg.Note
}
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(response)
}
func main() {
internal.LoadConfig()
internal.InitLogger()
if err := internal.GetTokenManager().Start(); err != nil {
internal.LogError("TokenManager 启动失败: %v", err)
}
internal.StartAnonymousTokenPool()
internal.StartVersionUpdater()
internal.StartModelFetcher()
internal.GetApiKeyManager().StartPeriodicSave()
http.HandleFunc("/", corsMiddleware(loggingMiddleware(handleRoot)))
http.HandleFunc("/v1/models", corsMiddleware(loggingMiddleware(internal.HandleModels)))
http.HandleFunc("/v1/chat/completions", corsMiddleware(loggingMiddleware(internal.HandleChatCompletions)))
// Admin Web UI
http.HandleFunc("/admin", corsMiddleware(loggingMiddleware(internal.HandleAdminUI)))
http.HandleFunc("/admin/", corsMiddleware(loggingMiddleware(internal.HandleAdminUI)))
http.HandleFunc("/admin/api/login", corsMiddleware(loggingMiddleware(internal.HandleAdminLogin)))
http.HandleFunc("/admin/api/logout", corsMiddleware(loggingMiddleware(internal.HandleAdminLogout)))
http.HandleFunc("/admin/api/overview", corsMiddleware(loggingMiddleware(internal.HandleAdminOverview)))
http.HandleFunc("/admin/api/config", corsMiddleware(loggingMiddleware(internal.HandleAdminConfig)))
// Tokens：GET 列表 / POST 添加 / DELETE 删除 / POST validate 验证
http.HandleFunc("/admin/api/tokens", corsMiddleware(loggingMiddleware(func(w http.ResponseWriter, r *http.Request) {
switch r.Method {
case http.MethodGet:
internal.HandleAdminTokens(w, r)
case http.MethodPost:
internal.HandleAdminTokenAdd(w, r)
case http.MethodDelete:
internal.HandleAdminTokenDelete(w, r)
default:
w.WriteHeader(http.StatusMethodNotAllowed)
}
})))
http.HandleFunc("/admin/api/tokens/validate", corsMiddleware(loggingMiddleware(internal.HandleAdminTokenValidate)))
// API Keys
http.HandleFunc("/admin/api/keys", corsMiddleware(loggingMiddleware(func(w http.ResponseWriter, r *http.Request) {
switch r.Method {
case http.MethodGet:
internal.HandleAdminKeysList(w, r)
case http.MethodPost:
internal.HandleAdminKeysCreate(w, r)
case http.MethodDelete:
internal.HandleAdminKeysDelete(w, r)
default:
w.WriteHeader(http.StatusMethodNotAllowed)
}
})))
http.HandleFunc("/admin/api/keys/toggle", corsMiddleware(loggingMiddleware(internal.HandleAdminKeysToggle)))
http.HandleFunc("/admin/api/models", corsMiddleware(loggingMiddleware(internal.HandleAdminModels)))
http.HandleFunc("/admin/api/test", corsMiddleware(loggingMiddleware(internal.HandleAdminTestModel)))
addr := ":" + internal.Cfg.Port
internal.LogInfo("Server starting on %s", addr)
internal.LogInfo("API docs available at http://localhost:%s/v1/models", internal.Cfg.Port)
if err := http.ListenAndServe(addr, nil); err != nil {
internal.LogError("Server failed: %v", err)
}
}
|