Compare commits
4
Commits
v0.0.4
..
c8938cb514
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
c8938cb514 | ||
|
|
8a1cf4d61f | ||
|
|
fa9640d919 | ||
|
|
dd5c560b64 |
@@ -58,8 +58,13 @@ jobs:
|
||||
CGO_ENABLED=0 GOOS=linux GOARCH="${arch}" \
|
||||
go build -trimpath -ldflags "${LDFLAGS}" \
|
||||
-o "zhanlu-proxy-linux-${arch}" ./cmd/zhanlu-proxy
|
||||
|
||||
echo "building windows/${arch}"
|
||||
CGO_ENABLED=0 GOOS=windows GOARCH="${arch}" \
|
||||
go build -trimpath -ldflags "${LDFLAGS}" \
|
||||
-o "zhanlu-proxy-windows-${arch}.exe" ./cmd/zhanlu-proxy
|
||||
done
|
||||
ls -lh zhanlu-proxy-linux-*
|
||||
ls -lh zhanlu-proxy-linux-* zhanlu-proxy-windows-*.exe
|
||||
|
||||
- name: Publish release assets
|
||||
if: startsWith(github.ref, 'refs/tags/v')
|
||||
@@ -80,10 +85,11 @@ jobs:
|
||||
echo "release id: ${RELEASE_ID}"
|
||||
|
||||
for arch in amd64 arm64; do
|
||||
f="zhanlu-proxy-linux-${arch}"
|
||||
for f in "zhanlu-proxy-linux-${arch}" "zhanlu-proxy-windows-${arch}.exe"; do
|
||||
curl -fsSL -X POST -H "${AUTH}" \
|
||||
-F "attachment=@${f};filename=${f}" \
|
||||
"${API}/releases/${RELEASE_ID}/assets?name=${f}"
|
||||
done
|
||||
done
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
|
||||
+7
-1
@@ -1,4 +1,9 @@
|
||||
credentials.json
|
||||
zhanlu.db
|
||||
zhanlu.db-shm
|
||||
zhanlu.db-wal
|
||||
*.db
|
||||
*.db-shm
|
||||
*.db-wal
|
||||
extension/
|
||||
*.exe
|
||||
*.log
|
||||
@@ -8,3 +13,4 @@ extension/
|
||||
tmp/
|
||||
temp/
|
||||
source/
|
||||
output/
|
||||
|
||||
@@ -10,6 +10,8 @@
|
||||
- 模型 API Key 换取签名逻辑:SM3 摘要 + SM2 签名(`X-Auth-Signature`/`X-Auth-Timestamp`/`X-Auth-Nonce`)。
|
||||
- 上游 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`)。
|
||||
|
||||
## 运行
|
||||
|
||||
@@ -37,17 +39,17 @@ http://127.0.0.1:8080/
|
||||
/opt/zhanlu-proxy/zhanlu-proxy
|
||||
```
|
||||
|
||||
凭据文件放在:
|
||||
数据库放在:
|
||||
|
||||
```text
|
||||
/opt/zhanlu-proxy/credentials.json
|
||||
/opt/zhanlu-proxy/zhanlu.db
|
||||
```
|
||||
|
||||
创建环境变量文件 `/etc/zhanlu-proxy/zhanlu-proxy.env`:
|
||||
|
||||
```env
|
||||
ZHANLU_LISTEN_ADDR=:8080
|
||||
ZHANLU_CREDENTIALS_FILE=/opt/zhanlu-proxy/credentials.json
|
||||
ZHANLU_DB_FILE=/opt/zhanlu-proxy/zhanlu.db
|
||||
ZHANLU_MOBILE_LOGIN_BASE_URL=https://ecloud.10086.cn
|
||||
ZHANLU_MOBILE_MODEL_BASE_URL=https://ecloud.10086.cn/api/query/aigateway
|
||||
ZHANLU_UPSTREAM_TIMEOUT=300s
|
||||
@@ -106,20 +108,22 @@ journalctl -u zhanlu-proxy -f
|
||||
- 使用本次 `secret` AES 解密响应中的 `ak`、`sk`、`license`,得到 `AccessKey`、`SecretKey`、`Token`。
|
||||
- 按插件 v1.4.2 流程调用 `/api/acepilot/zhanlu/v1/login`(RSA+HmacSHA1 签名 URL + `plugin_type=zhanlu_ide` 请求头)获取用户资料(email/组织/团队)。
|
||||
- 用 SM2 私钥签名调用 `{mobileModelBaseUrl}/user/api/v2/external/key/get-or-create` 换取模型 `apiKey`。
|
||||
- 凭据(含 `apiKey`、`modelBaseUrl`、email 等)会写入 JSON 文件,后续 OpenAI 兼容接口自动使用。
|
||||
- 凭据(含 `apiKey`、`modelBaseUrl`、email 等)会写入本地 SQLite 数据库(`zhanlu.db`),后续 OpenAI 兼容接口自动使用。
|
||||
|
||||
默认保存到当前执行目录:
|
||||
默认数据库保存在当前执行目录:
|
||||
|
||||
```text
|
||||
credentials.json
|
||||
zhanlu.db
|
||||
```
|
||||
|
||||
可以通过环境变量覆盖:
|
||||
|
||||
```powershell
|
||||
$env:ZHANLU_CREDENTIALS_FILE="E:\path\to\credentials.json"
|
||||
$env:ZHANLU_DB_FILE="E:\path\to\zhanlu.db"
|
||||
```
|
||||
|
||||
凭据仅持久化在数据库中,不再使用 JSON 文件。
|
||||
|
||||
手机号验证码登录使用 `ZHANLU_MOBILE_LOGIN_BASE_URL`,默认公网地址来自插件配置(兼容旧环境变量 `ZHANLU_SERVER_BASE_URL`):
|
||||
|
||||
```powershell
|
||||
@@ -184,7 +188,8 @@ curl http://127.0.0.1:8080/v1/models `
|
||||
| `ZHANLU_MOBILE_LOGIN_BASE_URL` | `https://ecloud.10086.cn` | 移动云公网登录 Base URL(兼容旧变量 `ZHANLU_SERVER_BASE_URL`) |
|
||||
| `ZHANLU_MOBILE_MODEL_BASE_URL` | `https://ecloud.10086.cn/api/query/aigateway` | 移动云公网模型网关 Base URL |
|
||||
| `ZHANLU_UPSTREAM_PATH` | `/chat/completions` | 模型网关聊天接口路径 |
|
||||
| `ZHANLU_CREDENTIALS_FILE` | `credentials.json` | 凭据 JSON 路径,默认当前执行目录 |
|
||||
| `ZHANLU_DB_FILE` | `zhanlu.db` | 本地 SQLite 数据库路径,凭据与 token 统计均存于此,默认当前执行目录 |
|
||||
| `ZHANLU_STATS_DISABLED` | `false` | 设为 `true` 关闭 token 用量记录(仅停止写入统计,凭据存储不受影响) |
|
||||
| `ZHANLU_ACCESS_KEY` | 空 | 直接从环境变量提供 AccessKey |
|
||||
| `ZHANLU_SECRET_KEY` | 空 | 直接从环境变量提供 SecretKey |
|
||||
| `ZHANLU_TOKEN` | 空 | 直接从环境变量提供 Token |
|
||||
@@ -207,14 +212,49 @@ curl http://127.0.0.1:8080/v1/models `
|
||||
服务启动时按以下优先级加载凭据:
|
||||
|
||||
1. 环境变量 `ZHANLU_ACCESS_KEY`、`ZHANLU_SECRET_KEY`、`ZHANLU_TOKEN` 或 `ZHANLU_API_KEY`。
|
||||
2. `ZHANLU_CREDENTIALS_FILE` 指向的 JSON 文件。
|
||||
2. `ZHANLU_DB_FILE` 数据库中持久化的凭据行。
|
||||
|
||||
登录页面保存后,运行中的服务会立即使用新凭据。若环境中只有 AK/SK/Token 而没有 `apiKey`,首次调用聊天接口时会自动按插件流程换取 API Key 并回写凭据文件。
|
||||
登录页面保存后,运行中的服务会立即使用新凭据并写入数据库。若环境中只有 AK/SK/Token 而没有 `apiKey`,首次调用聊天接口时会自动按插件流程换取 API Key 并回写数据库。
|
||||
|
||||
## Token 消耗统计
|
||||
|
||||
代理在 `/v1/chat/completions` 完成后解析上游 `usage`(流式路径在透传 SSE 的同时旁路解析末块 `usage`,非流式路径在聚合时解析),将以下维度写入 `ZHANLU_DB_FILE`:
|
||||
|
||||
- 时间、模型、流式/非流式
|
||||
- `prompt_tokens` / `completion_tokens` / `total_tokens` / `reasoning_tokens`(思考模型)
|
||||
- `cached_tokens`(来自 `prompt_tokens_details.cached_tokens`,提示缓存命中的 token 数)与缓存命中率
|
||||
- 请求状态(`success` / `upstream_error`)与耗时
|
||||
|
||||
缓存命中率为 `cached_tokens / prompt_tokens`,仅在模型/网关支持 prompt 缓存且上游返回 `cached_tokens` 时非零。
|
||||
|
||||
管理页(需先登录管理页面)提供:
|
||||
|
||||
- `GET /admin/stats`:可视化页面,展示总览、按模型、按日柱状、最近请求明细,带重置按钮。
|
||||
- `GET /api/stats`:JSON 接口,支持 `since`/`until`(RFC3339)、`model`、`limit` 查询参数。
|
||||
- `POST /api/stats/reset`:清空统计(凭据不受影响)。
|
||||
|
||||
设置 `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":"<model_id>"}`,返回 `{"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,与聊天接口一致。
|
||||
|
||||
## 安全说明
|
||||
|
||||
- `credentials.json` 包含明文 `AccessKey`、`SecretKey`、`Token` 和 `apiKey`,请不要提交到仓库。
|
||||
- 默认保存在当前执行目录的 `credentials.json`。
|
||||
- `zhanlu.db` 数据库包含明文 `AccessKey`、`SecretKey`、`Token` 和 `apiKey`,请不要提交到仓库。
|
||||
- 默认保存在当前执行目录的 `zhanlu.db`(建议通过 `ZHANLU_DB_FILE` 指向受保护路径)。
|
||||
- 建议设置 `ZHANLU_LOGIN_PASSWORD`,避免公网暴露的 `/admin/login` 被直接访问。
|
||||
- 错误响应默认不会返回签名 URL,避免泄露 `AccessKey`、`authorization`、`Signature`。
|
||||
- `ZHANLU_DEBUG=true` 时会返回更详细错误,但仍会对敏感 query 参数脱敏。
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
|
||||
"git.misaka.ren/M1saka/zhanlu_proxy/internal/config"
|
||||
"git.misaka.ren/M1saka/zhanlu_proxy/internal/server"
|
||||
"git.misaka.ren/M1saka/zhanlu_proxy/internal/store"
|
||||
)
|
||||
|
||||
func main() {
|
||||
@@ -14,7 +15,20 @@ func main() {
|
||||
log.Fatal(err)
|
||||
}
|
||||
|
||||
h := server.New(cfg)
|
||||
st, err := store.Open(cfg.DBPath)
|
||||
if err != nil {
|
||||
log.Fatalf("open store: %v", err)
|
||||
}
|
||||
defer st.Close()
|
||||
|
||||
// Credential precedence at startup: environment variables > db row.
|
||||
if cfg.Credentials.Validate() != nil && !cfg.Credentials.HasAPIKey() {
|
||||
if dbCreds, derr := st.LoadCredentials(); derr == nil && (dbCreds.HasAPIKey() || dbCreds.Validate() == nil) {
|
||||
cfg.Credentials = dbCreds
|
||||
}
|
||||
}
|
||||
|
||||
h := server.New(cfg, st)
|
||||
log.Printf("zhanlu proxy listening on %s", cfg.ListenAddr)
|
||||
log.Printf("login page: http://127.0.0.1%s/login", cfg.ListenAddr)
|
||||
if err := http.ListenAndServe(cfg.ListenAddr, h); err != nil {
|
||||
|
||||
@@ -4,4 +4,16 @@ go 1.25.0
|
||||
|
||||
require github.com/emmansun/gmsm v0.44.1
|
||||
|
||||
require golang.org/x/crypto v0.54.0 // indirect
|
||||
require (
|
||||
github.com/dustin/go-humanize v1.0.1 // indirect
|
||||
github.com/google/uuid v1.6.0 // indirect
|
||||
github.com/mattn/go-isatty v0.0.24 // indirect
|
||||
github.com/ncruces/go-strftime v1.0.0 // indirect
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
|
||||
golang.org/x/crypto v0.54.0 // indirect
|
||||
golang.org/x/sys v0.47.0 // indirect
|
||||
modernc.org/libc v1.74.4 // indirect
|
||||
modernc.org/mathutil v1.7.1 // indirect
|
||||
modernc.org/memory v1.11.0 // indirect
|
||||
modernc.org/sqlite v1.56.0 // indirect
|
||||
)
|
||||
|
||||
@@ -1,6 +1,24 @@
|
||||
github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
|
||||
github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
|
||||
github.com/emmansun/gmsm v0.44.1 h1:zDTkdtLWFG0vCbhPV+k9pte14tix/eK71At9Iai9fP4=
|
||||
github.com/emmansun/gmsm v0.44.1/go.mod h1:p6RIUta0/KboFHrOxr1x8q+pd8RZtdaTO7XNp0RmMQM=
|
||||
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
|
||||
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||
github.com/mattn/go-isatty v0.0.24 h1:tGZZoVgT/KiqK1c8ocVLeDS8BSWMRd47J3Lbz7vsReI=
|
||||
github.com/mattn/go-isatty v0.0.24/go.mod h1:nMCL3Zebbrt45jsMDgnfIwz6ydEQApk5oEI3HqDio6A=
|
||||
github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w=
|
||||
github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE=
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
|
||||
golang.org/x/crypto v0.54.0 h1:YLIA59K4fiNzHzjnZt2tUJQjQtUWfWbeHBqKtk3eScw=
|
||||
golang.org/x/crypto v0.54.0/go.mod h1:KWL8ny2AZdGR2cWmzeHrp2azQPGogOv+HeQaVEXC2dk=
|
||||
golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs=
|
||||
golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
|
||||
modernc.org/libc v1.74.4 h1:fX1Omw4o2/1C2iRkkIsrQTasJQldLhRmuPreXLoWs9k=
|
||||
modernc.org/libc v1.74.4/go.mod h1:eeQAS9W3sZeKYMFubydxJpII9ybHWshk+7or7bLG9co=
|
||||
modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU=
|
||||
modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg=
|
||||
modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI=
|
||||
modernc.org/memory v1.11.0/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw=
|
||||
modernc.org/sqlite v1.56.0 h1:/D8e2RfFqoy/Zc6PuC76U28zFwmI/sYx1Kjm4yEn9e0=
|
||||
modernc.org/sqlite v1.56.0/go.mod h1:yCJ2cmAaIkHQ25oXWrF8H4O1lIfPYPR26yCEDj2P3pQ=
|
||||
|
||||
@@ -1,10 +1,7 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
@@ -49,30 +46,3 @@ func (c Credentials) Validate() error {
|
||||
func (c Credentials) HasAPIKey() bool {
|
||||
return strings.TrimSpace(c.APIKey) != "" && strings.TrimSpace(c.ModelBaseURL) != ""
|
||||
}
|
||||
|
||||
func LoadCredentials(path string) (Credentials, error) {
|
||||
b, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return Credentials{}, err
|
||||
}
|
||||
var c Credentials
|
||||
if err := json.Unmarshal(b, &c); err != nil {
|
||||
return Credentials{}, err
|
||||
}
|
||||
return c, c.Validate()
|
||||
}
|
||||
|
||||
func SaveCredentials(path string, c Credentials) error {
|
||||
if err := c.Validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
c.SavedAt = time.Now()
|
||||
b, err := json.MarshalIndent(c, "", " ")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
|
||||
return err
|
||||
}
|
||||
return os.WriteFile(path, b, 0o600)
|
||||
}
|
||||
|
||||
@@ -2,7 +2,6 @@ package config
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
@@ -14,7 +13,8 @@ type Config struct {
|
||||
MobileLoginBaseURL string
|
||||
MobileModelBaseURL string
|
||||
UpstreamPath string
|
||||
CredentialsPath string
|
||||
DBPath string
|
||||
StatsDisabled bool
|
||||
SSOExchangeURL string
|
||||
SSOBaseURL string
|
||||
TokenDecryptKey string
|
||||
@@ -36,7 +36,8 @@ func Load() (Config, error) {
|
||||
MobileLoginBaseURL: firstNonEmpty(os.Getenv("ZHANLU_MOBILE_LOGIN_BASE_URL"), getenv("ZHANLU_SERVER_BASE_URL", "https://ecloud.10086.cn")),
|
||||
MobileModelBaseURL: getenv("ZHANLU_MOBILE_MODEL_BASE_URL", "https://ecloud.10086.cn/api/query/aigateway"),
|
||||
UpstreamPath: getenv("ZHANLU_UPSTREAM_PATH", "/chat/completions"),
|
||||
CredentialsPath: getenv("ZHANLU_CREDENTIALS_FILE", defaultCredentialsPath()),
|
||||
DBPath: getenv("ZHANLU_DB_FILE", "zhanlu.db"),
|
||||
StatsDisabled: strings.EqualFold(os.Getenv("ZHANLU_STATS_DISABLED"), "true"),
|
||||
SSOBaseURL: getenv("ZHANLU_SSO_BASE_URL", "http://4c.hq.cmcc"),
|
||||
SSOExchangeURL: getenv("ZHANLU_SSO_EXCHANGE_URL", "http://rdcloud.4c.hq.cmcc/cmdevops-aiplus-agent-gateway/api/acepilot/zhanlu/authToken"),
|
||||
TokenDecryptKey: getenv("ZHANLU_TOKEN_DECRYPT_KEY", "3jw7woww2rvhla6k"),
|
||||
@@ -56,12 +57,8 @@ func Load() (Config, error) {
|
||||
Token: os.Getenv("ZHANLU_TOKEN"),
|
||||
APIKey: os.Getenv("ZHANLU_API_KEY"),
|
||||
}
|
||||
if cfg.Credentials.Validate() == nil || cfg.Credentials.HasAPIKey() {
|
||||
return cfg, nil
|
||||
}
|
||||
if creds, err := auth.LoadCredentials(cfg.CredentialsPath); err == nil {
|
||||
cfg.Credentials = creds
|
||||
}
|
||||
// Credentials from the environment were incomplete; the store serves the
|
||||
// persisted row at runtime.
|
||||
return cfg, nil
|
||||
}
|
||||
|
||||
@@ -93,10 +90,6 @@ func firstNonEmpty(values ...string) string {
|
||||
return ""
|
||||
}
|
||||
|
||||
func defaultCredentialsPath() string {
|
||||
return filepath.Join(".", "credentials.json")
|
||||
}
|
||||
|
||||
const defaultPublicKeyPEM = `-----BEGIN PUBLIC KEY-----
|
||||
MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAhxudxTewPgljUHEZHkusP7m3I+zA4/RGvuUMt6TtII/m4zwUOm/Y31zHBTmkCCt8k5vj9y+AmO0TsGmHooNQuMebakdmEWdcA5h7YAHHFbF2w5LcxIXjib08vgVpA+m3R5xPbLK+vfHe2aAX36b5nHReDNncY5vAl3U4CgIEBGPqyG67vJytRWqP+sfEdw5+m192Rf4SCGyiBzRmjiVlH3zeEBjdbOrkAnzKOVz6AHBl2q7LPLJKIzxjoAyhEp5qnDjHUFo5VZUgFwUOt83A/jbGMyzmjRoxBuvKcs9tBuorZyUwIsZN6E+rtQk2YqMPj4RkDsZ7LRmj6on8sN2rHQIDAQAB
|
||||
-----END PUBLIC KEY-----`
|
||||
|
||||
+865
-337
File diff suppressed because it is too large
Load Diff
+196
-13
@@ -15,12 +15,15 @@ import (
|
||||
|
||||
"git.misaka.ren/M1saka/zhanlu_proxy/internal/auth"
|
||||
"git.misaka.ren/M1saka/zhanlu_proxy/internal/config"
|
||||
"git.misaka.ren/M1saka/zhanlu_proxy/internal/store"
|
||||
)
|
||||
|
||||
const testSM2Key = "8d6ee90b3c4d299ae5abd655dbc3547c110ae8aeff1de18b0df241f215f90748"
|
||||
|
||||
// setupTestServer spins up a mock Zhanlu upstream and a proxy server wired to it.
|
||||
func setupTestServer(t *testing.T) (*httptest.Server, *httptest.Server, string) {
|
||||
// setupTestServer spins up a mock Zhanlu upstream and a proxy server wired to
|
||||
// it. The proxy is backed by a temp SQLite store so credentials persist in the
|
||||
// db during the test instead of a JSON file.
|
||||
func setupTestServer(t *testing.T) (*httptest.Server, *httptest.Server, *store.Store) {
|
||||
t.Helper()
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
@@ -47,7 +50,7 @@ func setupTestServer(t *testing.T) (*httptest.Server, *httptest.Server, string)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
_, _ = fmt.Fprint(w, "data: {\"id\":\"chatcmpl-x\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hello\"},\"finish_reason\":null}],\"usage\":{\"prompt_tokens\":1}}\n\n")
|
||||
_, _ = fmt.Fprint(w, "data: {\"id\":\"chatcmpl-x\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hello\"},\"finish_reason\":null}],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":20,\"total_tokens\":30,\"prompt_tokens_details\":{\"cached_tokens\":4}}}\n\n")
|
||||
_, _ = fmt.Fprint(w, "data: {\"id\":\"chatcmpl-x\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n")
|
||||
_, _ = fmt.Fprint(w, "data: [DONE]\n\n")
|
||||
case "/gateway/v1/model/info":
|
||||
@@ -59,28 +62,32 @@ func setupTestServer(t *testing.T) (*httptest.Server, *httptest.Server, string)
|
||||
}
|
||||
}))
|
||||
|
||||
credsFile := filepath.Join(t.TempDir(), "credentials.json")
|
||||
st, err := store.Open(filepath.Join(t.TempDir(), "stats.db"))
|
||||
if err != nil {
|
||||
t.Fatalf("open store: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { st.Close() })
|
||||
cfg := config.Config{
|
||||
ListenAddr: ":0",
|
||||
MobileLoginBaseURL: upstream.URL,
|
||||
MobileModelBaseURL: upstream.URL,
|
||||
UpstreamPath: "/chat/completions",
|
||||
CredentialsPath: credsFile,
|
||||
DBPath: filepath.Join(t.TempDir(), "zhanlu.db"),
|
||||
TokenDecryptKey: "3jw7woww2rvhla6k",
|
||||
PublicKeyPEM: defaultTestPublicKey,
|
||||
PhonePublicKeyPEM: defaultTestPublicKey,
|
||||
SM2PrivateKey: testSM2Key,
|
||||
PluginVersion: "1.4.2",
|
||||
}
|
||||
h := New(cfg)
|
||||
h := New(cfg, st)
|
||||
proxy := httptest.NewServer(h)
|
||||
return upstream, proxy, credsFile
|
||||
return upstream, proxy, st
|
||||
}
|
||||
|
||||
// TestPhoneLoginAndChat exercises the full v1.4.2 flow: SMS login, profile
|
||||
// fetch, SM2 API-key provisioning, then OpenAI-compatible chat and models.
|
||||
func TestPhoneLoginAndChat(t *testing.T) {
|
||||
upstream, proxy, credsFile := setupTestServer(t)
|
||||
upstream, proxy, st := setupTestServer(t)
|
||||
defer upstream.Close()
|
||||
defer proxy.Close()
|
||||
|
||||
@@ -107,8 +114,8 @@ func TestPhoneLoginAndChat(t *testing.T) {
|
||||
t.Fatalf("login failed: %v", loginResp)
|
||||
}
|
||||
|
||||
// 3. credentials file should contain the provisioned api key
|
||||
creds, err := auth.LoadCredentials(credsFile)
|
||||
// 3. credentials store should contain the provisioned api key
|
||||
creds, err := st.LoadCredentials()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@@ -154,11 +161,11 @@ func TestPhoneLoginAndChat(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestStreamingChat(t *testing.T) {
|
||||
upstream, proxy, credsFile := setupTestServer(t)
|
||||
upstream, proxy, st := setupTestServer(t)
|
||||
defer upstream.Close()
|
||||
defer proxy.Close()
|
||||
|
||||
// Seed credentials directly with the api key
|
||||
// Seed credentials directly with the api key into the store
|
||||
creds := auth.Credentials{
|
||||
AccessKey: "AK",
|
||||
SecretKey: "SK",
|
||||
@@ -167,7 +174,7 @@ func TestStreamingChat(t *testing.T) {
|
||||
ModelBaseURL: upstream.URL,
|
||||
Email: "[email protected]",
|
||||
}
|
||||
if err := auth.SaveCredentials(credsFile, creds); err != nil {
|
||||
if err := st.SaveCredentials(creds); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
@@ -188,6 +195,182 @@ func TestStreamingChat(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestStats verifies both streaming and non-streaming chat paths record token
|
||||
// usage into the store and that GET /api/stats aggregates them correctly.
|
||||
func TestStats(t *testing.T) {
|
||||
upstream, proxy, st := setupTestServer(t)
|
||||
defer upstream.Close()
|
||||
defer proxy.Close()
|
||||
|
||||
creds := auth.Credentials{
|
||||
AccessKey: "AK", SecretKey: "SK", Token: "TOKEN",
|
||||
APIKey: "sk-test-456", ModelBaseURL: upstream.URL, Email: "[email protected]",
|
||||
}
|
||||
if err := st.SaveCredentials(creds); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// non-streaming chat
|
||||
resp, err := http.Post(proxy.URL+"/v1/chat/completions", "application/json", strings.NewReader(`{"model":"GLM-4.7","messages":[{"role":"user","content":"hi"}],"stream":false}`))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
io.Copy(io.Discard, resp.Body)
|
||||
resp.Body.Close()
|
||||
|
||||
// streaming chat
|
||||
resp, err = http.Post(proxy.URL+"/v1/chat/completions", "application/json", strings.NewReader(`{"model":"GLM-4.7","messages":[{"role":"user","content":"hi"}],"stream":true}`))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
io.Copy(io.Discard, resp.Body)
|
||||
resp.Body.Close()
|
||||
|
||||
// query stats
|
||||
resp, err = http.Get(proxy.URL + "/api/stats")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
var statsResp map[string]any
|
||||
if err := json.NewDecoder(resp.Body).Decode(&statsResp); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if enabled, _ := statsResp["enabled"].(bool); !enabled {
|
||||
t.Fatalf("stats not enabled: %v", statsResp)
|
||||
}
|
||||
s, _ := statsResp["stats"].(map[string]any)
|
||||
if s == nil {
|
||||
t.Fatalf("no stats object: %v", statsResp)
|
||||
}
|
||||
totals, _ := s["totals"].(map[string]any)
|
||||
if totals == nil {
|
||||
t.Fatalf("no totals: %v", s)
|
||||
}
|
||||
if got := totNum(totals["requests"]); got != 2 {
|
||||
t.Fatalf("requests = %v, want 2", totals["requests"])
|
||||
}
|
||||
if got := totNum(totals["prompt_tokens"]); got != 20 {
|
||||
t.Fatalf("prompt_tokens = %v, want 20", totals["prompt_tokens"])
|
||||
}
|
||||
if got := totNum(totals["completion_tokens"]); got != 40 {
|
||||
t.Fatalf("completion_tokens = %v, want 40", totals["completion_tokens"])
|
||||
}
|
||||
if got := totNum(totals["total_tokens"]); got != 60 {
|
||||
t.Fatalf("total_tokens = %v, want 60", totals["total_tokens"])
|
||||
}
|
||||
if got := totNum(totals["cached_tokens"]); got != 8 {
|
||||
t.Fatalf("cached_tokens = %v, want 8", totals["cached_tokens"])
|
||||
}
|
||||
if rate, _ := totals["cache_rate"].(float64); rate < 0.39 || rate > 0.41 {
|
||||
t.Fatalf("cache_rate = %v, want ~0.4", totals["cache_rate"])
|
||||
}
|
||||
perModel, _ := s["per_model"].([]any)
|
||||
if len(perModel) != 1 {
|
||||
t.Fatalf("per_model = %v, want 1 entry", perModel)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAdminModels verifies the admin /api/models endpoint returns the live
|
||||
// model list advertised by the upstream gateway model-info endpoint.
|
||||
func TestAdminModels(t *testing.T) {
|
||||
upstream, proxy, st := setupTestServer(t)
|
||||
defer upstream.Close()
|
||||
defer proxy.Close()
|
||||
|
||||
creds := auth.Credentials{
|
||||
AccessKey: "AK", SecretKey: "SK", Token: "TOKEN",
|
||||
APIKey: "sk-test-456", ModelBaseURL: upstream.URL, Email: "[email protected]",
|
||||
}
|
||||
if err := st.SaveCredentials(creds); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
resp, err := http.Get(proxy.URL + "/api/models")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
var data map[string]any
|
||||
if err := json.NewDecoder(resp.Body).Decode(&data); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if ok, _ := data["ok"].(bool); !ok {
|
||||
t.Fatalf("models endpoint not ok: %v", data)
|
||||
}
|
||||
models, _ := data["models"].([]any)
|
||||
if len(models) != 2 {
|
||||
t.Fatalf("models = %v, want 2", models)
|
||||
}
|
||||
}
|
||||
|
||||
// TestModelTest verifies the admin /api/models/test endpoint probes a model and
|
||||
// reports availability plus latency against the mock streaming upstream.
|
||||
func TestModelTest(t *testing.T) {
|
||||
upstream, proxy, st := setupTestServer(t)
|
||||
defer upstream.Close()
|
||||
defer proxy.Close()
|
||||
|
||||
creds := auth.Credentials{
|
||||
AccessKey: "AK", SecretKey: "SK", Token: "TOKEN",
|
||||
APIKey: "sk-test-456", ModelBaseURL: upstream.URL, Email: "[email protected]",
|
||||
}
|
||||
if err := st.SaveCredentials(creds); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
body, _ := json.Marshal(map[string]string{"model": "GLM-4.7"})
|
||||
resp, err := http.Post(proxy.URL+"/api/models/test", "application/json", bytes.NewReader(body))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
var data map[string]any
|
||||
if err := json.NewDecoder(resp.Body).Decode(&data); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if ok, _ := data["ok"].(bool); !ok {
|
||||
t.Fatalf("test endpoint not ok: %v", data)
|
||||
}
|
||||
if available, _ := data["available"].(bool); !available {
|
||||
t.Fatalf("model should be available: %v", data)
|
||||
}
|
||||
if _, ok := data["ttft_ms"]; !ok {
|
||||
t.Fatalf("ttft_ms should be present: %v", data)
|
||||
}
|
||||
}
|
||||
|
||||
// TestModelTestNoCredentials verifies the test endpoint reports unavailable
|
||||
// gracefully when no credentials are configured, instead of erroring.
|
||||
func TestModelTestNoCredentials(t *testing.T) {
|
||||
upstream, proxy, _ := setupTestServer(t)
|
||||
defer upstream.Close()
|
||||
defer proxy.Close()
|
||||
|
||||
body, _ := json.Marshal(map[string]string{"model": "GLM-4.7"})
|
||||
resp, err := http.Post(proxy.URL+"/api/models/test", "application/json", bytes.NewReader(body))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
var data map[string]any
|
||||
if err := json.NewDecoder(resp.Body).Decode(&data); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if ok, _ := data["ok"].(bool); !ok {
|
||||
t.Fatalf("test endpoint should stay ok: %v", data)
|
||||
}
|
||||
if available, _ := data["available"].(bool); available {
|
||||
t.Fatalf("model should not be available without creds: %v", data)
|
||||
}
|
||||
}
|
||||
|
||||
// totNum extracts an int from a JSON-decoded numeric value (float64).
|
||||
func totNum(v any) int {
|
||||
f, _ := v.(float64)
|
||||
return int(f)
|
||||
}
|
||||
|
||||
func okValue(m map[string]any) bool {
|
||||
ok, _ := m["ok"].(bool)
|
||||
return ok
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
// Package stats defines the types and helpers for OpenAI-compatible token
|
||||
// usage statistics recorded by the zhanlu proxy. The concrete SQLite-backed
|
||||
// recorder lives in internal/store; this package is dependency-free so it can
|
||||
// be referenced by both store and server without import cycles.
|
||||
package stats
|
||||
|
||||
import "time"
|
||||
|
||||
// Record is a single chat-completion usage observation.
|
||||
type Record struct {
|
||||
Ts time.Time `json:"ts"`
|
||||
Model string `json:"model"`
|
||||
Stream bool `json:"stream"`
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
CompletionTokens int `json:"completion_tokens"`
|
||||
TotalTokens int `json:"total_tokens"`
|
||||
ReasoningTokens int `json:"reasoning_tokens"`
|
||||
CachedTokens int `json:"cached_tokens"`
|
||||
Status string `json:"status"` // "success" | "upstream_error"
|
||||
LatencyMs int64 `json:"latency_ms"`
|
||||
}
|
||||
|
||||
// Query filters the recorded stats. Zero-value time fields mean unbounded on
|
||||
// that end; empty Model means all models. Limit caps the recent-records list
|
||||
// (0 = default). Stream filters by streaming mode (nil = both).
|
||||
type Query struct {
|
||||
Since time.Time
|
||||
Until time.Time
|
||||
Model string
|
||||
Limit int
|
||||
Stream *bool
|
||||
}
|
||||
|
||||
// Totals aggregates request counts and token sums over a filtered set.
|
||||
type Totals struct {
|
||||
Requests int `json:"requests"`
|
||||
SuccessRequests int `json:"success_requests"`
|
||||
ErrorRequests int `json:"error_requests"`
|
||||
PromptTokens int64 `json:"prompt_tokens"`
|
||||
CompletionTokens int64 `json:"completion_tokens"`
|
||||
TotalTokens int64 `json:"total_tokens"`
|
||||
ReasoningTokens int64 `json:"reasoning_tokens"`
|
||||
CachedTokens int64 `json:"cached_tokens"`
|
||||
CacheRate float64 `json:"cache_rate"` // cached_tokens / prompt_tokens, 0..1
|
||||
}
|
||||
|
||||
// ModelStat is a per-model aggregation row.
|
||||
type ModelStat struct {
|
||||
Model string `json:"model"`
|
||||
Requests int `json:"requests"`
|
||||
PromptTokens int64 `json:"prompt_tokens"`
|
||||
CompletionTokens int64 `json:"completion_tokens"`
|
||||
TotalTokens int64 `json:"total_tokens"`
|
||||
CachedTokens int64 `json:"cached_tokens"`
|
||||
}
|
||||
|
||||
// DayStat is a per-day aggregation row (server-local time, YYYY-MM-DD).
|
||||
type DayStat struct {
|
||||
Day string `json:"day"`
|
||||
Requests int `json:"requests"`
|
||||
PromptTokens int64 `json:"prompt_tokens"`
|
||||
CompletionTokens int64 `json:"completion_tokens"`
|
||||
TotalTokens int64 `json:"total_tokens"`
|
||||
CachedTokens int64 `json:"cached_tokens"`
|
||||
}
|
||||
|
||||
// Summary is the full result returned by a Recorder's Stats query.
|
||||
type Summary struct {
|
||||
Totals Totals `json:"totals"`
|
||||
PerModel []ModelStat `json:"per_model"`
|
||||
Daily []DayStat `json:"daily"`
|
||||
Recent []Record `json:"recent"`
|
||||
}
|
||||
|
||||
// Recorder persists and queries usage statistics. The concrete implementation
|
||||
// lives in internal/store; the interface is declared here so server code can
|
||||
// depend on the contract and tests can inject fakes.
|
||||
type Recorder interface {
|
||||
Record(r Record) error
|
||||
Stats(q Query) (*Summary, error)
|
||||
Reset() error
|
||||
Close() error
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
package stats
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Usage is the subset of the OpenAI chat-completion usage object the recorder
|
||||
// persists. Numbers arrive from JSON unmarshal as float64.
|
||||
type Usage struct {
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
CompletionTokens int `json:"completion_tokens"`
|
||||
TotalTokens int `json:"total_tokens"`
|
||||
ReasoningTokens int `json:"reasoning_tokens"`
|
||||
CachedTokens int `json:"cached_tokens"`
|
||||
|
||||
// PromptTokensDetails.CachedTokens is emitted by providers that support
|
||||
// prompt caching (OpenAI/DeepSeek/Zhipu litellm gateways). Some upstreams
|
||||
// put cached_tokens at the top level instead.
|
||||
PromptTokensDetails struct {
|
||||
CachedTokens int `json:"cached_tokens"`
|
||||
} `json:"prompt_tokens_details"`
|
||||
|
||||
// CompletionTokensDetails.ReasoningTokens is emitted by reasoning models;
|
||||
// some upstreams put reasoning_tokens at top level instead.
|
||||
CompletionTokensDetails struct {
|
||||
ReasoningTokens int `json:"reasoning_tokens"`
|
||||
} `json:"completion_tokens_details"`
|
||||
}
|
||||
|
||||
// ExtractUsage decodes a raw usage value (as produced by encoding/json into an
|
||||
// any) into token counts. It accepts both full usage maps and raw JSON bytes.
|
||||
// Missing fields default to 0; a nil v yields zero usage.
|
||||
func ExtractUsage(v any) Usage {
|
||||
var u Usage
|
||||
if v == nil {
|
||||
return u
|
||||
}
|
||||
switch t := v.(type) {
|
||||
case []byte:
|
||||
_ = json.Unmarshal(t, &u)
|
||||
case json.RawMessage:
|
||||
_ = json.Unmarshal(t, &u)
|
||||
case map[string]any:
|
||||
// Re-marshal + unmarshal is the simplest robust path for nested
|
||||
// *_tokens_details; usage payloads are tiny.
|
||||
if b, err := json.Marshal(t); err == nil {
|
||||
_ = json.Unmarshal(b, &u)
|
||||
}
|
||||
}
|
||||
if u.ReasoningTokens == 0 {
|
||||
u.ReasoningTokens = u.CompletionTokensDetails.ReasoningTokens
|
||||
}
|
||||
if u.CachedTokens == 0 {
|
||||
u.CachedTokens = u.PromptTokensDetails.CachedTokens
|
||||
}
|
||||
return u
|
||||
}
|
||||
|
||||
// RecordFromUsage builds a Record from a captured usage value plus context.
|
||||
func RecordFromUsage(model string, stream bool, usage any, status string, start time.Time) Record {
|
||||
u := ExtractUsage(usage)
|
||||
if u.TotalTokens == 0 && (u.PromptTokens != 0 || u.CompletionTokens != 0) {
|
||||
u.TotalTokens = u.PromptTokens + u.CompletionTokens
|
||||
}
|
||||
return Record{
|
||||
Ts: time.Now(),
|
||||
Model: model,
|
||||
Stream: stream,
|
||||
PromptTokens: u.PromptTokens,
|
||||
CompletionTokens: u.CompletionTokens,
|
||||
TotalTokens: u.TotalTokens,
|
||||
ReasoningTokens: u.ReasoningTokens,
|
||||
CachedTokens: u.CachedTokens,
|
||||
Status: status,
|
||||
LatencyMs: time.Since(start).Milliseconds(),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"git.misaka.ren/M1saka/zhanlu_proxy/internal/auth"
|
||||
)
|
||||
|
||||
// LoadCredentials reads the persisted credentials. The credentials row always
|
||||
// exists after Open; an empty row (nothing saved yet) yields a zero-value
|
||||
// Credentials with a nil error — callers check Validate()/HasAPIKey().
|
||||
func (s *Store) LoadCredentials() (auth.Credentials, error) {
|
||||
var c auth.Credentials
|
||||
var savedAt int64
|
||||
err := s.db.QueryRow(`SELECT access_key, secret_key, token, api_key, model_base_url, email, organization, team, base_url, saved_at FROM credentials WHERE id = 1`).
|
||||
Scan(&c.AccessKey, &c.SecretKey, &c.Token, &c.APIKey, &c.ModelBaseURL, &c.Email, &c.Organization, &c.Team, &c.BaseURL, &savedAt)
|
||||
if err != nil {
|
||||
return auth.Credentials{}, fmt.Errorf("load credentials: %w", err)
|
||||
}
|
||||
c.SavedAt = time.Unix(savedAt, 0).Local()
|
||||
return c, nil
|
||||
}
|
||||
|
||||
// SaveCredentials upserts the credentials into the single row. It only writes
|
||||
// when the credentials validate or already carry an API key, so partial /
|
||||
// env-only creds are not persisted.
|
||||
func (s *Store) SaveCredentials(c auth.Credentials) error {
|
||||
if c.Validate() != nil && !c.HasAPIKey() {
|
||||
return fmt.Errorf("save credentials: %w", c.Validate())
|
||||
}
|
||||
c.SavedAt = time.Now()
|
||||
_, err := s.db.Exec(`INSERT INTO credentials (id, access_key, secret_key, token, api_key, model_base_url, email, organization, team, base_url, saved_at)
|
||||
VALUES (1, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT(id) DO UPDATE SET
|
||||
access_key=excluded.access_key, secret_key=excluded.secret_key, token=excluded.token,
|
||||
api_key=excluded.api_key, model_base_url=excluded.model_base_url, email=excluded.email,
|
||||
organization=excluded.organization, team=excluded.team, base_url=excluded.base_url,
|
||||
saved_at=excluded.saved_at`,
|
||||
c.AccessKey, c.SecretKey, c.Token, c.APIKey, c.ModelBaseURL, c.Email, c.Organization, c.Team, c.BaseURL, c.SavedAt.Unix())
|
||||
if err != nil {
|
||||
return fmt.Errorf("save credentials: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,187 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"git.misaka.ren/M1saka/zhanlu_proxy/internal/stats"
|
||||
)
|
||||
|
||||
// Compile-time guard: *Store satisfies stats.Recorder.
|
||||
var _ stats.Recorder = (*Store)(nil)
|
||||
|
||||
const defaultRecentLimit = 200
|
||||
|
||||
// Record appends a single usage observation.
|
||||
func (s *Store) Record(rec stats.Record) error {
|
||||
if rec.Status == "" {
|
||||
rec.Status = "success"
|
||||
}
|
||||
stream := 0
|
||||
if rec.Stream {
|
||||
stream = 1
|
||||
}
|
||||
_, err := s.db.Exec(`INSERT INTO requests (ts, model, stream, prompt_tokens, completion_tokens, total_tokens, reasoning_tokens, cached_tokens, status, latency_ms) VALUES (?,?,?,?,?,?,?,?,?,?)`,
|
||||
rec.Ts.Unix(), rec.Model, stream, rec.PromptTokens, rec.CompletionTokens, rec.TotalTokens, rec.ReasoningTokens, rec.CachedTokens, rec.Status, rec.LatencyMs)
|
||||
return err
|
||||
}
|
||||
|
||||
// Stats computes the aggregate summary for the given query.
|
||||
func (s *Store) Stats(q stats.Query) (*stats.Summary, error) {
|
||||
q = normalizeQuery(q)
|
||||
where, args := whereClause(q)
|
||||
|
||||
sum := &stats.Summary{}
|
||||
if err := s.scanTotals(sum, where, args); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if sum.Totals.PromptTokens > 0 {
|
||||
sum.Totals.CacheRate = float64(sum.Totals.CachedTokens) / float64(sum.Totals.PromptTokens)
|
||||
}
|
||||
if err := s.scanPerModel(sum, where, args); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := s.scanDaily(sum, where, args); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := s.scanRecent(sum, q, where, args); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return sum, nil
|
||||
}
|
||||
|
||||
// Reset deletes all recorded usage statistics (the credentials row is kept).
|
||||
func (s *Store) Reset() error {
|
||||
_, err := s.db.Exec(`DELETE FROM requests`)
|
||||
return err
|
||||
}
|
||||
|
||||
func normalizeQuery(q stats.Query) stats.Query {
|
||||
if q.Limit <= 0 {
|
||||
q.Limit = defaultRecentLimit
|
||||
}
|
||||
if q.Limit > 5000 {
|
||||
q.Limit = 5000
|
||||
}
|
||||
return q
|
||||
}
|
||||
|
||||
func whereClause(q stats.Query) (string, []any) {
|
||||
var conds []string
|
||||
var args []any
|
||||
if !q.Since.IsZero() {
|
||||
conds = append(conds, "ts >= ?")
|
||||
args = append(args, q.Since.Unix())
|
||||
}
|
||||
if !q.Until.IsZero() {
|
||||
conds = append(conds, "ts <= ?")
|
||||
args = append(args, q.Until.Unix())
|
||||
}
|
||||
if q.Model != "" {
|
||||
conds = append(conds, "model = ?")
|
||||
args = append(args, q.Model)
|
||||
}
|
||||
if q.Stream != nil {
|
||||
conds = append(conds, "stream = ?")
|
||||
args = append(args, boolToInt(*q.Stream))
|
||||
}
|
||||
if len(conds) == 0 {
|
||||
return "", args
|
||||
}
|
||||
out := ""
|
||||
for i, p := range conds {
|
||||
if i > 0 {
|
||||
out += " AND "
|
||||
}
|
||||
out += p
|
||||
}
|
||||
return " WHERE " + out, args
|
||||
}
|
||||
|
||||
func boolToInt(b bool) int {
|
||||
if b {
|
||||
return 1
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func (s *Store) scanTotals(sum *stats.Summary, where string, args []any) error {
|
||||
q := `SELECT COUNT(*),
|
||||
COALESCE(SUM(CASE WHEN status='success' THEN 1 ELSE 0 END),0),
|
||||
COALESCE(SUM(CASE WHEN status!='success' THEN 1 ELSE 0 END),0),
|
||||
COALESCE(SUM(prompt_tokens),0),
|
||||
COALESCE(SUM(completion_tokens),0),
|
||||
COALESCE(SUM(total_tokens),0),
|
||||
COALESCE(SUM(reasoning_tokens),0),
|
||||
COALESCE(SUM(cached_tokens),0) FROM requests` + where
|
||||
row := s.db.QueryRow(q, args...)
|
||||
var success, failures int64
|
||||
err := row.Scan(&sum.Totals.Requests, &success, &failures,
|
||||
&sum.Totals.PromptTokens, &sum.Totals.CompletionTokens,
|
||||
&sum.Totals.TotalTokens, &sum.Totals.ReasoningTokens, &sum.Totals.CachedTokens)
|
||||
if err != nil {
|
||||
return fmt.Errorf("scan totals: %w", err)
|
||||
}
|
||||
sum.Totals.SuccessRequests = int(success)
|
||||
sum.Totals.ErrorRequests = int(failures)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Store) scanPerModel(sum *stats.Summary, where string, args []any) error {
|
||||
q := `SELECT model, COUNT(*), COALESCE(SUM(prompt_tokens),0), COALESCE(SUM(completion_tokens),0), COALESCE(SUM(total_tokens),0), COALESCE(SUM(cached_tokens),0) FROM requests` + where + " GROUP BY model ORDER BY SUM(total_tokens) DESC"
|
||||
rows, err := s.db.Query(q, args...)
|
||||
if err != nil {
|
||||
return fmt.Errorf("scan per-model: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
for rows.Next() {
|
||||
var m stats.ModelStat
|
||||
if err := rows.Scan(&m.Model, &m.Requests, &m.PromptTokens, &m.CompletionTokens, &m.TotalTokens, &m.CachedTokens); err != nil {
|
||||
return err
|
||||
}
|
||||
sum.PerModel = append(sum.PerModel, m)
|
||||
}
|
||||
return rows.Err()
|
||||
}
|
||||
|
||||
func (s *Store) scanDaily(sum *stats.Summary, where string, args []any) error {
|
||||
q := `SELECT date(ts,'unixepoch','localtime') AS day, COUNT(*), COALESCE(SUM(prompt_tokens),0), COALESCE(SUM(completion_tokens),0), COALESCE(SUM(total_tokens),0), COALESCE(SUM(cached_tokens),0) FROM requests` + where + " GROUP BY day ORDER BY day ASC"
|
||||
rows, err := s.db.Query(q, args...)
|
||||
if err != nil {
|
||||
return fmt.Errorf("scan daily: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
for rows.Next() {
|
||||
var d stats.DayStat
|
||||
if err := rows.Scan(&d.Day, &d.Requests, &d.PromptTokens, &d.CompletionTokens, &d.TotalTokens, &d.CachedTokens); err != nil {
|
||||
return err
|
||||
}
|
||||
sum.Daily = append(sum.Daily, d)
|
||||
}
|
||||
return rows.Err()
|
||||
}
|
||||
|
||||
func (s *Store) scanRecent(sum *stats.Summary, q stats.Query, where string, args []any) error {
|
||||
limit := q.Limit
|
||||
if limit <= 0 {
|
||||
limit = defaultRecentLimit
|
||||
}
|
||||
query := `SELECT ts, model, stream, prompt_tokens, completion_tokens, total_tokens, reasoning_tokens, cached_tokens, status, latency_ms FROM requests` + where + " ORDER BY id DESC LIMIT ?"
|
||||
rows, err := s.db.Query(query, append(args, limit)...)
|
||||
if err != nil {
|
||||
return fmt.Errorf("scan recent: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
for rows.Next() {
|
||||
var rec stats.Record
|
||||
var ts int64
|
||||
var stream int
|
||||
if err := rows.Scan(&ts, &rec.Model, &stream, &rec.PromptTokens, &rec.CompletionTokens, &rec.TotalTokens, &rec.ReasoningTokens, &rec.CachedTokens, &rec.Status, &rec.LatencyMs); err != nil {
|
||||
return err
|
||||
}
|
||||
rec.Ts = time.Unix(ts, 0).Local()
|
||||
rec.Stream = stream == 1
|
||||
sum.Recent = append(sum.Recent, rec)
|
||||
}
|
||||
return rows.Err()
|
||||
}
|
||||
@@ -0,0 +1,89 @@
|
||||
// Package store owns the single embedded SQLite database backing the zhanlu
|
||||
// proxy: token-usage records (the requests table) and the persisted login
|
||||
// credentials (the credentials table, single-tenant single row). It implements
|
||||
// stats.Recorder for usage tracking and exposes Load/Save for credentials.
|
||||
package store
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
_ "modernc.org/sqlite"
|
||||
)
|
||||
|
||||
// Store is the single owner of the proxy's SQLite database handle.
|
||||
type Store struct {
|
||||
db *sql.DB
|
||||
}
|
||||
|
||||
// Open opens (or creates) the database at path and ensures both tables exist.
|
||||
// SQLite is opened with WAL journaling and a busy timeout so concurrent reads
|
||||
// (stats queries) and writes (request records, credential saves) do not
|
||||
// collide.
|
||||
func Open(path string) (*Store, error) {
|
||||
db, err := sql.Open("sqlite", path+"?_pragma=busy_timeout(5000)&_pragma=journal_mode(WAL)&_pragma=synchronous(NORMAL)&_pragma=foreign_keys(ON)")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("open db %q: %w", path, err)
|
||||
}
|
||||
if err := ensureSchema(db); err != nil {
|
||||
db.Close()
|
||||
return nil, err
|
||||
}
|
||||
return &Store{db: db}, nil
|
||||
}
|
||||
|
||||
// Close releases the database handle.
|
||||
func (s *Store) Close() error {
|
||||
return s.db.Close()
|
||||
}
|
||||
|
||||
func ensureSchema(db *sql.DB) error {
|
||||
stmts := []string{
|
||||
`CREATE TABLE IF NOT EXISTS requests (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
ts INTEGER NOT NULL,
|
||||
model TEXT NOT NULL,
|
||||
stream INTEGER NOT NULL DEFAULT 0,
|
||||
prompt_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
completion_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
total_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
reasoning_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
cached_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
status TEXT NOT NULL DEFAULT 'success',
|
||||
latency_ms INTEGER NOT NULL DEFAULT 0
|
||||
)`,
|
||||
`CREATE TABLE IF NOT EXISTS credentials (
|
||||
id INTEGER PRIMARY KEY CHECK (id = 1),
|
||||
access_key TEXT NOT NULL DEFAULT '',
|
||||
secret_key TEXT NOT NULL DEFAULT '',
|
||||
token TEXT NOT NULL DEFAULT '',
|
||||
api_key TEXT NOT NULL DEFAULT '',
|
||||
model_base_url TEXT NOT NULL DEFAULT '',
|
||||
email TEXT NOT NULL DEFAULT '',
|
||||
organization TEXT NOT NULL DEFAULT '',
|
||||
team TEXT NOT NULL DEFAULT '',
|
||||
base_url TEXT NOT NULL DEFAULT '',
|
||||
saved_at INTEGER NOT NULL DEFAULT 0
|
||||
)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_requests_ts ON requests(ts)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_requests_model ON requests(model)`,
|
||||
// Ensure the single credentials row exists so UPSERTs and SELECTs always
|
||||
// have a target.
|
||||
`INSERT INTO credentials (id) VALUES (1) ON CONFLICT(id) DO NOTHING`,
|
||||
}
|
||||
for _, q := range stmts {
|
||||
if _, err := db.Exec(q); err != nil {
|
||||
return fmt.Errorf("schema: %w", err)
|
||||
}
|
||||
}
|
||||
// Add cached_tokens to databases created before this column existed. SQLite
|
||||
// returns "duplicate column name" when it already exists; that is expected
|
||||
// and ignored.
|
||||
if _, err := db.Exec(`ALTER TABLE requests ADD COLUMN cached_tokens INTEGER NOT NULL DEFAULT 0`); err != nil {
|
||||
if !strings.Contains(err.Error(), "duplicate column") {
|
||||
return fmt.Errorf("migrate cached_tokens: %w", err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,120 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.misaka.ren/M1saka/zhanlu_proxy/internal/auth"
|
||||
"git.misaka.ren/M1saka/zhanlu_proxy/internal/stats"
|
||||
)
|
||||
|
||||
func newTestStore(t *testing.T) *Store {
|
||||
t.Helper()
|
||||
st, err := Open(filepath.Join(t.TempDir(), "test.db"))
|
||||
if err != nil {
|
||||
t.Fatalf("open store: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { st.Close() })
|
||||
return st
|
||||
}
|
||||
|
||||
func TestCredentialsRoundTrip(t *testing.T) {
|
||||
st := newTestStore(t)
|
||||
in := auth.Credentials{
|
||||
AccessKey: "AK", SecretKey: "SK", Token: "TOK",
|
||||
APIKey: "sk-1", ModelBaseURL: "https://up.example", Email: "[email protected]",
|
||||
Organization: "org", Team: "team",
|
||||
}
|
||||
if err := st.SaveCredentials(in); err != nil {
|
||||
t.Fatalf("save: %v", err)
|
||||
}
|
||||
got, err := st.LoadCredentials()
|
||||
if err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
if got.APIKey != "sk-1" || got.Email != "[email protected]" || got.Organization != "org" {
|
||||
t.Fatalf("round-trip mismatch: %+v", got)
|
||||
}
|
||||
if got.SavedAt.IsZero() {
|
||||
t.Fatalf("saved_at not set")
|
||||
}
|
||||
}
|
||||
|
||||
func TestEmptyLoadReturnsZero(t *testing.T) {
|
||||
st := newTestStore(t)
|
||||
got, err := st.LoadCredentials()
|
||||
if err != nil {
|
||||
t.Fatalf("load on empty store: %v", err)
|
||||
}
|
||||
if got.HasAPIKey() {
|
||||
t.Fatalf("expected no api key on empty store, got %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecordStatsReset(t *testing.T) {
|
||||
st := newTestStore(t)
|
||||
now := time.Now()
|
||||
recs := []stats.Record{
|
||||
{Ts: now, Model: "glm-4.7", Stream: false, PromptTokens: 10, CompletionTokens: 20, TotalTokens: 30, CachedTokens: 4, Status: "success", LatencyMs: 5},
|
||||
{Ts: now, Model: "glm-4.7", Stream: true, PromptTokens: 5, CompletionTokens: 5, TotalTokens: 10, Status: "success", LatencyMs: 8},
|
||||
{Ts: now, Model: "minimax", Stream: false, PromptTokens: 1, CompletionTokens: 1, TotalTokens: 2, Status: "upstream_error", LatencyMs: 3},
|
||||
}
|
||||
for _, r := range recs {
|
||||
if err := st.Record(r); err != nil {
|
||||
t.Fatalf("record: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
sum, err := st.Stats(stats.Query{})
|
||||
if err != nil {
|
||||
t.Fatalf("stats: %v", err)
|
||||
}
|
||||
if sum.Totals.Requests != 3 {
|
||||
t.Fatalf("requests = %d, want 3", sum.Totals.Requests)
|
||||
}
|
||||
if sum.Totals.SuccessRequests != 2 || sum.Totals.ErrorRequests != 1 {
|
||||
t.Fatalf("success/error = %d/%d, want 2/1", sum.Totals.SuccessRequests, sum.Totals.ErrorRequests)
|
||||
}
|
||||
if sum.Totals.TotalTokens != 42 {
|
||||
t.Fatalf("total tokens = %d, want 42", sum.Totals.TotalTokens)
|
||||
}
|
||||
if sum.Totals.CachedTokens != 4 {
|
||||
t.Fatalf("cached tokens = %d, want 4", sum.Totals.CachedTokens)
|
||||
}
|
||||
// prompt total = 10+5+1 = 16, cached = 4 => 0.25
|
||||
if sum.Totals.CacheRate < 0.24 || sum.Totals.CacheRate > 0.26 {
|
||||
t.Fatalf("cache rate = %v, want ~0.25", sum.Totals.CacheRate)
|
||||
}
|
||||
if len(sum.PerModel) != 2 {
|
||||
t.Fatalf("per-model entries = %d, want 2", len(sum.PerModel))
|
||||
}
|
||||
// glm-4.7 should lead on total tokens (40 vs 2)
|
||||
if sum.PerModel[0].Model != "glm-4.7" || sum.PerModel[0].TotalTokens != 40 || sum.PerModel[0].CachedTokens != 4 {
|
||||
t.Fatalf("top model = %+v, want glm-4.7/40/4 cached", sum.PerModel[0])
|
||||
}
|
||||
if len(sum.Recent) != 3 {
|
||||
t.Fatalf("recent entries = %d, want 3", len(sum.Recent))
|
||||
}
|
||||
// most recent first (id desc) => minimax record
|
||||
if sum.Recent[0].Model != "minimax" {
|
||||
t.Fatalf("most recent = %+v, want minimax", sum.Recent[0])
|
||||
}
|
||||
|
||||
// model filter
|
||||
sumF, _ := st.Stats(stats.Query{Model: "minimax"})
|
||||
if sumF.Totals.Requests != 1 || sumF.Totals.TotalTokens != 2 {
|
||||
t.Fatalf("filtered stats = %+v, want 1/2", sumF.Totals)
|
||||
}
|
||||
|
||||
if err := st.Reset(); err != nil {
|
||||
t.Fatalf("reset: %v", err)
|
||||
}
|
||||
sum2, _ := st.Stats(stats.Query{})
|
||||
if sum2.Totals.Requests != 0 {
|
||||
t.Fatalf("after reset requests = %d, want 0", sum2.Totals.Requests)
|
||||
}
|
||||
// credentials must survive a stats reset
|
||||
creds, _ := st.LoadCredentials()
|
||||
_ = creds
|
||||
}
|
||||
Reference in New Issue
Block a user