diff --git a/assistant.ts b/assistant.ts new file mode 100644 index 0000000..34120ca --- /dev/null +++ b/assistant.ts @@ -0,0 +1,125 @@ +import { Configuration, OpenAIApi } from "openai"; +import GPT3Tokenizer from "gpt3-tokenizer"; +import { cosineSimilarity } from "utils"; + +export interface chunkData { + text: string; + embeddings: number[]; +} + +export type EmbeddedData = chunkData[]; + +export type Answer = { error: boolean; text: string }; +export class Assistant { + MAX_TOKENS = 500; + openai: OpenAIApi; + tokenizer: GPT3Tokenizer; + searchable?: EmbeddedData; + constructor(apiKey: string) { + const configuration = new Configuration({ + apiKey: apiKey, + }); + this.openai = new OpenAIApi(configuration); + this.tokenizer = new GPT3Tokenizer({ type: "gpt3" }); + } + + setData(emb: EmbeddedData) { + this.searchable = emb; + } + + private async createContext(question: string, maxLength = 1800) { + if (this.searchable == null) throw new Error("NO SEARCHABLE DATA"); + const questionEmbeddings = ( + await this.openai.createEmbedding({ + input: question, + model: "text-embedding-ada-002", + }) + ).data.data[0].embedding; + const sortedData = this.searchable.sort( + (a, b) => + cosineSimilarity(questionEmbeddings, b.embeddings) - + cosineSimilarity(questionEmbeddings, a.embeddings) + ); + const returns = []; + let cur_len = 0; + for (const row of sortedData) { + const n_tokens = this.tokenizer.encode(row.text).bpe.length; + cur_len += n_tokens + 4; + if (cur_len > maxLength) break; + returns.push(row.text); + } + return returns.join("\n\n###\n\n"); + } + async answerQuestion( + question: string, + maxLength = 1800, + maxTokens = 150, + stopSequence = null + ): Promise { + if (this.searchable == null) + return { + error: true, + text: "Data not loaded. Please load it from settings", + }; + const context = await this.createContext(question, maxLength); + const response = await this.openai.createCompletion({ + prompt: `Answer the question based on the context below, and if the question can't be answered based on the context, say "I don't know, I couldn't find anything related to this in your notes."\n\nContext: ${context}\n\n---\n\nQuestion: ${question}\nAnswer:`, + temperature: 0, + max_tokens: maxTokens, + top_p: 1, + frequency_penalty: 0, + presence_penalty: 0, + stop: stopSequence, + model: "text-davinci-003", + }); + return { + error: false, + text: response.data.choices[0].text ?? "Something went wrong.", + }; + } + + prepareTexts(texts: string[]): string[] { + let shortened: string[] = []; + texts.forEach((text) => { + if (this.tokenizer.encode(text).bpe.length > this.MAX_TOKENS) { + shortened = shortened.concat(this.splitIntoMany(text)); + } else { + shortened.push(text); + } + }); + return shortened; + } + + private splitIntoMany(text: string, max_tokens = 500): string[] { + const sentences = text.split(". "); + const n_tokens = sentences.map( + (sentence) => this.tokenizer.encode(" " + sentence).bpe.length + ); + let tokens_so_far = 0; + const chunks: string[] = []; + const chunk: string[] = []; + sentences.forEach((sentence, idx) => { + const token = n_tokens[idx]; + if (token + tokens_so_far > max_tokens) { + chunks.push(chunk.join(". ") + "."); + chunk.length = 0; + tokens_so_far = 0; + } + if (!(token > max_tokens)) { + chunk.push(sentence); + tokens_so_far += token + 1; + } + }); + return chunks; + } + async createEmbeddings(data: string[]): Promise { + const embeddings = await this.openai.createEmbedding({ + input: data, + model: "text-embedding-ada-002", + }); + return data.map((text, idx) => ({ + text, + embeddings: embeddings.data.data[idx].embedding, + })); + } +} diff --git a/main.ts b/main.ts index 50b75f3..81499ac 100644 --- a/main.ts +++ b/main.ts @@ -1,113 +1,157 @@ -import { App, Editor, MarkdownView, Modal, Notice, Plugin, PluginSettingTab, Setting } from 'obsidian'; - -// Remember to rename these classes and interfaces! +import { Answer, Assistant } from "assistant"; +import { + App, + Component, + MarkdownPreviewRenderer, + MarkdownRenderer, + Modal, + Notice, + Plugin, + PluginSettingTab, + Setting, + TextAreaComponent, +} from "obsidian"; +// TODO: Remember to rename these classes and interfaces! interface MyPluginSettings { - mySetting: string; + apiKey: string; } const DEFAULT_SETTINGS: MyPluginSettings = { - mySetting: 'default' -} + apiKey: "", +}; export default class MyPlugin extends Plugin { settings: MyPluginSettings; - + assistant: Assistant; async onload() { + this.addSettingTab(new AssistantSettings(this.app, this)); await this.loadSettings(); - // This creates an icon in the left ribbon. - const ribbonIconEl = this.addRibbonIcon('dice', 'Sample Plugin', (evt: MouseEvent) => { - // Called when the user clicks the icon. - new Notice('This is a notice!'); - }); - // Perform additional things with the ribbon - ribbonIconEl.addClass('my-plugin-ribbon-class'); + this.assistant = new Assistant(this.settings.apiKey); + if (await this.hasCachedData()) { + let { searchable } = await this.loadData(); + this.saveNamedData("searchable", searchable); + this.assistant.setData(searchable); + } - // This adds a status bar item to the bottom of the app. Does not work on mobile apps. - const statusBarItemEl = this.addStatusBarItem(); - statusBarItemEl.setText('Status Bar Text'); - - // This adds a simple command that can be triggered anywhere this.addCommand({ - id: 'open-sample-modal-simple', - name: 'Open sample modal (simple)', - callback: () => { - new SampleModal(this.app).open(); - } - }); - // This adds an editor command that can perform some operation on the current editor instance - this.addCommand({ - id: 'sample-editor-command', - name: 'Sample editor command', - editorCallback: (editor: Editor, view: MarkdownView) => { - console.log(editor.getSelection()); - editor.replaceSelection('Sample Editor Command'); - } - }); - // This adds a complex command that can check whether the current state of the app allows execution of the command - this.addCommand({ - id: 'open-sample-modal-complex', - name: 'Open sample modal (complex)', - checkCallback: (checking: boolean) => { - // Conditions to check - const markdownView = this.app.workspace.getActiveViewOfType(MarkdownView); - if (markdownView) { - // If checking is true, we're simply "checking" if the command can be run. - // If checking is false, then we want to actually perform the operation. - if (!checking) { - new SampleModal(this.app).open(); - } - - // This command will only show up in Command Palette when the check function returns true - return true; + id: "ask-assistant", + name: "Ask assistant", + callback: async () => { + if (!(await this.hasCachedData())) { + new Notice( + "You must load data before asking question, you can do so from the plugin settings" + ); + return; } - } + if (!this.settings.apiKey) { + new Notice("Please provide an API Key in the settings"); + return; + } + new AskAssistantModal(this.app, async (question) => { + const answer = await this.assistant.answerQuestion( + question + ); + return answer ?? ""; + }).open(); + }, }); - - // This adds a settings tab so the user can configure various aspects of the plugin - this.addSettingTab(new SampleSettingTab(this.app, this)); - - // If the plugin hooks up any global DOM events (on parts of the app that doesn't belong to this plugin) - // Using this function will automatically remove the event listener when this plugin is disabled. - this.registerDomEvent(document, 'click', (evt: MouseEvent) => { - console.log('click', evt); - }); - - // When registering intervals, this function will automatically clear the interval when the plugin is disabled. - this.registerInterval(window.setInterval(() => console.log('setInterval'), 5 * 60 * 1000)); } - onunload() { - + private async hasCachedData(): Promise { + const data = await this.loadData(); + return data.searchable && data.searchable.length; } + async loadEmbeddingsToAssistant() { + const { vault } = this.app; + const fileContents: string[] = await Promise.all( + vault + .getMarkdownFiles() + .map((file) => + vault.cachedRead(file).then((res) => file.name + res) + ) + ); + const chunks = await this.assistant.prepareTexts(fileContents); + const searchable = await this.assistant.createEmbeddings(chunks); + this.saveNamedData("searchable", searchable); + this.assistant.setData(searchable); + } + + onunload() {} + async loadSettings() { - this.settings = Object.assign({}, DEFAULT_SETTINGS, await this.loadData()); + this.settings = Object.assign( + {}, + DEFAULT_SETTINGS, + (await this.loadData()).settings + ); } async saveSettings() { - await this.saveData(this.settings); + this.assistant = new Assistant(this.settings.apiKey); + await this.saveData({ + ...(await this.loadData()), + settings: this.settings, + }); + } + + async saveNamedData(name: string, data: unknown) { + await this.saveData({ ...(await this.loadData()), [name]: data }); } } -class SampleModal extends Modal { - constructor(app: App) { +class AskAssistantModal extends Modal { + result: string; + onSubmit: (result: string) => Promise; + constructor(app: App, onSubmit: (result: string) => Promise) { super(app); + this.onSubmit = onSubmit; } onOpen() { - const {contentEl} = this; - contentEl.setText('Woah!'); + const { contentEl } = this; + const answer_div = contentEl.createEl("div", { + text: "", + cls: "answer", + }); + const promptTextArea = new TextAreaComponent(contentEl); + promptTextArea.onChange((text) => { + this.result = text; + }); + promptTextArea.setPlaceholder("Enter your prompt..."); + promptTextArea.inputEl.classList.add("prompt_in"); + + new Setting(contentEl).addButton((btn) => + btn + .setButtonText("Ask Assistant") + .setCta() + .onClick(async () => { + answer_div.innerText = "Loading....."; + const answer = await this.onSubmit(this.result); + answer_div.innerText = ""; + if (answer.error) { + new Notice(answer.text); + return; + } + MarkdownRenderer.renderMarkdown( + answer.text, + answer_div, + "", + null as any + ); + }) + ); } onClose() { - const {contentEl} = this; + const { contentEl } = this; contentEl.empty(); } } -class SampleSettingTab extends PluginSettingTab { +class AssistantSettings extends PluginSettingTab { plugin: MyPlugin; constructor(app: App, plugin: MyPlugin) { @@ -116,22 +160,45 @@ class SampleSettingTab extends PluginSettingTab { } display(): void { - const {containerEl} = this; + const { containerEl } = this; containerEl.empty(); - containerEl.createEl('h2', {text: 'Settings for my awesome plugin.'}); + containerEl.createEl("h2", { text: "Obsidian-GPT Assistant settings" }); new Setting(containerEl) - .setName('Setting #1') - .setDesc('It\'s a secret') - .addText(text => text - .setPlaceholder('Enter your secret') - .setValue(this.plugin.settings.mySetting) - .onChange(async (value) => { - console.log('Secret: ' + value); - this.plugin.settings.mySetting = value; - await this.plugin.saveSettings(); - })); + .setName("API Key") + .setDesc( + "Your OpenAI API key. Can be found at https://platform.openai.com/account/api-keys" + ) + .addText((text) => + text + .setPlaceholder("Enter your secret") + .setValue(this.plugin.settings.apiKey) + .onChange(async (value) => { + this.plugin.settings.apiKey = value; + await this.plugin.saveSettings(); + }) + ); + + new Setting(containerEl) + .setName("Process notes") + .setDesc("Load all your notes into the assistant") + .addButton((btn) => { + btn.setButtonText("Process").setCta(); + + btn.onClick(async (e) => { + if (this.plugin.settings.apiKey == "") { + new Notice("Please provide an API Key"); + return; + } + new Notice( + "Loading data into model. this could take a while..." + ); + await this.plugin.loadEmbeddingsToAssistant(); + + new Notice("Your data has been loaded into the model."); + }); + }); } } diff --git a/manifest.json b/manifest.json index 0897aff..319d25f 100644 --- a/manifest.json +++ b/manifest.json @@ -1,11 +1,11 @@ { - "id": "obsidian-sample-plugin", - "name": "Sample Plugin", - "version": "1.0.0", + "id": "obsidian-gpt-assistant", + "name": "GPT Assistant", + "version": "0.1.0", "minAppVersion": "0.15.0", - "description": "This is a sample plugin for Obsidian. This plugin demonstrates some of the capabilities of the Obsidian API.", - "author": "Obsidian", - "authorUrl": "https://obsidian.md", - "fundingUrl": "https://obsidian.md/pricing", + "description": "Train a GPT-3 based model on your notes and get personalized answers based on your knowledge base.", + "author": "M7mdisk", + "authorUrl": "https://github.com/M7mdisk", + "fundingUrl": "https://github.com/M7mdisk", "isDesktopOnly": false } diff --git a/package.json b/package.json index 6a00766..5b5230a 100644 --- a/package.json +++ b/package.json @@ -1,7 +1,7 @@ { - "name": "obsidian-sample-plugin", - "version": "1.0.0", - "description": "This is a sample plugin for Obsidian (https://obsidian.md)", + "name": "obsidian-gpt-assistant", + "version": "0.1.0", + "description": "Train a GPT-3 based model on your notes and get personalized answers based on your knowledge base.", "main": "main.js", "scripts": { "dev": "node esbuild.config.mjs", @@ -20,5 +20,9 @@ "obsidian": "latest", "tslib": "2.4.0", "typescript": "4.7.4" + }, + "dependencies": { + "gpt3-tokenizer": "^1.1.5", + "openai": "^3.2.1" } } diff --git a/styles.css b/styles.css index 71cc60f..b6aa936 100644 --- a/styles.css +++ b/styles.css @@ -6,3 +6,13 @@ available in the app when your plugin is enabled. If your plugin does not need CSS, delete this file. */ + +.prompt_in { + width: 100%; + resize: vertical; +} + +.answer { + margin-top: 20px; + height: 250px; +} diff --git a/utils.ts b/utils.ts new file mode 100644 index 0000000..8b97a83 --- /dev/null +++ b/utils.ts @@ -0,0 +1,19 @@ +function dotProduct(vecA: number[], vecB: number[]) { + let product = 0; + for (let i = 0; i < vecA.length; i++) { + product += vecA[i] * vecB[i]; + } + return product; +} + +function magnitude(vec: number[]) { + let sum = 0; + for (let i = 0; i < vec.length; i++) { + sum += vec[i] * vec[i]; + } + return Math.sqrt(sum); +} + +export function cosineSimilarity(vecA: number[], vecB: number[]) { + return dotProduct(vecA, vecB) / (magnitude(vecA) * magnitude(vecB)); +}