codersdk/aibridge.go
Rama main · réplica de lectura del repositorio original.
package codersdk
import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"strconv"
"strings"
"time"
"github.com/google/uuid"
"golang.org/x/xerrors"
"github.com/coder/coder/v2/coderd/util/slice"
)
// MaxAISpendLimitMicros is the highest AI spend limit that can be configured,
// $1,000,000 per member per budget period.
const MaxAISpendLimitMicros int64 = 1_000_000_000_000
// MaxAISpendPeriodDays bounds explicit AI spend reporting windows.
const MaxAISpendPeriodDays = 31
// AIBudgetLimitSource identifies which tier produced the user's
// effective budget limit.
type AIBudgetLimitSource string
const (
// AIBudgetLimitSourceUserOverride indicates the limit came from a
// per-user override.
AIBudgetLimitSourceUserOverride AIBudgetLimitSource = "user_override"
// AIBudgetLimitSourceGroup indicates the limit came from a group
// budget selected by the deployment budget policy.
AIBudgetLimitSourceGroup AIBudgetLimitSource = "group"
)
// AIBudgetLimit is an AI spend limit and the tier that produced it. Both
// fields are always populated together.
type AIBudgetLimit struct {
SpendLimitMicros int64 `json:"spend_limit_micros"`
LimitSource AIBudgetLimitSource `json:"limit_source"`
}
// UserAIBudgetSummary is the effective AI budget for a user. When no budget
// applies, the effective group falls back to the Everyone group with a null
// budget.
type UserAIBudgetSummary struct {
UserID uuid.UUID `json:"user_id" format:"uuid"`
// EffectiveGroupID is the group the spend is attributed to, falling back to
// the Everyone group when no budget applies. Null only when the user has no
// organization membership.
EffectiveGroupID *uuid.UUID `json:"effective_group_id" format:"uuid"`
// EffectiveBudget is the spend limit that applies to the user, whether it
// came from a group budget or a user override. Null when no budget
// applies, leaving the user's spend unlimited.
EffectiveBudget *AIBudgetLimit `json:"effective_budget"`
}
// AISpendPeriodWindow is the [Start, End) window over which AI spend is
// aggregated.
type AISpendPeriodWindow struct {
// PeriodStart is the inclusive lower bound of the current budget
// period.
PeriodStart time.Time `json:"period_start" format:"date-time"`
// PeriodEnd is the exclusive upper bound of the current budget
// period.
PeriodEnd time.Time `json:"period_end" format:"date-time"`
}
// UserAISpendStatus is the current AI spend snapshot for a user within
// the active budget period.
type UserAISpendStatus struct {
UserAIBudgetSummary
AISpendPeriodWindow
// CurrentSpendMicros is the user's spend on their effective group over
// the current budget period.
CurrentSpendMicros int64 `json:"current_spend_micros"`
}
// OrganizationGroupsAISpend reports AI spend for a set of groups in the
// active budget period.
type OrganizationGroupsAISpend struct {
AISpendPeriodWindow
Groups []OrganizationGroupAISpend `json:"groups"`
}
// OrganizationGroupAISpend is the current AI spend snapshot for a group
// within the active budget period.
type OrganizationGroupAISpend struct {
GroupID uuid.UUID `json:"group_id" format:"uuid"`
// SpendLimitMicros is the group's configured AI spend budget per member.
// Null when the group has no configured budget.
SpendLimitMicros *int64 `json:"spend_limit_micros"`
// TotalSpendLimitMicros is the currently configured combined budget of the
// members attributed to this group, with each member's override replacing
// their share. Null when the group has no budget, and zero when no members
// are attributed to it.
TotalSpendLimitMicros *int64 `json:"total_spend_limit_micros"`
// CurrentSpendMicros is the group's spend over the current budget period.
CurrentSpendMicros int64 `json:"current_spend_micros"`
}
// GroupAISpend is the current AI spend snapshot for a single group within
// the active budget period.
type GroupAISpend struct {
AISpendPeriodWindow
OrganizationGroupAISpend
}
// GroupMembersAISpend reports per-member AI spend attributed to a specific
// group in the active budget period.
type GroupMembersAISpend struct {
AISpendPeriodWindow
Members []GroupMemberAISpend `json:"members"`
}
// GroupMemberAISpend is a single member's AI spend attributed to the queried
// group in the current budget period.
type GroupMemberAISpend struct {
UserID uuid.UUID `json:"user_id" format:"uuid"`
// EffectiveGroupID is the user's effective budget group within the queried
// group's organization, falling back to the Everyone group when no budget
// applies. Null when the effective group belongs to a different organization
// than the queried group.
EffectiveGroupID *uuid.UUID `json:"effective_group_id" format:"uuid"`
// EffectiveBudget is the spend limit that currently applies to the user.
// Null when no budget applies or the effective group belongs to a different
// organization than the queried group.
EffectiveBudget *AIBudgetLimit `json:"effective_budget"`
// GroupBudget is the budget when the queried group is this user's
// effective budget source. When populated, it matches EffectiveBudget. Null
// when the user's budget resolves to another group or no budget applies.
// Deprecated: Use EffectiveBudget instead.
GroupBudget *AIBudgetLimit `json:"group_budget"`
// GroupSpendMicros is the user's spend attributed to the queried group
// over the current budget period.
GroupSpendMicros int64 `json:"group_spend_micros"`
}
type AIBridgeSession struct {
ID string `json:"id"`
Initiator MinimalUser `json:"initiator"`
Providers []string `json:"providers"`
Models []string `json:"models"`
Client *string `json:"client"`
Metadata map[string]any `json:"metadata"`
StartedAt time.Time `json:"started_at" format:"date-time"`
EndedAt *time.Time `json:"ended_at,omitempty" format:"date-time"`
Threads int64 `json:"threads"`
TokenUsageSummary AIBridgeSessionTokenUsageSummary `json:"token_usage_summary"`
// NetworkCalls summarizes the Agent Firewall network requests made during the
// session. A nil value means the session did not pass through Agent
// Firewall, so network call monitoring was not active, which the UI
// surfaces as "Disabled".
NetworkCalls *AIBridgeSessionNetworkCallSummary `json:"network_calls,omitempty"`
LastPrompt *string `json:"last_prompt,omitempty"`
LastActiveAt time.Time `json:"last_active_at" format:"date-time"`
}
type AIBridgeSessionTokenUsageSummary struct {
InputTokens int64 `json:"input_tokens"`
OutputTokens int64 `json:"output_tokens"`
CacheReadInputTokens int64 `json:"cache_read_input_tokens"`
CacheWriteInputTokens int64 `json:"cache_write_input_tokens"`
}
// AIBridgeSessionNetworkCallSummary aggregates the Agent Firewall network
// calls made during a session. Blocked counts calls denied by the firewall
// allow-list.
type AIBridgeSessionNetworkCallSummary struct {
Total int64 `json:"total"`
Blocked int64 `json:"blocked"`
}
// AIBridgeSessionNetworkDomain is one destination host contacted during a
// session, with the number of network calls made to it.
type AIBridgeSessionNetworkDomain struct {
Domain string `json:"domain"`
Count int64 `json:"count"`
}
type AIBridgeListSessionsResponse struct {
Count int64 `json:"count"`
Sessions []AIBridgeSession `json:"sessions"`
}
// AIBridgeSessionThreadsResponse is the response for GET
// /api/v2/ai-gateway/sessions/{session_id} which returns a single
// session with fully expanded threads.
type AIBridgeSessionThreadsResponse struct {
ID string `json:"id"`
Initiator MinimalUser `json:"initiator"`
Providers []string `json:"providers"`
Models []string `json:"models"`
Client *string `json:"client,omitempty"`
Metadata map[string]any `json:"metadata"`
PageStartedAt *time.Time `json:"page_started_at,omitempty" format:"date-time"`
PageEndedAt *time.Time `json:"page_ended_at,omitempty" format:"date-time"`
StartedAt time.Time `json:"started_at" format:"date-time"`
EndedAt *time.Time `json:"ended_at,omitempty" format:"date-time"`
TokenUsageSummary AIBridgeSessionThreadsTokenUsage `json:"token_usage_summary"`
// NetworkCalls summarizes the Agent Firewall network calls made during the
// session. A nil value means the session did not pass through Agent
// Firewall, so network call monitoring was not active, which the UI
// surfaces as "Disabled".
NetworkCalls *AIBridgeSessionNetworkCallSummary `json:"network_calls,omitempty"`
// NetworkTopDomains lists the most contacted destination hosts, ordered by
// call count descending. NetworkDomainCount is the total number of distinct
// domains, used to render a "+N more" overflow beyond the listed domains.
NetworkTopDomains []AIBridgeSessionNetworkDomain `json:"network_top_domains,omitempty"`
NetworkDomainCount int64 `json:"network_domain_count,omitempty"`
// NetworkCallLogs is the chronological list of individual network calls made
// during the session, holding the earliest calls up to a server-side cap.
// NetworkCalls remains authoritative for whole-session totals, so a shorter
// list than NetworkCalls.Total means the list was truncated. Empty when the
// session did not pass through Agent Firewall.
NetworkCallLogs []AgentFirewallLog `json:"network_call_logs,omitempty"`
Threads []AIBridgeThread `json:"threads"`
}
// AIBridgeSessionThreadsTokenUsage represents aggregated token usage
// with metadata containing provider-specific fields.
type AIBridgeSessionThreadsTokenUsage struct {
InputTokens int64 `json:"input_tokens"`
OutputTokens int64 `json:"output_tokens"`
CacheReadInputTokens int64 `json:"cache_read_input_tokens"`
CacheWriteInputTokens int64 `json:"cache_write_input_tokens"`
Metadata map[string]any `json:"metadata"`
// CostMicros is the summed cost of the priced usage, in micro-units
// (1000000 is $1.00).
CostMicros int64 `json:"cost_micros"`
// HasUnpricedUsage reports that some usage has no cost, so CostMicros
// is a lower bound.
HasUnpricedUsage bool `json:"has_unpriced_usage"`
}
// AIBridgePricedModel is the model and per-token prices used to compute the
// cost of an interception. Prices are integer micro-units per million
// tokens, so 1000000 is $1.00 per million tokens. A nil price means the
// price is not known and was treated as zero.
type AIBridgePricedModel struct {
// Model is the priced model, either the requested model or the model
// reported by the provider.
Model string `json:"model"`
InputPrice *int64 `json:"input_price"`
OutputPrice *int64 `json:"output_price"`
CacheReadPrice *int64 `json:"cache_read_price"`
CacheWritePrice *int64 `json:"cache_write_price"`
}
// AIBridgeAttribution contains the attribution fields recorded for one
// interception.
type AIBridgeAttribution map[string]string
// MarshalJSON encodes unknown attribution as an empty object.
func (a AIBridgeAttribution) MarshalJSON() ([]byte, error) {
if a == nil {
return []byte("{}"), nil
}
return json.Marshal(map[string]string(a))
}
// AIBridgeThread represents a single thread within a session.
// A thread groups interceptions by their thread_root_id.
type AIBridgeThread struct {
ID uuid.UUID `json:"id" format:"uuid"`
Prompt *string `json:"prompt,omitempty"`
Model string `json:"model"`
Provider string `json:"provider"`
CredentialKind string `json:"credential_kind"`
CredentialHint string `json:"credential_hint"`
StartedAt time.Time `json:"started_at" format:"date-time"`
EndedAt *time.Time `json:"ended_at,omitempty" format:"date-time"`
TokenUsage AIBridgeSessionThreadsTokenUsage `json:"token_usage"`
// PricedModel is the model and prices used for the root interception
// (the prompt). Other requests in the thread may use different models
// and prices, see agentic_actions[].priced_model. Nil when the root
// interception was not priced.
PricedModel *AIBridgePricedModel `json:"priced_model,omitempty"`
// Attribution contains attribution data from the root interception.
// Unknown attribution is serialized as an empty object.
Attribution AIBridgeAttribution `json:"attribution"`
AgenticActions []AIBridgeAgenticAction `json:"agentic_actions"`
// ErrorType is the categorized terminal upstream error from the root
// interception, or nil when the interception succeeded. See the
// aibridge_interception_error_type enum for possible values.
ErrorType *string `json:"error_type,omitempty"`
// ErrorMessage is the raw terminal upstream error message from the root
// interception. Nil when the interception succeeded.
ErrorMessage *string `json:"error_message,omitempty"`
// AgentFirewallSessionID links this thread to an agent firewall
// confinement session. Nil when the request did not pass through
// the agent firewall.
AgentFirewallSessionID *uuid.UUID `json:"agent_firewall_session_id,omitempty" format:"uuid"`
// AgentFirewallSequenceNumber is the firewall sequence number from
// the root interception. Used to determine the position of this
// LLM request in the firewall event stream. Nil when the request
// did not pass through the agent firewall.
AgentFirewallSequenceNumber *int32 `json:"agent_firewall_sequence_number,omitempty"`
}
// AIBridgeAgenticAction represents data from one interception, including
// tool calls, thinking blocks, and token usage. Tool-less child interceptions
// are represented as actions with an empty ToolCalls slice.
type AIBridgeAgenticAction struct {
InterceptionID uuid.UUID `json:"interception_id" format:"uuid"`
Model string `json:"model"`
// Attribution contains attribution data from this interception.
// Unknown attribution is serialized as an empty object.
Attribution AIBridgeAttribution `json:"attribution"`
TokenUsage AIBridgeSessionThreadsTokenUsage `json:"token_usage"`
// PricedModel is the model and prices used for this interception. Nil
// when it was not priced.
PricedModel *AIBridgePricedModel `json:"priced_model,omitempty"`
Thinking []AIBridgeModelThought `json:"thinking"`
ToolCalls []AIBridgeToolCall `json:"tool_calls"`
}
// AIBridgeModelThought represents a single thinking block from
// the model.
type AIBridgeModelThought struct {
Text string `json:"text"`
}
// AIBridgeToolCall represents a tool call recorded during an
// interception.
type AIBridgeToolCall struct {
ID uuid.UUID `json:"id" format:"uuid"`
InterceptionID uuid.UUID `json:"interception_id" format:"uuid"`
ProviderResponseID string `json:"provider_response_id"`
ServerURL string `json:"server_url"`
Tool string `json:"tool"`
Injected bool `json:"injected"`
Input string `json:"input"`
Metadata map[string]any `json:"metadata"`
CreatedAt time.Time `json:"created_at" format:"date-time"`
}
// @typescript-ignore AIBridgeListSessionsFilter
type AIBridgeListSessionsFilter struct {
// Limit defaults to 100, max is 1000.
Pagination Pagination `json:"pagination,omitempty"`
// Initiator is a user ID, username, or "me".
Initiator string `json:"initiator,omitempty"`
StartedBefore time.Time `json:"started_before,omitempty" format:"date-time"`
StartedAfter time.Time `json:"started_after,omitempty" format:"date-time"`
// Provider matches the runtime provider type column (openai,
// anthropic, copilot). The runtime type collapses the configured
// ai_provider_type: azure, google, openai-compat, openrouter, and
// vercel route through openai; bedrock routes through anthropic.
// Retained for backward compatibility; new clients should prefer
// ProviderName, which scopes to a specific configured row.
Provider string `json:"provider,omitempty"`
ProviderName string `json:"provider_name,omitempty"`
Model string `json:"model,omitempty"`
Client string `json:"client,omitempty"`
SessionID string `json:"session_id,omitempty"`
// AfterSessionID is a cursor for pagination. It is the session ID of the
// last session in the previous page.
AfterSessionID string `json:"after_session_id,omitempty"`
FilterQuery string `json:"q,omitempty"`
}
// asRequestOption returns a function that can be used in (*Client).Request.
func (f AIBridgeListSessionsFilter) asRequestOption() RequestOption {
return func(r *http.Request) {
var params []string
if f.Initiator != "" {
params = append(params, fmt.Sprintf("initiator:%q", f.Initiator))
}
if !f.StartedBefore.IsZero() {
params = append(params, fmt.Sprintf("started_before:%q", f.StartedBefore.Format(time.RFC3339Nano)))
}
if !f.StartedAfter.IsZero() {
params = append(params, fmt.Sprintf("started_after:%q", f.StartedAfter.Format(time.RFC3339Nano)))
}
if f.Provider != "" {
params = append(params, fmt.Sprintf("provider:%q", f.Provider))
}
if f.ProviderName != "" {
params = append(params, fmt.Sprintf("provider_name:%q", f.ProviderName))
}
if f.Model != "" {
params = append(params, fmt.Sprintf("model:%q", f.Model))
}
if f.Client != "" {
params = append(params, fmt.Sprintf("client:%q", f.Client))
}
if f.SessionID != "" {
params = append(params, fmt.Sprintf("session_id:%q", f.SessionID))
}
if f.FilterQuery != "" {
params = append(params, f.FilterQuery)
}
q := r.URL.Query()
q.Set("q", strings.Join(params, " "))
if f.AfterSessionID != "" {
q.Set("after_session_id", f.AfterSessionID)
}
r.URL.RawQuery = q.Encode()
}
}
// AIBridgeListSessions returns AI Bridge sessions with the given filter.
func (c *Client) AIBridgeListSessions(ctx context.Context, filter AIBridgeListSessionsFilter) (AIBridgeListSessionsResponse, error) {
res, err := c.Request(ctx, http.MethodGet, "/api/v2/ai-gateway/sessions", nil, filter.asRequestOption(), filter.Pagination.asRequestOption())
if err != nil {
return AIBridgeListSessionsResponse{}, err
}
defer res.Body.Close()
if res.StatusCode != http.StatusOK {
return AIBridgeListSessionsResponse{}, ReadBodyAsError(res)
}
var resp AIBridgeListSessionsResponse
return resp, ReadBodyAsJSON(res, &resp)
}
// AIBridgeGetSessionThreads returns a single session with expanded
// thread details including agentic actions and thinking blocks.
func (c *Client) AIBridgeGetSessionThreads(ctx context.Context, sessionID string, afterID, beforeID uuid.UUID, limit int32) (AIBridgeSessionThreadsResponse, error) {
res, err := c.Request(ctx, http.MethodGet, fmt.Sprintf("/api/v2/ai-gateway/sessions/%s", sessionID), nil, func(r *http.Request) {
q := r.URL.Query()
if afterID != uuid.Nil {
q.Set("after_id", afterID.String())
}
if beforeID != uuid.Nil {
q.Set("before_id", beforeID.String())
}
if limit > 0 {
q.Set("limit", fmt.Sprintf("%d", limit))
}
r.URL.RawQuery = q.Encode()
})
if err != nil {
return AIBridgeSessionThreadsResponse{}, err
}
defer res.Body.Close()
if res.StatusCode != http.StatusOK {
return AIBridgeSessionThreadsResponse{}, ReadBodyAsError(res)
}
var resp AIBridgeSessionThreadsResponse
return resp, ReadBodyAsJSON(res, &resp)
}
// AIBridgeListModelsFilter narrows the distinct models visible to the caller.
// @typescript-ignore AIBridgeListModelsFilter
type AIBridgeListModelsFilter struct {
// Limit defaults to 100, max is 1000.
Pagination Pagination
// Model keeps only identifiers starting with this literal prefix.
Model string
}
// AIBridgeListModels returns the distinct AI models visible to the caller.
func (c *Client) AIBridgeListModels(ctx context.Context, filter AIBridgeListModelsFilter) ([]string, error) {
res, err := c.Request(ctx, http.MethodGet, "/api/v2/ai-gateway/models", nil, filter.Pagination.asRequestOption(), func(r *http.Request) {
if filter.Model == "" {
return
}
q := r.URL.Query()
q.Set("model", filter.Model)
r.URL.RawQuery = q.Encode()
})
if err != nil {
return nil, err
}
defer res.Body.Close()
if res.StatusCode != http.StatusOK {
return nil, ReadBodyAsError(res)
}
var models []string
return models, ReadBodyAsJSON(res, &models)
}
// AIBridgeListClients returns the distinct AI clients visible to the caller.
func (c *Client) AIBridgeListClients(ctx context.Context) ([]string, error) {
res, err := c.Request(ctx, http.MethodGet, "/api/v2/ai-gateway/clients", nil)
if err != nil {
return nil, err
}
defer res.Body.Close()
if res.StatusCode != http.StatusOK {
return nil, ReadBodyAsError(res)
}
var clients []string
return clients, ReadBodyAsJSON(res, &clients)
}
// AIBridgeProvider is the display metadata for a configured AI provider,
// used to filter AI Gateway sessions by provider_name. It carries no
// configuration so it can be served to anyone who can read sessions.
type AIBridgeProvider struct {
Name string `json:"name"`
Type AIProviderType `json:"type"`
DisplayName string `json:"display_name"`
Icon string `json:"icon"`
}
// AIBridgeListProviders returns the providers available for filtering AI
// Gateway sessions, including disabled and deleted ones that past sessions
// may still reference.
func (c *Client) AIBridgeListProviders(ctx context.Context) ([]AIBridgeProvider, error) {
res, err := c.Request(ctx, http.MethodGet, "/api/v2/ai-gateway/providers", nil)
if err != nil {
return nil, err
}
defer res.Body.Close()
if res.StatusCode != http.StatusOK {
return nil, ReadBodyAsError(res)
}
var providers []AIBridgeProvider
return providers, ReadBodyAsJSON(res, &providers)
}
// OrganizationAISpendFilter narrows the organization per-user AI spend
// report. Zero values apply no filter: the period falls back to the current
// budget period on the server, and an empty dimension matches all usage.
type OrganizationAISpendFilter struct {
// PeriodStart and PeriodEnd bound the [PeriodStart, PeriodEnd) window and
// must be supplied together.
PeriodStart time.Time `json:"period_start,omitempty" format:"date-time"`
PeriodEnd time.Time `json:"period_end,omitempty" format:"date-time"`
// ProviderName matches the configured provider name recorded on the
// intercepted request.
ProviderName string `json:"provider_name,omitempty"`
Model string `json:"model,omitempty"`
// Client matches the client recorded on the intercepted request. Unknown
// matches usage without a recorded client.
Client string `json:"client,omitempty"`
}
// asRequestOption returns a function that can be used in (*Client).Request.
func (f OrganizationAISpendFilter) asRequestOption() RequestOption {
return func(r *http.Request) {
q := r.URL.Query()
if !f.PeriodStart.IsZero() {
q.Set("period_start", f.PeriodStart.UTC().Format(time.RFC3339Nano))
}
if !f.PeriodEnd.IsZero() {
q.Set("period_end", f.PeriodEnd.UTC().Format(time.RFC3339Nano))
}
if f.ProviderName != "" {
q.Set("provider_name", f.ProviderName)
}
if f.Model != "" {
q.Set("model", f.Model)
}
if f.Client != "" {
q.Set("client", f.Client)
}
r.URL.RawQuery = q.Encode()
}
}
// OrganizationAISpendUser is one user's AI spend within an organization
// report.
type OrganizationAISpendUser struct {
UserID uuid.UUID `json:"user_id" format:"uuid"`
Username string `json:"username"`
Name string `json:"name"`
AvatarURL string `json:"avatar_url"`
// CostMicros is the user's priced spend over the period.
CostMicros int64 `json:"cost_micros"`
// UnpricedUsageCount is the number of the user's token usage records that
// carry no cost because their model had no price when they were recorded.
UnpricedUsageCount int64 `json:"unpriced_usage_count"`
// Providers are the configured provider names the user spent through,
// sorted.
Providers []string `json:"providers"`
// Clients are the clients the user spent through, sorted. Usage without a
// recorded client is reported as Unknown.
Clients []string `json:"clients"`
// Models are the models the user spent through, sorted.
Models []string `json:"models"`
}
// OrganizationAISpendTotals aggregates every user matching the report's
// filter, not only the returned page.
type OrganizationAISpendTotals struct {
// CostMicros is the priced spend of every matching user.
CostMicros int64 `json:"cost_micros"`
// UnpricedUsageCount is the number of token usage records without a cost
// across every matching user.
UnpricedUsageCount int64 `json:"unpriced_usage_count"`
}
// OrganizationAISpendReport is one page of per-user AI spend for an
// organization over the applied period. Count and Totals cover every
// matching user, not only the returned page.
type OrganizationAISpendReport struct {
AISpendPeriodWindow
// RetentionStart is the oldest instant for which token usage is still
// retained. An explicit period must not start before it. Omitted when the
// deployment does not purge AI Gateway data.
RetentionStart *time.Time `json:"retention_start,omitempty" format:"date-time"`
// Count is the number of users with token usage matching the filter.
Count int64 `json:"count"`
Totals OrganizationAISpendTotals `json:"totals"`
// Users is the requested page, most expensive first.
Users []OrganizationAISpendUser `json:"users"`
}
// OrganizationAISpendDetailsFilter narrows organization AI spend.
type OrganizationAISpendDetailsFilter struct {
PeriodStart time.Time `json:"period_start,omitempty" format:"date-time"`
PeriodEnd time.Time `json:"period_end,omitempty" format:"date-time"`
UserID uuid.UUID `json:"user_id,omitempty" format:"uuid"`
GroupID uuid.UUID `json:"group_id,omitempty" format:"uuid"`
ProviderName string `json:"provider_name,omitempty"`
Model string `json:"model,omitempty"`
}
// asRequestOption returns a function that applies the filter's query
// parameters to a request.
func (f OrganizationAISpendDetailsFilter) asRequestOption() RequestOption {
return func(r *http.Request) {
q := r.URL.Query()
if !f.PeriodStart.IsZero() {
q.Set("period_start", f.PeriodStart.UTC().Format(time.RFC3339Nano))
}
if !f.PeriodEnd.IsZero() {
q.Set("period_end", f.PeriodEnd.UTC().Format(time.RFC3339Nano))
}
if f.UserID != uuid.Nil {
q.Set("user_id", f.UserID.String())
}
if f.GroupID != uuid.Nil {
q.Set("group_id", f.GroupID.String())
}
if f.ProviderName != "" {
q.Set("provider_name", f.ProviderName)
}
if f.Model != "" {
q.Set("model", f.Model)
}
r.URL.RawQuery = q.Encode()
}
}
// ExportOrganizationAISpend returns a CSV of per-user, per-group, per-model,
// per-provider AI spend for the organization over the requested period. Both
// bounds are optional and interpreted as UTC, and zero values fall back to the
// current budget period on the server. The caller is responsible for closing
// the returned ReadCloser.
func (c *Client) ExportOrganizationAISpend(ctx context.Context, organization uuid.UUID, opts AISpendPeriodWindow) (io.ReadCloser, error) {
return c.ExportOrganizationAISpendWithFilter(ctx, organization, OrganizationAISpendDetailsFilter{
PeriodStart: opts.PeriodStart,
PeriodEnd: opts.PeriodEnd,
})
}
// ExportOrganizationAISpendWithFilter returns organization AI spend as CSV,
// narrowed by the supplied filter. The caller is responsible for closing the
// returned ReadCloser.
func (c *Client) ExportOrganizationAISpendWithFilter(ctx context.Context, organization uuid.UUID, filter OrganizationAISpendDetailsFilter) (io.ReadCloser, error) {
res, err := c.Request(ctx, http.MethodGet,
fmt.Sprintf("/api/v2/organizations/%s/ai/spend/export", organization.String()),
nil,
filter.asRequestOption(),
)
if err != nil {
return nil, xerrors.Errorf("make request: %w", err)
}
if res.StatusCode != http.StatusOK {
defer res.Body.Close()
return nil, ReadBodyAsError(res)
}
return res.Body, nil
}
// OrganizationAISpendPage selects one page of the per-user report, which
// pages by offset only. A zero Limit uses the server default.
type OrganizationAISpendPage struct {
Limit int `json:"limit,omitempty"`
Offset int `json:"offset,omitempty"`
}
func (p OrganizationAISpendPage) asRequestOption() RequestOption {
return func(r *http.Request) {
q := r.URL.Query()
if p.Limit > 0 {
q.Set("limit", strconv.Itoa(p.Limit))
}
if p.Offset > 0 {
q.Set("offset", strconv.Itoa(p.Offset))
}
r.URL.RawQuery = q.Encode()
}
}
// OrganizationAISpendUsers returns one page of per-user AI spend for the
// organization matching the filter. It accounts for the same token usage as
// ExportOrganizationAISpend over the same period.
func (c *ExperimentalClient) OrganizationAISpendUsers(ctx context.Context, organization uuid.UUID, filter OrganizationAISpendFilter, page OrganizationAISpendPage) (OrganizationAISpendReport, error) {
res, err := c.Request(ctx, http.MethodGet,
fmt.Sprintf("/api/experimental/organizations/%s/ai/spend/users", organization.String()),
nil,
filter.asRequestOption(),
page.asRequestOption(),
)
if err != nil {
return OrganizationAISpendReport{}, xerrors.Errorf("make request: %w", err)
}
defer res.Body.Close()
if res.StatusCode != http.StatusOK {
return OrganizationAISpendReport{}, ReadBodyAsError(res)
}
var report OrganizationAISpendReport
return report, ReadBodyAsJSON(res, &report)
}
type GroupAIBudget struct {
GroupID uuid.UUID `json:"group_id" format:"uuid"`
SpendLimitMicros int64 `json:"spend_limit_micros"`
CreatedAt time.Time `json:"created_at" format:"date-time"`
UpdatedAt time.Time `json:"updated_at" format:"date-time"`
}
type UpsertGroupAIBudgetRequest struct {
// SpendLimitMicros must not exceed MaxAISpendLimitMicros.
SpendLimitMicros int64 `json:"spend_limit_micros" validate:"gte=0"`
}
// GroupAIBudget returns the AI spend budget configured for the given group.
func (c *Client) GroupAIBudget(ctx context.Context, group uuid.UUID) (GroupAIBudget, error) {
res, err := c.Request(ctx, http.MethodGet,
fmt.Sprintf("/api/v2/groups/%s/ai/budget", group.String()),
nil,
)
if err != nil {
return GroupAIBudget{}, xerrors.Errorf("make request: %w", err)
}
defer res.Body.Close()
if res.StatusCode != http.StatusOK {
return GroupAIBudget{}, ReadBodyAsError(res)
}
var resp GroupAIBudget
return resp, ReadBodyAsJSON(res, &resp)
}
// UpsertGroupAIBudget creates or updates the AI spend budget for the given group.
func (c *Client) UpsertGroupAIBudget(ctx context.Context, group uuid.UUID, req UpsertGroupAIBudgetRequest) (GroupAIBudget, error) {
res, err := c.Request(ctx, http.MethodPut,
fmt.Sprintf("/api/v2/groups/%s/ai/budget", group.String()),
req,
)
if err != nil {
return GroupAIBudget{}, xerrors.Errorf("make request: %w", err)
}
defer res.Body.Close()
if res.StatusCode != http.StatusOK {
return GroupAIBudget{}, ReadBodyAsError(res)
}
var resp GroupAIBudget
return resp, ReadBodyAsJSON(res, &resp)
}
// DeleteGroupAIBudget removes the AI spend budget for the given group.
func (c *Client) DeleteGroupAIBudget(ctx context.Context, group uuid.UUID) error {
res, err := c.Request(ctx, http.MethodDelete,
fmt.Sprintf("/api/v2/groups/%s/ai/budget", group.String()),
nil,
)
if err != nil {
return xerrors.Errorf("make request: %w", err)
}
defer res.Body.Close()
if res.StatusCode != http.StatusNoContent {
return ReadBodyAsError(res)
}
return nil
}
type UserAIBudgetOverride struct {
UserID uuid.UUID `json:"user_id" format:"uuid"`
GroupID uuid.UUID `json:"group_id" format:"uuid"`
SpendLimitMicros int64 `json:"spend_limit_micros"`
CreatedAt time.Time `json:"created_at" format:"date-time"`
UpdatedAt time.Time `json:"updated_at" format:"date-time"`
}
type UpsertUserAIBudgetOverrideRequest struct {
// GroupID is the group the user's spend is attributed to. The user must
// be a member of this group.
GroupID uuid.UUID `json:"group_id" format:"uuid" validate:"required"`
// SpendLimitMicros must not exceed MaxAISpendLimitMicros.
SpendLimitMicros int64 `json:"spend_limit_micros" validate:"gte=0"`
}
// UserAIBudgetOverride returns the AI spend budget override configured for the given user.
func (c *Client) UserAIBudgetOverride(ctx context.Context, user uuid.UUID) (UserAIBudgetOverride, error) {
res, err := c.Request(ctx, http.MethodGet,
fmt.Sprintf("/api/v2/users/%s/ai/budget/override", user.String()),
nil,
)
if err != nil {
return UserAIBudgetOverride{}, xerrors.Errorf("make request: %w", err)
}
defer res.Body.Close()
if res.StatusCode != http.StatusOK {
return UserAIBudgetOverride{}, ReadBodyAsError(res)
}
var resp UserAIBudgetOverride
return resp, ReadBodyAsJSON(res, &resp)
}
// UpsertUserAIBudgetOverride creates or updates the AI spend budget override for the given user.
func (c *Client) UpsertUserAIBudgetOverride(ctx context.Context, user uuid.UUID, req UpsertUserAIBudgetOverrideRequest) (UserAIBudgetOverride, error) {
res, err := c.Request(ctx, http.MethodPut,
fmt.Sprintf("/api/v2/users/%s/ai/budget/override", user.String()),
req,
)
if err != nil {
return UserAIBudgetOverride{}, xerrors.Errorf("make request: %w", err)
}
defer res.Body.Close()
if res.StatusCode != http.StatusOK {
return UserAIBudgetOverride{}, ReadBodyAsError(res)
}
var resp UserAIBudgetOverride
return resp, ReadBodyAsJSON(res, &resp)
}
// DeleteUserAIBudgetOverride removes the AI spend budget override for the given user.
func (c *Client) DeleteUserAIBudgetOverride(ctx context.Context, user uuid.UUID) error {
res, err := c.Request(ctx, http.MethodDelete,
fmt.Sprintf("/api/v2/users/%s/ai/budget/override", user.String()),
nil,
)
if err != nil {
return xerrors.Errorf("make request: %w", err)
}
defer res.Body.Close()
if res.StatusCode != http.StatusNoContent {
return ReadBodyAsError(res)
}
return nil
}
// UserAISpendStatus returns the current AI spend snapshot for the given user
// within the active budget period.
func (c *Client) UserAISpendStatus(ctx context.Context, user uuid.UUID) (UserAISpendStatus, error) {
res, err := c.Request(ctx, http.MethodGet,
fmt.Sprintf("/api/v2/users/%s/ai/spend", user.String()),
nil,
)
if err != nil {
return UserAISpendStatus{}, xerrors.Errorf("make request: %w", err)
}
defer res.Body.Close()
if res.StatusCode != http.StatusOK {
return UserAISpendStatus{}, ReadBodyAsError(res)
}
var resp UserAISpendStatus
return resp, ReadBodyAsJSON(res, &resp)
}
// OrganizationGroupsAISpend returns AI spend for the given groups within the
// organization for the active budget period. At most 100 group IDs may be
// requested per call, and callers with more groups are expected to batch
// across multiple requests.
func (c *Client) OrganizationGroupsAISpend(ctx context.Context, organization uuid.UUID, groupIDs []uuid.UUID) (OrganizationGroupsAISpend, error) {
ids := slice.List(groupIDs, func(id uuid.UUID) string { return id.String() })
res, err := c.Request(ctx, http.MethodGet,
fmt.Sprintf("/api/v2/organizations/%s/groups/ai/spend", organization.String()),
nil,
func(r *http.Request) {
q := r.URL.Query()
q.Set("group_ids", strings.Join(ids, ","))
r.URL.RawQuery = q.Encode()
},
)
if err != nil {
return OrganizationGroupsAISpend{}, xerrors.Errorf("make request: %w", err)
}
defer res.Body.Close()
if res.StatusCode != http.StatusOK {
return OrganizationGroupsAISpend{}, ReadBodyAsError(res)
}
var resp OrganizationGroupsAISpend
return resp, ReadBodyAsJSON(res, &resp)
}
// GroupAISpend returns AI spend for the given group within the active budget
// period.
func (c *Client) GroupAISpend(ctx context.Context, group uuid.UUID) (GroupAISpend, error) {
res, err := c.Request(ctx, http.MethodGet,
fmt.Sprintf("/api/v2/groups/%s/ai/spend", group.String()),
nil,
)
if err != nil {
return GroupAISpend{}, xerrors.Errorf("make request: %w", err)
}
defer res.Body.Close()
if res.StatusCode != http.StatusOK {
return GroupAISpend{}, ReadBodyAsError(res)
}
var resp GroupAISpend
return resp, ReadBodyAsJSON(res, &resp)
}
// GroupMembersAISpend returns AI spend attributed to the given group for the
// specified users within the active budget period. At most 100 user IDs may be
// requested per call, and callers with more members are expected to batch
// across multiple requests.
func (c *Client) GroupMembersAISpend(ctx context.Context, group uuid.UUID, userIDs []uuid.UUID) (GroupMembersAISpend, error) {
ids := slice.List(userIDs, func(id uuid.UUID) string { return id.String() })
res, err := c.Request(ctx, http.MethodGet,
fmt.Sprintf("/api/v2/groups/%s/members/ai/spend", group.String()),
nil,
func(r *http.Request) {
q := r.URL.Query()
q.Set("user_ids", strings.Join(ids, ","))
r.URL.RawQuery = q.Encode()
},
)
if err != nil {
return GroupMembersAISpend{}, xerrors.Errorf("make request: %w", err)
}
defer res.Body.Close()
if res.StatusCode != http.StatusOK {
return GroupMembersAISpend{}, ReadBodyAsError(res)
}
var resp GroupMembersAISpend
return resp, ReadBodyAsJSON(res, &resp)
}