mirror of
https://github.com/CJackHwang/ds2api.git
synced 2026-05-04 00:15:28 +08:00
311 lines
9.1 KiB
Go
311 lines
9.1 KiB
Go
package admin
|
||
|
||
import (
|
||
"bufio"
|
||
"bytes"
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"io"
|
||
"net/http"
|
||
"strings"
|
||
"time"
|
||
|
||
"github.com/go-chi/chi/v5"
|
||
|
||
authn "ds2api/internal/auth"
|
||
"ds2api/internal/config"
|
||
"ds2api/internal/sse"
|
||
)
|
||
|
||
func (h *Handler) listAccounts(w http.ResponseWriter, r *http.Request) {
|
||
page := intFromQuery(r, "page", 1)
|
||
pageSize := intFromQuery(r, "page_size", 10)
|
||
if page < 1 {
|
||
page = 1
|
||
}
|
||
if pageSize < 1 {
|
||
pageSize = 1
|
||
}
|
||
if pageSize > 100 {
|
||
pageSize = 100
|
||
}
|
||
accounts := h.Store.Snapshot().Accounts
|
||
total := len(accounts)
|
||
reverseAccounts(accounts)
|
||
totalPages := 1
|
||
if total > 0 {
|
||
totalPages = (total + pageSize - 1) / pageSize
|
||
}
|
||
start := (page - 1) * pageSize
|
||
if start > total {
|
||
start = total
|
||
}
|
||
end := start + pageSize
|
||
if end > total {
|
||
end = total
|
||
}
|
||
items := make([]map[string]any, 0, end-start)
|
||
for _, acc := range accounts[start:end] {
|
||
token := strings.TrimSpace(acc.Token)
|
||
preview := ""
|
||
if token != "" {
|
||
if len(token) > 20 {
|
||
preview = token[:20] + "..."
|
||
} else {
|
||
preview = token
|
||
}
|
||
}
|
||
items = append(items, map[string]any{"email": acc.Email, "mobile": acc.Mobile, "has_password": acc.Password != "", "has_token": token != "", "token_preview": preview})
|
||
}
|
||
writeJSON(w, http.StatusOK, map[string]any{"items": items, "total": total, "page": page, "page_size": pageSize, "total_pages": totalPages})
|
||
}
|
||
|
||
func (h *Handler) addAccount(w http.ResponseWriter, r *http.Request) {
|
||
var req map[string]any
|
||
_ = json.NewDecoder(r.Body).Decode(&req)
|
||
acc := toAccount(req)
|
||
if acc.Identifier() == "" {
|
||
writeJSON(w, http.StatusBadRequest, map[string]any{"detail": "需要 email 或 mobile"})
|
||
return
|
||
}
|
||
err := h.Store.Update(func(c *config.Config) error {
|
||
for _, a := range c.Accounts {
|
||
if acc.Email != "" && a.Email == acc.Email {
|
||
return fmt.Errorf("邮箱已存在")
|
||
}
|
||
if acc.Mobile != "" && a.Mobile == acc.Mobile {
|
||
return fmt.Errorf("手机号已存在")
|
||
}
|
||
}
|
||
c.Accounts = append(c.Accounts, acc)
|
||
return nil
|
||
})
|
||
if err != nil {
|
||
writeJSON(w, http.StatusBadRequest, map[string]any{"detail": err.Error()})
|
||
return
|
||
}
|
||
h.Pool.Reset()
|
||
writeJSON(w, http.StatusOK, map[string]any{"success": true, "total_accounts": len(h.Store.Snapshot().Accounts)})
|
||
}
|
||
|
||
func (h *Handler) deleteAccount(w http.ResponseWriter, r *http.Request) {
|
||
identifier := chi.URLParam(r, "identifier")
|
||
err := h.Store.Update(func(c *config.Config) error {
|
||
idx := -1
|
||
for i, a := range c.Accounts {
|
||
if a.Email == identifier || a.Mobile == identifier {
|
||
idx = i
|
||
break
|
||
}
|
||
}
|
||
if idx < 0 {
|
||
return fmt.Errorf("账号不存在")
|
||
}
|
||
c.Accounts = append(c.Accounts[:idx], c.Accounts[idx+1:]...)
|
||
return nil
|
||
})
|
||
if err != nil {
|
||
writeJSON(w, http.StatusNotFound, map[string]any{"detail": err.Error()})
|
||
return
|
||
}
|
||
h.Pool.Reset()
|
||
writeJSON(w, http.StatusOK, map[string]any{"success": true, "total_accounts": len(h.Store.Snapshot().Accounts)})
|
||
}
|
||
|
||
func (h *Handler) queueStatus(w http.ResponseWriter, _ *http.Request) {
|
||
writeJSON(w, http.StatusOK, h.Pool.Status())
|
||
}
|
||
|
||
func (h *Handler) testSingleAccount(w http.ResponseWriter, r *http.Request) {
|
||
var req map[string]any
|
||
_ = json.NewDecoder(r.Body).Decode(&req)
|
||
identifier, _ := req["identifier"].(string)
|
||
if strings.TrimSpace(identifier) == "" {
|
||
writeJSON(w, http.StatusBadRequest, map[string]any{"detail": "需要账号标识（email 或 mobile）"})
|
||
return
|
||
}
|
||
acc, ok := h.Store.FindAccount(identifier)
|
||
if !ok {
|
||
writeJSON(w, http.StatusNotFound, map[string]any{"detail": "账号不存在"})
|
||
return
|
||
}
|
||
model, _ := req["model"].(string)
|
||
if model == "" {
|
||
model = "deepseek-chat"
|
||
}
|
||
message, _ := req["message"].(string)
|
||
result := h.testAccount(r.Context(), acc, model, message)
|
||
writeJSON(w, http.StatusOK, result)
|
||
}
|
||
|
||
func (h *Handler) testAllAccounts(w http.ResponseWriter, r *http.Request) {
|
||
var req map[string]any
|
||
_ = json.NewDecoder(r.Body).Decode(&req)
|
||
model, _ := req["model"].(string)
|
||
if model == "" {
|
||
model = "deepseek-chat"
|
||
}
|
||
accounts := h.Store.Snapshot().Accounts
|
||
if len(accounts) == 0 {
|
||
writeJSON(w, http.StatusOK, map[string]any{"total": 0, "success": 0, "failed": 0, "results": []any{}})
|
||
return
|
||
}
|
||
results := make([]map[string]any, 0, len(accounts))
|
||
success := 0
|
||
for _, acc := range accounts {
|
||
res := h.testAccount(r.Context(), acc, model, "")
|
||
if ok, _ := res["success"].(bool); ok {
|
||
success++
|
||
}
|
||
results = append(results, res)
|
||
time.Sleep(time.Second)
|
||
}
|
||
writeJSON(w, http.StatusOK, map[string]any{"total": len(accounts), "success": success, "failed": len(accounts) - success, "results": results})
|
||
}
|
||
|
||
func (h *Handler) testAccount(ctx context.Context, acc config.Account, model, message string) map[string]any {
|
||
start := time.Now()
|
||
result := map[string]any{"account": acc.Identifier(), "success": false, "response_time": 0, "message": "", "model": model}
|
||
token := strings.TrimSpace(acc.Token)
|
||
if token == "" {
|
||
newToken, err := h.DS.Login(ctx, acc)
|
||
if err != nil {
|
||
result["message"] = "登录失败: " + err.Error()
|
||
return result
|
||
}
|
||
token = newToken
|
||
_ = h.Store.UpdateAccountToken(acc.Identifier(), token)
|
||
}
|
||
authCtx := &authn.RequestAuth{UseConfigToken: false, DeepSeekToken: token}
|
||
sessionID, err := h.DS.CreateSession(ctx, authCtx, 1)
|
||
if err != nil {
|
||
newToken, loginErr := h.DS.Login(ctx, acc)
|
||
if loginErr != nil {
|
||
result["message"] = "创建会话失败: " + err.Error()
|
||
return result
|
||
}
|
||
token = newToken
|
||
authCtx.DeepSeekToken = token
|
||
_ = h.Store.UpdateAccountToken(acc.Identifier(), token)
|
||
sessionID, err = h.DS.CreateSession(ctx, authCtx, 1)
|
||
if err != nil {
|
||
result["message"] = "创建会话失败: " + err.Error()
|
||
return result
|
||
}
|
||
}
|
||
if strings.TrimSpace(message) == "" {
|
||
result["success"] = true
|
||
result["message"] = "API 测试成功（仅会话创建）"
|
||
result["response_time"] = int(time.Since(start).Milliseconds())
|
||
return result
|
||
}
|
||
thinking, search, ok := config.GetModelConfig(model)
|
||
if !ok {
|
||
thinking, search = false, false
|
||
}
|
||
pow, err := h.DS.GetPow(ctx, authCtx, 1)
|
||
if err != nil {
|
||
result["message"] = "获取 PoW 失败: " + err.Error()
|
||
return result
|
||
}
|
||
payload := map[string]any{"chat_session_id": sessionID, "prompt": "<｜User｜>" + message, "ref_file_ids": []any{}, "thinking_enabled": thinking, "search_enabled": search}
|
||
resp, err := h.DS.CallCompletion(ctx, authCtx, payload, pow, 1)
|
||
if err != nil {
|
||
result["message"] = "请求失败: " + err.Error()
|
||
return result
|
||
}
|
||
defer resp.Body.Close()
|
||
if resp.StatusCode != http.StatusOK {
|
||
result["message"] = fmt.Sprintf("请求失败: HTTP %d", resp.StatusCode)
|
||
return result
|
||
}
|
||
text := strings.Builder{}
|
||
think := strings.Builder{}
|
||
currentType := "text"
|
||
if thinking {
|
||
currentType = "thinking"
|
||
}
|
||
scanner := bufio.NewScanner(resp.Body)
|
||
buf := make([]byte, 0, 64*1024)
|
||
scanner.Buffer(buf, 2*1024*1024)
|
||
for scanner.Scan() {
|
||
chunk, done, parsed := sse.ParseDeepSeekSSELine(scanner.Bytes())
|
||
if !parsed {
|
||
continue
|
||
}
|
||
if done {
|
||
break
|
||
}
|
||
parts, finished, newType := sse.ParseSSEChunkForContent(chunk, thinking, currentType)
|
||
currentType = newType
|
||
if finished {
|
||
break
|
||
}
|
||
for _, p := range parts {
|
||
if p.Type == "thinking" {
|
||
think.WriteString(p.Text)
|
||
} else {
|
||
text.WriteString(p.Text)
|
||
}
|
||
}
|
||
}
|
||
result["success"] = true
|
||
result["response_time"] = int(time.Since(start).Milliseconds())
|
||
if text.Len() > 0 {
|
||
result["message"] = text.String()
|
||
} else {
|
||
result["message"] = "（无回复内容）"
|
||
}
|
||
if think.Len() > 0 {
|
||
result["thinking"] = think.String()
|
||
}
|
||
return result
|
||
}
|
||
|
||
func (h *Handler) testAPI(w http.ResponseWriter, r *http.Request) {
|
||
var req map[string]any
|
||
_ = json.NewDecoder(r.Body).Decode(&req)
|
||
model, _ := req["model"].(string)
|
||
message, _ := req["message"].(string)
|
||
apiKey, _ := req["api_key"].(string)
|
||
if model == "" {
|
||
model = "deepseek-chat"
|
||
}
|
||
if message == "" {
|
||
message = "你好"
|
||
}
|
||
if apiKey == "" {
|
||
keys := h.Store.Snapshot().Keys
|
||
if len(keys) == 0 {
|
||
writeJSON(w, http.StatusBadRequest, map[string]any{"detail": "没有可用的 API Key"})
|
||
return
|
||
}
|
||
apiKey = keys[0]
|
||
}
|
||
host := r.Host
|
||
scheme := "http"
|
||
if strings.Contains(strings.ToLower(host), "vercel") || strings.Contains(strings.ToLower(r.Header.Get("X-Forwarded-Proto")), "https") {
|
||
scheme = "https"
|
||
}
|
||
payload := map[string]any{"model": model, "messages": []map[string]any{{"role": "user", "content": message}}, "stream": false}
|
||
b, _ := json.Marshal(payload)
|
||
request, _ := http.NewRequestWithContext(r.Context(), http.MethodPost, fmt.Sprintf("%s://%s/v1/chat/completions", scheme, host), bytes.NewReader(b))
|
||
request.Header.Set("Authorization", "Bearer "+apiKey)
|
||
request.Header.Set("Content-Type", "application/json")
|
||
resp, err := (&http.Client{Timeout: 60 * time.Second}).Do(request)
|
||
if err != nil {
|
||
writeJSON(w, http.StatusOK, map[string]any{"success": false, "error": err.Error()})
|
||
return
|
||
}
|
||
defer resp.Body.Close()
|
||
body, _ := io.ReadAll(resp.Body)
|
||
if resp.StatusCode == http.StatusOK {
|
||
var parsed any
|
||
_ = json.Unmarshal(body, &parsed)
|
||
writeJSON(w, http.StatusOK, map[string]any{"success": true, "status_code": resp.StatusCode, "response": parsed})
|
||
return
|
||
}
|
||
writeJSON(w, http.StatusOK, map[string]any{"success": false, "status_code": resp.StatusCode, "response": string(body)})
|
||
}
|