mirror of
https://github.com/memohai/Memoh.git
synced 2026-04-27 07:16:19 +09:00
Acbox Liu
272 lines
8.2 KiB
Go
272 lines
8.2 KiB
Go
package handlers
|
|
|
|
import (
|
|
"context"
|
|
"log/slog"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/jackc/pgx/v5/pgtype"
|
|
"github.com/labstack/echo/v4"
|
|
|
|
"github.com/memohai/memoh/internal/accounts"
|
|
"github.com/memohai/memoh/internal/bots"
|
|
"github.com/memohai/memoh/internal/db"
|
|
"github.com/memohai/memoh/internal/db/sqlc"
|
|
)
|
|
|
|
type TokenUsageHandler struct {
|
|
queries *sqlc.Queries
|
|
botService *bots.Service
|
|
accountService *accounts.Service
|
|
logger *slog.Logger
|
|
}
|
|
|
|
func NewTokenUsageHandler(log *slog.Logger, queries *sqlc.Queries, botService *bots.Service, accountService *accounts.Service) *TokenUsageHandler {
|
|
return &TokenUsageHandler{
|
|
queries: queries,
|
|
botService: botService,
|
|
accountService: accountService,
|
|
logger: log.With(slog.String("handler", "token_usage")),
|
|
}
|
|
}
|
|
|
|
func (h *TokenUsageHandler) Register(e *echo.Echo) {
|
|
e.GET("/bots/:bot_id/token-usage", h.GetTokenUsage)
|
|
}
|
|
|
|
// DailyTokenUsage represents aggregated token usage for a single day.
|
|
type DailyTokenUsage struct {
|
|
Day string `json:"day"`
|
|
InputTokens int64 `json:"input_tokens"`
|
|
OutputTokens int64 `json:"output_tokens"`
|
|
CacheReadTokens int64 `json:"cache_read_tokens"`
|
|
CacheWriteTokens int64 `json:"cache_write_tokens"`
|
|
ReasoningTokens int64 `json:"reasoning_tokens"`
|
|
}
|
|
|
|
// ModelTokenUsage represents aggregated token usage for a single model.
|
|
type ModelTokenUsage struct {
|
|
ModelID string `json:"model_id"`
|
|
ModelSlug string `json:"model_slug"`
|
|
ModelName string `json:"model_name"`
|
|
ProviderName string `json:"provider_name"`
|
|
InputTokens int64 `json:"input_tokens"`
|
|
OutputTokens int64 `json:"output_tokens"`
|
|
}
|
|
|
|
// TokenUsageResponse is the response body for GET /bots/:bot_id/token-usage.
|
|
type TokenUsageResponse struct {
|
|
Chat []DailyTokenUsage `json:"chat"`
|
|
Heartbeat []DailyTokenUsage `json:"heartbeat"`
|
|
ByModel []ModelTokenUsage `json:"by_model"`
|
|
}
|
|
|
|
// GetTokenUsage godoc
|
|
// @Summary Get token usage statistics
|
|
// @Description Get daily aggregated token usage for a bot, split by chat and heartbeat, with optional model filter and per-model breakdown
|
|
// @Tags token-usage
|
|
// @Param bot_id path string true "Bot ID"
|
|
// @Param from query string true "Start date (YYYY-MM-DD)"
|
|
// @Param to query string true "End date exclusive (YYYY-MM-DD)"
|
|
// @Param model_id query string false "Optional model UUID to filter by"
|
|
// @Success 200 {object} TokenUsageResponse
|
|
// @Failure 400 {object} ErrorResponse
|
|
// @Failure 403 {object} ErrorResponse
|
|
// @Failure 500 {object} ErrorResponse
|
|
// @Router /bots/{bot_id}/token-usage [get]
|
|
func (h *TokenUsageHandler) GetTokenUsage(c echo.Context) error {
|
|
userID, err := RequireChannelIdentityID(c)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
botID := strings.TrimSpace(c.Param("bot_id"))
|
|
if botID == "" {
|
|
return echo.NewHTTPError(http.StatusBadRequest, "bot id is required")
|
|
}
|
|
if _, err := AuthorizeBotAccess(c.Request().Context(), h.botService, h.accountService, userID, botID, bots.AccessPolicy{AllowPublicMember: false}); err != nil {
|
|
return err
|
|
}
|
|
|
|
fromStr := strings.TrimSpace(c.QueryParam("from"))
|
|
toStr := strings.TrimSpace(c.QueryParam("to"))
|
|
if fromStr == "" || toStr == "" {
|
|
return echo.NewHTTPError(http.StatusBadRequest, "from and to query parameters are required (YYYY-MM-DD)")
|
|
}
|
|
fromDate, err := time.Parse("2006-01-02", fromStr)
|
|
if err != nil {
|
|
return echo.NewHTTPError(http.StatusBadRequest, "invalid from date format, expected YYYY-MM-DD")
|
|
}
|
|
toDate, err := time.Parse("2006-01-02", toStr)
|
|
if err != nil {
|
|
return echo.NewHTTPError(http.StatusBadRequest, "invalid to date format, expected YYYY-MM-DD")
|
|
}
|
|
if !toDate.After(fromDate) {
|
|
return echo.NewHTTPError(http.StatusBadRequest, "to must be after from")
|
|
}
|
|
|
|
pgBotID, err := db.ParseUUID(botID)
|
|
if err != nil {
|
|
return echo.NewHTTPError(http.StatusBadRequest, "invalid bot id")
|
|
}
|
|
|
|
var pgModelID pgtype.UUID
|
|
if modelIDStr := strings.TrimSpace(c.QueryParam("model_id")); modelIDStr != "" {
|
|
pgModelID, err = db.ParseUUID(modelIDStr)
|
|
if err != nil {
|
|
return echo.NewHTTPError(http.StatusBadRequest, "invalid model_id")
|
|
}
|
|
}
|
|
|
|
fromTS := pgtype.Timestamptz{Time: fromDate, Valid: true}
|
|
toTS := pgtype.Timestamptz{Time: toDate, Valid: true}
|
|
|
|
ctx := c.Request().Context()
|
|
|
|
chatRows, hbRows, err := h.fetchUsage(ctx, pgBotID, fromTS, toTS, pgModelID)
|
|
if err != nil {
|
|
h.logger.Error("fetch token usage failed", slog.Any("error", err))
|
|
return echo.NewHTTPError(http.StatusInternalServerError, "failed to fetch token usage")
|
|
}
|
|
|
|
byModel, err := h.fetchUsageByModel(ctx, pgBotID, fromTS, toTS)
|
|
if err != nil {
|
|
h.logger.Error("fetch token usage by model failed", slog.Any("error", err))
|
|
return echo.NewHTTPError(http.StatusInternalServerError, "failed to fetch token usage by model")
|
|
}
|
|
|
|
resp := TokenUsageResponse{
|
|
Chat: convertMessageRows(chatRows),
|
|
Heartbeat: convertHeartbeatRows(hbRows),
|
|
ByModel: byModel,
|
|
}
|
|
return c.JSON(http.StatusOK, resp)
|
|
}
|
|
|
|
func (h *TokenUsageHandler) fetchUsage(ctx context.Context, botID pgtype.UUID, from, to pgtype.Timestamptz, modelID pgtype.UUID) ([]sqlc.GetMessageTokenUsageByDayRow, []sqlc.GetHeartbeatTokenUsageByDayRow, error) {
|
|
chatRows, err := h.queries.GetMessageTokenUsageByDay(ctx, sqlc.GetMessageTokenUsageByDayParams{
|
|
BotID: botID,
|
|
FromTime: from,
|
|
ToTime: to,
|
|
ModelID: modelID,
|
|
})
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
hbRows, err := h.queries.GetHeartbeatTokenUsageByDay(ctx, sqlc.GetHeartbeatTokenUsageByDayParams{
|
|
BotID: botID,
|
|
FromTime: from,
|
|
ToTime: to,
|
|
ModelID: modelID,
|
|
})
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
return chatRows, hbRows, nil
|
|
}
|
|
|
|
func (h *TokenUsageHandler) fetchUsageByModel(ctx context.Context, botID pgtype.UUID, from, to pgtype.Timestamptz) ([]ModelTokenUsage, error) {
|
|
merged := map[string]*ModelTokenUsage{}
|
|
|
|
chatRows, err := h.queries.GetMessageTokenUsageByModel(ctx, sqlc.GetMessageTokenUsageByModelParams{
|
|
BotID: botID,
|
|
FromTime: from,
|
|
ToTime: to,
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
for _, r := range chatRows {
|
|
key := r.ModelID.String()
|
|
if m, ok := merged[key]; ok {
|
|
m.InputTokens += r.InputTokens
|
|
m.OutputTokens += r.OutputTokens
|
|
} else {
|
|
merged[key] = &ModelTokenUsage{
|
|
ModelID: formatOptionalUUID(r.ModelID),
|
|
ModelSlug: r.ModelSlug,
|
|
ModelName: r.ModelName,
|
|
ProviderName: r.ProviderName,
|
|
InputTokens: r.InputTokens,
|
|
OutputTokens: r.OutputTokens,
|
|
}
|
|
}
|
|
}
|
|
|
|
hbRows, err := h.queries.GetHeartbeatTokenUsageByModel(ctx, sqlc.GetHeartbeatTokenUsageByModelParams{
|
|
BotID: botID,
|
|
FromTime: from,
|
|
ToTime: to,
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
for _, r := range hbRows {
|
|
key := r.ModelID.String()
|
|
if m, ok := merged[key]; ok {
|
|
m.InputTokens += r.InputTokens
|
|
m.OutputTokens += r.OutputTokens
|
|
} else {
|
|
merged[key] = &ModelTokenUsage{
|
|
ModelID: formatOptionalUUID(r.ModelID),
|
|
ModelSlug: r.ModelSlug,
|
|
ModelName: r.ModelName,
|
|
ProviderName: r.ProviderName,
|
|
InputTokens: r.InputTokens,
|
|
OutputTokens: r.OutputTokens,
|
|
}
|
|
}
|
|
}
|
|
|
|
result := make([]ModelTokenUsage, 0, len(merged))
|
|
for _, m := range merged {
|
|
result = append(result, *m)
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
func convertMessageRows(rows []sqlc.GetMessageTokenUsageByDayRow) []DailyTokenUsage {
|
|
out := make([]DailyTokenUsage, 0, len(rows))
|
|
for _, r := range rows {
|
|
out = append(out, DailyTokenUsage{
|
|
Day: formatPgDate(r.Day),
|
|
InputTokens: r.InputTokens,
|
|
OutputTokens: r.OutputTokens,
|
|
CacheReadTokens: r.CacheReadTokens,
|
|
CacheWriteTokens: r.CacheWriteTokens,
|
|
ReasoningTokens: r.ReasoningTokens,
|
|
})
|
|
}
|
|
return out
|
|
}
|
|
|
|
func convertHeartbeatRows(rows []sqlc.GetHeartbeatTokenUsageByDayRow) []DailyTokenUsage {
|
|
out := make([]DailyTokenUsage, 0, len(rows))
|
|
for _, r := range rows {
|
|
out = append(out, DailyTokenUsage{
|
|
Day: formatPgDate(r.Day),
|
|
InputTokens: r.InputTokens,
|
|
OutputTokens: r.OutputTokens,
|
|
CacheReadTokens: r.CacheReadTokens,
|
|
CacheWriteTokens: r.CacheWriteTokens,
|
|
ReasoningTokens: r.ReasoningTokens,
|
|
})
|
|
}
|
|
return out
|
|
}
|
|
|
|
func formatPgDate(d pgtype.Date) string {
|
|
if !d.Valid {
|
|
return ""
|
|
}
|
|
return d.Time.Format("2006-01-02")
|
|
}
|
|
|
|
func formatOptionalUUID(id pgtype.UUID) string {
|
|
if !id.Valid {
|
|
return ""
|
|
}
|
|
return id.String()
|
|
}
|