multica/server/pkg/db/generated/task_usage.sql.go

// Code generated by sqlc. DO NOT EDIT.
// versions:
//   sqlc v1.30.0
// source: task_usage.sql

package db

import (
	"context"

	"github.com/jackc/pgx/v5/pgtype"
)

const getIssueUsageSummary = `-- name: GetIssueUsageSummary :one
SELECT
    COALESCE(SUM(tu.input_tokens), 0)::bigint AS total_input_tokens,
    COALESCE(SUM(tu.output_tokens), 0)::bigint AS total_output_tokens,
    COALESCE(SUM(tu.cache_read_tokens), 0)::bigint AS total_cache_read_tokens,
    COALESCE(SUM(tu.cache_write_tokens), 0)::bigint AS total_cache_write_tokens,
    COUNT(DISTINCT tu.task_id)::int AS task_count
FROM task_usage tu
JOIN agent_task_queue atq ON atq.id = tu.task_id
WHERE atq.issue_id = $1
`

type GetIssueUsageSummaryRow struct {
	TotalInputTokens      int64 `json:"total_input_tokens"`
	TotalOutputTokens     int64 `json:"total_output_tokens"`
	TotalCacheReadTokens  int64 `json:"total_cache_read_tokens"`
	TotalCacheWriteTokens int64 `json:"total_cache_write_tokens"`
	TaskCount             int32 `json:"task_count"`
}

func (q *Queries) GetIssueUsageSummary(ctx context.Context, issueID pgtype.UUID) (GetIssueUsageSummaryRow, error) {
	row := q.db.QueryRow(ctx, getIssueUsageSummary, issueID)
	var i GetIssueUsageSummaryRow
	err := row.Scan(
		&i.TotalInputTokens,
		&i.TotalOutputTokens,
		&i.TotalCacheReadTokens,
		&i.TotalCacheWriteTokens,
		&i.TaskCount,
	)
	return i, err
}

const getTaskUsage = `-- name: GetTaskUsage :many
SELECT id, task_id, provider, model, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, created_at, updated_at FROM task_usage
WHERE task_id = $1
ORDER BY model
`

