From cd957acdd46802057f6ae36d52b3fe691d9ad60f Mon Sep 17 00:00:00 2001
From: Matic Jurglic <matic@jurglic.si>
Date: Mon, 18 Nov 2024 16:40:10 +0100
Subject: [PATCH 1/9] Track credit consumption in AI bot
---
packages/ai-bot/lib/send-response.ts | 2 +-
packages/ai-bot/main.ts | 113 +++++++++++++-
packages/ai-bot/package.json | 4 +-
packages/billing/billing-queries.ts | 60 ++++++++
packages/postgres/pg-config.ts | 4 +-
packages/realm-server/tests/billing-test.ts | 154 ++++++++++++++++++++
6 files changed, 329 insertions(+), 8 deletions(-)
diff --git a/packages/ai-bot/lib/send-response.ts b/packages/ai-bot/lib/send-response.ts
index 4b913c8987..494aa114e5 100644
--- a/packages/ai-bot/lib/send-response.ts
+++ b/packages/ai-bot/lib/send-response.ts
@@ -133,7 +133,7 @@ export class Responder {
}
}
- async onError(error: OpenAIError) {
+ async onError(error: OpenAIError | string) {
Sentry.captureException(error);
return await sendError(
this.client,
diff --git a/packages/ai-bot/main.ts b/packages/ai-bot/main.ts
index 910a3851ee..4d3f169dfb 100644
--- a/packages/ai-bot/main.ts
+++ b/packages/ai-bot/main.ts
@@ -7,7 +7,7 @@ import {
type MatrixEvent,
} from 'matrix-js-sdk';
import OpenAI from 'openai';
-import { logger, aiBotUsername } from '@cardstack/runtime-common';
+import { logger, aiBotUsername, retry } from '@cardstack/runtime-common';
import {
constructHistory,
getModifyPrompt,
@@ -24,21 +24,94 @@ import { handleDebugCommands } from './lib/debug';
import { MatrixClient } from './lib/matrix';
import type { MatrixEvent as DiscreteMatrixEvent } from 'https://cardstack.com/base/matrix-event';
import * as Sentry from '@sentry/node';
+import { PgAdapter, TransactionManager } from '@cardstack/postgres';
+import {
+ getUserByMatrixUserId,
+ spendCredits,
+} from '@cardstack/billing/billing-queries';
let log = logger('ai-bot');
+let trackAiUsageCostPromises = new Map<string, Promise<void>>();
+
class Assistant {
private openai: OpenAI;
private client: MatrixClient;
+ private pgAdapter: PgAdapter;
id: string;
constructor(client: MatrixClient, id: string) {
this.openai = new OpenAI({
- baseURL: 'https://openrouter.ai/api/v1', // We use openrouter so that we can track usage cost in $
+ baseURL: 'https://openrouter.ai/api/v1',
apiKey: process.env.OPENROUTER_API_KEY,
});
this.id = id;
this.client = client;
+ this.pgAdapter = new PgAdapter();
+ }
+
+ async fetchGenerationCost(generationId: string) {
+ let response = await (
+ await fetch(
+ `https://openrouter.ai/api/v1/generation?id=${generationId}`,
+ {
+ headers: {
+ Authorization: `Bearer ${process.env.OPENROUTER_API_KEY}`,
+ },
+ },
+ )
+ ).json();
+
+ if (response.error && response.error.includes('not found')) {
+ return null;
+ }
+
+ return response.data.total_cost;
+ }
+
+ async trackAiUsageCost(matrixUserId: string, generationId: string) {
+ try {
+ let costInUsd = await retry(
+ () => this.fetchGenerationCost(generationId),
+ {
+ retries: 10,
+ delayMs: 500,
+ },
+ );
+
+ let creditsConsumed = Math.round(costInUsd / 0.001);
+
+ let user = await getUserByMatrixUserId(this.pgAdapter, matrixUserId);
+ if (!user) {
+ throw new Error('should not happen: user with matrix id not found');
+ }
+
+ let txManager = new TransactionManager(this.pgAdapter);
+ await txManager.withTransaction(async () => {
+ await spendCredits(this.pgAdapter, user!.id, creditsConsumed);
+ });
+
+ log.info('Successfully tracked AI usage:', {
+ matrixUserId,
+ generationId,
+ creditsConsumed,
+ });
+ } catch (err) {
+ log.error('Failed to track AI usage:', err);
+ throw err;
+ }
+ }
+
+ async enqueueTrackAiUsageCost(matrixUserId: string, generationId: string) {
+ if (trackAiUsageCostPromises.has(matrixUserId)) {
+ return;
+ }
+ trackAiUsageCostPromises.set(
+ matrixUserId,
+ this.trackAiUsageCost(matrixUserId, generationId).finally(() => {
+ trackAiUsageCostPromises.delete(matrixUserId);
+ }),
+ );
}
getResponse(history: DiscreteMatrixEvent[]) {
@@ -177,6 +250,21 @@ Common issues are:
}
const responder = new Responder(client, room.roomId);
+
+ // Do not generate new responses if previous ones' cost is still being reported
+ let pendingCreditsConsumptionPromise = trackAiUsageCostPromises.get(
+ event.getSender()!,
+ );
+ if (pendingCreditsConsumptionPromise) {
+ try {
+ await pendingCreditsConsumptionPromise;
+ } catch (e) {
+ log.error(e);
+ return responder.onError(
+ 'There was an error reporting Boxel Credits usage. Try again or contact support if the problem persists.',
+ );
+ }
+ }
await responder.initialize();
if (historyError) {
@@ -186,9 +274,11 @@ Common issues are:
return;
}
+ let generationId: string | undefined;
const runner = assistant
.getResponse(history)
.on('chunk', async (chunk, _snapshot) => {
+ generationId = chunk.id;
await responder.onChunk(chunk);
})
.on('content', async (_delta, snapshot) => {
@@ -200,9 +290,24 @@ Common issues are:
.on('error', async (error) => {
await responder.onError(error);
});
+
// We also need to catch the error when getting the final content
- let finalContent = await runner.finalContent().catch(responder.onError);
- await responder.finalize(finalContent);
+ let finalContent;
+ try {
+ finalContent = await runner.finalContent();
+ await responder.finalize(finalContent);
+ } catch (error) {
+ await responder.onError(error);
+ } finally {
+ if (generationId) {
+ let userMatrixId = event.getSender()!;
+
+ assistant.enqueueTrackAiUsageCost(
+ userMatrixId,
+ generationId as string,
+ );
+ }
+ }
if (shouldSetRoomTitle(eventList, aiBotUserId, event)) {
return await assistant.setTitle(room.roomId, history, event);
diff --git a/packages/ai-bot/package.json b/packages/ai-bot/package.json
index 9578581ce1..a946eccbce 100644
--- a/packages/ai-bot/package.json
+++ b/packages/ai-bot/package.json
@@ -1,7 +1,9 @@
{
"name": "@cardstack/ai-bot",
"dependencies": {
- "@cardstack/runtime-common": "workspace:^",
+ "@cardstack/runtime-common": "workspace:*",
+ "@cardstack/postgres": "workspace:*",
+ "@cardstack/billing": "workspace:*",
"@sentry/node": "^8.31.0",
"@types/node": "^18.18.5",
"@types/stream-chain": "^2.0.1",
diff --git a/packages/billing/billing-queries.ts b/packages/billing/billing-queries.ts
index acc7e5e2ff..07f513f79d 100644
--- a/packages/billing/billing-queries.ts
+++ b/packages/billing/billing-queries.ts
@@ -490,3 +490,63 @@ export async function expireRemainingPlanAllowanceInSubscriptionCycle(
subscriptionCycleId,
});
}
+
+export async function spendCredits(
+ dbAdapter: DBAdapter,
+ userId: string,
+ creditsToSpend: number,
+) {
+ let subscription = await getMostRecentSubscription(dbAdapter, userId);
+ if (!subscription) {
+ throw new Error('subscription not found');
+ }
+ let subscriptionCycle = await getMostRecentSubscriptionCycle(
+ dbAdapter,
+ subscription.id,
+ );
+ if (!subscriptionCycle) {
+ throw new Error('subscription cycle not found');
+ }
+ let availablePlanAllowanceCredits = await sumUpCreditsLedger(dbAdapter, {
+ creditType: [
+ 'plan_allowance',
+ 'plan_allowance_used',
+ 'plan_allowance_expired',
+ ],
+ userId,
+ });
+
+ if (availablePlanAllowanceCredits >= creditsToSpend) {
+ await addToCreditsLedger(dbAdapter, {
+ userId,
+ creditAmount: -creditsToSpend,
+ creditType: 'plan_allowance_used',
+ subscriptionCycleId: subscriptionCycle.id,
+ });
+ } else {
+ // If user does not have enough plan allowance credits to cover the spend, use extra credits
+ let availableExtraCredits = await sumUpCreditsLedger(dbAdapter, {
+ creditType: ['extra_credit', 'extra_credit_used'],
+ userId,
+ });
+ let planAllowanceToSpend = availablePlanAllowanceCredits; // Spend all plan allowance credits first
+ let extraCreditsToSpend = creditsToSpend - planAllowanceToSpend;
+ if (extraCreditsToSpend > availableExtraCredits) {
+ extraCreditsToSpend = availableExtraCredits;
+ }
+
+ await addToCreditsLedger(dbAdapter, {
+ userId,
+ creditAmount: -planAllowanceToSpend,
+ creditType: 'plan_allowance_used',
+ subscriptionCycleId: subscriptionCycle.id,
+ });
+
+ await addToCreditsLedger(dbAdapter, {
+ userId,
+ creditAmount: -extraCreditsToSpend,
+ creditType: 'extra_credit_used',
+ subscriptionCycleId: subscriptionCycle.id,
+ });
+ }
+}
diff --git a/packages/postgres/pg-config.ts b/packages/postgres/pg-config.ts
index a2b7256275..d9d3939577 100644
--- a/packages/postgres/pg-config.ts
+++ b/packages/postgres/pg-config.ts
@@ -3,9 +3,9 @@ import { type ClientConfig } from 'pg';
export function postgresConfig(defaultConfig: ClientConfig = {}) {
return Object.assign({}, defaultConfig, {
host: process.env.PGHOST || 'localhost',
- port: process.env.PGPORT || '5432',
+ port: process.env.PGPORT || '5435',
user: process.env.PGUSER || 'postgres',
password: process.env.PGPASSWORD || undefined,
- database: process.env.PGDATABASE || 'postgres',
+ database: process.env.PGDATABASE || 'boxel',
});
}
diff --git a/packages/realm-server/tests/billing-test.ts b/packages/realm-server/tests/billing-test.ts
index c0cd26a1bc..ae91da67de 100644
--- a/packages/realm-server/tests/billing-test.ts
+++ b/packages/realm-server/tests/billing-test.ts
@@ -18,6 +18,9 @@ import {
addToCreditsLedger,
insertSubscription,
User,
+ spendCredits,
+ Plan,
+ Subscription,
} from '@cardstack/billing/billing-queries';
import {
@@ -25,6 +28,7 @@ import {
StripeSubscriptionDeletedWebhookEvent,
StripeCheckoutSessionCompletedWebhookEvent,
} from '@cardstack/billing/stripe-webhook-handlers';
+import { add } from 'date-fns';
async function fetchStripeEvents(dbAdapter: PgAdapter) {
return await query(dbAdapter, [`SELECT * FROM stripe_events`]);
@@ -796,4 +800,154 @@ module('billing', function (hooks) {
assert.strictEqual(availableExtraCredits, 25000);
});
});
+
+ // eslint-disable-next-line qunit/no-only
+ module.only('ai usage tracking', function (hooks) {
+ let user: User;
+ let creatorPlan: Plan;
+ let subscription: Subscription;
+ let subscriptionCycle: SubscriptionCycle;
+
+ hooks.beforeEach(async function () {
+ user = await insertUser(dbAdapter, 'testuser', 'cus_123');
+ creatorPlan = await insertPlan(
+ dbAdapter,
+ 'Creator',
+ 12,
+ 2500,
+ 'prod_creator',
+ );
+ subscription = await insertSubscription(dbAdapter, {
+ user_id: user.id,
+ plan_id: creatorPlan.id,
+ started_at: 1,
+ status: 'active',
+ stripe_subscription_id: 'sub_1234567890',
+ });
+ subscriptionCycle = await insertSubscriptionCycle(dbAdapter, {
+ subscriptionId: subscription.id,
+ periodStart: 1,
+ periodEnd: 2,
+ });
+ });
+
+ test('spends ai credits correctly when no extra credits are available', async function (assert) {
+ // User receives 2500 credits for the creator plan and spends 2490 credits
+ await addToCreditsLedger(dbAdapter, {
+ userId: user.id,
+ creditAmount: creatorPlan.creditsIncluded,
+ creditType: 'plan_allowance',
+ subscriptionCycleId: subscriptionCycle.id,
+ });
+
+ await addToCreditsLedger(dbAdapter, {
+ userId: user.id,
+ creditAmount: -2490,
+ creditType: 'plan_allowance_used',
+ subscriptionCycleId: subscriptionCycle.id,
+ });
+
+ assert.strictEqual(
+ await sumUpCreditsLedger(dbAdapter, {
+ userId: user.id,
+ }),
+ 10,
+ );
+
+ await spendCredits(dbAdapter, user.id, 2);
+
+ assert.strictEqual(
+ await sumUpCreditsLedger(dbAdapter, {
+ userId: user.id,
+ }),
+ 8,
+ );
+
+ await spendCredits(dbAdapter, user.id, 5);
+
+ assert.strictEqual(
+ await sumUpCreditsLedger(dbAdapter, {
+ userId: user.id,
+ }),
+ 3,
+ );
+
+ // Make sure that we can't spend more credits than the user has - in this case user has 3 credits left and we try to spend 5
+ await spendCredits(dbAdapter, user.id, 5);
+ assert.strictEqual(
+ await sumUpCreditsLedger(dbAdapter, {
+ userId: user.id,
+ }),
+ 0,
+ );
+ });
+
+ test('spends ai credits correctly when extra credits are available', async function (assert) {
+ // User receives 2500 credits for the creator plan and spends 2490 credits
+ await addToCreditsLedger(dbAdapter, {
+ userId: user.id,
+ creditAmount: creatorPlan.creditsIncluded,
+ creditType: 'plan_allowance',
+ subscriptionCycleId: subscriptionCycle.id,
+ });
+
+ await addToCreditsLedger(dbAdapter, {
+ userId: user.id,
+ creditAmount: -2490,
+ creditType: 'plan_allowance_used',
+ subscriptionCycleId: subscriptionCycle.id,
+ });
+
+ assert.strictEqual(
+ await sumUpCreditsLedger(dbAdapter, {
+ userId: user.id,
+ }),
+ 10,
+ );
+
+ // Add 5 extra credits
+ await addToCreditsLedger(dbAdapter, {
+ userId: user.id,
+ creditAmount: 5,
+ creditType: 'extra_credit',
+ subscriptionCycleId: null,
+ });
+
+ // User has 15 credits in total: 10 credits from the plan allowance and 5 extra credits
+ assert.strictEqual(
+ await sumUpCreditsLedger(dbAdapter, {
+ userId: user.id,
+ }),
+ 15,
+ );
+
+ // This should spend 10 credits from the plan allowance and 2 from the extra credits
+ await spendCredits(dbAdapter, user.id, 12);
+
+ // Plan allowance is now 0, 3 credits left from the extra credits
+ assert.strictEqual(
+ await sumUpCreditsLedger(dbAdapter, {
+ userId: user.id,
+ }),
+ 3,
+ );
+
+ // Make sure the available credits come from the extra credits and not the plan allowance
+ assert.strictEqual(
+ await sumUpCreditsLedger(dbAdapter, {
+ userId: user.id,
+ creditType: ['plan_allowance', 'plan_allowance_used'],
+ }),
+ 0,
+ );
+
+ assert.strictEqual(
+ await sumUpCreditsLedger(dbAdapter, {
+ userId: user.id,
+ creditType: ['extra_credit', 'extra_credit_used'],
+ }),
+ 3,
+ );
+ });
+ });
});
From 41b520ae97bbbd898ffc2917ef518b3a69156141 Mon Sep 17 00:00:00 2001
From: Matic Jurglic <matic@jurglic.si>
Date: Tue, 19 Nov 2024 10:59:37 +0100
Subject: [PATCH 2/9] AI bot saves credits usage for user in credits ledger
---
packages/ai-bot/lib/ai-cost.ts | 75 +++++++++++++++++++
packages/ai-bot/main.ts | 82 +++------------------
packages/billing/billing-queries.ts | 30 ++++----
packages/realm-server/tests/billing-test.ts | 4 +-
4 files changed, 105 insertions(+), 86 deletions(-)
create mode 100644 packages/ai-bot/lib/ai-cost.ts
diff --git a/packages/ai-bot/lib/ai-cost.ts b/packages/ai-bot/lib/ai-cost.ts
new file mode 100644
index 0000000000..0e1f7d0bd2
--- /dev/null
+++ b/packages/ai-bot/lib/ai-cost.ts
@@ -0,0 +1,75 @@
+import {
+ getMostRecentSubscription,
+ getUserByMatrixUserId,
+ spendCredits,
+} from '@cardstack/billing/billing-queries';
+import { PgAdapter, TransactionManager } from '@cardstack/postgres';
+import { logger, retry } from '@cardstack/runtime-common';
+import * as Sentry from '@sentry/node';
+
+let log = logger('ai-bot');
+
+export async function saveUsageCost(
+ pgAdapter: PgAdapter,
+ matrixUserId: string,
+ generationId: string,
+) {
+ try {
+ // Generation data is sometimes not immediately available, so we retry a couple of times until we are able to get the cost
+ let costInUsd = await retry(() => fetchGenerationCost(generationId), {
+ retries: 10,
+ delayMs: 500,
+ });
+
+ let creditsConsumed = Math.round(costInUsd / 0.001);
+
+ let user = await getUserByMatrixUserId(pgAdapter, matrixUserId);
+
+ // This check is for transition period where we don't have subscriptions have the subscription flow fully rolled out.
+ // When we have assurance that all users who use the bot have subscriptions, we can remove this subscription check.
+ let subscription = await getMostRecentSubscription(pgAdapter, user!.id);
+ if (!subscription) {
+ log.info(
+ `user ${matrixUserId} has no subscription, skipping credit usage tracking`,
+ );
+ return Promise.resolve();
+ }
+
+ if (!user) {
+ throw new Error(
+ `should not happen: user with matrix id ${matrixUserId} not found in the users table`,
+ );
+ }
+
+ let txManager = new TransactionManager(pgAdapter);
+
+ await txManager.withTransaction(async () => {
+ await spendCredits(pgAdapter, user!.id, creditsConsumed);
+
+ // TODO: send a signal to the host app to update credits balance displayed in the UI
+ });
+ } catch (err) {
+ log.error(
+ `Failed to track AI usage (matrixUserId: ${matrixUserId}, generationId: ${generationId}):`,
+ err,
+ );
+ Sentry.captureException(err);
+ // Don't throw, because we don't want to crash the bot over this
+ }
+}
+
+async function fetchGenerationCost(generationId: string) {
+ let response = await (
+ await fetch(`https://openrouter.ai/api/v1/generation?id=${generationId}`, {
+ headers: {
+ Authorization: `Bearer ${process.env.OPENROUTER_API_KEY}`,
+ },
+ })
+ ).json();
+
+ if (response.error && response.error.includes('not found')) {
+ return null;
+ }
+
+ return response.data.total_cost;
+}
diff --git a/packages/ai-bot/main.ts b/packages/ai-bot/main.ts
index 4d3f169dfb..a8ef6837a4 100644
--- a/packages/ai-bot/main.ts
+++ b/packages/ai-bot/main.ts
@@ -7,7 +7,7 @@ import {
type MatrixEvent,
} from 'matrix-js-sdk';
import OpenAI from 'openai';
-import { logger, aiBotUsername, retry } from '@cardstack/runtime-common';
+import { logger, aiBotUsername } from '@cardstack/runtime-common';
import {
constructHistory,
getModifyPrompt,
@@ -24,11 +24,9 @@ import { handleDebugCommands } from './lib/debug';
import { MatrixClient } from './lib/matrix';
import type { MatrixEvent as DiscreteMatrixEvent } from 'https://cardstack.com/base/matrix-event';
import * as Sentry from '@sentry/node';
-import { PgAdapter, TransactionManager } from '@cardstack/postgres';
-import {
- getUserByMatrixUserId,
- spendCredits,
-} from '@cardstack/billing/billing-queries';
+
+import { saveUsageCost } from './lib/ai-cost';
+import { PgAdapter } from '@cardstack/postgres';
let log = logger('ai-bot');
@@ -50,65 +48,13 @@ class Assistant {
this.pgAdapter = new PgAdapter();
}
- async fetchGenerationCost(generationId: string) {
- let response = await (
- await fetch(
- `https://openrouter.ai/api/v1/generation?id=${generationId}`,
- {
- headers: {
- Authorization: `Bearer ${process.env.OPENROUTER_API_KEY}`,
- },
- },
- )
- ).json();
-
- if (response.error && response.error.includes('not found')) {
- return null;
- }
-
- return response.data.total_cost;
- }
-
async trackAiUsageCost(matrixUserId: string, generationId: string) {
- try {
- let costInUsd = await retry(
- () => this.fetchGenerationCost(generationId),
- {
- retries: 10,
- delayMs: 500,
- },
- );
-
- let creditsConsumed = Math.round(costInUsd / 0.001);
-
- let user = await getUserByMatrixUserId(this.pgAdapter, matrixUserId);
- if (!user) {
- throw new Error('should not happen: user with matrix id not found');
- }
-
- let txManager = new TransactionManager(this.pgAdapter);
- await txManager.withTransaction(async () => {
- await spendCredits(this.pgAdapter, user!.id, creditsConsumed);
- });
-
- log.info('Successfully tracked AI usage:', {
- matrixUserId,
- generationId,
- creditsConsumed,
- });
- } catch (err) {
- log.error('Failed to track AI usage:', err);
- throw err;
- }
- }
-
- async enqueueTrackAiUsageCost(matrixUserId: string, generationId: string) {
if (trackAiUsageCostPromises.has(matrixUserId)) {
return;
}
trackAiUsageCostPromises.set(
matrixUserId,
- this.trackAiUsageCost(matrixUserId, generationId).finally(() => {
+ saveUsageCost(this.pgAdapter, matrixUserId, generationId).finally(() => {
trackAiUsageCostPromises.delete(matrixUserId);
}),
);
@@ -206,6 +152,7 @@ Common issues are:
async function (event, room, toStartOfTimeline) {
try {
let eventBody = event.getContent().body;
+ let senderMatrixUserId = event.getSender()!;
if (!room) {
return;
}
@@ -223,7 +170,7 @@ Common issues are:
return; // don't respond to card fragments, we just gather these in our history
}
- if (event.getSender() === aiBotUserId) {
+ if (senderMatrixUserId === aiBotUserId) {
return;
}
log.info(
@@ -231,7 +178,7 @@ Common issues are:
event.getType(),
room?.name,
room?.roomId,
- event.getSender(),
+ senderMatrixUserId,
eventBody,
);
@@ -250,10 +197,11 @@ Common issues are:
}
const responder = new Responder(client, room.roomId);
+ await responder.initialize();
// Do not generate new responses if previous ones' cost is still being reported
let pendingCreditsConsumptionPromise = trackAiUsageCostPromises.get(
- event.getSender()!,
+ senderMatrixUserId!,
);
if (pendingCreditsConsumptionPromise) {
try {
@@ -261,11 +209,10 @@ Common issues are:
} catch (e) {
log.error(e);
return responder.onError(
- 'There was an error reporting Boxel Credits usage. Try again or contact support if the problem persists.',
+ 'There was an error saving your Boxel credits usage. Try again or contact support if the problem persists.',
);
}
}
- await responder.initialize();
if (historyError) {
responder.finalize(
@@ -300,12 +247,7 @@ Common issues are:
await responder.onError(error);
} finally {
if (generationId) {
- let userMatrixId = event.getSender()!;
-
- assistant.enqueueTrackAiUsageCost(
- userMatrixId,
- generationId as string,
- );
+ assistant.trackAiUsageCost(senderMatrixUserId, generationId);
}
}
diff --git a/packages/billing/billing-queries.ts b/packages/billing/billing-queries.ts
index 07f513f79d..371968c77a 100644
--- a/packages/billing/billing-queries.ts
+++ b/packages/billing/billing-queries.ts
@@ -524,7 +524,7 @@ export async function spendCredits(
subscriptionCycleId: subscriptionCycle.id,
});
} else {
- // If user does not have enough plan allowance credits to cover the spend, use extra credits
+ // If user does not have enough plan allowance credits to cover the spend, try to also use extra credits
let availableExtraCredits = await sumUpCreditsLedger(dbAdapter, {
creditType: ['extra_credit', 'extra_credit_used'],
userId,
@@ -535,18 +535,22 @@ export async function spendCredits(
extraCreditsToSpend = availableExtraCredits;
}
- await addToCreditsLedger(dbAdapter, {
- userId,
- creditAmount: -planAllowanceToSpend,
- creditType: 'plan_allowance_used',
- subscriptionCycleId: subscriptionCycle.id,
- });
+ if (planAllowanceToSpend > 0) {
+ await addToCreditsLedger(dbAdapter, {
+ userId,
+ creditAmount: -planAllowanceToSpend,
+ creditType: 'plan_allowance_used',
+ subscriptionCycleId: subscriptionCycle.id,
+ });
+ }
- await addToCreditsLedger(dbAdapter, {
- userId,
- creditAmount: -extraCreditsToSpend,
- creditType: 'extra_credit_used',
- subscriptionCycleId: subscriptionCycle.id,
- });
+ if (extraCreditsToSpend > 0) {
+ await addToCreditsLedger(dbAdapter, {
+ userId,
+ creditAmount: -extraCreditsToSpend,
+ creditType: 'extra_credit_used',
+ subscriptionCycleId: subscriptionCycle.id,
+ });
+ }
}
}
diff --git a/packages/realm-server/tests/billing-test.ts b/packages/realm-server/tests/billing-test.ts
index ae91da67de..4a35f08f0d 100644
--- a/packages/realm-server/tests/billing-test.ts
+++ b/packages/realm-server/tests/billing-test.ts
@@ -28,7 +28,6 @@ import {
StripeSubscriptionDeletedWebhookEvent,
StripeCheckoutSessionCompletedWebhookEvent,
} from '@cardstack/billing/stripe-webhook-handlers';
-import { add } from 'date-fns';
async function fetchStripeEvents(dbAdapter: PgAdapter) {
return await query(dbAdapter, [`SELECT * FROM stripe_events`]);
@@ -801,8 +800,7 @@ module('billing', function (hooks) {
});
});
- // eslint-disable-next-line qunit/no-only
- module.only('ai usage tracking', function (hooks) {
+ module('AI usage tracking', function (hooks) {
let user: User;
let creatorPlan: Plan;
let subscription: Subscription;
From 08255dbe8945c2c40fb1b5dba145fcb122f8db52 Mon Sep 17 00:00:00 2001
From: Matic Jurglic <matic@jurglic.si>
Date: Tue, 19 Nov 2024 11:00:24 +0100
Subject: [PATCH 3/9] Save subscription cycle in extra credits ledger entries
---
.../checkout-session-completed.ts | 20 ++++++++++++++++++-
packages/realm-server/tests/billing-test.ts | 19 ++++++++++++++++++
2 files changed, 38 insertions(+), 1 deletion(-)
diff --git a/packages/billing/stripe-webhook-handlers/checkout-session-completed.ts b/packages/billing/stripe-webhook-handlers/checkout-session-completed.ts
index 9d7c108758..37a856bbf7 100644
--- a/packages/billing/stripe-webhook-handlers/checkout-session-completed.ts
+++ b/packages/billing/stripe-webhook-handlers/checkout-session-completed.ts
@@ -1,6 +1,8 @@
import { type DBAdapter } from '@cardstack/runtime-common';
import {
addToCreditsLedger,
+ getMostRecentSubscription,
+ getMostRecentSubscriptionCycle,
getUserByStripeId,
insertStripeEvent,
markStripeEventAsProcessed,
@@ -50,11 +52,27 @@ export async function handleCheckoutSessionCompleted(
);
}
+ let subscription = await getMostRecentSubscription(dbAdapter, user.id);
+ if (!subscription) {
+ throw new Error(
+ `User ${user.id} has no subscription, cannot add extra credits`,
+ );
+ }
+ let subscriptionCycle = await getMostRecentSubscriptionCycle(
+ dbAdapter,
+ subscription!.id,
+ );
+ if (!subscriptionCycle) {
+ throw new Error(
+ `User ${user.id} has no subscription cycle, cannot add extra credits`,
+ );
+ }
+
await addToCreditsLedger(dbAdapter, {
userId: user.id,
creditAmount: creditReloadAmount,
creditType: 'extra_credit',
- subscriptionCycleId: null,
+ subscriptionCycleId: subscriptionCycle.id,
});
}
});
diff --git a/packages/realm-server/tests/billing-test.ts b/packages/realm-server/tests/billing-test.ts
index 4a35f08f0d..60385f7998 100644
--- a/packages/realm-server/tests/billing-test.ts
+++ b/packages/realm-server/tests/billing-test.ts
@@ -771,6 +771,25 @@ module('billing', function (hooks) {
});
test('add extra credits to user ledger when checkout session completed', async function (assert) {
+ let creatorPlan = await insertPlan(
+ dbAdapter,
+ 'Creator',
+ 12,
+ 2500,
+ 'prod_creator',
+ );
+ let subscription = await insertSubscription(dbAdapter, {
+ user_id: user.id,
+ plan_id: creatorPlan.id,
+ started_at: 1,
+ status: 'active',
+ stripe_subscription_id: 'sub_1234567890',
+ });
+ await insertSubscriptionCycle(dbAdapter, {
+ subscriptionId: subscription.id,
+ periodStart: 1,
+ periodEnd: 2,
+ });
let stripeCheckoutSessionCompletedEvent = {
id: 'evt_1234567890',
object: 'event',
From 111020a24ea5bce64f148c3091f0b6dfcc4441e4 Mon Sep 17 00:00:00 2001
From: Matic Jurglic <matic@jurglic.si>
Date: Tue, 19 Nov 2024 11:02:58 +0100
Subject: [PATCH 4/9] Update pnpm lock
---
pnpm-lock.yaml | 8 +++++++-
1 file changed, 7 insertions(+), 1 deletion(-)
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index 4e9816ecdb..fac05f6088 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -117,8 +117,14 @@ importers:
packages/ai-bot:
dependencies:
+ '@cardstack/billing':
+ specifier: workspace:*
+ version: link:../billing
+ '@cardstack/postgres':
+ specifier: workspace:*
+ version: link:../postgres
'@cardstack/runtime-common':
- specifier: workspace:^
+ specifier: workspace:*
version: link:../runtime-common
'@sentry/node':
specifier: ^8.31.0
From 6a1713213a4134a8897b5dbfc5632e324b89a6db Mon Sep 17 00:00:00 2001
From: Matic Jurglic <matic@jurglic.si>
Date: Tue, 19 Nov 2024 11:34:18 +0100
Subject: [PATCH 5/9] Fixes after upstream merge
---
packages/billing/billing-queries.ts | 4 ++--
.../stripe-webhook-handlers/checkout-session-completed.ts | 4 ++--
2 files changed, 4 insertions(+), 4 deletions(-)
diff --git a/packages/billing/billing-queries.ts b/packages/billing/billing-queries.ts
index aaed46cc48..07ab2076ea 100644
--- a/packages/billing/billing-queries.ts
+++ b/packages/billing/billing-queries.ts
@@ -478,9 +478,9 @@ export async function spendCredits(
userId: string,
creditsToSpend: number,
) {
- let subscription = await getMostRecentSubscription(dbAdapter, userId);
+ let subscription = await getCurrentActiveSubscription(dbAdapter, userId);
if (!subscription) {
- throw new Error('subscription not found');
+ throw new Error('active subscription not found');
}
let subscriptionCycle = await getMostRecentSubscriptionCycle(
dbAdapter,
diff --git a/packages/billing/stripe-webhook-handlers/checkout-session-completed.ts b/packages/billing/stripe-webhook-handlers/checkout-session-completed.ts
index 37a856bbf7..d176fa5057 100644
--- a/packages/billing/stripe-webhook-handlers/checkout-session-completed.ts
+++ b/packages/billing/stripe-webhook-handlers/checkout-session-completed.ts
@@ -1,7 +1,7 @@
import { type DBAdapter } from '@cardstack/runtime-common';
import {
addToCreditsLedger,
- getMostRecentSubscription,
+ getCurrentActiveSubscription,
getMostRecentSubscriptionCycle,
getUserByStripeId,
insertStripeEvent,
@@ -52,7 +52,7 @@ export async function handleCheckoutSessionCompleted(
);
}
- let subscription = await getMostRecentSubscription(dbAdapter, user.id);
+ let subscription = await getCurrentActiveSubscription(dbAdapter, user.id);
if (!subscription) {
throw new Error(
`User ${user.id} has no subscription, cannot add extra credits`,
From 3e89ac8215d0a6f546e788b2796ff2f1e28e26ef Mon Sep 17 00:00:00 2001
From: Matic Jurglic <matic@jurglic.si>
Date: Tue, 19 Nov 2024 11:48:01 +0100
Subject: [PATCH 6/9] Fix query and comment
---
packages/ai-bot/lib/ai-cost.ts | 6 +++---
1 file changed, 3 insertions(+), 3 deletions(-)
diff --git a/packages/ai-bot/lib/ai-cost.ts b/packages/ai-bot/lib/ai-cost.ts
index 0e1f7d0bd2..bdbebf2219 100644
--- a/packages/ai-bot/lib/ai-cost.ts
+++ b/packages/ai-bot/lib/ai-cost.ts
@@ -1,5 +1,5 @@
import {
- getMostRecentSubscription,
+ getCurrentActiveSubscription,
getUserByMatrixUserId,
spendCredits,
} from '@cardstack/billing/billing-queries';
@@ -25,9 +25,9 @@ export async function saveUsageCost(
let user = await getUserByMatrixUserId(pgAdapter, matrixUserId);
- // This check is for transition period where we don't have subscriptions have the subscription flow fully rolled out.
+ // This check is for the transition period where we don't have subscriptions fully rolled out yet.
// When we have assurance that all users who use the bot have subscriptions, we can remove this subscription check.
- let subscription = await getMostRecentSubscription(pgAdapter, user!.id);
+ let subscription = await getCurrentActiveSubscription(pgAdapter, user!.id);
if (!subscription) {
log.info(
`user ${matrixUserId} has no subscription, skipping credit usage tracking`,
From bc5e24fecbf4916bf2fd338665d13fbc97b6fc55 Mon Sep 17 00:00:00 2001
From: Matic Jurglic <matic@jurglic.si>
Date: Tue, 19 Nov 2024 13:01:25 +0100
Subject: [PATCH 7/9] Revert default pg config, provide config with the start
command
---
packages/ai-bot/package.json | 2 +-
packages/postgres/pg-config.ts | 4 ++--
2 files changed, 3 insertions(+), 3 deletions(-)
diff --git a/packages/ai-bot/package.json b/packages/ai-bot/package.json
index a946eccbce..5eb5768711 100644
--- a/packages/ai-bot/package.json
+++ b/packages/ai-bot/package.json
@@ -23,7 +23,7 @@
},
"scripts": {
"lint": "eslint . --cache --ext ts",
- "start": "NODE_NO_WARNINGS=1 ts-node --transpileOnly main",
+ "start": "NODE_NO_WARNINGS=1 PGDATABASE=boxel PGPORT=5435 ts-node --transpileOnly main",
"test": "NODE_NO_WARNINGS=1 qunit --require ts-node/register/transpile-only tests/index.ts",
"get-chat": "NODE_NO_WARNINGS=1 ts-node --transpileOnly scripts/get_chat.ts"
},
diff --git a/packages/postgres/pg-config.ts b/packages/postgres/pg-config.ts
index d9d3939577..a2b7256275 100644
--- a/packages/postgres/pg-config.ts
+++ b/packages/postgres/pg-config.ts
@@ -3,9 +3,9 @@ import { type ClientConfig } from 'pg';
export function postgresConfig(defaultConfig: ClientConfig = {}) {
return Object.assign({}, defaultConfig, {
host: process.env.PGHOST || 'localhost',
- port: process.env.PGPORT || '5435',
+ port: process.env.PGPORT || '5432',
user: process.env.PGUSER || 'postgres',
password: process.env.PGPASSWORD || undefined,
- database: process.env.PGDATABASE || 'boxel',
+ database: process.env.PGDATABASE || 'postgres',
});
}
From 26e8ba0f644ba2d10fe29bed9848813a2b9c30a8 Mon Sep 17 00:00:00 2001
From: Matic Jurglic <matic@jurglic.si>
Date: Tue, 19 Nov 2024 14:22:18 +0100
Subject: [PATCH 8/9] Remove comment that is not relevant anymore
---
packages/billing/stripe-webhook-handlers/index.ts | 3 ---
1 file changed, 3 deletions(-)
diff --git a/packages/billing/stripe-webhook-handlers/index.ts b/packages/billing/stripe-webhook-handlers/index.ts
index 5f343e3d37..4b53bf6c3c 100644
--- a/packages/billing/stripe-webhook-handlers/index.ts
+++ b/packages/billing/stripe-webhook-handlers/index.ts
@@ -124,9 +124,6 @@ export default async function stripeWebhookHandler(
let type = event.type;
- // For adding extra credits, we should listen for charge.succeeded, and for
- // subsciptions, we should listen for invoice.payment_succeeded (I discovered this when I was
- // testing which webhooks arrive for both types of payments)
switch (type) {
// These handlers should eventually become jobs which workers will process asynchronously
case 'invoice.payment_succeeded':
From 39db7addaf7027bb31c92718d3b3f54ec1602302 Mon Sep 17 00:00:00 2001
From: Matic Jurglic <matic@jurglic.si>
Date: Wed, 20 Nov 2024 14:28:21 +0100
Subject: [PATCH 9/9] Check if user has enough credits before responding with
AI generated content
---
.../ai-bot/lib/{ai-cost.ts => ai-billing.ts} | 20 +++++++++++++
packages/ai-bot/lib/matrix.ts | 4 +--
packages/ai-bot/main.ts | 29 ++++++++++++++-----
3 files changed, 44 insertions(+), 9 deletions(-)
rename packages/ai-bot/lib/{ai-cost.ts => ai-billing.ts} (84%)
diff --git a/packages/ai-bot/lib/ai-cost.ts b/packages/ai-bot/lib/ai-billing.ts
similarity index 84%
rename from packages/ai-bot/lib/ai-cost.ts
rename to packages/ai-bot/lib/ai-billing.ts
index bdbebf2219..8ef38c4894 100644
--- a/packages/ai-bot/lib/ai-cost.ts
+++ b/packages/ai-bot/lib/ai-billing.ts
@@ -2,6 +2,7 @@ import {
getCurrentActiveSubscription,
getUserByMatrixUserId,
spendCredits,
+ sumUpCreditsLedger,
} from '@cardstack/billing/billing-queries';
import { PgAdapter, TransactionManager } from '@cardstack/postgres';
import { logger, retry } from '@cardstack/runtime-common';
@@ -58,6 +59,25 @@ export async function saveUsageCost(
}
}
+export async function getAvailableCredits(
+ pgAdapter: PgAdapter,
+ matrixUserId: string,
+) {
+ let user = await getUserByMatrixUserId(pgAdapter, matrixUserId);
+
+ if (!user) {
+ throw new Error(
+ `should not happen: user with matrix id ${matrixUserId} not found in the users table`,
+ );
+ }
+
+ let availableCredits = await sumUpCreditsLedger(pgAdapter, {
+ userId: user.id,
+ });
+
+ return availableCredits;
+}
+
async function fetchGenerationCost(generationId: string) {
let response = await (
await fetch(`https://openrouter.ai/api/v1/generation?id=${generationId}`, {
diff --git a/packages/ai-bot/lib/matrix.ts b/packages/ai-bot/lib/matrix.ts
index 5d129994e1..dc8399d696 100644
--- a/packages/ai-bot/lib/matrix.ts
+++ b/packages/ai-bot/lib/matrix.ts
@@ -145,7 +145,7 @@ function getErrorMessage(error: any): string {
return `OpenAI error: ${error.name} - ${error.message}`;
}
if (typeof error === 'string') {
- return `Unknown error: ${error}`;
+ return error;
}
- return `Unknown error`;
+ return 'Unknown error';
}
diff --git a/packages/ai-bot/main.ts b/packages/ai-bot/main.ts
index a8ef6837a4..a7bb0b6843 100644
--- a/packages/ai-bot/main.ts
+++ b/packages/ai-bot/main.ts
@@ -25,8 +25,12 @@ import { MatrixClient } from './lib/matrix';
import type { MatrixEvent as DiscreteMatrixEvent } from 'https://cardstack.com/base/matrix-event';
import * as Sentry from '@sentry/node';
-import { saveUsageCost } from './lib/ai-cost';
+import { getAvailableCredits, saveUsageCost } from './lib/ai-billing';
import { PgAdapter } from '@cardstack/postgres';
+import {
+ getUserByMatrixUserId,
+ sumUpCreditsLedger,
+} from '@cardstack/billing/billing-queries';
let log = logger('ai-bot');
@@ -35,7 +39,7 @@ let trackAiUsageCostPromises = new Map<string, Promise<void>>();
class Assistant {
private openai: OpenAI;
private client: MatrixClient;
- private pgAdapter: PgAdapter;
+ pgAdapter: PgAdapter;
id: string;
constructor(client: MatrixClient, id: string) {
@@ -197,6 +201,12 @@ Common issues are:
}
const responder = new Responder(client, room.roomId);
+ if (historyError) {
+ return responder.finalize(
+ 'There was an error processing chat history. Please open another session.',
+ );
+ }
+
await responder.initialize();
// Do not generate new responses if previous ones' cost is still being reported
@@ -214,11 +224,17 @@ Common issues are:
}
}
- if (historyError) {
- responder.finalize(
- 'There was an error processing chat history. Please open another session.',
+ let availableCredits = await getAvailableCredits(
+ assistant.pgAdapter,
+ senderMatrixUserId,
+ );
+
+ let minimumCredits = 10;
+
+ if (availableCredits < minimumCredits) {
+ return responder.onError(
+ `You need a minimum of ${minimumCredits} credits to continue using the AI bot. Please upgrade to a larger plan, or top up your account.`,
);
- return;
}
let generationId: string | undefined;
@@ -238,7 +254,6 @@ Common issues are:
await responder.onError(error);
});
- // We also need to catch the error when getting the final content
let finalContent;
try {
finalContent = await runner.finalContent();