kami
97 lines
2.2 KiB
Go
97 lines
2.2 KiB
Go
package herdr
|
|
|
|
import (
|
|
"bufio"
|
|
"encoding/json"
|
|
"os"
|
|
"strings"
|
|
)
|
|
|
|
type Usage struct{ Input, CacheRead, CacheWrite, Output int64 }
|
|
|
|
func (u Usage) Numerator() int64 { return u.Input + u.CacheRead + u.CacheWrite }
|
|
func Fraction(u Usage, w int64) float64 {
|
|
if w <= 0 {
|
|
return 0
|
|
}
|
|
f := float64(u.Numerator()) / float64(w)
|
|
if f < 0 {
|
|
return 0
|
|
}
|
|
if f > 1 {
|
|
return 1
|
|
}
|
|
return f
|
|
}
|
|
func ClaudeUsage(p string) (Usage, error) {
|
|
f, e := os.Open(p)
|
|
if e != nil {
|
|
return Usage{}, e
|
|
}
|
|
defer f.Close()
|
|
s := bufio.NewScanner(f)
|
|
var last Usage
|
|
for s.Scan() {
|
|
var x struct {
|
|
Message struct {
|
|
Usage struct {
|
|
Input int64 `json:"input_tokens"`
|
|
Read int64 `json:"cache_read_input_tokens"`
|
|
Write int64 `json:"cache_creation_input_tokens"`
|
|
Output int64 `json:"output_tokens"`
|
|
} `json:"usage"`
|
|
} `json:"message"`
|
|
}
|
|
if json.Unmarshal(s.Bytes(), &x) == nil && x.Message.Usage.Input > 0 {
|
|
last = Usage{x.Message.Usage.Input, x.Message.Usage.Read, x.Message.Usage.Write, x.Message.Usage.Output}
|
|
}
|
|
}
|
|
return last, s.Err()
|
|
}
|
|
func CodexUsage(p string) (Usage, error) {
|
|
f, e := os.Open(p)
|
|
if e != nil {
|
|
return Usage{}, e
|
|
}
|
|
defer f.Close()
|
|
s := bufio.NewScanner(f)
|
|
var u Usage
|
|
for s.Scan() {
|
|
var x struct {
|
|
Payload struct {
|
|
Type string `json:"type"`
|
|
Info struct {
|
|
Last struct {
|
|
Input int64 `json:"input"`
|
|
Read int64 `json:"cached_input"`
|
|
} `json:"last_token_usage"`
|
|
} `json:"info"`
|
|
} `json:"payload"`
|
|
}
|
|
if json.Unmarshal(s.Bytes(), &x) == nil && x.Payload.Type == "token_count" {
|
|
u = Usage{x.Payload.Info.Last.Input, x.Payload.Info.Last.Read, 0, 0}
|
|
}
|
|
}
|
|
return u, s.Err()
|
|
}
|
|
func OpenCodeUsage(p string) (Usage, error) {
|
|
f, e := os.Open(p)
|
|
if e != nil {
|
|
return Usage{}, e
|
|
}
|
|
defer f.Close()
|
|
var x struct {
|
|
Tokens struct {
|
|
Input int64 `json:"input"`
|
|
Output int64 `json:"output"`
|
|
Cache struct {
|
|
Read int64 `json:"read"`
|
|
Write int64 `json:"write"`
|
|
} `json:"cache"`
|
|
} `json:"tokens"`
|
|
}
|
|
e = json.NewDecoder(f).Decode(&x)
|
|
return Usage{x.Tokens.Input, x.Tokens.Cache.Read, x.Tokens.Cache.Write, x.Tokens.Output}, e
|
|
}
|
|
func IsBusy(s string) bool { return strings.EqualFold(s, "busy") }
|