2.17.3: feat: switch from internal embeddings to open source another-…

…ai/embedbase
different-ai · Feb 10, 2023 · b95ea13 · b95ea13
1 parent ce23b7c
commit b95ea13
Show file tree

Hide file tree

Showing 7 changed files with 159 additions and 208 deletions.
diff --git a/manifest.json b/manifest.json
@@ -1,7 +1,7 @@
 {
 	"id": "ava",
 	"name": "🧙 AVA",
-	"version": "2.17.2",
+	"version": "2.17.3",
 	"minAppVersion": "0.12.0",
 	"description": "AI assistant for Obsidian",
 	"author": "louis030195",

diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "ava",
-  "version": "2.17.2",
+  "version": "2.17.3",
   "description": "AI assistant for Obsidian",
   "main": "main.js",
   "scripts": {

diff --git a/src/LinkComponent.tsx b/src/LinkComponent.tsx
@@ -1,15 +1,16 @@
 import { ChevronDownIcon, ChevronRightIcon } from '@heroicons/react/24/outline';
+import { Notice } from 'obsidian';
 import { posthog } from 'posthog-js';
 import * as React from 'react';
 import { useForm } from 'react-hook-form';
 import { ReactMarkdown } from 'react-markdown/lib/react-markdown';
 import { PrimaryButton } from './Button';
-import { API_HOST } from './constants';
+import { EMBEDBASE_URL } from './constants';
 import { CopyToClipboardButton } from './CopyToClipboard';
 import { useApp } from './hooks';
 import { InsertButton } from './InsertButton';
 import { store } from './store';
-import { camelize, ISearchBody } from './utils';
+import { camelize } from './utils';
 
 export interface ButtonProps extends React.HTMLAttributes<HTMLButtonElement> {
   children: React.ReactNode;
@@ -90,34 +91,40 @@ const ControlForm = () => {
   const state = React.useSyncExternalStore(store.subscribe, store.getState);
 
   const onSubmit = async (data: { limit: number; useNoteTitle: boolean }) => {
+    if (!state.settings.useLinks) {
+      new Notice('🧙 Link - You need to enable links in settings', 3000);
+      return;
+    }
+    if (!state.settings.token) {
+      new Notice('🧙 You need to login to use this feature', 3000);
+      return;
+    }
     posthog.capture('use-feature', {
       feature: 'search',
       limit: data.limit,
     });
     state.setEmbedsLoading(true);
 
-    const body: ISearchBody = {
-      // parenthese are needed for the ternary operator to work
-      query:
-        (data.useNoteTitle ? `File:\n${state.currentFilePath}\n` : '') +
-        `Content:\n${state.currentFileContent}`,
-      vault_id: state.settings.vaultId,
-      top_k: Number(data.limit),
-    };
-    const url = posthog.isFeatureEnabled('new-search') ?
-      `https://obsidian-search-dev-c6txy76x2q-uc.a.run.app/v1/search` :
-      `${API_HOST}/v1/search`;
-    const response = await fetch(url, {
-      method: 'POST',
-      headers: {
-        'Content-Type': 'application/json',
-        Authorization: `Bearer ${state.settings.token}`,
-        'X-Client-Version': state.version,
-      },
-      body: JSON.stringify(body),
-    }).then((res) => res.json());
+    const response = await fetch(`${EMBEDBASE_URL}/v1/${state.settings.vaultId}/search`, {
+        method: 'POST',
+        headers: {
+          'Content-Type': 'application/json',
+          Authorization: `Bearer ${state.settings.token}`,
+          'X-Client-Version': state.version,
+        },
+        body: JSON.stringify({
+          // parenthese are needed for the ternary operator to work
+          query:
+            (data.useNoteTitle ? `File:\n${state.currentFilePath}\n` : '') +
+            `Content:\n${state.currentFileContent}`,
+          top_k: Number(data.limit),
+        }),
+      }).then((res) => res.json());
 
-    const embeds = camelize(response.similarities);
+    const embeds = camelize(response.similarities.filter(
+      (similarity: any) =>
+        similarity.id !== state.currentFilePath
+    ));
 
     state.setEmbedsLoading(false);
     /* @ts-expect-error need to work on better types */
@@ -194,7 +201,7 @@ export function LinkComponent() {
     // get the top value
     const topValue = embeds[0].score;
     const results = embeds.map((embed) => {
-      const [path, similarity] = [embed.notePath, embed.score];
+      const [path, similarity] = [embed.id, embed.score];
       const opacity = sCurve(similarity, threshold, topValue);
 
       return {
@@ -206,7 +213,6 @@ export function LinkComponent() {
     });
     setResults(results);
   }, [embeds, threshold]);
-
   const textToInsert = `${results
     .map((similarity) => '- [[' + similarity?.path?.replace('.md', '') + ']]')
     .join('\n')}`;

diff --git a/src/constants.ts b/src/constants.ts
@@ -1,5 +1,6 @@
 // TODO: firebase hosting custom domain name does not work with SSE somehow?
 export const API_HOST = process.env.API_HOST || "https://obsidian-ai-c6txy76x2q-uc.a.run.app";
+export const EMBEDBASE_URL = "https://embedbase-internal-c6txy76x2q-uc.a.run.app";
 
 export const buildHeaders = (token: string, version: string) => ({
     'Content-Type': 'application/json',

diff --git a/src/main.tsx b/src/main.tsx
@@ -8,7 +8,7 @@ import {
   Plugin,
   PluginSettingTab,
   TFile,
-  WorkspaceLeaf,
+  WorkspaceLeaf
 } from 'obsidian';
 
 import * as React from 'react';
@@ -17,26 +17,24 @@ import { CustomSettings } from './Settings';
 import {
   createImage,
   RequestImageCreate,
-  ResponseImageCreate,
+  ResponseImageCreate
 } from './stableDiffusion';
 import { StatusBar } from './StatusBar';
 import {
   clearIndex,
   complete,
   createParagraph,
   createSemanticLinks,
+  deleteFromIndex,
   EMBED_CHAR_LIMIT,
   getCompleteFiles,
   getLinkData,
   getVaultId,
   ICompletion,
-  ISearchData,
-  ISearchRequest,
-  refreshSemanticSearch,
-  rewrite,
+  ISearchRequest, ISearchResponse, rewrite,
   REWRITE_CHAR_LIMIT as TEXT_CREATE_CHAR_LIMIT,
   search,
-  suggestTags,
+  suggestTags, syncIndex
 } from './utils';
 
 import posthog from 'posthog-js';
@@ -135,7 +133,7 @@ export default class AvaPlugin extends Plugin {
   ```
   */
 
-  public search: (request: ISearchRequest) => Promise<ISearchData>;
+  public search: (request: ISearchRequest) => Promise<ISearchResponse>;
   public clearIndex: () => Promise<any>;
   private eventRefRenamed: EventRef;
   private eventRefDeleted: EventRef;
@@ -148,7 +146,7 @@ export default class AvaPlugin extends Plugin {
     this.streamingSource = source;
   }
 
-  private async link(currentText: string, tags: string[], path: string) {
+  private async link(currentText: string, path: string) {
     if (currentText.length > EMBED_CHAR_LIMIT) {
       new Notice(
         'Link - Note is too long. 🧙 AVA  only supports notes that are up to 25k characters'
@@ -160,10 +158,9 @@ export default class AvaPlugin extends Plugin {
       completion = await createSemanticLinks(
         path,
         currentText,
-        tags,
         this.settings?.token,
         this.settings.vaultId,
-        this.manifest.version
+        this.manifest.version,
       );
 
       this.statusBarItem.render(<StatusBar status="disabled" />);
@@ -238,11 +235,10 @@ export default class AvaPlugin extends Plugin {
             return acc;
           }, [])
           .map((batch) =>
-            refreshSemanticSearch(
+            syncIndex(
               batch.map((file: any) => ({
-                notePath: file.path,
-                noteTags: file.tags,
-                noteContent: file.content,
+                id: file.path,
+                data: `File:\n${file.path}\nContent:\n${file.content}`
               })),
               this.settings?.token,
               this.settings?.vaultId,
@@ -312,12 +308,11 @@ export default class AvaPlugin extends Plugin {
             if (!cache) return setLastFile(leaf);
             this.app.vault.adapter.read(this.lastFile.path).then((data) => {
               if (!this.lastFile) return setLastFile(leaf);
-              refreshSemanticSearch(
+              syncIndex(
                 [
                   {
-                    notePath: this.lastFile.path,
-                    noteTags: cache.tags?.map((tag) => tag.tag) || [],
-                    noteContent: data,
+                    id: this.lastFile.path,
+                    data: `File:\n${this.lastFile.path}\nContent:\n${data}`
                   },
                 ],
                 this.settings?.token,
@@ -344,18 +339,20 @@ export default class AvaPlugin extends Plugin {
         const f = file as TFile;
         // if file in ignored folder, ignore
         if (isIgnored(this.settings?.ignoredFolders, f.path)) return;
+
         try {
+          if (oldPath) {
+            deleteFromIndex([oldPath], this.settings?.token, this.settings?.vaultId, this.manifest.version);
+          }
           if (!this.settings.useLinks) {
             this.unlistenToNoteEvents();
             return;
           }
-          refreshSemanticSearch(
+          syncIndex(
             [
               {
-                notePath: f.path,
-                noteTags: cache.tags?.map((tag) => tag.tag) || [],
-                noteContent: data,
-                pathToDelete: oldPath,
+                id: f.path,
+                data: `File:\n${f.path}\nContent:\n${data}`
               },
             ],
             this.settings?.token,
@@ -377,16 +374,7 @@ export default class AvaPlugin extends Plugin {
         }
         // if file in ignored folder, ignore
         if (isIgnored(this.settings?.ignoredFolders, file.path)) return;
-        refreshSemanticSearch(
-          [
-            {
-              pathToDelete: file.path,
-            },
-          ],
-          this.settings?.token,
-          this.settings.vaultId,
-          this.manifest.version
-        );
+        deleteFromIndex([file.path], this.settings?.token, this.settings?.vaultId, this.manifest.version);
       } catch (e) {
         onGeneralError(e);
         this.unlistenToNoteEvents();
@@ -442,18 +430,15 @@ export default class AvaPlugin extends Plugin {
 
     const file = this.app.workspace.getActiveFile();
     const currentText = await this.app.vault.read(file);
-    const tags = this.app.metadataCache.getFileCache(file).tags || [];
-    const tagsArray = tags.map((tag) => tag.tag);
     const path = file.path;
 
     // we need to do this so we can fire /search inside of the sidebar later
     store.setState({
       currentFileContent: currentText,
       currentFilePath: path,
-      currentFileTags: tagsArray,
     });
 
-    const results = await this.link(currentText, tagsArray, path);
+    const results = await this.link(currentText, path);
     if (results) {
       store.setState({ embeds: results });
     }
@@ -523,7 +508,7 @@ export default class AvaPlugin extends Plugin {
           req,
           this.settings.token,
           this.settings.vaultId,
-          this.manifest.version
+          this.manifest.version,
         );
       this.clearIndex = () =>
         clearIndex(
@@ -877,6 +862,8 @@ export default class AvaPlugin extends Plugin {
   onunload(): void {
     this.app.workspace.detachLeavesOfType(VIEW_TYPE_LINK);
     this.app.workspace.detachLeavesOfType(VIEW_TYPE_WRITE);
+    // TODO: skip when local development (annoying have to index every time I change a line of code)
+    // TODO: careful not using node stuff for mobile?
     this.unlistenToNoteEvents();
   }
 }