Merge pull request #2 from jmilldotdev/chatgpt

Chatgpt

main.ts

@@ -1,5 +1,10 @@
-import { Editor, Plugin, WorkspaceLeaf } from "obsidian";
+import { Editor, MarkdownView, Plugin, WorkspaceLeaf } from "obsidian";
 import { getAI21Completion } from "src/models/ai21";
+import {
+  ChatMessage,
+  ChatRole,
+  getChatGPTCompletion,
+} from "src/models/chatGPT";
 import { getCohereCompletion } from "src/models/cohere";
 import { getGPT3Completion } from "src/models/gpt3";
 import {
@@ -38,6 +43,11 @@ export default class GPTPlugin extends Plugin {
     return currentLineContents;
   }
 
+  getNoteContents(editor: Editor) {
+    const noteContents = editor.getValue();
+    return noteContents;
+  }
+
   getSuffix(selection: string) {
     if (selection.includes(this.settings.insertToken)) {
       const prompt = selection.split(this.settings.insertToken)[0];
@@ -48,7 +58,7 @@ export default class GPTPlugin extends Plugin {
   }
 
   async getCompletion(selection: string): Promise<string | null> {
-    const { ai21, gpt3, cohere } = this.settings.models;
+    const { ai21, chatgpt, gpt3, cohere } = this.settings.models;
     let completion: string;
     const notice = gettingCompletionNotice(this.settings.activeModel);
     if (this.settings.activeModel === SupportedModels.AI21) {
@@ -69,16 +79,60 @@ export default class GPTPlugin extends Plugin {
         selection,
         cohere.settings
       );
+    } else if (this.settings.activeModel === SupportedModels.CHATGPT) {
+      const messages: ChatMessage[] = [
+        {
+          role: "system",
+          content: chatgpt.settings.systemMessage,
+        },
+        {
+          role: "user",
+          content: selection,
+        },
+      ];
+      const message = await getChatGPTCompletion(
+        chatgpt.apiKey,
+        messages,
+        chatgpt.settings
+      );
+      completion = "\n\n" + message;
     }
     notice.hide();
     return completion;
   }
 
+  async getChatCompletion(selection: string) {
+    const { chatgpt } = this.settings.models;
+    const messagesText = selection.split(this.settings.chatSeparator);
+    let messages: ChatMessage[] = [
+      {
+        role: "system",
+        content: chatgpt.settings.systemMessage,
+      },
+      ...messagesText.map((message, idx) => {
+        return {
+          role: idx % 2 === 0 ? "user" : ("assistant" as ChatRole),
+          content: message.trim(),
+        };
+      }),
+    ];
+    const completion = await getChatGPTCompletion(
+      chatgpt.apiKey,
+      messages,
+      chatgpt.settings
+    );
+    return completion;
+  }
+
   handleGetCompletionError() {
     errorGettingCompletionNotice();
   }
 
-  formatCompletion(prompt: string, completion: string) {
+  formatCompletion(
+    prompt: string,
+    completion: string,
+    isChatCompletion = false
+  ) {
     const {
       tagCompletions,
       tagCompletionsHandlerTags,
@@ -94,6 +148,11 @@ export default class GPTPlugin extends Plugin {
       prompt = `${tagPromptsHandlerTags.openingTag}${prompt}${tagPromptsHandlerTags.closingTag}`;
     }
 
+    if (isChatCompletion) {
+      prompt += "\n\n" + this.settings.chatSeparator + "\n\n";
+      completion += "\n\n" + this.settings.chatSeparator + "\n\n";
+    }
+
     return prompt + completion;
   }
 
@@ -126,6 +185,29 @@ export default class GPTPlugin extends Plugin {
     }
   }
 
+  async chatCompletionHandler(editor: Editor) {
+    const selection: string = this.getSelectedText(editor);
+    if (selection) {
+      const completion = await this.getChatCompletion(selection);
+      if (!completion) {
+        this.handleGetCompletionError();
+        return;
+      }
+      editor.replaceSelection(
+        this.formatCompletion(selection, completion, true)
+      );
+      return;
+    } else {
+      const noteContents = this.getNoteContents(editor);
+      const completion = await this.getChatCompletion(noteContents);
+      if (!completion) {
+        this.handleGetCompletionError();
+        return;
+      }
+      editor.setValue(this.formatCompletion(noteContents, completion, true));
+    }
+  }
+
   initLeaf(): void {
     if (this.app.workspace.getLeavesOfType(VIEW_TYPE_MODEL_SETTINGS).length) {
       return;
@@ -161,8 +243,28 @@ export default class GPTPlugin extends Plugin {
     }
   }
 
+  async populateSettingDefaults() {
+    // ensure that each model's default settings are populated
+    const settings = this.settings;
+    console.log(settings);
+    Object.values(SupportedModels).forEach((model) => {
+      if (!settings.models[model]) {
+        console.log("populating default settings for", model);
+        settings.models[model] = {
+          settings: DEFAULT_SETTINGS.models[model].settings as never,
+          apiKey: "",
+        };
+      }
+    });
+    if (!settings.chatSeparator) {
+      settings.chatSeparator = DEFAULT_SETTINGS.chatSeparator;
+    }
+    await this.saveData(settings);
+  }
+
   async onload() {
     await this.loadSettings();
+    await this.populateSettingDefaults();
 
     this.registerView(
       VIEW_TYPE_MODEL_SETTINGS,
@@ -175,6 +277,12 @@ export default class GPTPlugin extends Plugin {
       editorCallback: (editor: Editor) => this.getCompletionHandler(editor),
     });
 
+    this.addCommand({
+      id: "chat-completion",
+      name: "Chat Completion",
+      editorCallback: (editor: Editor) => this.chatCompletionHandler(editor),
+    });
+
     this.addCommand({
       id: "show-model-settings",
       name: "Show Model Settings",

manifest.json

@@ -1,7 +1,7 @@
 {
   "id": "obsidian-gpt",
   "name": "GPT",
-  "version": "1.0.2",
+  "version": "1.1.0",
   "minAppVersion": "0.9.12",
   "description": "GPT & Large Language Model completions in Obsidian editor via API",
   "author": "Jonathan Miller",

package.json

@@ -1,6 +1,6 @@
 {
   "name": "obsidian-gpt",
-  "version": "1.0.2",
+  "version": "1.1.0",
   "description": "GPT & Large Language Model completions in Obsidian editor via API",
   "main": "main.js",
   "scripts": {

src/SettingsTab.ts

@@ -11,7 +11,7 @@ class GPTSettingTab extends PluginSettingTab {
 
   display(): void {
     let { containerEl } = this;
-    let { gpt3, ai21, cohere } = this.plugin.settings.models;
+    let { gpt3, chatgpt, ai21, cohere } = this.plugin.settings.models;
 
     containerEl.empty();
 
@@ -32,6 +32,19 @@ class GPTSettingTab extends PluginSettingTab {
           })
       );
 
+    new Setting(containerEl)
+      .setName("ChatGPT API Key")
+      .setDesc("Enter your OpenAI API Key (to use with ChatGPT)")
+      .addText((text) =>
+        text
+          .setPlaceholder("API Key")
+          .setValue(chatgpt.apiKey)
+          .onChange(async (value) => {
+            chatgpt.apiKey = value;
+            await this.plugin.saveSettings();
+          })
+      );
+
     new Setting(containerEl)
       .setName("AI21 API Key")
       .setDesc("Enter your AI21 API Key")
@@ -125,6 +138,17 @@ class GPTSettingTab extends PluginSettingTab {
           await this.plugin.saveSettings();
         })
     );
+
+    new Setting(containerEl)
+      .setName("Chat Completion Separator")
+      .addText((text) =>
+        text
+          .setValue(this.plugin.settings.chatSeparator)
+          .onChange(async (value) => {
+            this.plugin.settings.chatSeparator = value;
+            await this.plugin.saveSettings();
+          })
+      );
   }
 }
 

src/models/chatGPT.ts

@@ -0,0 +1,72 @@
+import { request, RequestParam } from "obsidian";
+import { pythonifyKeys } from "src/util";
+
+export enum ChatGPTModelType {
+  Default = "gpt-3.5-turbo",
+}
+
+export type ChatRole = "user" | "system" | "assistant";
+
+export interface ChatMessage {
+  role: ChatRole;
+  content: string;
+}
+
+export interface ChatGPTSettings {
+  modelType: ChatGPTModelType;
+  systemMessage: string;
+  maxTokens: number;
+  temperature: number;
+  topP: number;
+  presencePenalty: number;
+  frequencyPenalty: number;
+  stop: string[];
+}
+
+export const defaultChatGPTSettings: ChatGPTSettings = {
+  modelType: ChatGPTModelType.Default,
+  systemMessage:
+    "You are ChatGPT, a large language model trained by OpenAI. Answer as concisely as possible.",
+  maxTokens: 200,
+  temperature: 1.0,
+  topP: 1.0,
+  presencePenalty: 0,
+  frequencyPenalty: 0,
+  stop: [],
+};
+
+export const getChatGPTCompletion = async (
+  apiKey: string,
+  messages: ChatMessage[],
+  settings: ChatGPTSettings,
+  suffix?: string
+): Promise<string> => {
+  const apiUrl = `https://api.openai.com/v1/chat/completions`;
+  const headers = {
+    Authorization: `Bearer ${apiKey}`,
+    "Content-Type": "application/json",
+  };
+  const { modelType, systemMessage, ...params } = settings;
+  let body = {
+    messages,
+    model: modelType,
+    ...pythonifyKeys(params),
+    stop: settings.stop.length > 0 ? settings.stop : undefined,
+    suffix: suffix ? suffix : undefined,
+  };
+  const requestParam: RequestParam = {
+    url: apiUrl,
+    method: "POST",
+    contentType: "application/json",
+    body: JSON.stringify(body),
+    headers,
+  };
+  const res: any = await request(requestParam)
+    .then((response) => {
+      return JSON.parse(response);
+    })
+    .catch((err) => {
+      console.error(err);
+    });
+  return res?.choices?.[0]?.message?.content ?? null;
+};

src/models/gpt3.ts

@@ -6,6 +6,7 @@ export enum GPT3ModelType {
   Babbage = "text-babbage-001",
   Curie = "text-curie-001",
   TextDaVinci = "text-davinci-003",
+  CodeDaVinci = "code-davinci-002",
   DaVinci = "davinci",
 }
 

src/types.ts

@@ -1,16 +1,22 @@
 import { AI21Settings, defaultAI21Settings } from "src/models/ai21";
+import { ChatGPTSettings, defaultChatGPTSettings } from "src/models/chatGPT";
 import { CohereSettings, defaultCohereSettings } from "src/models/cohere";
 import { GPT3Settings, defaultGPT3Settings } from "src/models/gpt3";
 
 export const VIEW_TYPE_MODEL_SETTINGS = "gptModelSettings";
 
 export enum SupportedModels {
-  GPT3 = "GPT-3",
-  AI21 = "AI21",
-  COHERE = "Cohere",
+  CHATGPT = "chatgpt",
+  GPT3 = "gpt3",
+  AI21 = "ai21",
+  COHERE = "cohere",
 }
 
 export interface Models {
+  chatgpt: {
+    apiKey: string;
+    settings: ChatGPTSettings;
+  };
   gpt3: {
     apiKey: string;
     settings: GPT3Settings;
@@ -38,11 +44,16 @@ export interface GPTPluginSettings {
   tagPrompts: boolean;
   tagPromptsHandlerTags: HandlerTags;
   insertToken: string;
+  chatSeparator: string;
 }
 
 export const DEFAULT_SETTINGS: GPTPluginSettings = {
   activeModel: SupportedModels.GPT3,
   models: {
+    chatgpt: {
+      apiKey: "",
+      settings: defaultChatGPTSettings,
+    },
     gpt3: {
       apiKey: "",
       settings: defaultGPT3Settings,
@@ -67,6 +78,7 @@ export const DEFAULT_SETTINGS: GPTPluginSettings = {
     closingTag: "</Prompt>",
   },
   insertToken: "[insert]",
+  chatSeparator: "|||",
 };
 
 // Utils