From 0589820878e2d76cb069c2d052a397b53a19db73 Mon Sep 17 00:00:00 2001
From: celeste <parafactual@gmail.com>
Date: Mon, 20 Mar 2023 07:14:12 -0700
Subject: [PATCH] add support for chat models

---
 main.ts           | 31 ++++++++++++++++++++++++++-----
 package-lock.json |  8 ++++----
 package.json      |  2 +-
 3 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/main.ts b/main.ts
index 736add9..9d3e1f6 100644
--- a/main.ts
+++ b/main.ts
@@ -49,6 +49,15 @@ const DEFAULT_SETTINGS: LoomSettings = {
   showExport: false,
 };
 
+const CHAT_MODELS = [
+  "gpt-3.5-turbo",
+  "gpt-3.5-turbo-0301",
+  "gpt-4",
+  "gpt-4-0314",
+  "gpt-4-32k",
+  "gpt-4-32k-0314",
+];
+
 type Color = "red" | "orange" | "yellow" | "green" | "blue" | "purple" | null;
 
 interface Node {
@@ -847,15 +856,27 @@ export default class LoomPlugin extends Plugin {
     // complete, or visually display an error and return if that fails
     let completions;
     try {
-      completions = (
-        await this.openai.createCompletion({
+      if (CHAT_MODELS.contains(this.settings.model)) {
+        completions = (await this.openai.createChatCompletion({
           model: this.settings.model,
-          prompt,
+          messages: [
+            { role: "assistant", content: prompt },
+          ],
           max_tokens: this.settings.maxTokens,
           n: this.settings.n,
           temperature: this.settings.temperature,
-        })
-      ).data.choices.map((choice) => choice.text);
+        })).data.choices.map((choice) => choice.message?.content);
+      } else {
+        completions = (
+          await this.openai.createCompletion({
+            model: this.settings.model,
+            prompt,
+            max_tokens: this.settings.maxTokens,
+            n: this.settings.n,
+            temperature: this.settings.temperature,
+          })
+        ).data.choices.map((choice) => choice.text);
+      }
     } catch (e) {
       if (e.response.status === 401)
         new Notice(
diff --git a/package-lock.json b/package-lock.json
index cd0806a..f3c9b5a 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -12,7 +12,7 @@
 				"@types/lodash": "^4.14.191",
 				"gpt3-tokenizer": "^1.1.5",
 				"lodash": "^4.17.21",
-				"openai": "^3.1.0",
+				"openai": "^3.2.0",
 				"untildify": "^4.0.0",
 				"uuid": "^9.0.0"
 			},
@@ -1451,9 +1451,9 @@
 			}
 		},
 		"node_modules/openai": {
-			"version": "3.1.0",
-			"resolved": "https://registry.npmjs.org/openai/-/openai-3.1.0.tgz",
-			"integrity": "sha512-v5kKFH5o+8ld+t0arudj833Mgm3GcgBnbyN9946bj6u7bvel4Yg6YFz2A4HLIYDzmMjIo0s6vSG9x73kOwvdCg==",
+			"version": "3.2.1",
+			"resolved": "https://registry.npmjs.org/openai/-/openai-3.2.1.tgz",
+			"integrity": "sha512-762C9BNlJPbjjlWZi4WYK9iM2tAVAv0uUp1UmI34vb0CN5T2mjB/qM6RYBmNKMh/dN9fC+bxqPwWJZUTWW052A==",
 			"dependencies": {
 				"axios": "^0.26.0",
 				"form-data": "^4.0.0"
diff --git a/package.json b/package.json
index 5a3a3f0..ee437a2 100644
--- a/package.json
+++ b/package.json
@@ -27,7 +27,7 @@
 		"@types/lodash": "^4.14.191",
 		"gpt3-tokenizer": "^1.1.5",
 		"lodash": "^4.17.21",
-		"openai": "^3.1.0",
+		"openai": "^3.2.0",
 		"untildify": "^4.0.0",
 		"uuid": "^9.0.0"
 	}