Skip to content

Commit

Permalink
feat: ability to select model
Browse files Browse the repository at this point in the history
  • Loading branch information
ravenwits committed May 18, 2024
1 parent 7503cc6 commit 069f33d
Show file tree
Hide file tree
Showing 3 changed files with 19 additions and 4 deletions.
4 changes: 2 additions & 2 deletions src/plugin/gpt.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
const url = 'https://api.openai.com/v1/chat/completions';

export async function* prompt(prompt: string, apiKey: string, maxTokens: number) {
export async function* prompt(prompt: string, apiKey: string, maxTokens: number, model: string = 'gpt-3.5-turbo') {
const requestOptions = {
method: 'POST',
headers: {
Expand All @@ -9,7 +9,7 @@ export async function* prompt(prompt: string, apiKey: string, maxTokens: number)
},
body: JSON.stringify({
messages: [{ role: 'system', content: prompt }],
model: 'gpt-3.5-turbo',
model: model,
temperature: 0.7,
max_tokens: maxTokens,
top_p: 1,
Expand Down
4 changes: 2 additions & 2 deletions src/plugin/main.ts
Original file line number Diff line number Diff line change
Expand Up @@ -29,11 +29,11 @@ export default class AISummarizePlugin extends Plugin {
const promptText = `${this.settings.defaultPrompt} ${title ? 'title of the note is: ' + title + '\n' : ''} \n\n${selectedText}`;

if (!!this.settings.apiKey) {
Notify('Generating summary...');
Notify(`Generating summary...`);
let message = 'Summary updated successfully.';
(async () => {
let summary = '';
for await (const summaryChunk of prompt(promptText, this.settings.apiKey, this.settings.maxTokens)) {
for await (const summaryChunk of prompt(promptText, this.settings.apiKey, this.settings.maxTokens, this.settings.model)) {
if (!this.settings.putSummaryInFrontmatter) editor.replaceSelection(summaryChunk);
summary += summaryChunk;
}
Expand Down
15 changes: 15 additions & 0 deletions src/settings/settings.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ export interface AiSummarizePluginSettings {
maxTokens: number;
defaultPrompt: string;
putSummaryInFrontmatter: boolean;
model?: string;
}

const defaultMaxTokens = 1000;
Expand All @@ -14,6 +15,7 @@ export const default_settings: AiSummarizePluginSettings = {
maxTokens: defaultMaxTokens,
defaultPrompt: 'Generate a very short summary of the following note in 3-4 sentences:\n\n',
putSummaryInFrontmatter: false,
model: 'gpt-3.5-turbo',
};

export default class AiSummarizeSettingTab extends PluginSettingTab {
Expand Down Expand Up @@ -57,6 +59,19 @@ export default class AiSummarizeSettingTab extends PluginSettingTab {
}),
);

new Setting(containerEl)
.setName('AI Model')
.setDesc('The AI model to use for generating the summary. Default is gpt-3.5-turbo.')
.addDropdown((dropdown) =>
dropdown
.addOptions({ 'gpt-3.5-turbo': 'gpt-3.5-turbo', 'gpt-4': 'gpt-4', 'gpt-4-turbo': 'gpt-4-turbo', 'gpt-4o': 'gpt-4o' })
.setValue(this.plugin.settings.model)
.onChange(async (value) => {
this.plugin.settings.model = value;
await this.plugin.saveSettings();
}),
);

new Setting(containerEl)
.setName('Default prompt')
.setDesc('Default prompt for the AI to generate a summary.')
Expand Down

0 comments on commit 069f33d

Please sign in to comment.