diff --git a/messages/en-US.json b/messages/en-US.json index 5da27a834..43625ee73 100644 --- a/messages/en-US.json +++ b/messages/en-US.json @@ -1989,6 +1989,14 @@ "aiUsageUserCost": "User Cost", "aiUsageUserTokenUsage": "User Token Usage", "aiUsageUnknownUser": "Unknown user", + "aiUsageVirtualApiKeysTab": "Virtual API Keys", + "aiUsageFilterVirtualApiKey": "Virtual API Key", + "aiUsageFilterAllVirtualApiKeys": "All Virtual API Keys", + "aiUsageTopVirtualApiKeys": "Top Virtual API Keys", + "aiUsageVirtualApiKeyCost": "Virtual API Key Cost", + "aiUsageVirtualApiKeyTokenUsage": "Virtual API Key Token Usage", + "aiUsageUnknownVirtualApiKey": "No virtual API key", + "aiUsageUnnamedVirtualApiKey": "Unnamed key", "aiUsageLoading": "Loading...", "aiUsageNoData": "No data", "resourceBudgetSettings": "Budget", @@ -3457,6 +3465,8 @@ "provider": "Provider", "capability": "Capability", "model": "Model", + "virtualApiKey": "Virtual API Key", + "noVirtualApiKey": "No virtual API key", "stream": "Stream", "streaming": "Streaming", "nonStreaming": "Non-streaming", diff --git a/server/db/pg/schema/schema.ts b/server/db/pg/schema/schema.ts index afaccc480..a5ebfa530 100644 --- a/server/db/pg/schema/schema.ts +++ b/server/db/pg/schema/schema.ts @@ -1819,18 +1819,22 @@ export const aiUsageRecords = pgTable( .references(() => orgs.orgId, { onDelete: "cascade" }), providerId: integer("providerId") .notNull() - .references(() => aiProviders.providerId, { onDelete: "cascade" }), + .references(() => aiProviders.providerId, { onDelete: "set null" }), resourceId: integer("resourceId").references( () => resources.resourceId, - { onDelete: "cascade" } + { onDelete: "set null" } ), siteResourceId: integer("siteResourceId").references( () => siteResources.siteResourceId, - { onDelete: "cascade" } + { onDelete: "set null" } ), userId: varchar("userId").references(() => users.userId, { onDelete: "set null" }), + virtualApiKeyId: varchar("virtualApiKeyId").references( + () => virtualApiKeys.virtualApiKeyId, + { onDelete: "set null" } + ), // Links this usage record back to the aiSessionLog row for the same // request (aiSessionLog.sessionId), so token/cost usage can be shown // alongside the session transcript. Not a DB-level FK - aiSessionLog @@ -1870,6 +1874,11 @@ export const aiUsageRecords = pgTable( t.userId, t.createdAt ), + index("idx_ai_usage_records_org_virtual_api_key_created").on( + t.orgId, + t.virtualApiKeyId, + t.createdAt + ), index("idx_ai_usage_records_session").on(t.sessionId) ] ); @@ -1927,19 +1936,23 @@ export const aiSessionLog = pgTable( }), providerId: integer("providerId") .notNull() - .references(() => aiProviders.providerId, { onDelete: "cascade" }), + .references(() => aiProviders.providerId, { onDelete: "set null" }), capability: varchar("capability").notNull(), resourceId: integer("resourceId").references( () => resources.resourceId, - { onDelete: "cascade" } + { onDelete: "set null" } ), siteResourceId: integer("siteResourceId").references( () => siteResources.siteResourceId, - { onDelete: "cascade" } + { onDelete: "set null" } ), userId: varchar("userId").references(() => users.userId, { onDelete: "set null" }), + virtualApiKeyId: varchar("virtualApiKeyId").references( + () => virtualApiKeys.virtualApiKeyId, + { onDelete: "set null" } + ), requestedModel: varchar("requestedModel"), isStream: boolean("isStream").notNull().default(false), requestBody: text("requestBody"), @@ -1979,6 +1992,11 @@ export const aiSessionLog = pgTable( t.userId, t.createdAt ), + index("idx_ai_session_log_org_virtual_api_key_created").on( + t.orgId, + t.virtualApiKeyId, + t.createdAt + ), index("idx_ai_session_log_session").on(t.sessionId) ] ); diff --git a/server/db/sqlite/schema/schema.ts b/server/db/sqlite/schema/schema.ts index bf3be9f87..450b99fce 100644 --- a/server/db/sqlite/schema/schema.ts +++ b/server/db/sqlite/schema/schema.ts @@ -1807,18 +1807,22 @@ export const aiUsageRecords = sqliteTable( .references(() => orgs.orgId, { onDelete: "cascade" }), providerId: integer("providerId") .notNull() - .references(() => aiProviders.providerId, { onDelete: "cascade" }), + .references(() => aiProviders.providerId, { onDelete: "set null" }), resourceId: integer("resourceId").references( () => resources.resourceId, - { onDelete: "cascade" } + { onDelete: "set null" } ), siteResourceId: integer("siteResourceId").references( () => siteResources.siteResourceId, - { onDelete: "cascade" } + { onDelete: "set null" } ), userId: text("userId").references(() => users.userId, { onDelete: "set null" }), + virtualApiKeyId: text("virtualApiKeyId").references( + () => virtualApiKeys.virtualApiKeyId, + { onDelete: "set null" } + ), // Links this usage record back to the aiSessionLog row for the same // request (aiSessionLog.sessionId), so token/cost usage can be shown // alongside the session transcript. Not a DB-level FK - aiSessionLog @@ -1860,6 +1864,11 @@ export const aiUsageRecords = sqliteTable( t.userId, t.createdAt ), + index("idx_ai_usage_records_org_virtual_api_key_created").on( + t.orgId, + t.virtualApiKeyId, + t.createdAt + ), index("idx_ai_usage_records_session").on(t.sessionId) ] ); @@ -1917,19 +1926,23 @@ export const aiSessionLog = sqliteTable( }), providerId: integer("providerId") .notNull() - .references(() => aiProviders.providerId, { onDelete: "cascade" }), + .references(() => aiProviders.providerId, { onDelete: "set null" }), capability: text("capability").notNull(), resourceId: integer("resourceId").references( () => resources.resourceId, - { onDelete: "cascade" } + { onDelete: "set null" } ), siteResourceId: integer("siteResourceId").references( () => siteResources.siteResourceId, - { onDelete: "cascade" } + { onDelete: "set null" } ), userId: text("userId").references(() => users.userId, { onDelete: "set null" }), + virtualApiKeyId: text("virtualApiKeyId").references( + () => virtualApiKeys.virtualApiKeyId, + { onDelete: "set null" } + ), requestedModel: text("requestedModel"), isStream: integer("isStream", { mode: "boolean" }) .notNull() @@ -1973,6 +1986,11 @@ export const aiSessionLog = sqliteTable( t.userId, t.createdAt ), + index("idx_ai_session_log_org_virtual_api_key_created").on( + t.orgId, + t.virtualApiKeyId, + t.createdAt + ), index("idx_ai_session_log_session").on(t.sessionId) ] ); diff --git a/server/lib/aiBudgetEnforcement.ts b/server/lib/aiBudgetEnforcement.ts index d73f5ba59..f6daf9694 100644 --- a/server/lib/aiBudgetEnforcement.ts +++ b/server/lib/aiBudgetEnforcement.ts @@ -1,4 +1,14 @@ -import { and, eq, gte, inArray, isNull, or, sql, SQL, type InferInsertModel } from "drizzle-orm"; +import { + and, + eq, + gte, + inArray, + isNull, + or, + sql, + SQL, + type InferInsertModel +} from "drizzle-orm"; import { AiBudget, aiBudgetBreachEvents, @@ -424,6 +434,7 @@ export type UsageRecordInput = { resourceId: number | null; siteResourceId: number | null; userId: string | null; + virtualApiKeyId: string | null; requestedModel: string; usage: AiUsage; costUsd: number | null; @@ -460,7 +471,10 @@ async function flushUsageRecords() { isUsageFlushInProgress = true; - const recordsToWrite = usageRecordBuffer.splice(0, usageRecordBuffer.length); + const recordsToWrite = usageRecordBuffer.splice( + 0, + usageRecordBuffer.length + ); try { // Use a transaction to ensure all inserts succeed or fail together @@ -472,16 +486,25 @@ async function flushUsageRecords() { await tx.insert(aiUsageRecords).values(batch); } }); - logger.debug(`Flushed ${recordsToWrite.length} AI usage records to database`); + logger.debug( + `Flushed ${recordsToWrite.length} AI usage records to database` + ); } catch (error) { logger.error("Error flushing AI usage records:", error); // On transaction error, put records back at the front of the buffer // to retry, but only if the buffer isn't too large - if (usageRecordBuffer.length < USAGE_MAX_BUFFER_SIZE - recordsToWrite.length) { + if ( + usageRecordBuffer.length < + USAGE_MAX_BUFFER_SIZE - recordsToWrite.length + ) { usageRecordBuffer.unshift(...recordsToWrite); - logger.info(`Re-queued ${recordsToWrite.length} AI usage records for retry`); + logger.info( + `Re-queued ${recordsToWrite.length} AI usage records for retry` + ); } else { - logger.error(`Buffer full, dropped ${recordsToWrite.length} AI usage records`); + logger.error( + `Buffer full, dropped ${recordsToWrite.length} AI usage records` + ); } } finally { isUsageFlushInProgress = false; @@ -544,6 +567,7 @@ export async function recordUsage(input: UsageRecordInput): Promise { resourceId: input.resourceId, siteResourceId: input.siteResourceId, userId: input.userId, + virtualApiKeyId: input.virtualApiKeyId, sessionId: input.sessionId, requestedModel: input.requestedModel, promptTokens: usage.promptTokens, diff --git a/server/routers/aiGateway/logAiSession.ts b/server/routers/aiGateway/logAiSession.ts index ee05517f5..5a2e29bdb 100644 --- a/server/routers/aiGateway/logAiSession.ts +++ b/server/routers/aiGateway/logAiSession.ts @@ -180,6 +180,7 @@ export function logAiSession(data: { resourceId: number | null; siteResourceId: number | null; requestUserId: string | null; + virtualApiKeyId: string | null; }): void { (async () => { try { @@ -237,6 +238,9 @@ export function logAiSession(data: { resourceId: data.resourceId ?? undefined, siteResourceId: data.siteResourceId ?? undefined, userId: sanitizeString(data.requestUserId ?? undefined), + virtualApiKeyId: sanitizeString( + data.virtualApiKeyId ?? undefined + ), requestedModel: sanitizeString(data.requestedModel), isStream: data.isStream, requestBody: sanitizeString(requestBodyText.value), diff --git a/server/routers/aiGateway/pipeline.ts b/server/routers/aiGateway/pipeline.ts index 5aa69d7fc..2be7a2592 100644 --- a/server/routers/aiGateway/pipeline.ts +++ b/server/routers/aiGateway/pipeline.ts @@ -161,6 +161,15 @@ export type RequestUser = { roleIds: number[]; }; +// Identity resolved for a gateway request: the app/session or virtual-API-key +// user (if any) plus the virtual API key that authenticated the request (if +// any) - a manual virtual API key with no associated user has a +// virtualApiKeyId but no user. +export type RequestIdentity = { + user: RequestUser | null; + virtualApiKeyId: string | null; +}; + // Identity headers forwarded to the upstream inference endpoint when the // requesting user is known. Omitted entirely (not sent empty) when we // couldn't resolve a user for the request. @@ -223,10 +232,12 @@ async function resolveRequestUser( req: Request, resourceId: number | null, orgId: string | null -): Promise { +): Promise { // Public inference: identity comes from Badger via Remote-* only when the // Traefik trust header proves the request passed verify-session (VAK). if (isAiGatewayTrustHeaderValid(req.headers as Record)) { + const virtualApiKeyId = + getRequestHeader(req, "remote-virtual-api-key-id") || null; const userId = getRequestHeader(req, "remote-user-id"); if (userId) { const username = getRequestHeader(req, "remote-user") || userId; @@ -236,32 +247,41 @@ async function resolveRequestUser( const orgRoles = orgId ? await getUserOrgRoles(userId, orgId) : []; return { - userId, - username, - email: email || null, - name: name || null, - role: - role || orgRoles.map((r) => r.roleName).join(", ") || null, - roleIds: orgRoles.map((r) => r.roleId) + user: { + userId, + username, + email: email || null, + name: name || null, + role: + role || + orgRoles.map((r) => r.roleName).join(", ") || + null, + roleIds: orgRoles.map((r) => r.roleId) + }, + virtualApiKeyId }; } - // Trusted request with no associated user (manual key without userId). + // Trusted request with no associated user (manual key without + // userId) - still attribute usage to the virtual API key itself. if (resourceId != null) { - return null; + return { user: null, virtualApiKeyId }; } } // Public inference must come through Badger; do not authorize via app session. if (resourceId != null) { - return null; + return { user: null, virtualApiKeyId: null }; } const sessionToken = req.cookies?.[SESSION_COOKIE_NAME]; if (sessionToken) { const { session, user } = await validateSessionToken(sessionToken); if (session && user) { - return buildRequestUser(user.userId, orgId); + return { + user: await buildRequestUser(user.userId, orgId), + virtualApiKeyId: null + }; } } @@ -269,7 +289,7 @@ async function resolveRequestUser( const ip = req.ip; if (!ip) { - return null; + return { user: null, virtualApiKeyId: null }; } const exitNodeRanges = await getExitNodeRanges(); @@ -277,15 +297,18 @@ async function resolveRequestUser( isIpInCidr(ip, range) ); if (!inExitNodeRange) { - return null; + return { user: null, virtualApiKeyId: null }; } const client = await findClientByIp(ip); if (!client || !client.userId) { - return null; + return { user: null, virtualApiKeyId: null }; } - return buildRequestUser(client.userId, orgId); + return { + user: await buildRequestUser(client.userId, orgId), + virtualApiKeyId: null + }; } function getRequestHeader(req: Request, name: string): string | undefined { @@ -608,6 +631,7 @@ export function recordAiGatewayCompletion(args: { resourceId: number | null; siteResourceId: number | null; requestUserId: string | null; + virtualApiKeyId: string | null; budgets: AiBudget[]; }): void { const { @@ -623,6 +647,7 @@ export function recordAiGatewayCompletion(args: { resourceId, siteResourceId, requestUserId, + virtualApiKeyId, budgets } = args; @@ -668,6 +693,7 @@ export function recordAiGatewayCompletion(args: { resourceId, siteResourceId, userId: requestUserId, + virtualApiKeyId, requestedModel: model ?? "unknown", usage, costUsd: cost?.totalCost ?? null, @@ -699,7 +725,8 @@ export function recordAiGatewayCompletion(args: { orgId, resourceId, siteResourceId, - requestUserId + requestUserId, + virtualApiKeyId }); } @@ -770,7 +797,7 @@ export async function handleAiGatewayProxy( const requestedModel = def.extractModel(req); - const [requestUser, selection] = await Promise.all([ + const [identity, selection] = await Promise.all([ resolveRequestUser(req, resourceId, orgId), selectProvider( capableAttachments, @@ -778,6 +805,7 @@ export async function handleAiGatewayProxy( requestedModel ) ]); + const requestUser = identity.user; if (requestUser) { logger.debug( @@ -836,7 +864,8 @@ export async function handleAiGatewayProxy( resourceId, siteResourceId, requestedModel, - budgets: appliedBudgets + budgets: appliedBudgets, + virtualApiKeyId: identity.virtualApiKeyId } ); } @@ -987,6 +1016,7 @@ export async function handleAiGatewayProxy( resourceId, siteResourceId, requestUserId: requestUser?.userId ?? null, + virtualApiKeyId: identity.virtualApiKeyId, budgets: appliedBudgets }); } diff --git a/server/routers/aiGateway/targetRouting.ts b/server/routers/aiGateway/targetRouting.ts index 9ffbf1d1b..8c4f244aa 100644 --- a/server/routers/aiGateway/targetRouting.ts +++ b/server/routers/aiGateway/targetRouting.ts @@ -179,6 +179,7 @@ export async function proxyAiGatewayToSiteTarget( siteResourceId: number | null; requestedModel: string | undefined; budgets: AiBudget[]; + virtualApiKeyId: string | null; } ): Promise { const providerTargets = await getProviderTargets(provider.providerId); @@ -319,6 +320,7 @@ export async function proxyAiGatewayToSiteTarget( resourceId: ctx.resourceId, siteResourceId: ctx.siteResourceId, requestUserId: requestUser?.userId ?? null, + virtualApiKeyId: ctx.virtualApiKeyId, budgets: ctx.budgets }); } diff --git a/server/routers/auditLogs/aiUsageAnalyticsShared.ts b/server/routers/auditLogs/aiUsageAnalyticsShared.ts index 95a6f825b..ada94a9b9 100644 --- a/server/routers/auditLogs/aiUsageAnalyticsShared.ts +++ b/server/routers/auditLogs/aiUsageAnalyticsShared.ts @@ -58,7 +58,8 @@ export const aiUsageAnalyticsFiltersQuery = z.object({ .transform(Number) .pipe(z.int().positive()) .optional(), - userId: z.string().optional() + userId: z.string().optional(), + virtualApiKeyId: z.string().optional() }); export const aiUsageAnalyticsParams = z.object({ @@ -114,6 +115,9 @@ export function buildAiUsageWhere( ) : undefined, data.userId ? eq(aiUsageRecords.userId, data.userId) : undefined, + data.virtualApiKeyId + ? eq(aiUsageRecords.virtualApiKeyId, data.virtualApiKeyId) + : undefined, roleUserIds ? inArray(aiUsageRecords.userId, roleUserIds) : undefined ); } diff --git a/server/routers/auditLogs/index.ts b/server/routers/auditLogs/index.ts index 12735d650..59969134f 100644 --- a/server/routers/auditLogs/index.ts +++ b/server/routers/auditLogs/index.ts @@ -8,3 +8,4 @@ export * from "./queryAiUsageOverview"; export * from "./queryAiUsageProviders"; export * from "./queryAiUsageResources"; export * from "./queryAiUsageUsersRoles"; +export * from "./queryAiUsageVirtualApiKeys"; diff --git a/server/routers/auditLogs/queryAiSessionLog.ts b/server/routers/auditLogs/queryAiSessionLog.ts index d61bdf5c8..437d74550 100644 --- a/server/routers/auditLogs/queryAiSessionLog.ts +++ b/server/routers/auditLogs/queryAiSessionLog.ts @@ -6,6 +6,7 @@ import { resources, siteResources, users, + virtualApiKeys, db, primaryDb } from "@server/db"; @@ -67,6 +68,7 @@ export const queryAiSessionLogsQuery = z.strictObject({ .pipe(z.int().positive()) .optional(), actor: z.string().optional(), + virtualApiKeyId: z.string().optional(), model: z.string().optional(), isStream: z .union([z.boolean(), z.string()]) @@ -117,7 +119,9 @@ function getWhere(data: Q) { data.providerId ? eq(aiSessionLog.providerId, data.providerId) : undefined, - data.capability ? eq(aiSessionLog.capability, data.capability) : undefined, + data.capability + ? eq(aiSessionLog.capability, data.capability) + : undefined, data.resourceId ? or( eq(aiSessionLog.resourceId, data.resourceId), @@ -125,9 +129,10 @@ function getWhere(data: Q) { ) : undefined, data.actor ? eq(aiSessionLog.userId, data.actor) : undefined, - data.model - ? eq(aiSessionLog.requestedModel, data.model) + data.virtualApiKeyId + ? eq(aiSessionLog.virtualApiKeyId, data.virtualApiKeyId) : undefined, + data.model ? eq(aiSessionLog.requestedModel, data.model) : undefined, data.isStream !== undefined ? eq(aiSessionLog.isStream, data.isStream) : undefined @@ -145,6 +150,7 @@ export function queryAiSession(data: Q) { resourceId: aiSessionLog.resourceId, siteResourceId: aiSessionLog.siteResourceId, userId: aiSessionLog.userId, + virtualApiKeyId: aiSessionLog.virtualApiKeyId, requestedModel: aiSessionLog.requestedModel, isStream: aiSessionLog.isStream, requestBody: aiSessionLog.requestBody, @@ -182,6 +188,14 @@ async function enrichWithDetails( ) ]; + const virtualApiKeyIds = [ + ...new Set( + logs + .map((log) => log.virtualApiKeyId) + .filter((id): id is string => id !== null && id !== undefined) + ) + ]; + const providerMap = new Map< number, { name: string | null; type: string | null } @@ -254,6 +268,28 @@ async function enrichWithDetails( } } + const virtualApiKeyMap = new Map< + string, + { name: string | null; lastChars: string } + >(); + if (virtualApiKeyIds.length > 0) { + const virtualApiKeyDetails = await db + .select({ + virtualApiKeyId: virtualApiKeys.virtualApiKeyId, + name: virtualApiKeys.name, + lastChars: virtualApiKeys.lastChars + }) + .from(virtualApiKeys) + .where(inArray(virtualApiKeys.virtualApiKeyId, virtualApiKeyIds)); + + for (const k of virtualApiKeyDetails) { + virtualApiKeyMap.set(k.virtualApiKeyId, { + name: k.name, + lastChars: k.lastChars + }); + } + } + const usageMap = new Map< string, { @@ -328,6 +364,12 @@ async function enrichWithDetails( resourceName, resourceNiceId, userEmail: log.userId ? (userMap.get(log.userId) ?? null) : null, + virtualApiKeyName: log.virtualApiKeyId + ? (virtualApiKeyMap.get(log.virtualApiKeyId)?.name ?? null) + : null, + virtualApiKeyLastChars: log.virtualApiKeyId + ? (virtualApiKeyMap.get(log.virtualApiKeyId)?.lastChars ?? null) + : null, usage: usageMap.get(log.sessionId) ?? null }; }); @@ -385,7 +427,8 @@ async function queryUniqueFilterAttributes( uniqueUsers, uniqueResources, uniqueSiteResources, - uniqueModels + uniqueModels, + uniqueVirtualApiKeys ] = await Promise.all([ logsDb .selectDistinct({ id: aiSessionLog.providerId }) @@ -411,6 +454,11 @@ async function queryUniqueFilterAttributes( .selectDistinct({ model: aiSessionLog.requestedModel }) .from(aiSessionLog) .where(baseConditions) + .limit(DISTINCT_LIMIT + 1), + logsDb + .selectDistinct({ id: aiSessionLog.virtualApiKeyId }) + .from(aiSessionLog) + .where(baseConditions) .limit(DISTINCT_LIMIT + 1) ]); @@ -498,10 +546,37 @@ async function queryUniqueFilterAttributes( ]; } + const virtualApiKeyIds = uniqueVirtualApiKeys + .map((row) => row.id) + .filter((id): id is string => id !== null); + + let virtualApiKeyList: Array<{ + id: string; + name: string | null; + lastChars: string | null; + }> = []; + if (virtualApiKeyIds.length > 0) { + const virtualApiKeyDetails = await primaryDb + .select({ + virtualApiKeyId: virtualApiKeys.virtualApiKeyId, + name: virtualApiKeys.name, + lastChars: virtualApiKeys.lastChars + }) + .from(virtualApiKeys) + .where(inArray(virtualApiKeys.virtualApiKeyId, virtualApiKeyIds)); + + virtualApiKeyList = virtualApiKeyDetails.map((k) => ({ + id: k.virtualApiKeyId, + name: k.name, + lastChars: k.lastChars + })); + } + return { providers: sortNamedFilterOptions(providers), resources: sortNamedFilterOptions(resourcesWithNames), users: userList, + virtualApiKeys: virtualApiKeyList, models: models.sort() }; } diff --git a/server/routers/auditLogs/queryAiUsageFilterOptions.ts b/server/routers/auditLogs/queryAiUsageFilterOptions.ts index 46cd3199e..6d7191187 100644 --- a/server/routers/auditLogs/queryAiUsageFilterOptions.ts +++ b/server/routers/auditLogs/queryAiUsageFilterOptions.ts @@ -6,7 +6,8 @@ import { siteResources, users, roles, - userOrgRoles + userOrgRoles, + virtualApiKeys } from "@server/db"; import { registry } from "@server/openApi"; import { NextFunction } from "express"; @@ -43,8 +44,9 @@ const queryAiUsageFilterOptionsParams = z.object({ orgId: z.string() }); -const queryAiUsageFilterOptionsCombined = - queryAiUsageFilterOptionsQuery.merge(queryAiUsageFilterOptionsParams); +const queryAiUsageFilterOptionsCombined = queryAiUsageFilterOptionsQuery.merge( + queryAiUsageFilterOptionsParams +); type Q = z.infer; function sortNamedFilterOptions( @@ -73,7 +75,8 @@ async function query(data: Q) { uniqueModels, uniqueResources, uniqueSiteResources, - uniqueUsers + uniqueUsers, + uniqueVirtualApiKeys ] = await Promise.all([ db .selectDistinct({ id: aiUsageRecords.providerId }) @@ -105,6 +108,13 @@ async function query(data: Q) { .selectDistinct({ userId: aiUsageRecords.userId }) .from(aiUsageRecords) .where(and(baseConditions, not(isNull(aiUsageRecords.userId)))) + .limit(DISTINCT_LIMIT + 1), + db + .selectDistinct({ id: aiUsageRecords.virtualApiKeyId }) + .from(aiUsageRecords) + .where( + and(baseConditions, not(isNull(aiUsageRecords.virtualApiKeyId))) + ) .limit(DISTINCT_LIMIT + 1) ]); @@ -120,11 +130,17 @@ async function query(data: Q) { let providers: Array<{ id: number; name: string | null }> = []; if (providerIds.length > 0) { const providerDetails = await db - .select({ providerId: aiProviders.providerId, name: aiProviders.name }) + .select({ + providerId: aiProviders.providerId, + name: aiProviders.name + }) .from(aiProviders) .where(inArray(aiProviders.providerId, providerIds)); - providers = providerDetails.map((p) => ({ id: p.providerId, name: p.name })); + providers = providerDetails.map((p) => ({ + id: p.providerId, + name: p.name + })); } const resourceIds = uniqueResources @@ -155,7 +171,10 @@ async function query(data: Q) { .where(inArray(siteResources.siteResourceId, siteResourceIds)); resourcesWithNames = resourcesWithNames.concat( - siteResourceDetails.map((r) => ({ id: r.siteResourceId, name: r.name })) + siteResourceDetails.map((r) => ({ + id: r.siteResourceId, + name: r.name + })) ); } @@ -190,11 +209,37 @@ async function query(data: Q) { roleList = [...roleMap.entries()].map(([id, name]) => ({ id, name })); } + const virtualApiKeyIds = uniqueVirtualApiKeys + .map((row) => row.id) + .filter((id): id is string => id !== null); + + let virtualApiKeyList: Array<{ + id: string; + name: string | null; + lastChars: string; + }> = []; + if (virtualApiKeyIds.length > 0) { + const virtualApiKeyDetails = await db + .select({ + virtualApiKeyId: virtualApiKeys.virtualApiKeyId, + name: virtualApiKeys.name, + lastChars: virtualApiKeys.lastChars + }) + .from(virtualApiKeys) + .where(inArray(virtualApiKeys.virtualApiKeyId, virtualApiKeyIds)); + virtualApiKeyList = virtualApiKeyDetails.map((k) => ({ + id: k.virtualApiKeyId, + name: k.name, + lastChars: k.lastChars + })); + } + return { providers: sortNamedFilterOptions(providers), resources: sortNamedFilterOptions(resourcesWithNames), roles: sortNamedFilterOptions(roleList), users: userList, + virtualApiKeys: virtualApiKeyList, models }; } @@ -227,7 +272,9 @@ registry.registerPath({ } }); -export type QueryAiUsageFilterOptionsResponse = Awaited>; +export type QueryAiUsageFilterOptionsResponse = Awaited< + ReturnType +>; export async function queryAiUsageFilterOptions( req: Request, @@ -238,7 +285,10 @@ export async function queryAiUsageFilterOptions( const parsedQuery = queryAiUsageFilterOptionsQuery.safeParse(req.query); if (!parsedQuery.success) { return next( - createHttpError(HttpCode.BAD_REQUEST, fromError(parsedQuery.error)) + createHttpError( + HttpCode.BAD_REQUEST, + fromError(parsedQuery.error) + ) ); } @@ -247,7 +297,10 @@ export async function queryAiUsageFilterOptions( ); if (!parsedParams.success) { return next( - createHttpError(HttpCode.BAD_REQUEST, fromError(parsedParams.error)) + createHttpError( + HttpCode.BAD_REQUEST, + fromError(parsedParams.error) + ) ); } diff --git a/server/routers/auditLogs/queryAiUsageVirtualApiKeys.ts b/server/routers/auditLogs/queryAiUsageVirtualApiKeys.ts new file mode 100644 index 000000000..9d1b52ad2 --- /dev/null +++ b/server/routers/auditLogs/queryAiUsageVirtualApiKeys.ts @@ -0,0 +1,222 @@ +import { db, aiUsageRecords, virtualApiKeys } from "@server/db"; +import { registry } from "@server/openApi"; +import { NextFunction } from "express"; +import { Request, Response } from "express"; +import { count, desc, inArray, sql } from "drizzle-orm"; +import { OpenAPITags } from "@server/openApi"; +import { z } from "zod"; +import createHttpError from "http-errors"; +import HttpCode from "@server/types/HttpCode"; +import { fromError } from "zod-validation-error"; +import response from "@server/lib/response"; +import logger from "@server/logger"; +import { + aiUsageAnalyticsFiltersQuery, + aiUsageAnalyticsParams, + buildAiUsageWhere, + resolveRoleUserIds, + dayBucketExpr, + pickTopNKeys, + bucketTopNPerDay, + DISTINCT_LIMIT, + type AiUsageAnalyticsQuery +} from "./aiUsageAnalyticsShared"; + +type Q = AiUsageAnalyticsQuery; + +const UNKNOWN_VIRTUAL_API_KEY_KEY = "unknown"; + +async function query(data: Q) { + const roleUserIds = await resolveRoleUserIds(data.orgId, data.roleId); + const baseConditions = buildAiUsageWhere(data, roleUserIds); + const dayExpr = dayBucketExpr(); + + const virtualApiKeyByDay = await db + .select({ + day: dayExpr.as("day"), + virtualApiKeyId: aiUsageRecords.virtualApiKeyId, + cost: sql`COALESCE(SUM(${aiUsageRecords.costUsd}), 0)`, + tokens: sql`COALESCE(SUM(${aiUsageRecords.totalTokens}), 0)` + }) + .from(aiUsageRecords) + .where(baseConditions) + .groupBy(dayExpr, aiUsageRecords.virtualApiKeyId) + .orderBy(dayExpr); + + const virtualApiKeyTotalsRaw = await db + .select({ + virtualApiKeyId: aiUsageRecords.virtualApiKeyId, + requests: count(), + totalTokens: sql`COALESCE(SUM(${aiUsageRecords.totalTokens}), 0)`, + costUsd: sql`COALESCE(SUM(${aiUsageRecords.costUsd}), 0)` + }) + .from(aiUsageRecords) + .where(baseConditions) + .groupBy(aiUsageRecords.virtualApiKeyId) + .orderBy(desc(sql`COALESCE(SUM(${aiUsageRecords.costUsd}), 0)`)) + .limit(DISTINCT_LIMIT + 1); + + if (virtualApiKeyTotalsRaw.length > DISTINCT_LIMIT) { + throw createHttpError( + HttpCode.BAD_REQUEST, + "Too many distinct virtual API keys. Please narrow your query." + ); + } + + const virtualApiKeyIds = virtualApiKeyTotalsRaw + .map((r) => r.virtualApiKeyId) + .filter((id): id is string => id !== null); + + const detailsMap = new Map< + string, + { name: string | null; lastChars: string; kind: "user" | "manual" } + >(); + if (virtualApiKeyIds.length > 0) { + const details = await db + .select({ + virtualApiKeyId: virtualApiKeys.virtualApiKeyId, + name: virtualApiKeys.name, + lastChars: virtualApiKeys.lastChars, + kind: virtualApiKeys.kind + }) + .from(virtualApiKeys) + .where(inArray(virtualApiKeys.virtualApiKeyId, virtualApiKeyIds)); + for (const k of details) { + detailsMap.set(k.virtualApiKeyId, { + name: k.name, + lastChars: k.lastChars, + kind: k.kind + }); + } + } + + const topVirtualApiKeys = virtualApiKeyTotalsRaw.map((r) => { + const details = r.virtualApiKeyId + ? detailsMap.get(r.virtualApiKeyId) + : undefined; + return { + virtualApiKeyId: r.virtualApiKeyId, + name: details?.name ?? null, + lastChars: details?.lastChars ?? null, + kind: details?.kind ?? null, + requests: r.requests, + totalTokens: r.totalTokens, + costUsd: r.costUsd + }; + }); + + const virtualApiKeyCostTotals = new Map(); + const virtualApiKeyTokenTotals = new Map(); + for (const row of virtualApiKeyByDay) { + const key = row.virtualApiKeyId ?? UNKNOWN_VIRTUAL_API_KEY_KEY; + virtualApiKeyCostTotals.set( + key, + (virtualApiKeyCostTotals.get(key) ?? 0) + row.cost + ); + virtualApiKeyTokenTotals.set( + key, + (virtualApiKeyTokenTotals.get(key) ?? 0) + row.tokens + ); + } + const topVirtualApiKeysByCost = pickTopNKeys(virtualApiKeyCostTotals); + const topVirtualApiKeysByTokens = pickTopNKeys(virtualApiKeyTokenTotals); + + const virtualApiKeyCostPerDay = bucketTopNPerDay( + virtualApiKeyByDay.map((r) => ({ + day: r.day, + key: r.virtualApiKeyId ?? UNKNOWN_VIRTUAL_API_KEY_KEY, + value: r.cost + })), + topVirtualApiKeysByCost + ); + const virtualApiKeyTokensPerDay = bucketTopNPerDay( + virtualApiKeyByDay.map((r) => ({ + day: r.day, + key: r.virtualApiKeyId ?? UNKNOWN_VIRTUAL_API_KEY_KEY, + value: r.tokens + })), + topVirtualApiKeysByTokens + ); + + return { + topVirtualApiKeys, + virtualApiKeyCostPerDay, + virtualApiKeyTokensPerDay + }; +} + +registry.registerPath({ + method: "get", + path: "/org/{orgId}/logs/ai/usage/virtual-api-keys", + description: + "Query the AI usage analytics virtual API key breakdown for an organization", + tags: [OpenAPITags.Logs], + request: { + query: aiUsageAnalyticsFiltersQuery, + params: aiUsageAnalyticsParams + }, + responses: { + 200: { + description: "Successful response", + content: { + "application/json": { + schema: z.object({ + data: z.record(z.string(), z.any()).nullable(), + success: z.boolean(), + error: z.boolean(), + message: z.string(), + status: z.number() + }) + } + } + } + } +}); + +export type QueryAiUsageVirtualApiKeysResponse = Awaited< + ReturnType +>; + +export async function queryAiUsageVirtualApiKeys( + req: Request, + res: Response, + next: NextFunction +): Promise { + try { + const parsedQuery = aiUsageAnalyticsFiltersQuery.safeParse(req.query); + if (!parsedQuery.success) { + return next( + createHttpError( + HttpCode.BAD_REQUEST, + fromError(parsedQuery.error) + ) + ); + } + + const parsedParams = aiUsageAnalyticsParams.safeParse(req.params); + if (!parsedParams.success) { + return next( + createHttpError( + HttpCode.BAD_REQUEST, + fromError(parsedParams.error) + ) + ); + } + + const data = await query({ ...parsedQuery.data, ...parsedParams.data }); + + return response(res, { + data, + success: true, + error: false, + message: + "AI usage virtual API key breakdown retrieved successfully", + status: HttpCode.OK + }); + } catch (error) { + logger.error(error); + return next( + createHttpError(HttpCode.INTERNAL_SERVER_ERROR, "An error occurred") + ); + } +} diff --git a/server/routers/auditLogs/types.ts b/server/routers/auditLogs/types.ts index b63acbf30..dc8d4015f 100644 --- a/server/routers/auditLogs/types.ts +++ b/server/routers/auditLogs/types.ts @@ -110,6 +110,9 @@ export type QueryAiSessionLogResponse = { resourceType: "public" | "site" | null; userId: string | null; userEmail: string | null; + virtualApiKeyId: string | null; + virtualApiKeyName: string | null; + virtualApiKeyLastChars: string | null; requestedModel: string | null; isStream: boolean; requestBody: string | null; @@ -148,6 +151,11 @@ export type QueryAiSessionLogResponse = { id: string; email: string | null; }[]; + virtualApiKeys: { + id: string; + name: string | null; + lastChars: string | null; + }[]; models: string[]; }; }; diff --git a/server/routers/badger/verifySession.ts b/server/routers/badger/verifySession.ts index ee69ce163..525128fc8 100644 --- a/server/routers/badger/verifySession.ts +++ b/server/routers/badger/verifySession.ts @@ -96,6 +96,10 @@ export type VerifyUserResponse = { userData?: BasicUserData; pangolinVersion?: string; dontStripSession?: boolean; + // Set independently of userData so a manual virtual API key with no + // associated user still gets attributed to the key that authenticated + // the request (see the mode === "inference" branch below). + virtualApiKeyId?: string; }; export async function verifyResourceSession( @@ -401,7 +405,12 @@ export async function verifyResourceSession( parsedBody.data ); - return allowed(res, vakUserData, dontStripSession); + return allowed( + res, + vakUserData, + dontStripSession, + key.virtualApiKeyId + ); } } @@ -987,16 +996,20 @@ async function notAllowed( function allowed( res: Response, userData?: BasicUserData, - dontStripSession?: boolean + dontStripSession?: boolean, + virtualApiKeyId?: string ) { const baseData = userData !== undefined && userData !== null ? { valid: true, ...userData, pangolinVersion: APP_VERSION } : { valid: true, pangolinVersion: APP_VERSION }; + const withVirtualApiKey = virtualApiKeyId + ? { ...baseData, virtualApiKeyId } + : baseData; const data = { data: dontStripSession - ? { ...baseData, dontStripSession: true } - : baseData, + ? { ...withVirtualApiKey, dontStripSession: true } + : withVirtualApiKey, success: true, error: false, message: "Access allowed", diff --git a/server/routers/external.ts b/server/routers/external.ts index b50de6bea..6f44c1efe 100644 --- a/server/routers/external.ts +++ b/server/routers/external.ts @@ -1537,6 +1537,13 @@ authenticated.get( logs.queryAiUsageUsersRoles ); +authenticated.get( + "/org/:orgId/logs/ai/usage/virtual-api-keys", + verifyOrgAccess, + verifyUserHasAction(ActionsEnum.viewLogs), + logs.queryAiUsageVirtualApiKeys +); + authenticated.get( "/org/:orgId/blueprints", verifyOrgAccess, diff --git a/server/routers/integration.ts b/server/routers/integration.ts index 02e1ab987..cc283990e 100644 --- a/server/routers/integration.ts +++ b/server/routers/integration.ts @@ -1580,6 +1580,13 @@ authenticated.get( logs.queryAiUsageUsersRoles ); +authenticated.get( + "/org/:orgId/logs/ai/usage/virtual-api-keys", + verifyApiKeyOrgAccess, + verifyApiKeyHasAction(ActionsEnum.viewLogs), + logs.queryAiUsageVirtualApiKeys +); + authenticated.get( "/org/:orgId/logs/analytics", verifyApiKeyOrgAccess, diff --git a/src/app/[orgId]/settings/logs/ai/page.tsx b/src/app/[orgId]/settings/logs/ai/page.tsx index 5d51c118b..15c573337 100644 --- a/src/app/[orgId]/settings/logs/ai/page.tsx +++ b/src/app/[orgId]/settings/logs/ai/page.tsx @@ -12,6 +12,7 @@ import { useTranslations } from "next-intl"; import { getSevenDaysAgo } from "@app/lib/getSevenDaysAgo"; import { getPrivateResourceSettingsHref } from "@app/lib/launcherResourceAdminHref"; import { logQueries } from "@app/lib/queries"; +import { formatVirtualApiKeyPreview } from "@app/lib/virtualApiKeyFormat"; import { ColumnDef } from "@tanstack/react-table"; import { useQuery } from "@tanstack/react-query"; import axios from "axios"; @@ -50,6 +51,7 @@ export default function AiSessionLogsPage() { capability?: string; resourceId?: string; actor?: string; + virtualApiKeyId?: string; model?: string; isStream?: string; }>({ @@ -57,6 +59,7 @@ export default function AiSessionLogsPage() { capability: searchParams.get("capability") || undefined, resourceId: searchParams.get("resourceId") || undefined, actor: searchParams.get("actor") || undefined, + virtualApiKeyId: searchParams.get("virtualApiKeyId") || undefined, model: searchParams.get("model") || undefined, isStream: searchParams.get("isStream") || undefined }); @@ -135,6 +138,7 @@ export default function AiSessionLogsPage() { providers: [], resources: [], users: [], + virtualApiKeys: [], models: [] }; @@ -358,7 +362,9 @@ export default function AiSessionLogsPage() { }, cell: ({ row }) => { if (!row.original.resourceNiceId) { - return -; + return ( + - + ); } return ( ); } + }, + { + accessorKey: "virtualApiKeyId", + header: ({ column }) => { + return ( +
+ ({ + value: key.id, + label: + key.name ?? + (key.lastChars + ? formatVirtualApiKeyPreview( + key.id, + key.lastChars + ) + : key.id) + }) + )} + selectedValue={filters.virtualApiKeyId} + onValueChange={(value) => + handleFilterChange("virtualApiKeyId", value) + } + label={t("virtualApiKey")} + searchPlaceholder={t("searchPlaceholder")} + emptyMessage={t("emptySearchOptions")} + /> +
+ ); + }, + cell: ({ row }) => { + if (!row.original.virtualApiKeyId) { + return ( + - + ); + } + return ( +
+ + {row.original.virtualApiKeyName ?? + t("aiUsageUnnamedVirtualApiKey")} + + {row.original.virtualApiKeyLastChars && ( + + {formatVirtualApiKeyPreview( + row.original.virtualApiKeyId, + row.original.virtualApiKeyLastChars + )} + + )} +
+ ); + } } ]; const renderExpandedRow = (row: any) => { return (
-
+
{t("aiSessionId")}

@@ -490,6 +550,30 @@ export default function AiSessionLogsPage() { : "N/A"}

+
+ {t("virtualApiKey")} +

+ {row.virtualApiKeyId ? ( + <> + {row.virtualApiKeyName ?? + t("aiUsageUnnamedVirtualApiKey")} + {row.virtualApiKeyLastChars && ( + <> + {" "} + ( + {formatVirtualApiKeyPreview( + row.virtualApiKeyId, + row.virtualApiKeyLastChars + )} + ) + + )} + + ) : ( + t("noVirtualApiKey") + )} +

+
{row.usage && (
@@ -599,15 +683,25 @@ function generateSampleAiSessionLogs(): QueryAiSessionLogResponse["log"] { ]; const actors = ["alice@example.com", "bob@example.com", null]; const models = ["gpt-4o", "claude-sonnet-5", "gemini-2.5-pro"]; + const virtualApiKeysSample = [ + { id: "vak00001", name: "CI pipeline", lastChars: "ab12" }, + { id: "vak00002", name: null, lastChars: "cd34" }, + null + ]; const now = Date.now(); const sevenDaysAgoMs = now - 7 * 24 * 60 * 60 * 1000; return Array.from({ length: 10 }, (_, i) => { - const provider = providers[Math.floor(Math.random() * providers.length)]; + const provider = + providers[Math.floor(Math.random() * providers.length)]; const resource = resourcesSample[Math.floor(Math.random() * resourcesSample.length)]; const actor = actors[Math.floor(Math.random() * actors.length)]; + const virtualApiKey = + virtualApiKeysSample[ + Math.floor(Math.random() * virtualApiKeysSample.length) + ]; return { id: i, @@ -625,6 +719,9 @@ function generateSampleAiSessionLogs(): QueryAiSessionLogResponse["log"] { resourceType: "public", userId: actor ? `user-${i}` : null, userEmail: actor, + virtualApiKeyId: virtualApiKey?.id ?? null, + virtualApiKeyName: virtualApiKey?.name ?? null, + virtualApiKeyLastChars: virtualApiKey?.lastChars ?? null, requestedModel: models[Math.floor(Math.random() * models.length)], isStream: Math.random() > 0.5, requestBody: null, diff --git a/src/components/AiUsageAnalyticsData.tsx b/src/components/AiUsageAnalyticsData.tsx index c19ca2346..a5d3db2b8 100644 --- a/src/components/AiUsageAnalyticsData.tsx +++ b/src/components/AiUsageAnalyticsData.tsx @@ -29,6 +29,8 @@ import { ProvidersTab } from "./ai-usage-analytics/ProvidersTab"; import { ResourcesTab } from "./ai-usage-analytics/ResourcesTab"; import { RolesTab } from "./ai-usage-analytics/RolesTab"; import { UsersTab } from "./ai-usage-analytics/UsersTab"; +import { VirtualApiKeysTab } from "./ai-usage-analytics/VirtualApiKeysTab"; +import { formatVirtualApiKeyPreview } from "@app/lib/virtualApiKeyFormat"; export type AiUsageAnalyticsDataProps = { orgId: string; @@ -140,13 +142,20 @@ export function AiUsageAnalyticsData(props: AiUsageAnalyticsDataProps) { value: u.id, label: u.email ?? u.id })); + const virtualApiKeyOptions = (filterOptions?.virtualApiKeys ?? []).map( + (k) => ({ + value: k.id, + label: k.name ?? formatVirtualApiKeyPreview(k.id, k.lastChars) + }) + ); const tabs: TabItem[] = [ { title: t("aiUsageTabOverview"), href: "#" }, { title: t("aiUsageTabProviders"), href: "#" }, { title: t("aiUsageTabResources"), href: "#" }, { title: t("aiUsageRolesTab"), href: "#" }, - { title: t("aiUsageUsersTab"), href: "#" } + { title: t("aiUsageUsersTab"), href: "#" }, + { title: t("aiUsageVirtualApiKeysTab"), href: "#" } ]; return ( @@ -180,7 +189,9 @@ export function AiUsageAnalyticsData(props: AiUsageAnalyticsDataProps) { value={filters.providerId?.toString()} options={providerOptions} placeholder={t("aiUsageFilterAllProviders")} - onValueChange={(v) => setFilter("providerId", v)} + onValueChange={(v) => + setFilter("providerId", v) + } /> setFilter("resourceId", v)} + onValueChange={(v) => + setFilter("resourceId", v) + } /> setFilter("userId", v)} /> + + setFilter("virtualApiKeyId", v) + } + /> {!isEmptySearchParams && (
); diff --git a/src/components/ai-usage-analytics/VirtualApiKeysTab.tsx b/src/components/ai-usage-analytics/VirtualApiKeysTab.tsx new file mode 100644 index 000000000..9ed8edb3d --- /dev/null +++ b/src/components/ai-usage-analytics/VirtualApiKeysTab.tsx @@ -0,0 +1,108 @@ +"use client"; + +import { useQuery } from "@tanstack/react-query"; +import { aiUsageAnalyticsQueries } from "@app/lib/queries"; +import type { AiUsageAnalyticsFilters } from "@app/lib/queries"; +import { formatVirtualApiKeyPreview } from "@app/lib/virtualApiKeyFormat"; +import { useTranslations } from "next-intl"; +import { Card, CardContent, CardHeader } from "@app/components/ui/card"; +import { ToggleableTrendChart } from "./ToggleableTrendChart"; +import { TopEntitiesList, type TopEntity } from "./TopEntitiesList"; +import { buildSeriesFromData, formatCost } from "./shared"; + +type VirtualApiKeysTabProps = { + orgId: string; + filters: AiUsageAnalyticsFilters; +}; + +const UNKNOWN_VIRTUAL_API_KEY_KEY = "unknown"; + +export function VirtualApiKeysTab(props: VirtualApiKeysTabProps) { + const t = useTranslations(); + const { data, isLoading } = useQuery( + aiUsageAnalyticsQueries.virtualApiKeys({ + orgId: props.orgId, + filters: props.filters + }) + ); + + const labelByKey = new Map(); + for (const k of data?.topVirtualApiKeys ?? []) { + if (k.virtualApiKeyId) { + labelByKey.set(k.virtualApiKeyId, k.name ?? k.virtualApiKeyId); + } + } + const virtualApiKeyLabelFor = (key: string) => + key === UNKNOWN_VIRTUAL_API_KEY_KEY + ? t("aiUsageUnknownVirtualApiKey") + : (labelByKey.get(key) ?? key); + + const virtualApiKeyCostSeries = buildSeriesFromData( + data?.virtualApiKeyCostPerDay ?? [], + virtualApiKeyLabelFor, + t("aiUsageOther") + ); + const virtualApiKeyTokensSeries = buildSeriesFromData( + data?.virtualApiKeyTokensPerDay ?? [], + virtualApiKeyLabelFor, + t("aiUsageOther") + ); + + const topVirtualApiKeys: TopEntity[] = (data?.topVirtualApiKeys ?? []).map( + (k) => ({ + key: k.virtualApiKeyId ?? UNKNOWN_VIRTUAL_API_KEY_KEY, + label: k.virtualApiKeyId + ? (k.name ?? t("aiUsageUnnamedVirtualApiKey")) + : t("aiUsageUnknownVirtualApiKey"), + sublabel: + k.virtualApiKeyId && k.lastChars + ? formatVirtualApiKeyPreview(k.virtualApiKeyId, k.lastChars) + : undefined, + requests: k.requests, + totalTokens: k.totalTokens, + costUsd: k.costUsd + }) + ); + + return ( +
+ + +

+ {t("aiUsageTopVirtualApiKeys")} +

+ + + + + +
+ + + formatCost(v)} + /> + + + + + + + +
+
+ ); +} diff --git a/src/lib/queries.ts b/src/lib/queries.ts index bf7f69b8f..56a41cee8 100644 --- a/src/lib/queries.ts +++ b/src/lib/queries.ts @@ -12,7 +12,8 @@ import type { QueryAiUsageOverviewResponse, QueryAiUsageProvidersResponse, QueryAiUsageResourcesResponse, - QueryAiUsageUsersRolesResponse + QueryAiUsageUsersRolesResponse, + QueryAiUsageVirtualApiKeysResponse } from "@server/routers/auditLogs"; import type { QueryAccessAuditLogResponse, @@ -954,7 +955,8 @@ export const aiUsageAnalyticsFiltersSchema = z.object({ model: z.string().optional().catch(undefined), resourceId: z.coerce.number().optional().catch(undefined), roleId: z.coerce.number().optional().catch(undefined), - userId: z.string().optional().catch(undefined) + userId: z.string().optional().catch(undefined), + virtualApiKeyId: z.string().optional().catch(undefined) }); export type AiUsageAnalyticsFilters = z.output< @@ -1086,6 +1088,7 @@ export const aiSessionLogsFiltersSchema = z.object({ capability: z.string().optional().catch(undefined), resourceId: z.string().optional().catch(undefined), actor: z.string().optional().catch(undefined), + virtualApiKeyId: z.string().optional().catch(undefined), model: z.string().optional().catch(undefined), isStream: z.string().optional().catch(undefined) }); @@ -1284,7 +1287,12 @@ export const aiUsageAnalyticsQueries = { filters: Pick; }) => queryOptions({ - queryKey: ["AI_USAGE_ANALYTICS", orgId, "FILTERS", filters] as const, + queryKey: [ + "AI_USAGE_ANALYTICS", + orgId, + "FILTERS", + filters + ] as const, queryFn: async ({ signal, meta }) => { const res = await meta!.api.get< AxiosResponse @@ -1304,7 +1312,12 @@ export const aiUsageAnalyticsQueries = { filters: AiUsageAnalyticsFilters; }) => queryOptions({ - queryKey: ["AI_USAGE_ANALYTICS", orgId, "OVERVIEW", filters] as const, + queryKey: [ + "AI_USAGE_ANALYTICS", + orgId, + "OVERVIEW", + filters + ] as const, queryFn: async ({ signal, meta }) => { const res = await meta!.api.get< AxiosResponse @@ -1330,7 +1343,12 @@ export const aiUsageAnalyticsQueries = { filters: AiUsageAnalyticsFilters; }) => queryOptions({ - queryKey: ["AI_USAGE_ANALYTICS", orgId, "PROVIDERS", filters] as const, + queryKey: [ + "AI_USAGE_ANALYTICS", + orgId, + "PROVIDERS", + filters + ] as const, queryFn: async ({ signal, meta }) => { const res = await meta!.api.get< AxiosResponse @@ -1356,7 +1374,12 @@ export const aiUsageAnalyticsQueries = { filters: AiUsageAnalyticsFilters; }) => queryOptions({ - queryKey: ["AI_USAGE_ANALYTICS", orgId, "RESOURCES", filters] as const, + queryKey: [ + "AI_USAGE_ANALYTICS", + orgId, + "RESOURCES", + filters + ] as const, queryFn: async ({ signal, meta }) => { const res = await meta!.api.get< AxiosResponse @@ -1403,6 +1426,37 @@ export const aiUsageAnalyticsQueries = { } return false; } + }), + + virtualApiKeys: ({ + orgId, + filters + }: { + orgId: string; + filters: AiUsageAnalyticsFilters; + }) => + queryOptions({ + queryKey: [ + "AI_USAGE_ANALYTICS", + orgId, + "VIRTUAL_API_KEYS", + filters + ] as const, + queryFn: async ({ signal, meta }) => { + const res = await meta!.api.get< + AxiosResponse + >(`/org/${orgId}/logs/ai/usage/virtual-api-keys`, { + params: filters, + signal + }); + return res.data.data; + }, + refetchInterval: (query) => { + if (query.state.data) { + return durationToMs(30, "seconds"); + } + return false; + } }) };