feat: add support for inferring benefit and category with OpenAI when

submitting multiple claims from a CSV Fixes #6.
timrogers · Aug 1, 2023 · f83dbad · f83dbad
1 parent 26356e2
commit f83dbad
Show file tree

Hide file tree

Showing 3 changed files with 137 additions and 104 deletions.
diff --git a/src/commands/submit-claim.ts b/src/commands/submit-claim.ts
@@ -1,16 +1,12 @@
 import * as commander from 'commander';
 import chalk from 'chalk';
-import { Configuration, OpenAIApi } from 'openai';
 
-import { actionRunner, prompt } from '../utils.js';
+import { actionRunner } from '../utils.js';
 import { getAccessToken } from '../config.js';
-import {
-  type BenefitWithCategories,
-  createClaim,
-  getBenefitsWithCategories,
-} from '../forma.js';
+import { createClaim, getBenefitsWithCategories } from '../forma.js';
 import { claimParamsToCreateClaimOptions } from '../claims.js';
 import VERSION from '../version.js';
+import { attemptToInferCategoryAndBenefit } from '../openai.js';
 
 const command = new commander.Command();
 
@@ -26,93 +22,6 @@ interface Arguments {
   openaiApiKey?: string;
 }
 
-const generateOpenaiPrompt = (opts: {
-  validCategories: string[];
-  merchant: string;
-  description: string;
-}): string => {
-  const { description, merchant, validCategories } = opts;
-
-  return `Your job is to predict the category for an expense claim based on the name of the merchant and a description of what was purchased. You should give a single, specific answer without any extra words or punctuation.
-
-Here are the possible categories:
-
-${validCategories.join('\n')}
-
-Please predict the category for the following claim:
-
-Merchant: ${merchant}
-Description: ${description}`;
-};
-
-const attemptToinferCategoryAndBenefit = async (opts: {
-  merchant: string;
-  description: string;
-  benefitsWithCategories: BenefitWithCategories[];
-  openaiApiKey: string;
-}): Promise<{ category: string; benefit: string }> => {
-  const { merchant, description, benefitsWithCategories, openaiApiKey: apiKey } = opts;
-
-  const configuration = new Configuration({
-    apiKey,
-  });
-
-  const openai = new OpenAIApi(configuration);
-
-  const categoriesWithBenefits = benefitsWithCategories.flatMap((benefit) =>
-    benefit.categories.map((category) => ({ ...category, benefit })),
-  );
-
-  const validCategories = categoriesWithBenefits.flatMap(
-    (category) => category.subcategory_alias ?? category.subcategory_name,
-  );
-
-  const content = generateOpenaiPrompt({ validCategories, merchant, description });
-
-  const chatCompletion = await openai.createChatCompletion({
-    model: 'gpt-3.5-turbo',
-    messages: [
-      {
-        role: 'user',
-        content,
-      },
-    ],
-  });
-
-  const returnedCategoryAsString = chatCompletion.data.choices[0].message?.content;
-
-  if (!returnedCategoryAsString) {
-    throw new Error(
-      `Something went wrong while inferring the benefit and category for your claim. OpenAI returned an unexpected response: ${JSON.stringify(
-        chatCompletion.data,
-      )}`,
-    );
-  }
-
-  const returnedCategory = categoriesWithBenefits.find(
-    (category) =>
-      category.subcategory_alias === returnedCategoryAsString ||
-      category.subcategory_name === returnedCategoryAsString,
-  );
-
-  if (!returnedCategory) {
-    throw new Error(
-      `Something went wrong while inferring the benefit and category for your claim. OpenAI returned a response that wasn't a valid category: ${returnedCategoryAsString}`,
-    );
-  }
-
-  console.log(
-    `OpenAI inferred that you should claim using the ${chalk.magenta(
-      returnedCategory.benefit.name,
-    )} benefit and ${chalk.magenta(
-      returnedCategoryAsString,
-    )} category. If that seems right, hit Enter. If not, press Ctrl + C to end your session.`,
-  );
-  prompt('> ');
-
-  return { category: returnedCategoryAsString, benefit: returnedCategory.benefit.name };
-};
-
 command
   .name('submit-claim')
   .version(VERSION)
