descrybe/apps/api/internal/httpapi/admin_ai_costs_handlers.go

package httpapi

import (
	"context"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/descrybe/descrybe-v2/apps/api/internal/aiaudit"
	"github.com/google/uuid"
)

// Admin AI cost report: what each tenant and each user costs, plus the total.
//
// Reads ai_usage_daily, the durable rollup written alongside every captured call.
// ai_call_logs (full prompts) is pruned after a week; this survives for
// aiaudit.UsageRetentionDays (3 years), so quarter- and year-scale questions work.

const (
	adminAICostsTimeout = 30 * time.Second
	adminAICostsMaxDays = 1200 // slightly over the 3-year retention
	adminAICostsMaxRows = 500
	adminAICostsDefDays = 30
)

// handleAdminAICosts returns spend grouped by company and by user, with totals.
// GET /api/admin/ai-costs?days=30&from=&to=&company_id=
func (s *Server) handleAdminAICosts(w http.ResponseWriter, r *http.Request) {
	ctx, cancel := context.WithTimeout(r.Context(), adminAICostsTimeout)
	defer cancel()
	if s.Pool == nil {
		Error(w, http.StatusServiceUnavailable, "database unavailable")
		return
	}

	q := r.URL.Query()
	from, to, err := adminCostRange(q.Get("from"), q.Get("to"), q.Get("days"))
	if err != nil {
		Error(w, http.StatusBadRequest, err.Error())
		return
	}

	where := []string{"day >= $1", "day <= $2"}
	args := []any{from, to}
	if raw := strings.TrimSpace(q.Get("company_id")); raw != "" {
		id, perr := uuid.Parse(raw)
		if perr != nil {
			Error(w, http.StatusBadRequest, "invalid company_id")
			return
		}
		args = append(args, id)
		where = append(where, "company_id = $"+strconv.Itoa(len(args)))
	}
	whereSQL := strings.Join(where, " AND ")

	totals, err := s.scanCostTotals(ctx, whereSQL, args)
	if err != nil {
		Error(w, http.StatusInternalServerError, "load ai costs failed")
		return
	}
	byCompany, err := s.scanCostGroup(ctx, costGroupCompany, whereSQL, args)
	if err != nil {
		Error(w, http.StatusInternalServerError, "load ai costs by company failed")
		return
	}
	byUser, err := s.scanCostGroup(ctx, costGroupUser, whereSQL, args)
	if err != nil {
		Error(w, http.StatusInternalServerError, "load ai costs by user failed")
		return
	}
	byModel, err := s.scanCostGroup(ctx, costGroupModel, whereSQL, args)
	if err != nil {
		Error(w, http.StatusInternalServerError, "load ai costs by model failed")
		return
	}

	JSON(w, http.StatusOK, map[string]any{
		"from":                  from.Format("2006-01-02"),
		"to":                    to.Format("2006-01-02"),
		"currency":              "USD",
		"totals":                totals,
		"by_company":            byCompany,
		"by_user":               byUser,
		"by_model":              byModel,
		"usage_retention_days":  aiaudit.UsageRetentionDays,
		"prompt_retention_days": aiaudit.RetentionDays,
	})
}

func adminCostRange(fromRaw, toRaw, daysRaw string) (time.Time, time.Time, error) {
	today := time.Now().UTC().Truncate(24 * time.Hour)
	to := today
	if s := strings.TrimSpace(toRaw); s != "" {
		t, err := time.Parse("2006-01-02", s)
		if err != nil {
			return time.Time{}, time.Time{}, errBadRequest("invalid to date (YYYY-MM-DD)")
		}
		to = t.UTC()
	}
	if s := strings.TrimSpace(fromRaw); s != "" {
		f, err := time.Parse("2006-01-02", s)
		if err != nil {
			return time.Time{}, time.Time{}, errBadRequest("invalid from date (YYYY-MM-DD)")
		}
		from := f.UTC()
		if from.After(to) {
			return time.Time{}, time.Time{}, errBadRequest("from must not be after to")
		}
		if to.Sub(from) > time.Duration(adminAICostsMaxDays)*24*time.Hour {
			return time.Time{}, time.Time{}, errBadRequest("range too large")
		}
		return from, to, nil
	}
	days := adminAICostsDefDays
	if s := strings.TrimSpace(daysRaw); s != "" {
		n, err := strconv.Atoi(s)
		if err != nil || n <= 0 {
			return time.Time{}, time.Time{}, errBadRequest("invalid days")
		}
		days = n
	}
	if days > adminAICostsMaxDays {
		days = adminAICostsMaxDays
	}
	return to.AddDate(0, 0, -(days - 1)), to, nil
}

type badRequestErr string

func (e badRequestErr) Error() string { return string(e) }

func errBadRequest(msg string) error { return badRequestErr(msg) }

