diff --git a/README.md b/README.md index e7fd84b..a49359a 100644 --- a/README.md +++ b/README.md @@ -11,6 +11,7 @@ - 上游 SSE 直接透传为 OpenAI SSE;非流式请求在本地聚合为 OpenAI Chat Completion JSON。 - OpenAI 函数/工具调用:支持 `tools`、`tool_choice`、流式 `delta.tool_calls`、非流式 `message.tool_calls` 以及 `role: tool` 结果续传。 - Token 消耗统计:流式与非流式请求均解析上游 `usage`,按模型/按日/最近明细写入本地 SQLite(`zhanlu.db`),凭据也一并持久化在同一个库中。管理页提供 `GET /admin/stats` 可视化与 `GET /api/stats` JSON 接口。 +- 可用模型:管理后台新增「可用模型」tab,实时拉取上游 `/gateway/v1/model/info` 返回的模型列表,并提供单模型可用性及首字延时(TTFT)探测(`GET /api/models`、`POST /api/models/test`)。 ## 运行 @@ -234,6 +235,22 @@ curl http://127.0.0.1:8080/v1/models ` 设置 `ZHANLU_STATS_DISABLED=true` 可停止写入统计。 +## 可用模型 + +管理后台「可用模型」tab(位于「Token 统计」之后)实时展示当前上游接口返回的模型列表,并提供单模型可用性与延时探测: + +- 进入 tab 时自动拉取一次,也可随时点击「刷新」重新获取。 +- 每个模型行可单独「测试」,或点击「全部测试」依次探测所有模型。 +- 探测向上游 `/chat/completions` 发送一条极简流式请求(`hi`),测量首字延时(TTFT,首个 SSE 数据块到达时间)与总耗时,并在收到首个内容块后立即关闭连接,避免消耗额外 token。 +- 探测请求不计入 Token 统计。 + +对应 JSON 接口(需先登录管理页面): + +- `GET /api/models`:实时返回上游模型列表 `{"ok":true,"models":["GLM-4.7",...]}`。 +- `POST /api/models/test`:请求体 `{"model":""}`,返回 `{"ok":true,"model":"...","available":true,"ttft_ms":123,"total_ms":456}` 或 `{"ok":true,"model":"...","available":false,"error":"...","total_ms":789}`。 + +探测超时上限为 30 秒(`modelTestTimeout`);凭据未配置或 API Key 缺失时 `GET /api/models` 与 `POST /api/models/test` 会按需自动换取 API Key,与聊天接口一致。 + ## 安全说明 - `zhanlu.db` 数据库包含明文 `AccessKey`、`SecretKey`、`Token` 和 `apiKey`,请不要提交到仓库。 diff --git a/internal/server/server.go b/internal/server/server.go index 3653c78..5498af9 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -61,6 +61,8 @@ func (s *Server) routes() { s.mux.HandleFunc("POST /api/sso/exchange", s.withLoginSession(s.exchangeSSOCode)) s.mux.HandleFunc("GET /api/stats", s.withLoginSession(s.getStats)) s.mux.HandleFunc("POST /api/stats/reset", s.withLoginSession(s.resetStats)) + s.mux.HandleFunc("GET /api/models", s.withLoginSession(s.getModels)) + s.mux.HandleFunc("POST /api/models/test", s.withLoginSession(s.testModel)) s.mux.HandleFunc("GET /v1/models", s.withAPIKey(s.models)) s.mux.HandleFunc("POST /v1/chat/completions", s.withAPIKey(s.chatCompletions)) } @@ -593,6 +595,159 @@ func (s *Server) resetStats(w http.ResponseWriter, r *http.Request) { writeJSON(w, http.StatusOK, map[string]any{"ok": true}) } +// modelTestTimeout caps how long a single model probe may wait for the first +// token. The full upstream timeout (default 300s) is far too long for a probe. +const modelTestTimeout = 30 * time.Second + +// getModels returns the live model list advertised by the upstream gateway +// model-info endpoint, fetched on every request so the admin console reflects +// newly published models without a restart. It provisions an API key on demand +// if the stored credentials lack one, mirroring the chat path. +func (s *Server) getModels(w http.ResponseWriter, r *http.Request) { + creds, err := s.currentCredentials() + if err != nil { + writeJSON(w, http.StatusOK, map[string]any{"ok": false, "error": "凭据未配置,请先在凭据登录页登录"}) + return + } + if !creds.HasAPIKey() { + creds, err = s.provisionCredentials(r.Context(), creds) + if err != nil { + writeJSON(w, http.StatusBadGateway, map[string]any{"ok": false, "error": err.Error()}) + return + } + s.cfg.Credentials = creds + _ = s.st.SaveCredentials(creds) + } + modelBaseURL := firstNonEmpty(creds.ModelBaseURL, s.cfg.MobileModelBaseURL) + client, err := s.zhanluClientWithBase(modelBaseURL) + if err != nil { + writeJSON(w, http.StatusInternalServerError, map[string]any{"ok": false, "error": err.Error()}) + return + } + models, err := client.Models(r.Context(), creds.APIKey) + if err != nil { + writeJSON(w, http.StatusBadGateway, map[string]any{"ok": false, "error": err.Error()}) + return + } + writeJSON(w, http.StatusOK, map[string]any{"ok": true, "models": models}) +} + +// testModel sends a minimal streaming chat completion to the upstream gateway +// for the requested model and measures the time to first token (TTFT) and the +// total probe duration, so the admin console can report availability and +// latency. The stream is closed as soon as the first content chunk arrives to +// avoid consuming tokens beyond what the probe needs. Probe requests are not +// recorded in token statistics. +func (s *Server) testModel(w http.ResponseWriter, r *http.Request) { + var in struct { + Model string `json:"model"` + } + if err := json.NewDecoder(r.Body).Decode(&in); err != nil { + writeJSON(w, http.StatusBadRequest, map[string]any{"ok": false, "error": err.Error()}) + return + } + model := strings.TrimSpace(in.Model) + if model == "" { + writeJSON(w, http.StatusBadRequest, map[string]any{"ok": false, "error": "model 不能为空"}) + return + } + creds, err := s.currentCredentials() + if err != nil { + writeJSON(w, http.StatusOK, map[string]any{"ok": true, "model": model, "available": false, "error": "凭据未配置"}) + return + } + if !creds.HasAPIKey() { + creds, err = s.provisionCredentials(r.Context(), creds) + if err != nil { + writeJSON(w, http.StatusOK, map[string]any{"ok": true, "model": model, "available": false, "error": err.Error()}) + return + } + s.cfg.Credentials = creds + _ = s.st.SaveCredentials(creds) + } + modelBaseURL := firstNonEmpty(creds.ModelBaseURL, s.cfg.MobileModelBaseURL) + client, err := s.zhanluClientWithBase(modelBaseURL) + if err != nil { + writeJSON(w, http.StatusOK, map[string]any{"ok": true, "model": model, "available": false, "error": err.Error()}) + return + } + body, _ := json.Marshal(map[string]any{ + "model": model, + "messages": []map[string]any{{"role": "user", "content": "hi"}}, + "stream": true, + "stream_options": map[string]any{"include_usage": true}, + }) + ctx, cancel := context.WithTimeout(r.Context(), modelTestTimeout) + defer cancel() + start := time.Now() + resp, err := client.ChatCompletions(ctx, creds.APIKey, body) + if err != nil { + writeJSON(w, http.StatusOK, map[string]any{"ok": true, "model": model, "available": false, "error": err.Error(), "total_ms": time.Since(start).Milliseconds()}) + return + } + defer resp.Body.Close() + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + b, _ := io.ReadAll(io.LimitReader(resp.Body, 1024)) + writeJSON(w, http.StatusOK, map[string]any{"ok": true, "model": model, "available": false, "error": fmt.Sprintf("上游返回 %d: %s", resp.StatusCode, string(b)), "total_ms": time.Since(start).Milliseconds()}) + return + } + // Read the SSE stream until the first content chunk arrives (or an error), + // then close the connection so the probe consumes at most one token. + reader := bufio.NewReader(resp.Body) + var ttft int64 + gotStream := false + available := false + errMsg := "" + for { + line, err := reader.ReadString('\n') + if line != "" { + payload := sseDataPayload(line) + if payload != "" && payload != "[DONE]" { + var evt struct { + State string `json:"state"` + ErrorMessage string `json:"errorMessage"` + Choices []struct { + Delta struct { + Content string `json:"content"` + } `json:"delta"` + } `json:"choices"` + } + if json.Unmarshal([]byte(payload), &evt) == nil { + if evt.State == "ERROR" { + errMsg = firstNonEmpty(evt.ErrorMessage, "上游返回错误") + break + } + if !gotStream { + gotStream = true + ttft = time.Since(start).Milliseconds() + } + if len(evt.Choices) > 0 && evt.Choices[0].Delta.Content != "" { + available = true + break + } + } + } + } + if err != nil { + if !gotStream && errMsg == "" { + errMsg = err.Error() + } + break + } + } + if gotStream && errMsg == "" { + available = true + } + result := map[string]any{"ok": true, "model": model, "available": available, "total_ms": time.Since(start).Milliseconds()} + if gotStream { + result["ttft_ms"] = ttft + } + if !available { + result["error"] = firstNonEmpty(errMsg, "未收到响应内容") + } + writeJSON(w, http.StatusOK, result) +} + func parseLimit(s string) int { n := 0 for _, c := range s { @@ -1262,6 +1417,7 @@ var adminTemplate = template.Must(template.New("admin").Funcs(template.FuncMap{
+
@@ -1314,6 +1470,23 @@ var adminTemplate = template.Must(template.New("admin").Funcs(template.FuncMap{ +
+
+ + +
+
+

可用模型

+

点击刷新获取当前上游接口返回的模型列表。

+
+ + + +
模型状态首字延时总耗时操作
+
+
+
+