Update the usage records to use logs db

This commit is contained in:
Owen
2026-08-19 12:02:22 -04:00
parent a7bbafd2f0
commit 437ec50019
8 changed files with 38 additions and 29 deletions
+4 -3
View File
@@ -16,6 +16,7 @@ import {
aiModels,
aiUsageRecords,
db,
logsDb,
userOrgRoles
} from "@server/db";
import { modelKeyMatches } from "@server/lib/aiModelKeyMatch";
@@ -168,7 +169,7 @@ async function sumUsageAmount(
): Promise<number> {
const column =
unit === "usd" ? aiUsageRecords.costUsd : aiUsageRecords.totalTokens;
const [row] = await db
const [row] = await logsDb
.select({ total: sql<number>`coalesce(sum(${column}), 0)` })
.from(aiUsageRecords)
.where(where);
@@ -200,7 +201,7 @@ export async function sumUsageForBudget(
if (!model) {
return 0;
}
const rows = await db
const rows = await logsDb
.select({
requestedModel: aiUsageRecords.requestedModel,
costUsd: aiUsageRecords.costUsd,
@@ -496,7 +497,7 @@ async function flushUsageRecords() {
try {
// Use a transaction to ensure all inserts succeed or fail together
await db.transaction(async (tx) => {
await logsDb.transaction(async (tx) => {
// Batch insert in groups to avoid overwhelming the database
const DB_BATCH_SIZE = 25;
for (let i = 0; i < recordsToWrite.length; i += DB_BATCH_SIZE) {
@@ -312,7 +312,7 @@ async function enrichWithDetails(
>();
const sessionIds = logs.map((log) => log.sessionId);
if (sessionIds.length > 0) {
const usageDetails = await primaryDb
const usageDetails = await logsDb
.select({
sessionId: aiUsageRecords.sessionId,
promptTokens: aiUsageRecords.promptTokens,
@@ -1,5 +1,6 @@
import {
db,
logsDb,
aiUsageRecords,
aiProviders,
resources,
@@ -78,22 +79,22 @@ async function query(data: Q) {
uniqueUsers,
uniqueVirtualApiKeys
] = await Promise.all([
db
logsDb
.selectDistinct({ id: aiUsageRecords.providerId })
.from(aiUsageRecords)
.where(baseConditions)
.limit(DISTINCT_LIMIT + 1),
db
logsDb
.selectDistinct({ model: aiUsageRecords.requestedModel })
.from(aiUsageRecords)
.where(baseConditions)
.limit(DISTINCT_LIMIT + 1),
db
logsDb
.selectDistinct({ id: aiUsageRecords.resourceId })
.from(aiUsageRecords)
.where(and(baseConditions, not(isNull(aiUsageRecords.resourceId))))
.limit(DISTINCT_LIMIT + 1),
db
logsDb
.selectDistinct({ id: aiUsageRecords.siteResourceId })
.from(aiUsageRecords)
.where(
@@ -104,12 +105,12 @@ async function query(data: Q) {
)
)
.limit(DISTINCT_LIMIT + 1),
db
logsDb
.selectDistinct({ userId: aiUsageRecords.userId })
.from(aiUsageRecords)
.where(and(baseConditions, not(isNull(aiUsageRecords.userId))))
.limit(DISTINCT_LIMIT + 1),
db
logsDb
.selectDistinct({ id: aiUsageRecords.virtualApiKeyId })
.from(aiUsageRecords)
.where(
@@ -1,4 +1,4 @@
import { db, aiUsageRecords } from "@server/db";
import { logsDb, aiUsageRecords } from "@server/db";
import { registry } from "@server/openApi";
import { NextFunction } from "express";
import { Request, Response } from "express";
@@ -29,7 +29,7 @@ async function query(data: Q) {
const roleUserIds = await resolveRoleUserIds(data.orgId, data.roleId);
const baseConditions = buildAiUsageWhere(data, roleUserIds);
const [totalsRow] = await db
const [totalsRow] = await logsDb
.select({
requests: count(),
promptTokens: sql<number>`COALESCE(SUM(${aiUsageRecords.promptTokens}), 0)`,
@@ -46,7 +46,7 @@ async function query(data: Q) {
const dayExpr = dayBucketExpr();
const requestsPerDay = await db
const requestsPerDay = await logsDb
.select({
day: dayExpr.as("day"),
requests: count()
@@ -56,7 +56,7 @@ async function query(data: Q) {
.groupBy(dayExpr)
.orderBy(dayExpr);
const tokensPerDay = await db
const tokensPerDay = await logsDb
.select({
day: dayExpr.as("day"),
promptTokens: sql<number>`COALESCE(SUM(${aiUsageRecords.promptTokens}), 0)`,
@@ -70,7 +70,7 @@ async function query(data: Q) {
.groupBy(dayExpr)
.orderBy(dayExpr);
const costPerDay = await db
const costPerDay = await logsDb
.select({
day: dayExpr.as("day"),
cost: sql<number>`COALESCE(SUM(${aiUsageRecords.costUsd}), 0)`
@@ -80,7 +80,7 @@ async function query(data: Q) {
.groupBy(dayExpr)
.orderBy(dayExpr);
const modelByDay = await db
const modelByDay = await logsDb
.select({
day: dayExpr.as("day"),
model: aiUsageRecords.requestedModel,
@@ -117,7 +117,7 @@ async function query(data: Q) {
topModelsByTokens
);
const topModelsRaw = await db
const topModelsRaw = await logsDb
.select({
model: aiUsageRecords.requestedModel,
requests: count(),
@@ -1,4 +1,4 @@
import { db, aiUsageRecords, aiProviders } from "@server/db";
import { db, logsDb, aiUsageRecords, aiProviders } from "@server/db";
import { registry } from "@server/openApi";
import { NextFunction } from "express";
import { Request, Response } from "express";
@@ -29,7 +29,7 @@ async function query(data: Q) {
const baseConditions = buildAiUsageWhere(data, roleUserIds);
const dayExpr = dayBucketExpr();
const providerByDay = await db
const providerByDay = await logsDb
.select({
day: dayExpr.as("day"),
providerId: aiUsageRecords.providerId,
@@ -69,7 +69,7 @@ async function query(data: Q) {
topByTokens
);
const topProvidersRaw = await db
const topProvidersRaw = await logsDb
.select({
providerId: aiUsageRecords.providerId,
requests: count(),
@@ -1,4 +1,4 @@
import { db, aiUsageRecords, resources, siteResources } from "@server/db";
import { db, logsDb, aiUsageRecords, resources, siteResources } from "@server/db";
import { registry } from "@server/openApi";
import { NextFunction } from "express";
import { Request, Response } from "express";
@@ -39,7 +39,7 @@ async function query(data: Q) {
const baseConditions = buildAiUsageWhere(data, roleUserIds);
const dayExpr = dayBucketExpr();
const resourceByDay = await db
const resourceByDay = await logsDb
.select({
day: dayExpr.as("day"),
resourceId: aiUsageRecords.resourceId,
@@ -80,7 +80,7 @@ async function query(data: Q) {
topByTokens
);
const topResourcesRaw = await db
const topResourcesRaw = await logsDb
.select({
resourceId: aiUsageRecords.resourceId,
siteResourceId: aiUsageRecords.siteResourceId,
@@ -1,4 +1,11 @@
import { db, aiUsageRecords, users, roles, userOrgRoles } from "@server/db";
import {
db,
logsDb,
aiUsageRecords,
users,
roles,
userOrgRoles
} from "@server/db";
import { registry } from "@server/openApi";
import { NextFunction } from "express";
import { Request, Response } from "express";
@@ -36,7 +43,7 @@ async function query(data: Q) {
// userId, so role totals are derived by expanding each user's usage into
// every role they hold in the org (per-role double counting for
// multi-role users is expected/accepted).
const userByDay = await db
const userByDay = await logsDb
.select({
day: dayExpr.as("day"),
userId: aiUsageRecords.userId,
@@ -48,7 +55,7 @@ async function query(data: Q) {
.groupBy(dayExpr, aiUsageRecords.userId)
.orderBy(dayExpr);
const userTotalsRaw = await db
const userTotalsRaw = await logsDb
.select({
userId: aiUsageRecords.userId,
requests: count(),
@@ -1,4 +1,4 @@
import { db, aiUsageRecords, virtualApiKeys } from "@server/db";
import { db, logsDb, aiUsageRecords, virtualApiKeys } from "@server/db";
import { registry } from "@server/openApi";
import { NextFunction } from "express";
import { Request, Response } from "express";
@@ -31,7 +31,7 @@ async function query(data: Q) {
const baseConditions = buildAiUsageWhere(data, roleUserIds);
const dayExpr = dayBucketExpr();
const virtualApiKeyByDay = await db
const virtualApiKeyByDay = await logsDb
.select({
day: dayExpr.as("day"),
virtualApiKeyId: aiUsageRecords.virtualApiKeyId,
@@ -43,7 +43,7 @@ async function query(data: Q) {
.groupBy(dayExpr, aiUsageRecords.virtualApiKeyId)
.orderBy(dayExpr);
const virtualApiKeyTotalsRaw = await db
const virtualApiKeyTotalsRaw = await logsDb
.select({
virtualApiKeyId: aiUsageRecords.virtualApiKeyId,
requests: count(),