func (s *Server) scanCostTotals(ctx context.Context, whereSQL string, args []any) (map[string]any, error) {
	var calls, failed, promptTok, cachedTok, outTok, totalTok, costMicros int64
	var companies, users int64
	err := s.Pool.QueryRow(ctx, `
		SELECT COALESCE(sum(calls),0), COALESCE(sum(failed_calls),0),
			COALESCE(sum(prompt_tokens),0), COALESCE(sum(cached_prompt_tokens),0),
			COALESCE(sum(completion_tokens),0), COALESCE(sum(total_tokens),0),
			COALESCE(sum(cost_micros),0),
			count(DISTINCT company_id), count(DISTINCT user_id)
		FROM ai_usage_daily WHERE `+whereSQL, args...).
		Scan(&calls, &failed, &promptTok, &cachedTok, &outTok, &totalTok, &costMicros, &companies, &users)
	if err != nil {
		return nil, err
	}
	return costRow(map[string]any{
		"companies": companies,
		"users":     users,
	}, calls, failed, promptTok, cachedTok, outTok, totalTok, costMicros), nil
}

type costGroup string

const (
	costGroupCompany costGroup = "company"
	costGroupUser    costGroup = "user"
	costGroupModel   costGroup = "model"
)

func (s *Server) scanCostGroup(ctx context.Context, group costGroup, whereSQL string, args []any) ([]map[string]any, error) {
	var sql string
	switch group {
	case costGroupCompany:
		sql = `
			SELECT u.company_id::text, COALESCE(c.name, ''), '', '',
				COALESCE(sum(u.calls),0), COALESCE(sum(u.failed_calls),0),
				COALESCE(sum(u.prompt_tokens),0), COALESCE(sum(u.cached_prompt_tokens),0),
				COALESCE(sum(u.completion_tokens),0), COALESCE(sum(u.total_tokens),0),
				COALESCE(sum(u.cost_micros),0)
			FROM ai_usage_daily u
			LEFT JOIN companies c ON c.id = u.company_id
			WHERE ` + whereSQL + `
			GROUP BY u.company_id, c.name
			ORDER BY COALESCE(sum(u.cost_micros),0) DESC
			LIMIT ` + strconv.Itoa(adminAICostsMaxRows)
	case costGroupUser:
		sql = `
			SELECT u.user_id::text, COALESCE(usr.email, ''), u.company_id::text, COALESCE(c.name, ''),
				COALESCE(sum(u.calls),0), COALESCE(sum(u.failed_calls),0),
				COALESCE(sum(u.prompt_tokens),0), COALESCE(sum(u.cached_prompt_tokens),0),
				COALESCE(sum(u.completion_tokens),0), COALESCE(sum(u.total_tokens),0),
				COALESCE(sum(u.cost_micros),0)
			FROM ai_usage_daily u
			LEFT JOIN users usr ON usr.id = u.user_id
			LEFT JOIN companies c ON c.id = u.company_id
			WHERE ` + whereSQL + `
			GROUP BY u.user_id, usr.email, u.company_id, c.name
			ORDER BY COALESCE(sum(u.cost_micros),0) DESC
			LIMIT ` + strconv.Itoa(adminAICostsMaxRows)
	default:
		sql = `
			SELECT u.model, u.model, '', '',
				COALESCE(sum(u.calls),0), COALESCE(sum(u.failed_calls),0),
				COALESCE(sum(u.prompt_tokens),0), COALESCE(sum(u.cached_prompt_tokens),0),
				COALESCE(sum(u.completion_tokens),0), COALESCE(sum(u.total_tokens),0),
				COALESCE(sum(u.cost_micros),0)
			FROM ai_usage_daily u
			WHERE ` + whereSQL + `
			GROUP BY u.model
			ORDER BY COALESCE(sum(u.cost_micros),0) DESC
			LIMIT ` + strconv.Itoa(adminAICostsMaxRows)
	}

	rows, err := s.Pool.Query(ctx, sql, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := []map[string]any{}
	for rows.Next() {
		var id, label, companyID, companyName string
		var calls, failed, promptTok, cachedTok, outTok, totalTok, costMicros int64
		if err := rows.Scan(&id, &label, &companyID, &companyName,
			&calls, &failed, &promptTok, &cachedTok, &outTok, &totalTok, &costMicros); err != nil {
			return nil, err
		}
		base := map[string]any{"id": id, "label": label}
		if group == costGroupUser {
			base["company_id"] = companyID
			base["company_name"] = companyName
		}
		out = append(out, costRow(base, calls, failed, promptTok, cachedTok, outTok, totalTok, costMicros))
	}
	return out, rows.Err()
}

// costRow returns cost as integer micros AND a display-ready USD float. Callers
// must sum micros, never the float — that is the whole reason cost is stored as an
// integer.
func costRow(base map[string]any, calls, failed, promptTok, cachedTok, outTok, totalTok, costMicros int64) map[string]any {
	base["calls"] = calls
	base["failed_calls"] = failed
	base["prompt_tokens"] = promptTok
	base["cached_prompt_tokens"] = cachedTok
	base["completion_tokens"] = outTok
	base["total_tokens"] = totalTok
	base["cost_micros"] = costMicros
	base["cost_usd"] = float64(costMicros) / float64(aiaudit.MicrosPerUSD)
	return base
}