mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: record cost on aibridge token usages (#26229)
Implements https://linear.app/codercom/issue/AIGOV-286/add-interception-cost-calculation-to-aibridge-token-usages Adds spend attribution to AI Gateway. After the upstream response, each token-usage record now captures the user's effective group, the per-token prices in effect at that moment, and a computed cost — so spend is recorded as an immutable, point-in-time snapshot. Concretely, `aibridge_token_usages` gains `effective_group_id`, `input_price_micros`, `output_price_micros`, `cache_read_price_micros`, `cache_write_price_micros`, and `cost_micros`. When a usage record is written, the effective group is resolved (per-user override, else the deployment budget policy), the `(provider, model)` price is looked up and snapshotted onto the row, and cost is computed from the provider-reported token counts. A model that isn't in the price table records its tokens with a `NULL` cost; any *other* resolution failure fails the write, so a `NULL` cost unambiguously means "model not priced" rather than "lookup errored." All values are stored in micro-units (1 unit = 1,000,000 micro-units; Phase 1 assumes USD, so 1 micro-unit = $0.000001). Prices are quoted per million tokens. This also grants the AI Bridge RBAC subject `read` on `ai_model_prices` (the per-interception price lookup needs it; it previously only had `update` for the startup seeder). ## Cost precision Cost is computed per token category as `tokens × price / 1_000_000` with integer division, then the four categories are summed. The division is done **per category** (not once over the summed numerator) on purpose: it keeps the per-category line items summing exactly to the stored total — no "the parts don't add up to the whole" in reporting). Integer division truncates sub-micro-unit fractions. For example, a cheap model at $0.10 per million tokens is a price of `100_000`; 9 tokens cost `9 × 100_000 / 1_000_000 = 900_000 / 1_000_000 = 0` (the true 0.9 micro-units floors to 0). At real list prices this rarely bites — $3/M input is a price of `3_000_000`, so even a single token is 3 micro-units. The per-record under-count is bounded below 1 micro-unit per category, so under $0.000004 total across the four categories, which is acceptable for list-price-based cost approximation. ## Overflow safety `cost_micros` is a `BIGINT` (int64), and the largest intermediate value is a single category's `tokens × price` before division. int64's ceiling is ≈ `9.223e18`. - At a steep $75/M model (price `75_000_000`), overflow would require ~123 billion tokens in one response: `123e9 × 75e6 = 9.225e18`, just over the limit. `122e9` stays under at `9.15e18`. - A realistically maxed-out Opus 4.8 response (≈1M input + 128K output at list prices) costs about $15, with a numerator around `1.5e13` — roughly six orders of magnitude below the ceiling. So overflow is unreachable from real token counts. ### Multi-currency support In the future, we may encounter issues with multi-currency support, especially when dealing with currencies that have very large exchange rates relative to USD, for example: IRR: ~1,300,000 IRR ≈ 1 USD VND: ~26,000 VND ≈ 1 USD For currencies with such large denominations, numeric overflow is technically possible, considering that we have only about six orders of magnitude of headroom before reaching the limit (see above). ## `effective_group_id` has no foreign key `effective_group_id` records the group a spend was attributed to, as an immutable historical fact. It is intentionally **not** a foreign key, so the record survives deletion of the group. Alternatives were considered and rejected: - **`ON DELETE SET NULL`** would mutate an "immutable" record — deleting a group silently erases that interception's attribution and under-counts the group's historical spend. - **`RESTRICT` / `NO ACTION`** would block group deletion entirely (groups are hard-deleted). - **`CASCADE`** would delete spend history when a group is deleted — the worst outcome for an audit record. There is also no insert-time check that the group still exists: the id comes from a budget that was just resolved, meaning it was valid at some point. ## Open question: group name snapshotting Should we also snapshot the group *name* onto each record? Two options: - **Denormalize it now** — readable in historical reports even after a group is deleted, but the snapshot can drift from the current name on rename, raising a "show point-in-time vs. current name" question. - **Postpone until needed** — it's a purely additive column later, and the name is display-only (not correctness-bearing like the price). The cost: names of groups deleted before the column is added can't be backfilled. Leaning toward postponing until a concrete reporting need settles the drift question.
This commit is contained in:
Yevhenii Shcherbina
Generated
+52
-47
@@ -6,51 +6,56 @@ type CheckConstraint string
|
||||
|
||||
// CheckConstraint enums.
|
||||
const (
|
||||
CheckAIGatewayKeysHashedSecretCheck CheckConstraint = "ai_gateway_keys_hashed_secret_check" // ai_gateway_keys
|
||||
CheckAIGatewayKeysNameCheck CheckConstraint = "ai_gateway_keys_name_check" // ai_gateway_keys
|
||||
CheckAIGatewayKeysSecretPrefixCheck CheckConstraint = "ai_gateway_keys_secret_prefix_check" // ai_gateway_keys
|
||||
CheckAIModelPricesCacheReadPriceCheck CheckConstraint = "ai_model_prices_cache_read_price_check" // ai_model_prices
|
||||
CheckAIModelPricesCacheWritePriceCheck CheckConstraint = "ai_model_prices_cache_write_price_check" // ai_model_prices
|
||||
CheckAIModelPricesInputPriceCheck CheckConstraint = "ai_model_prices_input_price_check" // ai_model_prices
|
||||
CheckAIModelPricesOutputPriceCheck CheckConstraint = "ai_model_prices_output_price_check" // ai_model_prices
|
||||
CheckAIProvidersNameCheck CheckConstraint = "ai_providers_name_check" // ai_providers
|
||||
CheckAPIKeysAllowListNotEmpty CheckConstraint = "api_keys_allow_list_not_empty" // api_keys
|
||||
CheckBoundaryLogsSequenceNumberCheck CheckConstraint = "boundary_logs_sequence_number_check" // boundary_logs
|
||||
CheckChatModelConfigsAIProviderRequiredWhenActive CheckConstraint = "chat_model_configs_ai_provider_required_when_active" // chat_model_configs
|
||||
CheckChatModelConfigsCompressionThresholdCheck CheckConstraint = "chat_model_configs_compression_threshold_check" // chat_model_configs
|
||||
CheckChatModelConfigsContextLimitCheck CheckConstraint = "chat_model_configs_context_limit_check" // chat_model_configs
|
||||
CheckChatUsageLimitConfigDefaultLimitMicrosCheck CheckConstraint = "chat_usage_limit_config_default_limit_micros_check" // chat_usage_limit_config
|
||||
CheckChatUsageLimitConfigPeriodCheck CheckConstraint = "chat_usage_limit_config_period_check" // chat_usage_limit_config
|
||||
CheckChatUsageLimitConfigSingletonCheck CheckConstraint = "chat_usage_limit_config_singleton_check" // chat_usage_limit_config
|
||||
CheckChatAclOnlyOnRootChats CheckConstraint = "chat_acl_only_on_root_chats" // chats
|
||||
CheckChatGroupAclNotNullJsonb CheckConstraint = "chat_group_acl_not_null_jsonb" // chats
|
||||
CheckChatUserAclNotNullJsonb CheckConstraint = "chat_user_acl_not_null_jsonb" // chats
|
||||
CheckChatsPinOrderArchivedCheck CheckConstraint = "chats_pin_order_archived_check" // chats
|
||||
CheckChatsPinOrderParentCheck CheckConstraint = "chats_pin_order_parent_check" // chats
|
||||
CheckOneTimePasscodeSet CheckConstraint = "one_time_passcode_set" // users
|
||||
CheckUsersChatSpendLimitMicrosCheck CheckConstraint = "users_chat_spend_limit_micros_check" // users
|
||||
CheckUsersEmailNotEmpty CheckConstraint = "users_email_not_empty" // users
|
||||
CheckUsersServiceAccountLoginType CheckConstraint = "users_service_account_login_type" // users
|
||||
CheckUsersUsernameMinLength CheckConstraint = "users_username_min_length" // users
|
||||
CheckOrganizationIDNotZero CheckConstraint = "organization_id_not_zero" // custom_roles
|
||||
CheckGroupAIBudgetsSpendLimitMicrosCheck CheckConstraint = "group_ai_budgets_spend_limit_micros_check" // group_ai_budgets
|
||||
CheckGroupsChatSpendLimitMicrosCheck CheckConstraint = "groups_chat_spend_limit_micros_check" // groups
|
||||
CheckMcpServerConfigsAuthTypeCheck CheckConstraint = "mcp_server_configs_auth_type_check" // mcp_server_configs
|
||||
CheckMcpServerConfigsAvailabilityCheck CheckConstraint = "mcp_server_configs_availability_check" // mcp_server_configs
|
||||
CheckMcpServerConfigsTransportCheck CheckConstraint = "mcp_server_configs_transport_check" // mcp_server_configs
|
||||
CheckMaxProvisionerLogsLength CheckConstraint = "max_provisioner_logs_length" // provisioner_jobs
|
||||
CheckMaxLogsLength CheckConstraint = "max_logs_length" // workspace_agents
|
||||
CheckSubsystemsNotNone CheckConstraint = "subsystems_not_none" // workspace_agents
|
||||
CheckWorkspaceBuildsDeadlineBelowMaxDeadline CheckConstraint = "workspace_builds_deadline_below_max_deadline" // workspace_builds
|
||||
CheckGroupAclIsObject CheckConstraint = "group_acl_is_object" // workspaces
|
||||
CheckUserAclIsObject CheckConstraint = "user_acl_is_object" // workspaces
|
||||
CheckTelemetryLockEventTypeConstraint CheckConstraint = "telemetry_lock_event_type_constraint" // telemetry_locks
|
||||
CheckValidationMonotonicOrder CheckConstraint = "validation_monotonic_order" // template_version_parameters
|
||||
CheckUsageEventTypeCheck CheckConstraint = "usage_event_type_check" // usage_events
|
||||
CheckUserAIBudgetOverridesSpendLimitMicrosCheck CheckConstraint = "user_ai_budget_overrides_spend_limit_micros_check" // user_ai_budget_overrides
|
||||
CheckUserAIProviderKeysAPIKeyCheck CheckConstraint = "user_ai_provider_keys_api_key_check" // user_ai_provider_keys
|
||||
CheckUserSkillsContentSize CheckConstraint = "user_skills_content_size" // user_skills
|
||||
CheckUserSkillsDescriptionSize CheckConstraint = "user_skills_description_size" // user_skills
|
||||
CheckUserSkillsNameFormat CheckConstraint = "user_skills_name_format" // user_skills
|
||||
CheckUserSkillsNameSize CheckConstraint = "user_skills_name_size" // user_skills
|
||||
CheckAIGatewayKeysHashedSecretCheck CheckConstraint = "ai_gateway_keys_hashed_secret_check" // ai_gateway_keys
|
||||
CheckAIGatewayKeysNameCheck CheckConstraint = "ai_gateway_keys_name_check" // ai_gateway_keys
|
||||
CheckAIGatewayKeysSecretPrefixCheck CheckConstraint = "ai_gateway_keys_secret_prefix_check" // ai_gateway_keys
|
||||
CheckAIModelPricesCacheReadPriceCheck CheckConstraint = "ai_model_prices_cache_read_price_check" // ai_model_prices
|
||||
CheckAIModelPricesCacheWritePriceCheck CheckConstraint = "ai_model_prices_cache_write_price_check" // ai_model_prices
|
||||
CheckAIModelPricesInputPriceCheck CheckConstraint = "ai_model_prices_input_price_check" // ai_model_prices
|
||||
CheckAIModelPricesOutputPriceCheck CheckConstraint = "ai_model_prices_output_price_check" // ai_model_prices
|
||||
CheckAIProvidersNameCheck CheckConstraint = "ai_providers_name_check" // ai_providers
|
||||
CheckAibridgeTokenUsagesCacheReadPriceMicrosCheck CheckConstraint = "aibridge_token_usages_cache_read_price_micros_check" // aibridge_token_usages
|
||||
CheckAibridgeTokenUsagesCacheWritePriceMicrosCheck CheckConstraint = "aibridge_token_usages_cache_write_price_micros_check" // aibridge_token_usages
|
||||
CheckAibridgeTokenUsagesCostMicrosCheck CheckConstraint = "aibridge_token_usages_cost_micros_check" // aibridge_token_usages
|
||||
CheckAibridgeTokenUsagesInputPriceMicrosCheck CheckConstraint = "aibridge_token_usages_input_price_micros_check" // aibridge_token_usages
|
||||
CheckAibridgeTokenUsagesOutputPriceMicrosCheck CheckConstraint = "aibridge_token_usages_output_price_micros_check" // aibridge_token_usages
|
||||
CheckAPIKeysAllowListNotEmpty CheckConstraint = "api_keys_allow_list_not_empty" // api_keys
|
||||
CheckBoundaryLogsSequenceNumberCheck CheckConstraint = "boundary_logs_sequence_number_check" // boundary_logs
|
||||
CheckChatModelConfigsAIProviderRequiredWhenActive CheckConstraint = "chat_model_configs_ai_provider_required_when_active" // chat_model_configs
|
||||
CheckChatModelConfigsCompressionThresholdCheck CheckConstraint = "chat_model_configs_compression_threshold_check" // chat_model_configs
|
||||
CheckChatModelConfigsContextLimitCheck CheckConstraint = "chat_model_configs_context_limit_check" // chat_model_configs
|
||||
CheckChatUsageLimitConfigDefaultLimitMicrosCheck CheckConstraint = "chat_usage_limit_config_default_limit_micros_check" // chat_usage_limit_config
|
||||
CheckChatUsageLimitConfigPeriodCheck CheckConstraint = "chat_usage_limit_config_period_check" // chat_usage_limit_config
|
||||
CheckChatUsageLimitConfigSingletonCheck CheckConstraint = "chat_usage_limit_config_singleton_check" // chat_usage_limit_config
|
||||
CheckChatAclOnlyOnRootChats CheckConstraint = "chat_acl_only_on_root_chats" // chats
|
||||
CheckChatGroupAclNotNullJsonb CheckConstraint = "chat_group_acl_not_null_jsonb" // chats
|
||||
CheckChatUserAclNotNullJsonb CheckConstraint = "chat_user_acl_not_null_jsonb" // chats
|
||||
CheckChatsPinOrderArchivedCheck CheckConstraint = "chats_pin_order_archived_check" // chats
|
||||
CheckChatsPinOrderParentCheck CheckConstraint = "chats_pin_order_parent_check" // chats
|
||||
CheckOneTimePasscodeSet CheckConstraint = "one_time_passcode_set" // users
|
||||
CheckUsersChatSpendLimitMicrosCheck CheckConstraint = "users_chat_spend_limit_micros_check" // users
|
||||
CheckUsersEmailNotEmpty CheckConstraint = "users_email_not_empty" // users
|
||||
CheckUsersServiceAccountLoginType CheckConstraint = "users_service_account_login_type" // users
|
||||
CheckUsersUsernameMinLength CheckConstraint = "users_username_min_length" // users
|
||||
CheckOrganizationIDNotZero CheckConstraint = "organization_id_not_zero" // custom_roles
|
||||
CheckGroupAIBudgetsSpendLimitMicrosCheck CheckConstraint = "group_ai_budgets_spend_limit_micros_check" // group_ai_budgets
|
||||
CheckGroupsChatSpendLimitMicrosCheck CheckConstraint = "groups_chat_spend_limit_micros_check" // groups
|
||||
CheckMcpServerConfigsAuthTypeCheck CheckConstraint = "mcp_server_configs_auth_type_check" // mcp_server_configs
|
||||
CheckMcpServerConfigsAvailabilityCheck CheckConstraint = "mcp_server_configs_availability_check" // mcp_server_configs
|
||||
CheckMcpServerConfigsTransportCheck CheckConstraint = "mcp_server_configs_transport_check" // mcp_server_configs
|
||||
CheckMaxProvisionerLogsLength CheckConstraint = "max_provisioner_logs_length" // provisioner_jobs
|
||||
CheckMaxLogsLength CheckConstraint = "max_logs_length" // workspace_agents
|
||||
CheckSubsystemsNotNone CheckConstraint = "subsystems_not_none" // workspace_agents
|
||||
CheckWorkspaceBuildsDeadlineBelowMaxDeadline CheckConstraint = "workspace_builds_deadline_below_max_deadline" // workspace_builds
|
||||
CheckGroupAclIsObject CheckConstraint = "group_acl_is_object" // workspaces
|
||||
CheckUserAclIsObject CheckConstraint = "user_acl_is_object" // workspaces
|
||||
CheckTelemetryLockEventTypeConstraint CheckConstraint = "telemetry_lock_event_type_constraint" // telemetry_locks
|
||||
CheckValidationMonotonicOrder CheckConstraint = "validation_monotonic_order" // template_version_parameters
|
||||
CheckUsageEventTypeCheck CheckConstraint = "usage_event_type_check" // usage_events
|
||||
CheckUserAIBudgetOverridesSpendLimitMicrosCheck CheckConstraint = "user_ai_budget_overrides_spend_limit_micros_check" // user_ai_budget_overrides
|
||||
CheckUserAIProviderKeysAPIKeyCheck CheckConstraint = "user_ai_provider_keys_api_key_check" // user_ai_provider_keys
|
||||
CheckUserSkillsContentSize CheckConstraint = "user_skills_content_size" // user_skills
|
||||
CheckUserSkillsDescriptionSize CheckConstraint = "user_skills_description_size" // user_skills
|
||||
CheckUserSkillsNameFormat CheckConstraint = "user_skills_name_format" // user_skills
|
||||
CheckUserSkillsNameSize CheckConstraint = "user_skills_name_size" // user_skills
|
||||
)
|
||||
|
||||
@@ -650,9 +650,9 @@ var (
|
||||
},
|
||||
rbac.ResourceApiKey.Type: {policy.ActionRead}, // Validate API keys.
|
||||
rbac.ResourceAibridgeInterception.Type: {policy.ActionCreate, policy.ActionRead, policy.ActionUpdate, policy.ActionDelete},
|
||||
rbac.ResourceAiModelPrice.Type: {policy.ActionUpdate}, // Required for the startup price seeder.
|
||||
rbac.ResourceAiSeat.Type: {policy.ActionCreate}, // Required for UpsertAISeatState.
|
||||
rbac.ResourceAIProvider.Type: {policy.ActionRead}, // Required to load the provider snapshot (and per-provider keys) at startup.
|
||||
rbac.ResourceAiModelPrice.Type: {policy.ActionRead, policy.ActionUpdate}, // Read: per-interception cost lookup. Update: startup price seeder.
|
||||
rbac.ResourceAiSeat.Type: {policy.ActionCreate}, // Required for UpsertAISeatState.
|
||||
rbac.ResourceAIProvider.Type: {policy.ActionRead}, // Required to load the provider snapshot (and per-provider keys) at startup.
|
||||
}),
|
||||
User: []rbac.Permission{},
|
||||
ByOrgID: map[string]rbac.OrgPermissions{},
|
||||
|
||||
@@ -2037,6 +2037,12 @@ func AIBridgeTokenUsage(t testing.TB, db database.Store, seed database.InsertAIB
|
||||
CacheWriteInputTokens: seed.CacheWriteInputTokens,
|
||||
Metadata: takeFirstSlice(seed.Metadata, json.RawMessage("{}")),
|
||||
CreatedAt: takeFirst(seed.CreatedAt, dbtime.Now()),
|
||||
EffectiveGroupID: seed.EffectiveGroupID,
|
||||
InputPriceMicros: seed.InputPriceMicros,
|
||||
OutputPriceMicros: seed.OutputPriceMicros,
|
||||
CacheReadPriceMicros: seed.CacheReadPriceMicros,
|
||||
CacheWritePriceMicros: seed.CacheWritePriceMicros,
|
||||
CostMicros: seed.CostMicros,
|
||||
})
|
||||
require.NoError(t, err, "insert aibridge token usage")
|
||||
return usage
|
||||
|
||||
Generated
+12
-1
@@ -1576,7 +1576,18 @@ CREATE TABLE aibridge_token_usages (
|
||||
metadata jsonb,
|
||||
created_at timestamp with time zone NOT NULL,
|
||||
cache_read_input_tokens bigint DEFAULT 0 NOT NULL,
|
||||
cache_write_input_tokens bigint DEFAULT 0 NOT NULL
|
||||
cache_write_input_tokens bigint DEFAULT 0 NOT NULL,
|
||||
effective_group_id uuid,
|
||||
input_price_micros bigint,
|
||||
output_price_micros bigint,
|
||||
cache_read_price_micros bigint,
|
||||
cache_write_price_micros bigint,
|
||||
cost_micros bigint,
|
||||
CONSTRAINT aibridge_token_usages_cache_read_price_micros_check CHECK ((cache_read_price_micros >= 0)),
|
||||
CONSTRAINT aibridge_token_usages_cache_write_price_micros_check CHECK ((cache_write_price_micros >= 0)),
|
||||
CONSTRAINT aibridge_token_usages_cost_micros_check CHECK ((cost_micros >= 0)),
|
||||
CONSTRAINT aibridge_token_usages_input_price_micros_check CHECK ((input_price_micros >= 0)),
|
||||
CONSTRAINT aibridge_token_usages_output_price_micros_check CHECK ((output_price_micros >= 0))
|
||||
);
|
||||
|
||||
COMMENT ON TABLE aibridge_token_usages IS 'Audit log of tokens used by intercepted requests in AI Bridge';
|
||||
|
||||
@@ -0,0 +1,7 @@
|
||||
ALTER TABLE aibridge_token_usages
|
||||
DROP COLUMN effective_group_id,
|
||||
DROP COLUMN input_price_micros,
|
||||
DROP COLUMN output_price_micros,
|
||||
DROP COLUMN cache_read_price_micros,
|
||||
DROP COLUMN cache_write_price_micros,
|
||||
DROP COLUMN cost_micros;
|
||||
@@ -0,0 +1,15 @@
|
||||
ALTER TABLE aibridge_token_usages
|
||||
-- Effective group this interception's spend is attributed to. NULL if the
|
||||
-- user has no effective group (no budget configured). Intentionally not a
|
||||
-- foreign key: this is an immutable historical attribution that must
|
||||
-- survive group deletion, so the id is retained even after the group is gone.
|
||||
ADD COLUMN effective_group_id UUID,
|
||||
-- Snapshotted prices at interception time, in micro-units per million
|
||||
-- tokens. NULL if the model is not present in ai_model_prices.
|
||||
ADD COLUMN input_price_micros BIGINT CHECK (input_price_micros >= 0),
|
||||
ADD COLUMN output_price_micros BIGINT CHECK (output_price_micros >= 0),
|
||||
ADD COLUMN cache_read_price_micros BIGINT CHECK (cache_read_price_micros >= 0),
|
||||
ADD COLUMN cache_write_price_micros BIGINT CHECK (cache_write_price_micros >= 0),
|
||||
-- Computed cost in micro-units at interception time. NULL if the model is
|
||||
-- not present in ai_model_prices.
|
||||
ADD COLUMN cost_micros BIGINT CHECK (cost_micros >= 0);
|
||||
Generated
+6
@@ -4571,6 +4571,12 @@ type AIBridgeTokenUsage struct {
|
||||
CreatedAt time.Time `db:"created_at" json:"created_at"`
|
||||
CacheReadInputTokens int64 `db:"cache_read_input_tokens" json:"cache_read_input_tokens"`
|
||||
CacheWriteInputTokens int64 `db:"cache_write_input_tokens" json:"cache_write_input_tokens"`
|
||||
EffectiveGroupID uuid.NullUUID `db:"effective_group_id" json:"effective_group_id"`
|
||||
InputPriceMicros sql.NullInt64 `db:"input_price_micros" json:"input_price_micros"`
|
||||
OutputPriceMicros sql.NullInt64 `db:"output_price_micros" json:"output_price_micros"`
|
||||
CacheReadPriceMicros sql.NullInt64 `db:"cache_read_price_micros" json:"cache_read_price_micros"`
|
||||
CacheWritePriceMicros sql.NullInt64 `db:"cache_write_price_micros" json:"cache_write_price_micros"`
|
||||
CostMicros sql.NullInt64 `db:"cost_micros" json:"cost_micros"`
|
||||
}
|
||||
|
||||
// Audit log of tool calls in intercepted requests in AI Bridge
|
||||
|
||||
Generated
+37
-5
@@ -1216,7 +1216,7 @@ func (q *sqlQuerier) GetAIBridgeInterceptions(ctx context.Context) ([]AIBridgeIn
|
||||
|
||||
const getAIBridgeTokenUsagesByInterceptionID = `-- name: GetAIBridgeTokenUsagesByInterceptionID :many
|
||||
SELECT
|
||||
id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at, cache_read_input_tokens, cache_write_input_tokens
|
||||
id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at, cache_read_input_tokens, cache_write_input_tokens, effective_group_id, input_price_micros, output_price_micros, cache_read_price_micros, cache_write_price_micros, cost_micros
|
||||
FROM
|
||||
aibridge_token_usages WHERE interception_id = $1::uuid
|
||||
ORDER BY
|
||||
@@ -1243,6 +1243,12 @@ func (q *sqlQuerier) GetAIBridgeTokenUsagesByInterceptionID(ctx context.Context,
|
||||
&i.CreatedAt,
|
||||
&i.CacheReadInputTokens,
|
||||
&i.CacheWriteInputTokens,
|
||||
&i.EffectiveGroupID,
|
||||
&i.InputPriceMicros,
|
||||
&i.OutputPriceMicros,
|
||||
&i.CacheReadPriceMicros,
|
||||
&i.CacheWritePriceMicros,
|
||||
&i.CostMicros,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -1452,11 +1458,13 @@ func (q *sqlQuerier) InsertAIBridgeModelThought(ctx context.Context, arg InsertA
|
||||
|
||||
const insertAIBridgeTokenUsage = `-- name: InsertAIBridgeTokenUsage :one
|
||||
INSERT INTO aibridge_token_usages (
|
||||
id, interception_id, provider_response_id, input_tokens, output_tokens, cache_read_input_tokens, cache_write_input_tokens, metadata, created_at
|
||||
id, interception_id, provider_response_id, input_tokens, output_tokens, cache_read_input_tokens, cache_write_input_tokens, metadata, created_at,
|
||||
effective_group_id, input_price_micros, output_price_micros, cache_read_price_micros, cache_write_price_micros, cost_micros
|
||||
) VALUES (
|
||||
$1, $2, $3, $4, $5, $6, $7, COALESCE($8::jsonb, '{}'::jsonb), $9
|
||||
$1, $2, $3, $4, $5, $6, $7, COALESCE($8::jsonb, '{}'::jsonb), $9,
|
||||
$10, $11, $12, $13, $14, $15
|
||||
)
|
||||
RETURNING id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at, cache_read_input_tokens, cache_write_input_tokens
|
||||
RETURNING id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at, cache_read_input_tokens, cache_write_input_tokens, effective_group_id, input_price_micros, output_price_micros, cache_read_price_micros, cache_write_price_micros, cost_micros
|
||||
`
|
||||
|
||||
type InsertAIBridgeTokenUsageParams struct {
|
||||
@@ -1469,6 +1477,12 @@ type InsertAIBridgeTokenUsageParams struct {
|
||||
CacheWriteInputTokens int64 `db:"cache_write_input_tokens" json:"cache_write_input_tokens"`
|
||||
Metadata json.RawMessage `db:"metadata" json:"metadata"`
|
||||
CreatedAt time.Time `db:"created_at" json:"created_at"`
|
||||
EffectiveGroupID uuid.NullUUID `db:"effective_group_id" json:"effective_group_id"`
|
||||
InputPriceMicros sql.NullInt64 `db:"input_price_micros" json:"input_price_micros"`
|
||||
OutputPriceMicros sql.NullInt64 `db:"output_price_micros" json:"output_price_micros"`
|
||||
CacheReadPriceMicros sql.NullInt64 `db:"cache_read_price_micros" json:"cache_read_price_micros"`
|
||||
CacheWritePriceMicros sql.NullInt64 `db:"cache_write_price_micros" json:"cache_write_price_micros"`
|
||||
CostMicros sql.NullInt64 `db:"cost_micros" json:"cost_micros"`
|
||||
}
|
||||
|
||||
func (q *sqlQuerier) InsertAIBridgeTokenUsage(ctx context.Context, arg InsertAIBridgeTokenUsageParams) (AIBridgeTokenUsage, error) {
|
||||
@@ -1482,6 +1496,12 @@ func (q *sqlQuerier) InsertAIBridgeTokenUsage(ctx context.Context, arg InsertAIB
|
||||
arg.CacheWriteInputTokens,
|
||||
arg.Metadata,
|
||||
arg.CreatedAt,
|
||||
arg.EffectiveGroupID,
|
||||
arg.InputPriceMicros,
|
||||
arg.OutputPriceMicros,
|
||||
arg.CacheReadPriceMicros,
|
||||
arg.CacheWritePriceMicros,
|
||||
arg.CostMicros,
|
||||
)
|
||||
var i AIBridgeTokenUsage
|
||||
err := row.Scan(
|
||||
@@ -1494,6 +1514,12 @@ func (q *sqlQuerier) InsertAIBridgeTokenUsage(ctx context.Context, arg InsertAIB
|
||||
&i.CreatedAt,
|
||||
&i.CacheReadInputTokens,
|
||||
&i.CacheWriteInputTokens,
|
||||
&i.EffectiveGroupID,
|
||||
&i.InputPriceMicros,
|
||||
&i.OutputPriceMicros,
|
||||
&i.CacheReadPriceMicros,
|
||||
&i.CacheWritePriceMicros,
|
||||
&i.CostMicros,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
@@ -2162,7 +2188,7 @@ func (q *sqlQuerier) ListAIBridgeSessions(ctx context.Context, arg ListAIBridgeS
|
||||
|
||||
const listAIBridgeTokenUsagesByInterceptionIDs = `-- name: ListAIBridgeTokenUsagesByInterceptionIDs :many
|
||||
SELECT
|
||||
id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at, cache_read_input_tokens, cache_write_input_tokens
|
||||
id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at, cache_read_input_tokens, cache_write_input_tokens, effective_group_id, input_price_micros, output_price_micros, cache_read_price_micros, cache_write_price_micros, cost_micros
|
||||
FROM
|
||||
aibridge_token_usages
|
||||
WHERE
|
||||
@@ -2191,6 +2217,12 @@ func (q *sqlQuerier) ListAIBridgeTokenUsagesByInterceptionIDs(ctx context.Contex
|
||||
&i.CreatedAt,
|
||||
&i.CacheReadInputTokens,
|
||||
&i.CacheWriteInputTokens,
|
||||
&i.EffectiveGroupID,
|
||||
&i.InputPriceMicros,
|
||||
&i.OutputPriceMicros,
|
||||
&i.CacheReadPriceMicros,
|
||||
&i.CacheWritePriceMicros,
|
||||
&i.CostMicros,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
@@ -38,9 +38,11 @@ WHERE aibridge_interceptions.id = (
|
||||
|
||||
-- name: InsertAIBridgeTokenUsage :one
|
||||
INSERT INTO aibridge_token_usages (
|
||||
id, interception_id, provider_response_id, input_tokens, output_tokens, cache_read_input_tokens, cache_write_input_tokens, metadata, created_at
|
||||
id, interception_id, provider_response_id, input_tokens, output_tokens, cache_read_input_tokens, cache_write_input_tokens, metadata, created_at,
|
||||
effective_group_id, input_price_micros, output_price_micros, cache_read_price_micros, cache_write_price_micros, cost_micros
|
||||
) VALUES (
|
||||
@id, @interception_id, @provider_response_id, @input_tokens, @output_tokens, @cache_read_input_tokens, @cache_write_input_tokens, COALESCE(@metadata::jsonb, '{}'::jsonb), @created_at
|
||||
@id, @interception_id, @provider_response_id, @input_tokens, @output_tokens, @cache_read_input_tokens, @cache_write_input_tokens, COALESCE(@metadata::jsonb, '{}'::jsonb), @created_at,
|
||||
@effective_group_id, @input_price_micros, @output_price_micros, @cache_read_price_micros, @cache_write_price_micros, @cost_micros
|
||||
)
|
||||
RETURNING *;
|
||||
|
||||
|
||||
Reference in New Issue
Block a user