Memoh/internal/command/usage.go

package command

import (
	"fmt"
	"strings"
	"time"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5/pgtype"

	dbsqlc "github.com/memohai/memoh/internal/db/sqlc"
)

func (h *Handler) buildUsageGroup() *CommandGroup {
	g := newCommandGroup("usage", "View token usage")
	g.DefaultAction = "summary"
	g.Register(SubCommand{
		Name:  "summary",
		Usage: "summary - Token usage summary (last 7 days)",
		Handler: func(cc CommandContext) (string, error) {
			botUUID, err := parseBotUUID(cc.BotID)
			if err != nil {
				return "", err
			}
			now := time.Now().UTC()
			from := now.AddDate(0, 0, -7)
			fromTS := pgtype.Timestamptz{Time: from, Valid: true}
			toTS := pgtype.Timestamptz{Time: now, Valid: true}
			nullModel := pgtype.UUID{Valid: false}

			chatRows, err := h.queries.GetMessageTokenUsageByDay(cc.Ctx, dbsqlc.GetMessageTokenUsageByDayParams{
				BotID: botUUID, FromTime: fromTS, ToTime: toTS, ModelID: nullModel,
			})
			if err != nil {
				return "", err
			}

			hbRows, err := h.queries.GetHeartbeatTokenUsageByDay(cc.Ctx, dbsqlc.GetHeartbeatTokenUsageByDayParams{
				BotID: botUUID, FromTime: fromTS, ToTime: toTS, ModelID: nullModel,
			})
			if err != nil {
				return "", err
			}

			if len(chatRows) == 0 && len(hbRows) == 0 {
				return "No token usage in the last 7 days.", nil
			}

			var b strings.Builder
			b.WriteString("Token usage (last 7 days):\n\n")

			if len(chatRows) > 0 {
				b.WriteString("Chat:\n")
				var totalIn, totalOut int64
				for _, r := range chatRows {
					day := r.Day.Time.Format("01-02")
					fmt.Fprintf(&b, "  %s: in=%d out=%d\n", day, r.InputTokens, r.OutputTokens)
					totalIn += r.InputTokens
					totalOut += r.OutputTokens
				}
				fmt.Fprintf(&b, "  Total: in=%d out=%d\n", totalIn, totalOut)
			}

			if len(hbRows) > 0 {
				if len(chatRows) > 0 {
					b.WriteByte('\n')
				}
				b.WriteString("Heartbeat:\n")
				var totalIn, totalOut int64
				for _, r := range hbRows {
					day := r.Day.Time.Format("01-02")
					fmt.Fprintf(&b, "  %s: in=%d out=%d\n", day, r.InputTokens, r.OutputTokens)
					totalIn += r.InputTokens
					totalOut += r.OutputTokens
				}
				fmt.Fprintf(&b, "  Total: in=%d out=%d\n", totalIn, totalOut)
			}

			return strings.TrimRight(b.String(), "\n"), nil
		},
	})
	g.Register(SubCommand{
		Name:  "by-model",
		Usage: "by-model - Token usage grouped by model",
		Handler: func(cc CommandContext) (string, error) {
			botUUID, err := parseBotUUID(cc.BotID)
			if err != nil {
				return "", err
			}
			now := time.Now().UTC()
			from := now.AddDate(0, 0, -7)
			fromTS := pgtype.Timestamptz{Time: from, Valid: true}
			toTS := pgtype.Timestamptz{Time: now, Valid: true}

			chatRows, err := h.queries.GetMessageTokenUsageByModel(cc.Ctx, dbsqlc.GetMessageTokenUsageByModelParams{
				BotID: botUUID, FromTime: fromTS, ToTime: toTS,
			})
			if err != nil {
				return "", err
			}
			hbRows, err := h.queries.GetHeartbeatTokenUsageByModel(cc.Ctx, dbsqlc.GetHeartbeatTokenUsageByModelParams{
				BotID: botUUID, FromTime: fromTS, ToTime: toTS,
			})
			if err != nil {
				return "", err
			}

			if len(chatRows) == 0 && len(hbRows) == 0 {
				return "No token usage in the last 7 days.", nil
			}

			var b strings.Builder
			b.WriteString("Token usage by model (last 7 days):\n\n")

			if len(chatRows) > 0 {
				b.WriteString("Chat:\n")
				for _, r := range chatRows {
					fmt.Fprintf(&b, "  %s (%s): in=%d out=%d\n", r.ModelName, r.ProviderName, r.InputTokens, r.OutputTokens)
				}
			}

			if len(hbRows) > 0 {
				if len(chatRows) > 0 {
					b.WriteByte('\n')
				}
				b.WriteString("Heartbeat:\n")
				for _, r := range hbRows {
					fmt.Fprintf(&b, "  %s (%s): in=%d out=%d\n", r.ModelName, r.ProviderName, r.InputTokens, r.OutputTokens)
				}
			}

			return strings.TrimRight(b.String(), "\n"), nil
		},
	})
	return g
}

func parseBotUUID(botID string) (pgtype.UUID, error) {
	parsed, err := uuid.Parse(botID)
	if err != nil {
		return pgtype.UUID{}, fmt.Errorf("invalid bot ID: %w", err)
	}
	return pgtype.UUID{Bytes: parsed, Valid: true}, nil
}