Merge pull request #3614 from fosrl/dev

Dev

server/lib/aiBudgetEnforcement.ts

@@ -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) {

server/routers/auditLogs/queryAiSessionLog.ts

@@ -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,

server/routers/auditLogs/queryAiUsageFilterOptions.ts

@@ -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(

server/routers/auditLogs/queryAiUsageOverview.ts

@@ -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(),

server/routers/auditLogs/queryAiUsageProviders.ts

@@ -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(),

server/routers/auditLogs/queryAiUsageResources.ts

@@ -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,

server/routers/auditLogs/queryAiUsageUsersRoles.ts

@@ -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(),

server/routers/auditLogs/queryAiUsageVirtualApiKeys.ts

@@ -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(),

server/setup/scriptsPg/1.22.0.ts

@@ -18,9 +18,6 @@ export default async function migration() {
     try {
         await db.execute(sql`BEGIN`);
 
-        await db.execute(sql`
-        `);
-
         await db.execute(sql`
             CREATE TABLE "aiBudgetBreachEvents" (
                 "id" serial PRIMARY KEY NOT NULL,
@@ -454,14 +451,14 @@ export default async function migration() {
             throw new Error(fromZodError(parsedConfig.error).toString());
         }
 
-        traefikConfig.experimental.plugins.badger.version = "v1.6.0";
+        traefikConfig.experimental.plugins.badger.version = "v1.6.1";
 
         const updatedTraefikYaml = yaml.dump(traefikConfig);
 
         fs.writeFileSync(traefikPath, updatedTraefikYaml, "utf8");
 
         console.log(
-            "Updated the version of Badger in your Traefik configuration to v1.6.0"
+            "Updated the version of Badger in your Traefik configuration to v1.6.1"
         );
     } catch (e) {
         console.log(

server/setup/scriptsSqlite/1.22.0.ts

@@ -456,14 +456,14 @@ export default async function migration() {
             throw new Error(fromZodError(parsedConfig.error).toString());
         }
 
-        traefikConfig.experimental.plugins.badger.version = "v1.6.0";
+        traefikConfig.experimental.plugins.badger.version = "v1.6.1";
 
         const updatedTraefikYaml = yaml.dump(traefikConfig);
 
         fs.writeFileSync(traefikPath, updatedTraefikYaml, "utf8");
 
         console.log(
-            "Updated the version of Badger in your Traefik configuration to v1.6.0"
+            "Updated the version of Badger in your Traefik configuration to v1.6.1"
         );
     } catch (e) {
         console.log(