send, process, store, display virtual api key ai information in usage and sessions

This commit is contained in:
Owen
2026-08-12 11:20:25 -04:00
parent 355294a2d4
commit bed5817da3
20 changed files with 849 additions and 68 deletions
@@ -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
);
}
+1
View File
@@ -8,3 +8,4 @@ export * from "./queryAiUsageOverview";
export * from "./queryAiUsageProviders";
export * from "./queryAiUsageResources";
export * from "./queryAiUsageUsersRoles";
export * from "./queryAiUsageVirtualApiKeys";
+79 -4
View File
@@ -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()
};
}
@@ -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<typeof queryAiUsageFilterOptionsCombined>;
function sortNamedFilterOptions<T extends { id: number; name: string | null }>(
@@ -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<ReturnType<typeof query>>;
export type QueryAiUsageFilterOptionsResponse = Awaited<
ReturnType<typeof query>
>;
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)
)
);
}
@@ -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<number>`COALESCE(SUM(${aiUsageRecords.costUsd}), 0)`,
tokens: sql<number>`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<number>`COALESCE(SUM(${aiUsageRecords.totalTokens}), 0)`,
costUsd: sql<number>`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<string, number>();
const virtualApiKeyTokenTotals = new Map<string, number>();
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<typeof query>
>;
export async function queryAiUsageVirtualApiKeys(
req: Request,
res: Response,
next: NextFunction
): Promise<any> {
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<QueryAiUsageVirtualApiKeysResponse>(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")
);
}
}
+8
View File
@@ -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[];
};
};