240 lines
7.3 KiB
Go
240 lines
7.3 KiB
Go
package gui
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"net/url"
|
|
"os"
|
|
"path/filepath"
|
|
"strconv"
|
|
"time"
|
|
|
|
"github.com/yetone/magpie/internal/agent"
|
|
"github.com/yetone/magpie/internal/edit"
|
|
"github.com/yetone/magpie/internal/provider"
|
|
"github.com/yetone/magpie/internal/usage"
|
|
)
|
|
|
|
// usageGroup is a usage.Group with what the UI needs to draw it.
|
|
type usageGroup struct {
|
|
usage.Group
|
|
Name string `json:"name"`
|
|
Sub string `json:"sub,omitempty"` // models: the provider's name
|
|
Icon string `json:"icon,omitempty"` // a real logo, or "generic" for an unknown client
|
|
}
|
|
|
|
type usageJSON struct {
|
|
usage.Summary
|
|
Agents []usageGroup `json:"agents"`
|
|
Models []usageGroup `json:"models"`
|
|
Path string `json:"path"`
|
|
}
|
|
|
|
func usageState(p usage.Period) usageJSON {
|
|
s := usage.Summarize(p)
|
|
out := usageJSON{Summary: s, Agents: []usageGroup{}, Models: []usageGroup{}, Path: tilde(usage.Path())}
|
|
agents := map[string]*agent.Agent{}
|
|
for _, a := range agent.Clients() {
|
|
agents[a.ID] = a
|
|
}
|
|
for _, g := range s.Agents {
|
|
ug := usageGroup{Group: g, Name: g.ID, Icon: "generic"}
|
|
if a := agents[g.ID]; a != nil {
|
|
ug.Name, ug.Icon = a.Name, a.Icon
|
|
}
|
|
out.Agents = append(out.Agents, ug)
|
|
}
|
|
providers := map[string]provider.Provider{}
|
|
for _, p := range provider.All() {
|
|
providers[p.ID] = p
|
|
}
|
|
for _, g := range s.Models {
|
|
ug := usageGroup{Group: g, Name: g.Model, Sub: g.Provider, Icon: "generic"}
|
|
if p, ok := providers[g.Provider]; ok {
|
|
ug.Sub = p.Name
|
|
if p.Icon != "" {
|
|
ug.Icon = p.Icon
|
|
}
|
|
}
|
|
if g.Host != "" {
|
|
ug.Sub += " · " + g.Host
|
|
}
|
|
out.Models = append(out.Models, ug)
|
|
}
|
|
return out
|
|
}
|
|
|
|
func periodOf(s string) usage.Period {
|
|
switch p := usage.Period(s); p {
|
|
case usage.Today, usage.Week, usage.Month, usage.All:
|
|
return p
|
|
}
|
|
return usage.Month
|
|
}
|
|
|
|
func ledgerFilter(q url.Values) usage.Filter {
|
|
return usage.Filter{Agent: q.Get("agent"), Failed: q.Get("failed") == "1", Query: q.Get("q")}
|
|
}
|
|
|
|
// ledgerRow is a usage.Row with the names the page shows it by.
|
|
type ledgerRow struct {
|
|
usage.Row
|
|
AgentName string `json:"agentName"`
|
|
Icon string `json:"icon"` // the agent's
|
|
ProviderName string `json:"providerName"`
|
|
}
|
|
|
|
type ledgerJSON struct {
|
|
Period usage.Period `json:"period"`
|
|
Rows []ledgerRow `json:"rows"`
|
|
Offset int `json:"offset"`
|
|
Total int `json:"total"` // the rows the filter keeps, on every page
|
|
usage.Totals
|
|
// Agents: the agents with calls in the period, for the filter
|
|
Agents []ledgerAgent `json:"agents"`
|
|
}
|
|
|
|
type ledgerAgent struct {
|
|
ID string `json:"id"`
|
|
Name string `json:"name"`
|
|
Icon string `json:"icon"`
|
|
}
|
|
|
|
// ledgerPage is one page of the ledger: limit rows (100 when none is
|
|
// given, 500 at most) from offset.
|
|
func ledgerPage(p usage.Period, f usage.Filter, offset, limit int) ledgerJSON {
|
|
rows, sum, ids := usage.Ledger(p, f)
|
|
if limit <= 0 {
|
|
limit = 100
|
|
}
|
|
limit = min(limit, 500)
|
|
offset = max(0, min(offset, len(rows)))
|
|
page := rows[offset:min(len(rows), offset+limit)]
|
|
agents := map[string]*agent.Agent{}
|
|
for _, a := range agent.Clients() {
|
|
agents[a.ID] = a
|
|
}
|
|
names := map[string]string{}
|
|
for _, pr := range provider.All() {
|
|
names[pr.ID] = pr.Name
|
|
}
|
|
who := func(id string) ledgerAgent {
|
|
if a := agents[id]; a != nil {
|
|
return ledgerAgent{ID: id, Name: a.Name, Icon: a.Icon}
|
|
}
|
|
return ledgerAgent{ID: id, Name: id, Icon: "generic"}
|
|
}
|
|
out := ledgerJSON{Period: p, Rows: make([]ledgerRow, 0, len(page)), Offset: offset, Total: len(rows), Totals: sum, Agents: []ledgerAgent{}}
|
|
for _, r := range page {
|
|
a := who(r.Agent)
|
|
lr := ledgerRow{Row: r, AgentName: a.Name, Icon: a.Icon, ProviderName: names[r.Provider]}
|
|
if lr.ProviderName == "" {
|
|
lr.ProviderName = r.Provider
|
|
}
|
|
out.Rows = append(out.Rows, lr)
|
|
}
|
|
for _, id := range ids {
|
|
out.Agents = append(out.Agents, who(id))
|
|
}
|
|
return out
|
|
}
|
|
|
|
func usageRoutes(mux *http.ServeMux, w Windows) {
|
|
mux.HandleFunc("GET /api/usage", func(rw http.ResponseWriter, r *http.Request) {
|
|
p := usage.Period(r.URL.Query().Get("period"))
|
|
switch p {
|
|
case usage.Today, usage.Week, usage.Month, usage.All:
|
|
default:
|
|
p = usage.Month
|
|
}
|
|
writeJSON(rw, usageState(p))
|
|
})
|
|
// the ledger: the period's calls, newest first, a page at a time
|
|
mux.HandleFunc("GET /api/usage/requests", func(rw http.ResponseWriter, r *http.Request) {
|
|
q := r.URL.Query()
|
|
offset, _ := strconv.Atoi(q.Get("offset"))
|
|
limit, _ := strconv.Atoi(q.Get("limit"))
|
|
writeJSON(rw, ledgerPage(periodOf(q.Get("period")), ledgerFilter(q), offset, limit))
|
|
})
|
|
// the same CSV to a browser (magpie web), which saves it itself
|
|
mux.HandleFunc("GET /api/usage/requests.csv", func(rw http.ResponseWriter, r *http.Request) {
|
|
q := r.URL.Query()
|
|
p := periodOf(q.Get("period"))
|
|
rows, _, _ := usage.Ledger(p, ledgerFilter(q))
|
|
rw.Header().Set("Content-Type", "text/csv; charset=utf-8")
|
|
rw.Header().Set("Content-Disposition", `attachment; filename="magpie-requests-`+string(p)+"-"+time.Now().Format("2006-01-02")+`.csv"`)
|
|
usage.WriteCSV(rw, rows)
|
|
})
|
|
// the rows the ledger shows, all its pages, as a CSV in Downloads
|
|
mux.HandleFunc("POST /api/usage/requests/export", func(rw http.ResponseWriter, r *http.Request) {
|
|
q := r.URL.Query()
|
|
p := periodOf(q.Get("period"))
|
|
rows, _, _ := usage.Ledger(p, ledgerFilter(q))
|
|
var b bytes.Buffer
|
|
if err := usage.WriteCSV(&b, rows); err != nil {
|
|
fail(rw, err)
|
|
return
|
|
}
|
|
dir := downloads()
|
|
stamp := "magpie-requests-" + string(p) + "-" + time.Now().Format("2006-01-02")
|
|
name := filepath.Join(dir, stamp+".csv")
|
|
for i := 2; ; i++ { // never over an earlier one
|
|
if _, err := os.Stat(name); err != nil {
|
|
break
|
|
}
|
|
name = filepath.Join(dir, fmt.Sprintf("%s-%d.csv", stamp, i))
|
|
}
|
|
if err := edit.WriteAtomic(name, b.Bytes()); err != nil {
|
|
fail(rw, err)
|
|
return
|
|
}
|
|
_ = w.OpenFolder(dir) // saved either way; the path is in the answer
|
|
writeJSON(rw, map[string]any{"path": tilde(name), "rows": len(rows)})
|
|
})
|
|
// The subscriptions' quotas, the plans' bought with a key, and the
|
|
// keys' balances come from the
|
|
// vendors, which can be slow or unreachable, so the page asks for them
|
|
// apart from the local log.
|
|
mux.HandleFunc("GET /api/usage/quotas", func(rw http.ResponseWriter, r *http.Request) {
|
|
ctx, cancel := context.WithTimeout(r.Context(), 12*time.Second)
|
|
defer cancel()
|
|
writeJSON(rw, provider.Quotas(ctx))
|
|
})
|
|
// spends one of a Codex account's rate-limit resets, which the page
|
|
// has asked the user about first; what it did comes back
|
|
mux.HandleFunc("POST /api/usage/codex-reset", func(rw http.ResponseWriter, r *http.Request) {
|
|
var in struct{ User string }
|
|
if err := json.NewDecoder(r.Body).Decode(&in); err != nil {
|
|
fail(rw, err)
|
|
return
|
|
}
|
|
ctx, cancel := context.WithTimeout(r.Context(), 20*time.Second)
|
|
defer cancel()
|
|
out, err := provider.UseCodexReset(ctx, in.User)
|
|
if err != nil {
|
|
fail(rw, err)
|
|
return
|
|
}
|
|
writeJSON(rw, out)
|
|
})
|
|
// and one of a Claude account's usage-limit resets, the same way
|
|
mux.HandleFunc("POST /api/usage/claude-reset", func(rw http.ResponseWriter, r *http.Request) {
|
|
var in struct{ User string }
|
|
if err := json.NewDecoder(r.Body).Decode(&in); err != nil {
|
|
fail(rw, err)
|
|
return
|
|
}
|
|
ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second)
|
|
defer cancel()
|
|
out, err := provider.UseClaudeReset(ctx, in.User)
|
|
if err != nil {
|
|
fail(rw, err)
|
|
return
|
|
}
|
|
writeJSON(rw, out)
|
|
})
|
|
}
|