feat: adding additional settings to configure semantic and generative…

… search parameters
solderneer · Jul 18, 2023 · ff7c541 · ff7c541
1 parent 1073468
commit ff7c541
Show file tree

Hide file tree

Showing 6 changed files with 253 additions and 68 deletions.
diff --git a/README.md b/README.md
@@ -10,9 +10,9 @@ Powered by [Supabase Vector](https://supabase.com/vector) and the [OpenAI API](h
 - ✅ Semantic search over your notes
 - ✅ Talk to your notes
 - ✅ Simple unified UI
+- ✅ Public endpoint to allow others to talk to your notes
 
 ## Wishlist
-- Public endpoint for others to talk to selected notes **(working on it!)**
 - Suggest related notes to link to the active note
 - Suggest tags for note
 - Hybrid search with keyword and semantic matching

diff --git a/src/main.ts b/src/main.ts
@@ -6,26 +6,40 @@ import * as path from "path";
 import { oneLine } from "common-tags";
 
 import { createClient, SupabaseClient } from "@supabase/supabase-js";
-import { Configuration, OpenAIApi } from "openai-edge";
+import {
+	Configuration,
+	OpenAIApi,
+	CreateModerationResponse,
+} from "openai-edge";
 
 // Local things
 import { generateEmbeddings } from "./generate-embeddings";
 import { generativeSearch, semanticSearch } from "./search";
 import { truncateString, removeMarkdown } from "./utils";
 
+interface SearchSettings {
+	matchThreshold: number;
+	matchCount: number;
+	minContentLength: number;
+}
+
 interface ObsidianAISettings {
 	supabaseUrl: string;
 	supabaseKey: string;
 	openaiKey: string;
 	indexOnOpen: boolean;
 
+	// Directory settings
 	excludedDirs: string;
 	excludedDirsList: string[];
-
 	publicDirs: string;
 	publicDirsList: string[];
 
+	// Prompt injection
 	prompt: string;
+
+	semanticSearch: SearchSettings;
+	generativeSearch: SearchSettings;
 }
 
 const DEFAULT_SETTINGS: ObsidianAISettings = {
@@ -45,6 +59,18 @@ const DEFAULT_SETTINGS: ObsidianAISettings = {
 		the following sections from your notes, answer the question with the provided information. If
 		you are unsure, and the notes don't include relevant information, you may also say
 		"Sorry, I don't know the answer to this question :("`,
+
+	semanticSearch: {
+		matchThreshold: 0.78,
+		matchCount: 10,
+		minContentLength: 50,
+	},
+
+	generativeSearch: {
+		matchThreshold: 0.78,
+		matchCount: 10,
+		minContentLength: 50,
+	},
 };
 
 export default class ObsidianAIPlugin extends Plugin {
@@ -145,7 +171,9 @@ export default class ObsidianAIPlugin extends Plugin {
 							this.app,
 							this.supabaseClient,
 							this.openai,
-							this.settings.prompt
+							this.settings.prompt,
+							this.settings.semanticSearch,
+							this.settings.generativeSearch
 						).open();
 					}
 
@@ -212,24 +240,32 @@ interface SearchResult {
 	similarity: string;
 }
 
+type KeyListener = (event: KeyboardEvent) => void;
+
 class AISearchModal extends SuggestModal<SearchResult> {
-	private keyListener: any;
+	private keyListener: KeyListener;
 	private supabaseClient: SupabaseClient;
 	private prompt: string;
 	private typedInstance: Typed;
 	private openai: OpenAIApi;
+	private semanticSearchSettings: SearchSettings;
+	private generativeSearchSettings: SearchSettings;
 
 	constructor(
 		app: App,
 		supabaseClient: SupabaseClient,
 		openai: OpenAIApi,
-		prompt: string
+		prompt: string,
+		semanticSearchSettings: SearchSettings,
+		generativeSearchSettings: SearchSettings
 	) {
 		super(app);
 
 		this.supabaseClient = supabaseClient;
 		this.openai = openai;
 		this.prompt = prompt;
+		this.semanticSearchSettings = semanticSearchSettings;
+		this.generativeSearchSettings = generativeSearchSettings;
 
 		// Adding the instructions
 		const instructions = [
@@ -294,11 +330,28 @@ class AISearchModal extends SuggestModal<SearchResult> {
 					".prompt-input"
 				) as HTMLInputElement;
 
+				const query = inputEl.value;
+				// Sanitize input query
+				// Moderate the content to comply with OpenAI T&C
+				const moderationResponse: CreateModerationResponse =
+					await this.openai
+						.createModeration({ input: query.trim() })
+						.then((res) => res.json());
+
+				const [moderationRes] = moderationResponse.results;
+
+				if (moderationRes.flagged) {
+					throw new Error("Flagged content");
+				}
+
 				const answer = await generativeSearch(
 					this.supabaseClient,
 					this.openai,
-					inputEl.value,
-					this.prompt
+					query,
+					this.prompt,
+					this.generativeSearchSettings.matchThreshold,
+					this.generativeSearchSettings.matchCount,
+					this.generativeSearchSettings.minContentLength
 				);
 
 				if (this.typedInstance) {
@@ -325,10 +378,25 @@ class AISearchModal extends SuggestModal<SearchResult> {
 
 	// Returns all available suggestions.
 	async getSuggestions(query: string): Promise<SearchResult[]> {
+		// Sanitize input query
+		// Moderate the content to comply with OpenAI T&C
+		const moderationResponse: CreateModerationResponse = await this.openai
+			.createModeration({ input: query.trim() })
+			.then((res) => res.json());
+
+		const [moderationRes] = moderationResponse.results;
+
+		if (moderationRes.flagged) {
+			throw new Error("Flagged content");
+		}
+
 		const results: SearchResult[] = await semanticSearch(
 			this.supabaseClient,
 			this.openai,
-			query
+			query,
+			this.semanticSearchSettings.matchThreshold,
+			this.semanticSearchSettings.matchCount,
+			this.semanticSearchSettings.minContentLength
 		);
 		return results;
 	}
@@ -471,5 +539,127 @@ class SettingTab extends PluginSettingTab {
 					this.plugin.setupOpenai();
 				})
 		);
+
+		containerEl.createEl("div", {
+			cls: "setting-item setting-item-heading",
+			text: "Semantic Search Settings",
+		});
+
+		new Setting(containerEl)
+			.setName("Match Threshold")
+			.setDesc(
+				"The minimum similarity score to return a match (between 0 and 1)"
+			)
+			.addText((text) =>
+				text
+					.setPlaceholder("Enter number")
+					.setValue(
+						this.plugin.settings.semanticSearch.matchThreshold.toString()
+					)
+					.onChange(async (value) => {
+						if (!isNaN(parseInt(value))) {
+							this.plugin.settings.semanticSearch.matchThreshold =
+								parseFloat(value);
+							await this.plugin.saveSettings();
+						}
+					})
+			);
+
+		new Setting(containerEl)
+			.setName("Match Count")
+			.setDesc("The maximum number of results to return")
+			.addText((text) =>
+				text
+					.setPlaceholder("Enter integer")
+					.setValue(
+						this.plugin.settings.semanticSearch.matchCount.toString()
+					)
+					.onChange(async (value) => {
+						if (!isNaN(parseInt(value))) {
+							this.plugin.settings.semanticSearch.matchCount =
+								parseInt(value);
+							await this.plugin.saveSettings();
+						}
+					})
+			);
+
+		new Setting(containerEl)
+			.setName("Min Content Length")
+			.setDesc("The minimum length for valid result string")
+			.addText((text) =>
+				text
+					.setPlaceholder("Enter integer")
+					.setValue(
+						this.plugin.settings.semanticSearch.minContentLength.toString()
+					)
+					.onChange(async (value) => {
+						if (!isNaN(parseInt(value))) {
+							this.plugin.settings.semanticSearch.minContentLength =
+								parseInt(value);
+							await this.plugin.saveSettings();
+						}
+					})
+			);
+
+		containerEl.createEl("div", {
+			cls: "setting-item setting-item-heading",
+			text: "Generative Search Settings",
+		});
+
+		new Setting(containerEl)
+			.setName("Match Threshold")
+			.setDesc(
+				"The minimum similarity score to return a match (between 0 and 1)"
+			)
+			.addText((text) =>
+				text
+					.setPlaceholder("Enter number")
+					.setValue(
+						this.plugin.settings.generativeSearch.matchThreshold.toString()
+					)
+					.onChange(async (value) => {
+						if (!isNaN(parseInt(value))) {
+							this.plugin.settings.generativeSearch.matchThreshold =
+								parseFloat(value);
+							await this.plugin.saveSettings();
+						}
+					})
+			);
+
+		new Setting(containerEl)
+			.setName("Match Count")
+			.setDesc("The maximum number of results to return")
+			.addText((text) =>
+				text
+					.setPlaceholder("Enter integer")
+					.setValue(
+						this.plugin.settings.generativeSearch.matchCount.toString()
+					)
+					.onChange(async (value) => {
+						if (!isNaN(parseInt(value))) {
+							this.plugin.settings.generativeSearch.matchCount =
+								parseInt(value);
+							await this.plugin.saveSettings();
+						}
+					})
+			);
+
+		new Setting(containerEl)
+			.setName("Min Content Length")
+			.setDesc("The minimum length for valid result string")
+			.addText((text) =>
+				text
+					.setPlaceholder("Enter integer")
+					.setValue(
+						this.plugin.settings.generativeSearch.minContentLength.toString()
+					)
+					.onChange(async (value) => {
+						if (!isNaN(parseInt(value))) {
+							this.plugin.settings.generativeSearch.minContentLength =
+								parseInt(value);
+							await this.plugin.saveSettings();
+						}
+					})
+			);
 	}
 }