func (q *Queries) GetTaskUsage(ctx context.Context, taskID pgtype.UUID) ([]TaskUsage, error) {
	rows, err := q.db.Query(ctx, getTaskUsage, taskID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	items := []TaskUsage{}
	for rows.Next() {
		var i TaskUsage
		if err := rows.Scan(
			&i.ID,
			&i.TaskID,
			&i.Provider,
			&i.Model,
			&i.InputTokens,
			&i.OutputTokens,
			&i.CacheReadTokens,
			&i.CacheWriteTokens,
			&i.CreatedAt,
			&i.UpdatedAt,
		); err != nil {
			return nil, err
		}
		items = append(items, i)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return items, nil
}

const getWorkspaceUsageByDay = `-- name: GetWorkspaceUsageByDay :many
SELECT
    DATE(tu.created_at) AS date,
    tu.model,
    SUM(tu.input_tokens)::bigint AS total_input_tokens,
    SUM(tu.output_tokens)::bigint AS total_output_tokens,
    SUM(tu.cache_read_tokens)::bigint AS total_cache_read_tokens,
    SUM(tu.cache_write_tokens)::bigint AS total_cache_write_tokens,
    COUNT(DISTINCT tu.task_id)::int AS task_count
FROM task_usage tu
JOIN agent_task_queue atq ON atq.id = tu.task_id
JOIN agent a ON a.id = atq.agent_id
WHERE a.workspace_id = $1
  AND tu.created_at >= DATE_TRUNC('day', $2::timestamptz)
GROUP BY DATE(tu.created_at), tu.model
ORDER BY DATE(tu.created_at) DESC, tu.model
`

type GetWorkspaceUsageByDayParams struct {
	WorkspaceID pgtype.UUID        `json:"workspace_id"`
	Since       pgtype.Timestamptz `json:"since"`
}

type GetWorkspaceUsageByDayRow struct {
	Date                  pgtype.Date `json:"date"`
	Model                 string      `json:"model"`
	TotalInputTokens      int64       `json:"total_input_tokens"`
	TotalOutputTokens     int64       `json:"total_output_tokens"`
	TotalCacheReadTokens  int64       `json:"total_cache_read_tokens"`
	TotalCacheWriteTokens int64       `json:"total_cache_write_tokens"`
	TaskCount             int32       `json:"task_count"`
}

// Bucket by tu.created_at (usage report time, ~= task completion time), not
// atq.created_at (task enqueue time), so tasks that queue one day and execute
// the next are attributed to the day tokens were actually produced. The since
// cutoff is truncated to start-of-day so `days=N` yields full calendar days.
func (q *Queries) GetWorkspaceUsageByDay(ctx context.Context, arg GetWorkspaceUsageByDayParams) ([]GetWorkspaceUsageByDayRow, error) {
	rows, err := q.db.Query(ctx, getWorkspaceUsageByDay, arg.WorkspaceID, arg.Since)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	items := []GetWorkspaceUsageByDayRow{}
	for rows.Next() {
		var i GetWorkspaceUsageByDayRow
		if err := rows.Scan(
			&i.Date,
			&i.Model,
			&i.TotalInputTokens,
			&i.TotalOutputTokens,
			&i.TotalCacheReadTokens,
			&i.TotalCacheWriteTokens,
			&i.TaskCount,
		); err != nil {
			return nil, err
		}
		items = append(items, i)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return items, nil
}

const getWorkspaceUsageSummary = `-- name: GetWorkspaceUsageSummary :many
SELECT
    tu.model,
    SUM(tu.input_tokens)::bigint AS total_input_tokens,
    SUM(tu.output_tokens)::bigint AS total_output_tokens,
    SUM(tu.cache_read_tokens)::bigint AS total_cache_read_tokens,
    SUM(tu.cache_write_tokens)::bigint AS total_cache_write_tokens,
    COUNT(DISTINCT tu.task_id)::int AS task_count
FROM task_usage tu
JOIN agent_task_queue atq ON atq.id = tu.task_id
JOIN agent a ON a.id = atq.agent_id
WHERE a.workspace_id = $1
  AND tu.created_at >= DATE_TRUNC('day', $2::timestamptz)
GROUP BY tu.model
ORDER BY (SUM(tu.input_tokens) + SUM(tu.output_tokens)) DESC
`

type GetWorkspaceUsageSummaryParams struct {
	WorkspaceID pgtype.UUID        `json:"workspace_id"`
	Since       pgtype.Timestamptz `json:"since"`
}

type GetWorkspaceUsageSummaryRow struct {
	Model                 string `json:"model"`
	TotalInputTokens      int64  `json:"total_input_tokens"`
	TotalOutputTokens     int64  `json:"total_output_tokens"`
	TotalCacheReadTokens  int64  `json:"total_cache_read_tokens"`
	TotalCacheWriteTokens int64  `json:"total_cache_write_tokens"`
	TaskCount             int32  `json:"task_count"`
}

// Filter by tu.created_at (usage report time), aligned to start-of-day, so
// `days=N` is interpreted as N full calendar days like the other usage queries.
func (q *Queries) GetWorkspaceUsageSummary(ctx context.Context, arg GetWorkspaceUsageSummaryParams) ([]GetWorkspaceUsageSummaryRow, error) {
	rows, err := q.db.Query(ctx, getWorkspaceUsageSummary, arg.WorkspaceID, arg.Since)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	items := []GetWorkspaceUsageSummaryRow{}
	for rows.Next() {
		var i GetWorkspaceUsageSummaryRow
		if err := rows.Scan(
			&i.Model,
			&i.TotalInputTokens,
			&i.TotalOutputTokens,
			&i.TotalCacheReadTokens,
			&i.TotalCacheWriteTokens,
			&i.TaskCount,
		); err != nil {
			return nil, err
		}
		items = append(items, i)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return items, nil
}

const upsertTaskUsage = `-- name: UpsertTaskUsage :exec
INSERT INTO task_usage (task_id, provider, model, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, updated_at)
VALUES ($1, $2, $3, $4, $5, $6, $7, now())
ON CONFLICT (task_id, provider, model)
DO UPDATE SET
    input_tokens = EXCLUDED.input_tokens,
    output_tokens = EXCLUDED.output_tokens,
    cache_read_tokens = EXCLUDED.cache_read_tokens,
    cache_write_tokens = EXCLUDED.cache_write_tokens,
    updated_at = now()
`

type UpsertTaskUsageParams struct {
	TaskID           pgtype.UUID `json:"task_id"`
	Provider         string      `json:"provider"`
	Model            string      `json:"model"`
	InputTokens      int64       `json:"input_tokens"`
	OutputTokens     int64       `json:"output_tokens"`
	CacheReadTokens  int64       `json:"cache_read_tokens"`
	CacheWriteTokens int64       `json:"cache_write_tokens"`
}

// Bumps `updated_at` on INSERT and on conflict so the daily-rollup worker
// (migration 073) detects the row as dirty and re-aggregates its bucket.
// Without the conflict-side bump, a correction to historical token counts
// would never propagate to the rollup.
func (q *Queries) UpsertTaskUsage(ctx context.Context, arg UpsertTaskUsageParams) error {
	_, err := q.db.Exec(ctx, upsertTaskUsage,
		arg.TaskID,
		arg.Provider,
		arg.Model,
		arg.InputTokens,
		arg.OutputTokens,
		arg.CacheReadTokens,
		arg.CacheWriteTokens,
	)
	return err
}