@@ -162,7 +71,7 @@ command
         await createClaim(createClaimOptions);
       } else if (openaiApiKey) {
         const benefitsWithCategories = await getBenefitsWithCategories(accessToken);
-        const { benefit, category } = await attemptToinferCategoryAndBenefit({
+        const { benefit, category } = await attemptToInferCategoryAndBenefit({
           merchant: opts.merchant,
           description: opts.description,
           benefitsWithCategories,

diff --git a/src/commands/submit-claims-from-csv.ts b/src/commands/submit-claims-from-csv.ts
@@ -5,15 +5,17 @@ import chalk from 'chalk';
 
 import { actionRunner, serializeError } from '../utils.js';
 import { getAccessToken } from '../config.js';
-import { createClaim } from '../forma.js';
+import { createClaim, getBenefitsWithCategories } from '../forma.js';
 import { type Claim, claimParamsToCreateClaimOptions } from '../claims.js';
 import VERSION from '../version.js';
+import { attemptToInferCategoryAndBenefit } from '../openai.js';
 
 const command = new commander.Command();
 
 interface Arguments {
   accessToken?: string;
   inputPath: string;
+  openaiApiKey: string;
 }
 
 const EXPECTED_HEADERS = [
@@ -66,8 +68,15 @@ command
   )
   .requiredOption('--input-path <input_path>', 'The path to the CSV to read claims from')
   .option('--access-token <access_token>', 'Access token used to authenticate with Forma')
+  .option(
+    '--openai-api-key <openai_token>',
+    'An optional OpenAI API key used to infer the benefit and category based on the merchant and description. If this is set, you may omit the `--benefit` and `--category` options. This can also be configured using the `OPENAI_API_KEY` environment variable.',
+    process.env.OPENAI_API_KEY,
+  )
   .action(
     actionRunner(async (opts: Arguments) => {
+      const { openaiApiKey } = opts;
+
       const accessToken = opts.accessToken ?? getAccessToken();
 
       if (!accessToken) {
@@ -92,14 +101,37 @@ command
         console.log(`Submitting claim ${index + 1}/${claims.length}`);
 
         try {
-          const createClaimOptions = await claimParamsToCreateClaimOptions(
-            claim,
-            accessToken,
-          );
-          await createClaim(createClaimOptions);
-          console.log(
-            chalk.green(`Successfully submitted claim ${index + 1}/${claims.length}`),
-          );
+          if (claim.benefit !== '' && claim.category !== '') {
+            const createClaimOptions = await claimParamsToCreateClaimOptions(
+              claim,
+              accessToken,
+            );
+            await createClaim(createClaimOptions);
+            console.log(
+              chalk.green(`Successfully submitted claim ${index + 1}/${claims.length}`),
+            );
+          } else if (openaiApiKey) {
+            const benefitsWithCategories = await getBenefitsWithCategories(accessToken);
+            const { benefit, category } = await attemptToInferCategoryAndBenefit({
+              merchant: claim.merchant,
+              description: claim.description,
+              benefitsWithCategories,
+              openaiApiKey,
+            });
+
+            const createClaimOptions = await claimParamsToCreateClaimOptions(
+              { ...claim, benefit, category },
+              accessToken,
+            );
+            await createClaim(createClaimOptions);
+            console.log(
+              chalk.green(`Successfully submitted claim ${index + 1}/${claims.length}`),
+            );
+          } else {
+            throw new Error(
+              'You must either fill out the `benefit` and `category` columns, or specify an OpenAI API key.',
+            );
+          }
         } catch (e) {
           console.error(
             chalk.red(

diff --git a/src/openai.ts b/src/openai.ts
@@ -0,0 +1,92 @@
+import { Configuration, OpenAIApi } from 'openai';
+import chalk from 'chalk';
+
+import { prompt } from './utils.js';
+import { type BenefitWithCategories } from './forma.js';
+
+export const attemptToInferCategoryAndBenefit = async (opts: {
+  merchant: string;
+  description: string;
+  benefitsWithCategories: BenefitWithCategories[];
+  openaiApiKey: string;
+}): Promise<{ category: string; benefit: string }> => {
+  const { merchant, description, benefitsWithCategories, openaiApiKey: apiKey } = opts;
+
+  const configuration = new Configuration({
+    apiKey,
+  });
+
+  const openai = new OpenAIApi(configuration);
+
+  const categoriesWithBenefits = benefitsWithCategories.flatMap((benefit) =>
+    benefit.categories.map((category) => ({ ...category, benefit })),
+  );
+
+  const validCategories = categoriesWithBenefits.flatMap(
+    (category) => category.subcategory_alias ?? category.subcategory_name,
+  );
+
+  const content = generateOpenaiPrompt({ validCategories, merchant, description });
+
+  const chatCompletion = await openai.createChatCompletion({
+    model: 'gpt-3.5-turbo',
+    messages: [
+      {
+        role: 'user',
+        content,
+      },
+    ],
+  });
+
+  const returnedCategoryAsString = chatCompletion.data.choices[0].message?.content;
+
+  if (!returnedCategoryAsString) {
+    throw new Error(
+      `Something went wrong while inferring the benefit and category for your claim. OpenAI returned an unexpected response: ${JSON.stringify(
+        chatCompletion.data,
+      )}`,
+    );
+  }
+
+  const returnedCategory = categoriesWithBenefits.find(
+    (category) =>
+      category.subcategory_alias === returnedCategoryAsString ||
+      category.subcategory_name === returnedCategoryAsString,
+  );
+
+  if (!returnedCategory) {
+    throw new Error(
+      `Something went wrong while inferring the benefit and category for your claim. OpenAI returned a response that wasn't a valid category: ${returnedCategoryAsString}`,
+    );
+  }
+
+  console.log(
+    `OpenAI inferred that you should claim using the ${chalk.magenta(
+      returnedCategory.benefit.name,
+    )} benefit and ${chalk.magenta(
+      returnedCategoryAsString,
+    )} category. If that seems right, hit Enter. If not, press Ctrl + C to end your session.`,
+  );
+  prompt('> ');
+
+  return { category: returnedCategoryAsString, benefit: returnedCategory.benefit.name };
+};
+
+const generateOpenaiPrompt = (opts: {
+  validCategories: string[];
+  merchant: string;
+  description: string;
+}): string => {
+  const { description, merchant, validCategories } = opts;
+
+  return `Your job is to predict the category for an expense claim based on the name of the merchant and a description of what was purchased. You should give a single, specific answer without any extra words or punctuation.
+
+Here are the possible categories:
+
+${validCategories.join('\n')}
+
+Please predict the category for the following claim:
+
+Merchant: ${merchant}
+Description: ${description}`;
+};