307 lines
9.2 KiB
Go
307 lines
9.2 KiB
Go
package server
|
|
|
|
import (
|
|
"ai-router/internal/protocol"
|
|
"context"
|
|
"database/sql"
|
|
"database/sql/driver"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"sort"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
_ "time/tzdata"
|
|
|
|
"modernc.org/sqlite"
|
|
)
|
|
|
|
func init() {
|
|
var locationMu sync.Mutex
|
|
var cachedZone string
|
|
var cachedLocation *time.Location
|
|
if err := sqlite.RegisterDeterministicScalarFunction("router_log_day", 2, func(_ *sqlite.FunctionContext, args []driver.Value) (driver.Value, error) {
|
|
stampText, ok := args[0].(string)
|
|
zone, zoneOK := args[1].(string)
|
|
if !ok || !zoneOK {
|
|
return nil, fmt.Errorf("invalid log timestamp or timezone")
|
|
}
|
|
stamp, err := time.Parse(time.RFC3339Nano, stampText)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
locationMu.Lock()
|
|
defer locationMu.Unlock()
|
|
if cachedLocation == nil || cachedZone != zone {
|
|
loc, err := time.LoadLocation(zone)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
cachedZone, cachedLocation = zone, loc
|
|
}
|
|
return stamp.In(cachedLocation).Format("2006-01-02"), nil
|
|
}); err != nil {
|
|
panic(err)
|
|
}
|
|
}
|
|
|
|
type requestLog struct {
|
|
ID int64 `json:"id"`
|
|
CreatedAt string `json:"created_at"`
|
|
KeyID string `json:"key_id"`
|
|
KeyName string `json:"key_name"`
|
|
Channel string `json:"channel"`
|
|
RequestModel string `json:"request_model"`
|
|
ActualModel string `json:"actual_model"`
|
|
Downstream string `json:"downstream"`
|
|
Upstream string `json:"upstream"`
|
|
Status int `json:"status"`
|
|
DurationMS int64 `json:"duration_ms"`
|
|
Stream bool `json:"stream"`
|
|
Input *int64 `json:"input_tokens"`
|
|
Output *int64 `json:"output_tokens"`
|
|
CacheRead *int64 `json:"cache_read_tokens"`
|
|
CacheWrite *int64 `json:"cache_write_tokens"`
|
|
UsageKnown bool `json:"usage_known"`
|
|
Error string `json:"error"`
|
|
RawUsage json.RawMessage `json:"raw_usage"`
|
|
Usage protocol.Usage `json:"-"`
|
|
}
|
|
|
|
func insertLog(ctx context.Context, tx *sql.Tx, l requestLog) error {
|
|
l = boundedLog(l)
|
|
var in, out, cr, cw any
|
|
if l.Usage.Known || len(l.Usage.Raw) > 0 {
|
|
in = l.Usage.Input
|
|
out = l.Usage.Output
|
|
cr = l.Usage.CacheRead
|
|
cw = l.Usage.CacheWrite
|
|
}
|
|
_, err := tx.ExecContext(ctx, "INSERT INTO logs(created_at,key_id,key_name,channel,request_model,actual_model,downstream,upstream,status,duration_ms,stream,input_tokens,output_tokens,cache_read_tokens,cache_write_tokens,usage_known,error,raw_usage) VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)", l.CreatedAt, l.KeyID, l.KeyName, l.Channel, l.RequestModel, l.ActualModel, l.Downstream, l.Upstream, l.Status, l.DurationMS, l.Stream, in, out, cr, cw, l.Usage.Known, l.Error, string(l.Usage.Raw))
|
|
return err
|
|
}
|
|
func filter(r *http.Request) (string, []any, error) {
|
|
parts := []string{"1=1"}
|
|
args := []any{}
|
|
for _, name := range []string{"channel", "model"} {
|
|
if v := r.URL.Query().Get(name); v != "" {
|
|
col := name
|
|
if name == "model" {
|
|
col = "actual_model"
|
|
}
|
|
parts = append(parts, col+"=?")
|
|
args = append(args, v)
|
|
}
|
|
}
|
|
if k := r.URL.Query().Get("key_id"); k != "" {
|
|
parts = append(parts, "key_id=?")
|
|
args = append(args, k)
|
|
}
|
|
var from, to time.Time
|
|
for _, n := range []string{"from", "to"} {
|
|
if v := r.URL.Query().Get(n); v != "" {
|
|
t, e := time.Parse(time.RFC3339Nano, v)
|
|
if e != nil {
|
|
return "", nil, e
|
|
}
|
|
if n == "from" {
|
|
from = t
|
|
parts = append(parts, "julianday(created_at)>=julianday(?)")
|
|
} else {
|
|
to = t
|
|
parts = append(parts, "julianday(created_at)<=julianday(?)")
|
|
}
|
|
args = append(args, t.UTC().Format(time.RFC3339Nano))
|
|
}
|
|
}
|
|
if !from.IsZero() && !to.IsZero() && from.After(to) {
|
|
return "", nil, sql.ErrNoRows
|
|
}
|
|
return strings.Join(parts, " AND "), args, nil
|
|
}
|
|
func (s *Server) readLogs(r *http.Request, limited bool) ([]requestLog, error) {
|
|
where, args, e := filter(r)
|
|
if e != nil {
|
|
return nil, e
|
|
}
|
|
q := "SELECT id,created_at,key_id,key_name,channel,request_model,actual_model,downstream,upstream,status,duration_ms,stream,input_tokens,output_tokens,cache_read_tokens,cache_write_tokens,usage_known,error,raw_usage FROM logs WHERE " + where + " ORDER BY id DESC"
|
|
if limited {
|
|
q += " LIMIT 200"
|
|
}
|
|
rows, e := s.db.Query(q, args...)
|
|
if e != nil {
|
|
return nil, e
|
|
}
|
|
defer rows.Close()
|
|
out := []requestLog{}
|
|
for rows.Next() {
|
|
var l requestLog
|
|
var raw string
|
|
if e = rows.Scan(&l.ID, &l.CreatedAt, &l.KeyID, &l.KeyName, &l.Channel, &l.RequestModel, &l.ActualModel, &l.Downstream, &l.Upstream, &l.Status, &l.DurationMS, &l.Stream, &l.Input, &l.Output, &l.CacheRead, &l.CacheWrite, &l.UsageKnown, &l.Error, &raw); e != nil {
|
|
return nil, e
|
|
}
|
|
if json.Valid([]byte(raw)) {
|
|
l.RawUsage = json.RawMessage(raw)
|
|
}
|
|
out = append(out, l)
|
|
}
|
|
return out, rows.Err()
|
|
}
|
|
func (s *Server) logs(w http.ResponseWriter, r *http.Request) {
|
|
if _, _, e := filter(r); e != nil {
|
|
errJSON(w, 400, "invalid time filter")
|
|
return
|
|
}
|
|
ls, e := s.readLogs(r, true)
|
|
if e != nil {
|
|
errJSON(w, 500, "database error")
|
|
return
|
|
}
|
|
writeJSON(w, ls)
|
|
}
|
|
|
|
type summary struct {
|
|
Requests int64 `json:"requests"`
|
|
Success int64 `json:"success"`
|
|
Failed int64 `json:"failed"`
|
|
Input *int64 `json:"input_tokens"`
|
|
Output *int64 `json:"output_tokens"`
|
|
CacheRead *int64 `json:"cache_read_tokens"`
|
|
CacheWrite *int64 `json:"cache_write_tokens"`
|
|
Unknown int64 `json:"unknown_usage"`
|
|
}
|
|
|
|
func (a *summary) add(l requestLog) {
|
|
a.Requests++
|
|
if l.Status >= 200 && l.Status < 300 && l.Error == "" {
|
|
a.Success++
|
|
} else {
|
|
a.Failed++
|
|
}
|
|
if !l.UsageKnown {
|
|
a.Unknown++
|
|
return
|
|
}
|
|
if a.Input == nil {
|
|
a.Input = new(int64)
|
|
a.Output = new(int64)
|
|
a.CacheRead = new(int64)
|
|
a.CacheWrite = new(int64)
|
|
}
|
|
if l.Input != nil {
|
|
*a.Input += *l.Input
|
|
}
|
|
if l.Output != nil {
|
|
*a.Output += *l.Output
|
|
}
|
|
if l.CacheRead != nil {
|
|
*a.CacheRead += *l.CacheRead
|
|
}
|
|
if l.CacheWrite != nil {
|
|
*a.CacheWrite += *l.CacheWrite
|
|
}
|
|
}
|
|
func (s *Server) stats(w http.ResponseWriter, r *http.Request) {
|
|
zone := r.URL.Query().Get("timezone")
|
|
if zone == "" {
|
|
zone = "UTC"
|
|
}
|
|
_, err := time.LoadLocation(zone)
|
|
if err != nil || zone == "Local" {
|
|
errJSON(w, 400, "invalid timezone")
|
|
return
|
|
}
|
|
where, args, e := filter(r)
|
|
if e != nil {
|
|
errJSON(w, 400, "invalid time filter")
|
|
return
|
|
}
|
|
const aggregates = `COUNT(*), COALESCE(SUM(status>=200 AND status<300 AND error=''),0),
|
|
COALESCE(SUM(NOT(status>=200 AND status<300 AND error='')),0),
|
|
SUM(CASE WHEN usage_known THEN COALESCE(input_tokens,0) END),
|
|
SUM(CASE WHEN usage_known THEN COALESCE(output_tokens,0) END),
|
|
SUM(CASE WHEN usage_known THEN COALESCE(cache_read_tokens,0) END),
|
|
SUM(CASE WHEN usage_known THEN COALESCE(cache_write_tokens,0) END),
|
|
COALESCE(SUM(NOT usage_known),0), MAX(id)`
|
|
// NOT MATERIALIZED avoids copying the full log payload into a temporary table.
|
|
query := `WITH filtered AS NOT MATERIALIZED (SELECT * FROM logs WHERE ` + where + `)
|
|
SELECT 'total','','',` + aggregates + ` FROM filtered
|
|
UNION ALL SELECT 'key',key_id,key_name,` + aggregates + ` FROM filtered GROUP BY key_id
|
|
UNION ALL SELECT 'day',router_log_day(created_at,?), '',` + aggregates + ` FROM filtered GROUP BY router_log_day(created_at,?)
|
|
UNION ALL SELECT 'model',channel,actual_model,` + aggregates + ` FROM filtered GROUP BY channel,actual_model`
|
|
args = append(args, zone, zone)
|
|
ctx, cancel := context.WithTimeout(r.Context(), 10*time.Second)
|
|
defer cancel()
|
|
rows, e := s.db.QueryContext(ctx, query, args...)
|
|
if e != nil {
|
|
errJSON(w, 500, "database error")
|
|
return
|
|
}
|
|
defer rows.Close()
|
|
type byKey struct {
|
|
KeyID string `json:"key_id"`
|
|
Name string `json:"name"`
|
|
summary
|
|
}
|
|
type daily struct {
|
|
Day string `json:"day"`
|
|
summary
|
|
}
|
|
type byModel struct {
|
|
Channel string `json:"channel"`
|
|
Model string `json:"model"`
|
|
summary
|
|
}
|
|
total := summary{}
|
|
ks := map[string]*byKey{}
|
|
ds := map[string]*daily{}
|
|
ms := map[[2]string]*byModel{}
|
|
groups := 0
|
|
for rows.Next() {
|
|
groups++
|
|
if groups > 10000 {
|
|
errJSON(w, 422, "too many statistics groups; narrow the time range")
|
|
return
|
|
}
|
|
var kind, a, b string
|
|
var v summary
|
|
var latest sql.NullInt64
|
|
if err := rows.Scan(&kind, &a, &b, &v.Requests, &v.Success, &v.Failed, &v.Input, &v.Output, &v.CacheRead, &v.CacheWrite, &v.Unknown, &latest); err != nil {
|
|
errJSON(w, 500, "database error")
|
|
return
|
|
}
|
|
switch kind {
|
|
case "total":
|
|
total = v
|
|
case "key":
|
|
// SQLite's sole MAX selects key_name from the latest ID, not timestamp.
|
|
ks[a] = &byKey{KeyID: a, Name: b, summary: v}
|
|
case "day":
|
|
ds[a] = &daily{Day: a, summary: v}
|
|
case "model":
|
|
ms[[2]string{a, b}] = &byModel{Channel: a, Model: b, summary: v}
|
|
}
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
errJSON(w, 500, "database error")
|
|
return
|
|
}
|
|
keys := []*byKey{}
|
|
days := []*daily{}
|
|
mods := []*byModel{}
|
|
for _, v := range ks {
|
|
keys = append(keys, v)
|
|
}
|
|
for _, v := range ds {
|
|
days = append(days, v)
|
|
}
|
|
for _, v := range ms {
|
|
mods = append(mods, v)
|
|
}
|
|
sort.Slice(keys, func(i, j int) bool { return keys[i].KeyID < keys[j].KeyID })
|
|
sort.Slice(days, func(i, j int) bool { return days[i].Day < days[j].Day })
|
|
sort.Slice(mods, func(i, j int) bool { return mods[i].Channel+mods[i].Model < mods[j].Channel+mods[j].Model })
|
|
writeJSON(w, map[string]any{"summary": total, "by_key": keys, "daily": days, "by_model": mods, "timezone": zone})
|
|
}
|