usage.go

package main

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
	"time"

	_ "github.com/mattn/go-sqlite3"
)

// UsageTracker records per-request token usage in a SQLite database.
type UsageTracker struct {
	db *sql.DB
}

// NewUsageTracker opens (or creates) the SQLite database at dbPath and ensures
// the schema exists.
func NewUsageTracker(dbPath string) (*UsageTracker, error) {
	db, err := sql.Open("sqlite3", dbPath)
	if err != nil {
		return nil, fmt.Errorf("opening usage database: %w", err)
	}
	if err := db.Ping(); err != nil {
		db.Close()
		return nil, fmt.Errorf("pinging usage database: %w", err)
	}
	schema := `
CREATE TABLE IF NOT EXISTS usage (
    id               INTEGER PRIMARY KEY AUTOINCREMENT,
    created_at       DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
    model            TEXT     NOT NULL,
    prompt_tokens    INTEGER  NOT NULL DEFAULT 0,
    completion_tokens INTEGER NOT NULL DEFAULT 0,
    total_tokens     INTEGER  NOT NULL DEFAULT 0
);
CREATE INDEX IF NOT EXISTS idx_usage_created_at ON usage (created_at);
CREATE INDEX IF NOT EXISTS idx_usage_model ON usage (model);
`
	if _, err := db.Exec(schema); err != nil {
		db.Close()
		return nil, fmt.Errorf("creating usage schema: %w", err)
	}
	return &UsageTracker{db: db}, nil
}

// Close closes the underlying database.
func (t *UsageTracker) Close() error {
	return t.db.Close()
}

// Record inserts a new usage row.
func (t *UsageTracker) Record(model string, promptTokens, completionTokens int) error {
	_, err := t.db.Exec(
		`INSERT INTO usage (model, prompt_tokens, completion_tokens, total_tokens) VALUES (?, ?, ?, ?)`,
		model, promptTokens, completionTokens, promptTokens+completionTokens,
	)
	return err
}

// usageSummary is returned by the /usage endpoint.
type usageSummary struct {
	Period    string       `json:"period"`
	Model     string       `json:"model,omitempty"`
	Rows      []usageRow   `json:"rows"`
	Totals    usageTotals  `json:"totals"`
}

type usageRow struct {
	Date             string `json:"date"`
	Model            string `json:"model"`
	Requests         int    `json:"requests"`
	PromptTokens     int    `json:"prompt_tokens"`
	CompletionTokens int    `json:"completion_tokens"`
	TotalTokens      int    `json:"total_tokens"`
}

type usageTotals struct {
	Requests         int `json:"requests"`
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
	TotalTokens      int `json:"total_tokens"`
}

// HandleUsage handles GET /usage?days=7&model=xxx
func (t *UsageTracker) HandleUsage(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}

	days := 7
	if d := r.URL.Query().Get("days"); d != "" {
		if n, err := strconv.Atoi(d); err == nil && n > 0 {
			days = n
		}
	}
	modelFilter := r.URL.Query().Get("model")

	since := time.Now().AddDate(0, 0, -days).Format("2006-01-02")

	query := `
SELECT
    date(created_at) AS day,
    model,
    COUNT(*) AS requests,
    SUM(prompt_tokens) AS prompt_tokens,
    SUM(completion_tokens) AS completion_tokens,
    SUM(total_tokens) AS total_tokens
FROM usage
WHERE date(created_at) >= ?
`
	args := []interface{}{since}
	if modelFilter != "" {
		query += " AND model = ?"
		args = append(args, modelFilter)
	}
	query += " GROUP BY day, model ORDER BY day DESC, model"

	rows, err := t.db.Query(query, args...)
	if err != nil {
		logger.Printf("usage query error: %v", err)
		http.Error(w, "database error", http.StatusInternalServerError)
		return
	}
	defer rows.Close()

	var summary usageSummary
	summary.Period = fmt.Sprintf("%d days", days)
	summary.Model = modelFilter

	for rows.Next() {
		var row usageRow
		if err := rows.Scan(&row.Date, &row.Model, &row.Requests, &row.PromptTokens, &row.CompletionTokens, &row.TotalTokens); err != nil {
			logger.Printf("usage scan error: %v", err)
			continue
		}
		summary.Rows = append(summary.Rows, row)
		summary.Totals.Requests += row.Requests
		summary.Totals.PromptTokens += row.PromptTokens
		summary.Totals.CompletionTokens += row.CompletionTokens
		summary.Totals.TotalTokens += row.TotalTokens
	}

	if summary.Rows == nil {
		summary.Rows = []usageRow{}
	}

	w.Header().Set("Content-Type", "application/json")
	enc := json.NewEncoder(w)
	enc.SetIndent("", "  ")
	enc.Encode(summary)
}