Refactor generateGptResponse to include a transaction

sodic · sodic · commit ed6add742e0c · 2025-02-28T10:28:23.000+01:00
diff --git a/template/app/src/demo-ai-app/operations.ts b/template/app/src/demo-ai-app/operations.ts
@@ -1,5 +1,5 @@
 import * as z from 'zod';
-import type { Task, GptResponse } from 'wasp/entities';
+import type { Task, GptResponse, User } from 'wasp/entities';
 import type {
   GenerateGptResponse,
   CreateTask,
@@ -8,18 +8,18 @@ import type {
   GetGptResponses,
   GetAllTasksByUser,
 } from 'wasp/server/operations';
-import { HttpError } from 'wasp/server';
+import { HttpError, prisma } from 'wasp/server';
 import { GeneratedSchedule } from './schedule';
 import OpenAI from 'openai';
 import { SubscriptionStatus } from '../payment/plans';
 import { ensureArgsSchemaOrThrowHttpError } from '../server/validation';
 
-const openai = setupOpenAI();
-function setupOpenAI() {
-  if (!process.env.OPENAI_API_KEY) {
-    return new HttpError(500, 'OpenAI API key is not set');
+function getOpenAI(): OpenAI | null {
+  if (process.env.OPENAI_API_KEY) {
+    return new OpenAI({ apiKey: process.env.OPENAI_API_KEY });
+  } else {
+    return null;
   }
-  return new OpenAI({ apiKey: process.env.OPENAI_API_KEY });
 }
 
 //#region Actions
@@ -35,11 +35,15 @@ export const generateGptResponse: GenerateGptResponse<GenerateGptResponseInput,
   context
 ) => {
   if (!context.user) {
-    throw new HttpError(401);
+    throw new HttpError(401, 'Only authenticated users are allowed to perform this operation');
   }
 
   const { hours } = ensureArgsSchemaOrThrowHttpError(generateGptResponseInputSchema, rawArgs);
 
+  if (!isEligibleForResponse(context.user)) {
+    throw new HttpError(402, 'User has not paid or is out of credits');
+  }
+
   const tasks = await context.entities.Task.findMany({
     where: {
       user: {
@@ -48,150 +52,61 @@ export const generateGptResponse: GenerateGptResponse<GenerateGptResponseInput,
     },
   });
 
-  const parsedTasks = tasks.map(({ description, time }) => ({
-    description,
-    time,
-  }));
+  console.log('Calling open AI api');
+  const dailyPlanJson = await getDailyPlanFromGpt(tasks, hours);
+  if (dailyPlanJson === null) {
+    throw new HttpError(500, 'Bad response from OpenAI');
+  }
 
-  try {
-    // check if openai is initialized correctly with the API key
-    if (openai instanceof Error) {
-      throw openai;
-    }
-
-    const hasCredits = context.user.credits > 0;
-    const hasValidSubscription =
-      !!context.user.subscriptionStatus &&
-      context.user.subscriptionStatus !== SubscriptionStatus.Deleted &&
-      context.user.subscriptionStatus !== SubscriptionStatus.PastDue;
-    const canUserContinue = hasCredits || hasValidSubscription;
-
-    if (!canUserContinue) {
-      throw new HttpError(402, 'User has not paid or is out of credits');
-    } else {
-      console.log('decrementing credits');
-      await context.entities.User.update({
-        where: { id: context.user.id },
-        data: {
-          credits: {
-            decrement: 1,
-          },
-        },
-      });
-    }
-
-    const completion = await openai.chat.completions.create({
-      model: 'gpt-3.5-turbo', // you can use any model here, e.g. 'gpt-3.5-turbo', 'gpt-4', etc.
-      messages: [
-        {
-          role: 'system',
-          content:
-            'you are an expert daily planner. you will be given a list of main tasks and an estimated time to complete each task. You will also receive the total amount of hours to be worked that day. Your job is to return a detailed plan of how to achieve those tasks by breaking each task down into at least 3 subtasks each. MAKE SURE TO ALWAYS CREATE AT LEAST 3 SUBTASKS FOR EACH MAIN TASK PROVIDED BY THE USER! YOU WILL BE REWARDED IF YOU DO.',
-        },
-        {
-          role: 'user',
-          content: `I will work ${hours} hours today. Here are the tasks I have to complete: ${JSON.stringify(
-            parsedTasks
-          )}. Please help me plan my day by breaking the tasks down into actionable subtasks with time and priority status.`,
-        },
-      ],
-      tools: [
-        {
-          type: 'function',
-          function: {
-            name: 'parseTodaysSchedule',
-            description: 'parses the days tasks and returns a schedule',
-            parameters: {
-              type: 'object',
-              properties: {
-                mainTasks: {
-                  type: 'array',
-                  description: 'Name of main tasks provided by user, ordered by priority',
-                  items: {
-                    type: 'object',
-                    properties: {
-                      name: {
-                        type: 'string',
-                        description: 'Name of main task provided by user',
-                      },
-                      priority: {
-                        type: 'string',
-                        enum: ['low', 'medium', 'high'],
-                        description: 'task priority',
-                      },
-                    },
-                  },
-                },
-                subtasks: {
-                  type: 'array',
-                  items: {
-                    type: 'object',
-                    properties: {
-                      description: {
-                        type: 'string',
-                        description:
-                          'detailed breakdown and description of sub-task related to main task. e.g., "Prepare your learning session by first reading through the documentation"',
-                      },
-                      time: {
-                        type: 'number',
-                        description: 'time allocated for a given subtask in hours, e.g. 0.5',
-                      },
-                      mainTaskName: {
-                        type: 'string',
-                        description: 'name of main task related to subtask',
-                      },
-                    },
-                  },
-                },
-              },
-              required: ['mainTasks', 'subtasks', 'time', 'priority'],
-            },
-          },
-        },
-      ],
-      tool_choice: {
-        type: 'function',
-        function: {
-          name: 'parseTodaysSchedule',
-        },
-      },
-      temperature: 1,
-    });
+  // TODO: Do I need a try catch now that I'm saving a response and
+  // decrementing credits in a transaction?
 
-    const gptArgs = completion?.choices[0]?.message?.tool_calls?.[0]?.function.arguments;
+  // NOTE: I changed this up, first I do the request, and then I decrement
+  // credits. Is that dangerous? Protecting from it could be too complicated
+  // TODO: Potential issue here, user can lose credits between the point we
+  // inject it into the context and here
+  // Seems to me that there should be users in the database with negative credits
+  const decrementCredit = context.entities.User.update({
+    where: { id: context.user.id },
+    data: {
+      credits: {
+        decrement: 1,
+      },
+    },
+  });
 
-    if (!gptArgs) {
-      throw new HttpError(500, 'Bad response from OpenAI');
-    }
+  const createResponse = context.entities.GptResponse.create({
+    data: {
+      user: { connect: { id: context.user.id } },
+      content: dailyPlanJson,
+    },
+  });
 
-    console.log('gpt function call arguments: ', gptArgs);
+  // NOTE: Since these two are now brought together, I put them in a single transaction - so no rollback necessary
+  // TODO: But what if the server crashes after the transaction and before
+  // the response? I guess no big deal, since the response is in the db.
+  console.log('Decrementing credits and saving response');
+  prisma.$transaction([decrementCredit, createResponse]);
 
-    await context.entities.GptResponse.create({
-      data: {
-        user: { connect: { id: context.user.id } },
-        content: JSON.stringify(gptArgs),
-      },
-    });
-
-    return JSON.parse(gptArgs);
-  } catch (error: any) {
-    if (!context.user.subscriptionStatus && error?.statusCode != 402) {
-      await context.entities.User.update({
-        where: { id: context.user.id },
-        data: {
-          credits: {
-            increment: 1,
-          },
-        },
-      });
-    }
-    console.error(error);
-    const statusCode = error.statusCode || 500;
-    const errorMessage = error.message || 'Internal server error';
-    throw new HttpError(statusCode, errorMessage);
-  }
+  // TODO: Can this ever fail?
+  return JSON.parse(dailyPlanJson);
 };
 
+function isEligibleForResponse(user: User) {
+  // TODO: Why not check for allowed states?
+  const isUserSubscribed =
+    user.subscriptionStatus !== SubscriptionStatus.Deleted &&
+    user.subscriptionStatus !== SubscriptionStatus.PastDue;
+  const userHasCredits = user.credits > 0;
+  // TODO: If the user is subscribed (flat rate, use their subscription) -
+  // shouldn't take priority over credits since it makes no sece spending
+  // credits when a subscription is active?
+  // If they aren't subscribed, then it would make sense to spend credits.
+  // The old code always subtracted credits so I kept that, is this a bug?
+  // Also, no one checks for subscription plan (10credits or flat rate).
+  return isUserSubscribed || userHasCredits;
+}
+
 const createTaskInputSchema = z.object({
   description: z.string().nonempty(),
 });
@@ -296,3 +211,98 @@ export const getAllTasksByUser: GetAllTasksByUser<void, Task[]> = async (_args,
   });
 };
 //#endregion
+
+// TODO: Why is hours a string?
+async function getDailyPlanFromGpt(tasks: Task[], hours: string): Promise<string | null> {
+  // TODO: Why was this a singleton
+  const openai = getOpenAI();
+  if (openai === null) {
+    return null;
+  }
+
+  const parsedTasks = tasks.map(({ description, time }) => ({
+    description,
+    time,
+  }));
+
+  const completion = await openai.chat.completions.create({
+    model: 'gpt-3.5-turbo', // you can use any model here, e.g. 'gpt-3.5-turbo', 'gpt-4', etc.
+    messages: [
+      {
+        role: 'system',
+        content:
+          'you are an expert daily planner. you will be given a list of main tasks and an estimated time to complete each task. You will also receive the total amount of hours to be worked that day. Your job is to return a detailed plan of how to achieve those tasks by breaking each task down into at least 3 subtasks each. MAKE SURE TO ALWAYS CREATE AT LEAST 3 SUBTASKS FOR EACH MAIN TASK PROVIDED BY THE USER! YOU WILL BE REWARDED IF YOU DO.',
+      },
+      {
+        role: 'user',
+        content: `I will work ${hours} hours today. Here are the tasks I have to complete: ${JSON.stringify(
+          parsedTasks
+        )}. Please help me plan my day by breaking the tasks down into actionable subtasks with time and priority status.`,
+      },
+    ],
+    tools: [
+      {
+        type: 'function',
+        function: {
+          name: 'parseTodaysSchedule',
+          description: 'parses the days tasks and returns a schedule',
+          parameters: {
+            type: 'object',
+            properties: {
+              mainTasks: {
+                type: 'array',
+                description: 'Name of main tasks provided by user, ordered by priority',
+                items: {
+                  type: 'object',
+                  properties: {
+                    name: {
+                      type: 'string',
+                      description: 'Name of main task provided by user',
+                    },
+                    priority: {
+                      type: 'string',
+                      enum: ['low', 'medium', 'high'],
+                      description: 'task priority',
+                    },
+                  },
+                },
+              },
+              subtasks: {
+                type: 'array',
+                items: {
+                  type: 'object',
+                  properties: {
+                    description: {
+                      type: 'string',
+                      description:
+                        'detailed breakdown and description of sub-task related to main task. e.g., "Prepare your learning session by first reading through the documentation"',
+                    },
+                    time: {
+                      type: 'number',
+                      description: 'time allocated for a given subtask in hours, e.g. 0.5',
+                    },
+                    mainTaskName: {
+                      type: 'string',
+                      description: 'name of main task related to subtask',
+                    },
+                  },
+                },
+              },
+            },
+            required: ['mainTasks', 'subtasks', 'time', 'priority'],
+          },
+        },
+      },
+    ],
+    tool_choice: {
+      type: 'function',
+      function: {
+        name: 'parseTodaysSchedule',
+      },
+    },
+    temperature: 1,
+  });
+
+  const gptResponse = completion?.choices[0]?.message?.tool_calls?.[0]?.function.arguments;
+  return gptResponse ?? null;
+}
diff --git a/template/app/src/payment/plans.ts b/template/app/src/payment/plans.ts
@@ -1,4 +1,3 @@
-import * as z from 'zod';
 import { requireNodeEnvVar } from '../server/utils';
 
 export enum SubscriptionStatus {

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,3 @@`
`1`		`-import * as z from 'zod';`
`2`	`1`	`import { requireNodeEnvVar } from '../server/utils';`
`3`	`2`
`4`	`3`	`export enum SubscriptionStatus {`