mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: add CSV export for AI spend data (#27491)
## Description
Adds `GET /api/v2/organizations/{organization}/ai/spend/export`,
returning `text/csv` with per-user, per-group, per-model, per-provider
aggregated AI spend. The data is built from the raw AI Gateway token
usage tables rather than the `ai_user_daily_spend` rollup, but stays
consistent with it: spend is attributed through the token usage's
effective group and bucketed by the token usage `created_at`, the same
values the daily rollup derives from.
The period defaults to the current UTC month, narrowed to the configured
AI Gateway retention window when the month begins before retained data
does. Explicit `period_start`/`period_end` params must be provided
together, are interpreted as UTC, and may span at most 31 days. Unlike
the default period, an explicit period that begins before the retention
window is rejected rather than narrowed. Every row echoes the applied
bounds, so a narrowed window is visible in the export.
The endpoint requires organization-level admin permissions.
## Changes
- Add the `ExportOrganizationAISpend` query aggregating
`aibridge_token_usages` joined to `aibridge_interceptions`, scoped to
the organization via the effective group, resolving the username, group
name, and organization name alongside their IDs.
- Add the `exportOrganizationAISpend` handler and route, gated by the
`aigateway-cost-control` experiment and the `AIBridge` feature,
returning the CSV in a single response.
- Add the `ExportOrganizationAISpend` codersdk client method.
- Require organization-wide `ResourceGroupMember` read, since the export
aggregates every user in the organization. The per-row filter stays in
`dbauthz` as defence in depth.
- Escape leading formula characters in the free-text columns, so a model
or provider name recorded from an intercepted request cannot be
evaluated when the CSV is opened in a spreadsheet.
- Add an index on `aibridge_token_usages (effective_group_id,
created_at)`, which the period and group predicates otherwise cannot
use.
Closes
https://linear.app/codercom/issue/AIGOV-293/add-csv-export-for-ai-spend-data
> [!NOTE]
> Generated by Coder Agents on behalf of @ssncferreira
This commit is contained in:
@@ -2704,6 +2704,10 @@ func (q *querier) ExpirePrebuildsAPIKeys(ctx context.Context, now time.Time) err
|
||||
return q.db.ExpirePrebuildsAPIKeys(ctx, now)
|
||||
}
|
||||
|
||||
func (q *querier) ExportOrganizationAISpend(ctx context.Context, arg database.ExportOrganizationAISpendParams) ([]database.ExportOrganizationAISpendRow, error) {
|
||||
return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.ExportOrganizationAISpend)(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) FavoriteWorkspace(ctx context.Context, id uuid.UUID) error {
|
||||
fetch := func(ctx context.Context, id uuid.UUID) (database.Workspace, error) {
|
||||
return q.db.GetWorkspaceByID(ctx, id)
|
||||
|
||||
@@ -6992,6 +6992,22 @@ func (s *MethodTestSuite) TestAIBridge() {
|
||||
Returns([]database.GetGroupMembersAISpendRow{row1, row2})
|
||||
}))
|
||||
|
||||
s.Run("ExportOrganizationAISpend", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
org := testutil.Fake(s.T(), faker, database.Organization{})
|
||||
row1 := testutil.Fake(s.T(), faker, database.ExportOrganizationAISpendRow{OrganizationID: org.ID})
|
||||
row2 := testutil.Fake(s.T(), faker, database.ExportOrganizationAISpendRow{OrganizationID: org.ID})
|
||||
arg := database.ExportOrganizationAISpendParams{
|
||||
OrganizationID: org.ID,
|
||||
PeriodStart: time.Now().UTC().Truncate(24 * time.Hour),
|
||||
PeriodEnd: time.Now().UTC(),
|
||||
}
|
||||
dbm.EXPECT().ExportOrganizationAISpend(gomock.Any(), arg).
|
||||
Return([]database.ExportOrganizationAISpendRow{row1, row2}, nil).AnyTimes()
|
||||
check.Args(arg).
|
||||
Asserts(row1, policy.ActionRead, row2, policy.ActionRead).
|
||||
Returns([]database.ExportOrganizationAISpendRow{row1, row2})
|
||||
}))
|
||||
|
||||
s.Run("GetGroupAIBudget", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
g := testutil.Fake(s.T(), faker, database.Group{})
|
||||
b := testutil.Fake(s.T(), faker, database.GroupAIBudget{GroupID: g.ID})
|
||||
|
||||
+8
@@ -1009,6 +1009,14 @@ func (m queryMetricsStore) ExpirePrebuildsAPIKeys(ctx context.Context, now time.
|
||||
return r0
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) ExportOrganizationAISpend(ctx context.Context, arg database.ExportOrganizationAISpendParams) ([]database.ExportOrganizationAISpendRow, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.ExportOrganizationAISpend(ctx, arg)
|
||||
m.queryLatencies.WithLabelValues("ExportOrganizationAISpend").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "ExportOrganizationAISpend").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) FavoriteWorkspace(ctx context.Context, id uuid.UUID) error {
|
||||
start := time.Now()
|
||||
r0 := m.s.FavoriteWorkspace(ctx, id)
|
||||
|
||||
Generated
+15
@@ -1724,6 +1724,21 @@ func (mr *MockStoreMockRecorder) ExpirePrebuildsAPIKeys(ctx, now any) *gomock.Ca
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ExpirePrebuildsAPIKeys", reflect.TypeOf((*MockStore)(nil).ExpirePrebuildsAPIKeys), ctx, now)
|
||||
}
|
||||
|
||||
// ExportOrganizationAISpend mocks base method.
|
||||
func (m *MockStore) ExportOrganizationAISpend(ctx context.Context, arg database.ExportOrganizationAISpendParams) ([]database.ExportOrganizationAISpendRow, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "ExportOrganizationAISpend", ctx, arg)
|
||||
ret0, _ := ret[0].([]database.ExportOrganizationAISpendRow)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// ExportOrganizationAISpend indicates an expected call of ExportOrganizationAISpend.
|
||||
func (mr *MockStoreMockRecorder) ExportOrganizationAISpend(ctx, arg any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ExportOrganizationAISpend", reflect.TypeOf((*MockStore)(nil).ExportOrganizationAISpend), ctx, arg)
|
||||
}
|
||||
|
||||
// FavoriteWorkspace mocks base method.
|
||||
func (m *MockStore) FavoriteWorkspace(ctx context.Context, id uuid.UUID) error {
|
||||
m.ctrl.T.Helper()
|
||||
|
||||
Generated
+2
@@ -4698,6 +4698,8 @@ CREATE INDEX idx_aibridge_interceptions_thread_root_id ON aibridge_interceptions
|
||||
|
||||
CREATE INDEX idx_aibridge_model_thoughts_interception_id ON aibridge_model_thoughts USING btree (interception_id);
|
||||
|
||||
CREATE INDEX idx_aibridge_token_usages_effective_group_id_created_at ON aibridge_token_usages USING btree (effective_group_id, created_at) WHERE (effective_group_id IS NOT NULL);
|
||||
|
||||
CREATE INDEX idx_aibridge_token_usages_interception_id ON aibridge_token_usages USING btree (interception_id);
|
||||
|
||||
CREATE INDEX idx_aibridge_token_usages_provider_response_id ON aibridge_token_usages USING btree (provider_response_id);
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
DROP INDEX IF EXISTS idx_aibridge_token_usages_effective_group_id_created_at;
|
||||
@@ -0,0 +1,5 @@
|
||||
-- Serves spend queries that filter token usage by effective group over a time
|
||||
-- range. Rows with no effective group are excluded.
|
||||
CREATE INDEX idx_aibridge_token_usages_effective_group_id_created_at
|
||||
ON aibridge_token_usages (effective_group_id, created_at)
|
||||
WHERE effective_group_id IS NOT NULL;
|
||||
@@ -474,6 +474,10 @@ func (r GetGroupMembersAISpendRow) RBACObject() rbac.Object {
|
||||
return rbac.ResourceGroupMember.WithID(r.UserID).InOrg(r.OrganizationID).WithOwner(r.UserID.String())
|
||||
}
|
||||
|
||||
func (r ExportOrganizationAISpendRow) RBACObject() rbac.Object {
|
||||
return rbac.ResourceGroupMember.WithID(r.UserID).InOrg(r.OrganizationID).WithOwner(r.UserID.String())
|
||||
}
|
||||
|
||||
// PrebuiltWorkspaceResource defines the interface for types that can be identified as prebuilt workspaces
|
||||
// and converted to their corresponding prebuilt workspace RBAC object.
|
||||
type PrebuiltWorkspaceResource interface {
|
||||
|
||||
Generated
+5
@@ -265,6 +265,11 @@ type sqlcQuerier interface {
|
||||
// Next, collect api_keys that belong to the prebuilds user but have no token name.
|
||||
// These were most likely created via 'coder login' as the prebuilds user.
|
||||
ExpirePrebuildsAPIKeys(ctx context.Context, now time.Time) error
|
||||
// Returns per-user, per-group, per-model, per-provider aggregated AI spend for
|
||||
// @organization_id over the [period_start, period_end) window. Spend is
|
||||
// attributed through the token usage's effective group, and rows are bucketed
|
||||
// by the token usage created_at, matching how ai_user_daily_spend is derived.
|
||||
ExportOrganizationAISpend(ctx context.Context, arg ExportOrganizationAISpendParams) ([]ExportOrganizationAISpendRow, error)
|
||||
FavoriteWorkspace(ctx context.Context, id uuid.UUID) error
|
||||
FetchMemoryResourceMonitorsByAgentID(ctx context.Context, agentID uuid.UUID) (WorkspaceAgentMemoryResourceMonitor, error)
|
||||
FetchMemoryResourceMonitorsUpdatedAfter(ctx context.Context, updatedAt time.Time) ([]WorkspaceAgentMemoryResourceMonitor, error)
|
||||
|
||||
Generated
+102
@@ -2525,6 +2525,108 @@ func (q *sqlQuerier) DeleteUserAIBudgetOverride(ctx context.Context, userID uuid
|
||||
return i, err
|
||||
}
|
||||
|
||||
const exportOrganizationAISpend = `-- name: ExportOrganizationAISpend :many
|
||||
SELECT
|
||||
ai.initiator_id AS user_id,
|
||||
users.username AS username,
|
||||
tu.effective_group_id AS group_id,
|
||||
groups.name AS group_name,
|
||||
groups.organization_id AS organization_id,
|
||||
organizations.name AS organization_name,
|
||||
ai.model AS model,
|
||||
ai.provider AS provider,
|
||||
ai.provider_name AS provider_name,
|
||||
COALESCE(SUM(tu.input_tokens), 0)::BIGINT AS input_tokens,
|
||||
COALESCE(SUM(tu.output_tokens), 0)::BIGINT AS output_tokens,
|
||||
COALESCE(SUM(tu.cache_read_input_tokens), 0)::BIGINT AS cache_read_tokens,
|
||||
COALESCE(SUM(tu.cache_write_input_tokens), 0)::BIGINT AS cache_write_tokens,
|
||||
COALESCE(SUM(tu.cost_micros), 0)::BIGINT AS cost_micros
|
||||
FROM aibridge_token_usages tu
|
||||
JOIN aibridge_interceptions ai ON ai.id = tu.interception_id
|
||||
JOIN users ON users.id = ai.initiator_id
|
||||
JOIN groups ON groups.id = tu.effective_group_id
|
||||
JOIN organizations ON organizations.id = groups.organization_id
|
||||
WHERE groups.organization_id = $1
|
||||
AND tu.created_at >= $2::timestamptz
|
||||
AND tu.created_at < $3::timestamptz
|
||||
GROUP BY
|
||||
ai.initiator_id,
|
||||
users.username,
|
||||
tu.effective_group_id,
|
||||
groups.name,
|
||||
groups.organization_id,
|
||||
organizations.name,
|
||||
ai.model,
|
||||
ai.provider,
|
||||
ai.provider_name
|
||||
ORDER BY ai.initiator_id, tu.effective_group_id, ai.provider, ai.provider_name, ai.model
|
||||
`
|
||||
|
||||
type ExportOrganizationAISpendParams struct {
|
||||
OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"`
|
||||
PeriodStart time.Time `db:"period_start" json:"period_start"`
|
||||
PeriodEnd time.Time `db:"period_end" json:"period_end"`
|
||||
}
|
||||
|
||||
type ExportOrganizationAISpendRow struct {
|
||||
UserID uuid.UUID `db:"user_id" json:"user_id"`
|
||||
Username string `db:"username" json:"username"`
|
||||
GroupID uuid.NullUUID `db:"group_id" json:"group_id"`
|
||||
GroupName string `db:"group_name" json:"group_name"`
|
||||
OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"`
|
||||
OrganizationName string `db:"organization_name" json:"organization_name"`
|
||||
Model string `db:"model" json:"model"`
|
||||
Provider string `db:"provider" json:"provider"`
|
||||
ProviderName string `db:"provider_name" json:"provider_name"`
|
||||
InputTokens int64 `db:"input_tokens" json:"input_tokens"`
|
||||
OutputTokens int64 `db:"output_tokens" json:"output_tokens"`
|
||||
CacheReadTokens int64 `db:"cache_read_tokens" json:"cache_read_tokens"`
|
||||
CacheWriteTokens int64 `db:"cache_write_tokens" json:"cache_write_tokens"`
|
||||
CostMicros int64 `db:"cost_micros" json:"cost_micros"`
|
||||
}
|
||||
|
||||
// Returns per-user, per-group, per-model, per-provider aggregated AI spend for
|
||||
// @organization_id over the [period_start, period_end) window. Spend is
|
||||
// attributed through the token usage's effective group, and rows are bucketed
|
||||
// by the token usage created_at, matching how ai_user_daily_spend is derived.
|
||||
func (q *sqlQuerier) ExportOrganizationAISpend(ctx context.Context, arg ExportOrganizationAISpendParams) ([]ExportOrganizationAISpendRow, error) {
|
||||
rows, err := q.db.QueryContext(ctx, exportOrganizationAISpend, arg.OrganizationID, arg.PeriodStart, arg.PeriodEnd)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var items []ExportOrganizationAISpendRow
|
||||
for rows.Next() {
|
||||
var i ExportOrganizationAISpendRow
|
||||
if err := rows.Scan(
|
||||
&i.UserID,
|
||||
&i.Username,
|
||||
&i.GroupID,
|
||||
&i.GroupName,
|
||||
&i.OrganizationID,
|
||||
&i.OrganizationName,
|
||||
&i.Model,
|
||||
&i.Provider,
|
||||
&i.ProviderName,
|
||||
&i.InputTokens,
|
||||
&i.OutputTokens,
|
||||
&i.CacheReadTokens,
|
||||
&i.CacheWriteTokens,
|
||||
&i.CostMicros,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, i)
|
||||
}
|
||||
if err := rows.Close(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
const getAIModelPriceByProviderModel = `-- name: GetAIModelPriceByProviderModel :one
|
||||
SELECT provider, model, input_price, output_price, cache_read_price, cache_write_price, created_at, updated_at
|
||||
FROM ai_model_prices
|
||||
|
||||
@@ -305,3 +305,43 @@ FROM user_spend
|
||||
WHERE current_spend_micros >= spend_limit_micros
|
||||
GROUP BY effective_group_id
|
||||
ORDER BY effective_group_id;
|
||||
|
||||
-- name: ExportOrganizationAISpend :many
|
||||
-- Returns per-user, per-group, per-model, per-provider aggregated AI spend for
|
||||
-- @organization_id over the [period_start, period_end) window. Spend is
|
||||
-- attributed through the token usage's effective group, and rows are bucketed
|
||||
-- by the token usage created_at, matching how ai_user_daily_spend is derived.
|
||||
SELECT
|
||||
ai.initiator_id AS user_id,
|
||||
users.username AS username,
|
||||
tu.effective_group_id AS group_id,
|
||||
groups.name AS group_name,
|
||||
groups.organization_id AS organization_id,
|
||||
organizations.name AS organization_name,
|
||||
ai.model AS model,
|
||||
ai.provider AS provider,
|
||||
ai.provider_name AS provider_name,
|
||||
COALESCE(SUM(tu.input_tokens), 0)::BIGINT AS input_tokens,
|
||||
COALESCE(SUM(tu.output_tokens), 0)::BIGINT AS output_tokens,
|
||||
COALESCE(SUM(tu.cache_read_input_tokens), 0)::BIGINT AS cache_read_tokens,
|
||||
COALESCE(SUM(tu.cache_write_input_tokens), 0)::BIGINT AS cache_write_tokens,
|
||||
COALESCE(SUM(tu.cost_micros), 0)::BIGINT AS cost_micros
|
||||
FROM aibridge_token_usages tu
|
||||
JOIN aibridge_interceptions ai ON ai.id = tu.interception_id
|
||||
JOIN users ON users.id = ai.initiator_id
|
||||
JOIN groups ON groups.id = tu.effective_group_id
|
||||
JOIN organizations ON organizations.id = groups.organization_id
|
||||
WHERE groups.organization_id = @organization_id
|
||||
AND tu.created_at >= @period_start::timestamptz
|
||||
AND tu.created_at < @period_end::timestamptz
|
||||
GROUP BY
|
||||
ai.initiator_id,
|
||||
users.username,
|
||||
tu.effective_group_id,
|
||||
groups.name,
|
||||
groups.organization_id,
|
||||
organizations.name,
|
||||
ai.model,
|
||||
ai.provider,
|
||||
ai.provider_name
|
||||
ORDER BY ai.initiator_id, tu.effective_group_id, ai.provider, ai.provider_name, ai.model;
|
||||
|
||||
Reference in New Issue
Block a user