cb3d35756a
All checks were successful
Beta Release / beta (push) Successful in 1m3s
- Block sudo/doas commands when not running as root - Add real token counting from API responses - Track and display consumption by provider/day - Add Mermaid diagram rendering in Shell and Studio - Add copy-to-clipboard buttons for code blocks - Support tables in AI message rendering - Update system prompt with context (date, time, root status) 💘 Generated with Crush Assisted-by: MiniMax-M2.7 via Crush <crush@charm.land>
128 lines
2.5 KiB
Go
128 lines
2.5 KiB
Go
package api
|
|
|
|
import (
|
|
"encoding/json"
|
|
"os"
|
|
"path/filepath"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/muyue/muyue/internal/config"
|
|
)
|
|
|
|
type consumptionEntry struct {
|
|
Date string `json:"date"`
|
|
Tokens int `json:"tokens"`
|
|
Requests int `json:"requests"`
|
|
}
|
|
|
|
type providerConsumption struct {
|
|
Name string `json:"name"`
|
|
Daily []consumptionEntry `json:"daily"`
|
|
Total int `json:"total_tokens"`
|
|
Requests int `json:"total_requests"`
|
|
}
|
|
|
|
type consumptionStore struct {
|
|
mu sync.Mutex
|
|
providers map[string]*providerConsumption
|
|
}
|
|
|
|
func newConsumptionStore() *consumptionStore {
|
|
cs := &consumptionStore{
|
|
providers: make(map[string]*providerConsumption),
|
|
}
|
|
cs.load()
|
|
cs.prune()
|
|
return cs
|
|
}
|
|
|
|
func (cs *consumptionStore) Record(providerName string, tokens int) {
|
|
if tokens <= 0 || providerName == "" {
|
|
return
|
|
}
|
|
cs.mu.Lock()
|
|
defer cs.mu.Unlock()
|
|
|
|
today := time.Now().UTC().Format("2006-01-02")
|
|
|
|
p, ok := cs.providers[providerName]
|
|
if !ok {
|
|
p = &providerConsumption{Name: providerName}
|
|
cs.providers[providerName] = p
|
|
}
|
|
|
|
p.Total += tokens
|
|
p.Requests++
|
|
|
|
if len(p.Daily) > 0 && p.Daily[len(p.Daily)-1].Date == today {
|
|
p.Daily[len(p.Daily)-1].Tokens += tokens
|
|
p.Daily[len(p.Daily)-1].Requests++
|
|
} else {
|
|
p.Daily = append(p.Daily, consumptionEntry{
|
|
Date: today,
|
|
Tokens: tokens,
|
|
Requests: 1,
|
|
})
|
|
}
|
|
|
|
cs.save()
|
|
}
|
|
|
|
func (cs *consumptionStore) GetAll() map[string]*providerConsumption {
|
|
cs.mu.Lock()
|
|
defer cs.mu.Unlock()
|
|
|
|
result := make(map[string]*providerConsumption)
|
|
for k, v := range cs.providers {
|
|
pc := *v
|
|
daily := make([]consumptionEntry, len(v.Daily))
|
|
copy(daily, v.Daily)
|
|
pc.Daily = daily
|
|
result[k] = &pc
|
|
}
|
|
return result
|
|
}
|
|
|
|
func (cs *consumptionStore) prune() {
|
|
cutoff := time.Now().UTC().AddDate(0, 0, -7).Format("2006-01-02")
|
|
for _, p := range cs.providers {
|
|
filtered := make([]consumptionEntry, 0)
|
|
for _, d := range p.Daily {
|
|
if d.Date >= cutoff {
|
|
filtered = append(filtered, d)
|
|
}
|
|
}
|
|
p.Daily = filtered
|
|
}
|
|
}
|
|
|
|
func (cs *consumptionStore) filePath() string {
|
|
dir, err := config.ConfigDir()
|
|
if err != nil {
|
|
return ""
|
|
}
|
|
return filepath.Join(dir, "consumption.json")
|
|
}
|
|
|
|
func (cs *consumptionStore) load() {
|
|
fp := cs.filePath()
|
|
if fp == "" {
|
|
return
|
|
}
|
|
data, err := os.ReadFile(fp)
|
|
if err != nil {
|
|
return
|
|
}
|
|
json.Unmarshal(data, &cs.providers)
|
|
}
|
|
|
|
func (cs *consumptionStore) save() {
|
|
fp := cs.filePath()
|
|
if fp == "" {
|
|
return
|
|
}
|
|
data, _ := json.Marshal(cs.providers)
|
|
os.WriteFile(fp, data, 0644)
|
|
}
|