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 getTaskUsage = `-- name: GetTaskUsage :many
SELECT id, task_id, provider, model, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, created_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,
		); 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(atq.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 atq.created_at >= $2::timestamptz
GROUP BY DATE(atq.created_at), tu.model
ORDER BY DATE(atq.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"`
}

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 atq.created_at >= $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"`
}

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)
VALUES ($1, $2, $3, $4, $5, $6, $7)
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
`

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"`
}

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
}