From 9548e5007b76d74d9823f4bf0fbe707e72664c67 Mon Sep 17 00:00:00 2001
From: Robat Williams <robatwilliams@gmail.com>
Date: Wed, 20 Dec 2023 10:28:57 +0000
Subject: [PATCH] Make summing (or otherwise aggregating) costs Excel/user's
 responsibility

---
 README.md                        |  2 +-
 src/functions/functions.js       | 45 +++++++++++++++++---------------
 src/functions/functions.json     |  5 ++--
 src/functions/functions.test.mjs | 10 ++++---
 4 files changed, 35 insertions(+), 27 deletions(-)

diff --git a/README.md b/README.md
index 1e466ca..49cacd8 100644
--- a/README.md
+++ b/README.md
@@ -37,7 +37,7 @@ For other platforms, [clear the Office cache](https://learn.microsoft.com/en-us/
 
 `OAI.CHAT_COMPLETE()` calls the API to create completions.
 
-`OAI.COST()` calculates the billing cost of completion cell(s).
+`OAI.COST()` calculates the billing cost(s) of completion cell(s). Wrap it with `SUM()` to calculate a total.
 
 See the [function metadata](https://github.com/robatwilliams/openai-excel-formulas/blob/main/src/functions/functions.json) for full documentation of functions and parameters. Excel's presentation of custom function documentation varies by platform, but is best in the _Insert Function_ dialog and/or the desktop platform.
 
diff --git a/src/functions/functions.js b/src/functions/functions.js
index 3d9e2a5..ffe637b 100644
--- a/src/functions/functions.js
+++ b/src/functions/functions.js
@@ -91,32 +91,35 @@ async function chatComplete(messages, params, invocation) {
 // Terminology note: our _cost_ is driven by usage and OpenAI's _prices_.
 CustomFunctions.associate('COST', cost);
 function cost(completionsMatrix, pricesMatrix) {
-  const completions = completionsMatrix
-    .flat()
-    .filter((value) => value !== EMPTY_OR_ZERO);
-  completions.forEach(validateIsCompletion);
   const allPrices = Object.fromEntries(
     pricesMatrix.map((row) => [row[0], { input: row[1], output: row[2] }]),
   );
 
-  return completions.reduce((accumulator, completion) => {
-    const model = completion.properties.response.properties.model.basicValue;
-    const usage = completion.properties.response.properties.usage.properties;
-    const modelPrices = allPrices[model];
-
-    if (!modelPrices) {
-      throw new CustomFunctions.Error(
-        CustomFunctions.ErrorCode.invalidValue,
-        `No prices specified for model ${model}`,
+  return completionsMatrix.map((row) =>
+    row.map((cell) => {
+      if (cell === EMPTY_OR_ZERO) {
+        return 0;
+      } else {
+        validateIsCompletion(cell);
+      }
+
+      const model = cell.properties.response.properties.model.basicValue;
+      const usage = cell.properties.response.properties.usage.properties;
+      const modelPrices = allPrices[model];
+
+      if (!modelPrices) {
+        throw new CustomFunctions.Error(
+          CustomFunctions.ErrorCode.invalidValue,
+          `No prices specified for model ${model}`,
+        );
+      }
+
+      return (
+        (usage.prompt_tokens.basicValue / 1000) * modelPrices.input +
+        (usage.completion_tokens.basicValue / 1000) * modelPrices.output
       );
-    }
-
-    return (
-      accumulator +
-      (usage.prompt_tokens.basicValue / 1000) * modelPrices.input +
-      (usage.completion_tokens.basicValue / 1000) * modelPrices.output
-    );
-  }, 0);
+    }),
+  );
 }
 
 function toEntityProperty(value) {
diff --git a/src/functions/functions.json b/src/functions/functions.json
index 538bacc..f985b61 100644
--- a/src/functions/functions.json
+++ b/src/functions/functions.json
@@ -27,7 +27,7 @@
       }
     },
     {
-      "description": "Calculate the cost of the completions in the given cells",
+      "description": "Calculate the costs of the completions in the given cells",
       "id": "COST",
       "name": "COST",
       "parameters": [
@@ -45,7 +45,8 @@
         }
       ],
       "result": {
-        "type": "number"
+        "type": "number",
+        "dimensionality": "matrix"
       }
     }
   ]
diff --git a/src/functions/functions.test.mjs b/src/functions/functions.test.mjs
index 35e8a81..cfe21e4 100644
--- a/src/functions/functions.test.mjs
+++ b/src/functions/functions.test.mjs
@@ -181,10 +181,10 @@ describe('COST', () => {
     });
     const prices = [['gpt-4-0613', 0.03, 0.06]];
 
-    assert.strictEqual(cost([[completion]], prices), 0.24);
+    assert.deepStrictEqual(cost([[completion]], prices), [[0.24]]);
   });
 
-  it('calculates the cost for a range of completions including empty cells', () => {
+  it('calculates the costs for a range of completions including empty cells', () => {
     const completion = makeCompletionEntity({
       modelUsed: 'gpt-4-0613',
       tokensPrompt: 1000,
@@ -197,7 +197,11 @@ describe('COST', () => {
     ];
     const prices = [['gpt-4-0613', 0.03, 0.06]];
 
-    assert.strictEqual(cost(completions, prices), 0.27);
+    assert.deepStrictEqual(cost(completions, prices), [
+      [0.09, 0.09],
+      [0.09, 0],
+      [0, 0],
+    ]);
   });
 
   it('throws an error when no prices are specified for the model used', () => {