Local AI support for completions and embeddings (#93)

* Refactor AI context (1)

* Refactor AI (2) & local embeddings support

* cleanup for v0.9.8

* Performance tweaks and make embedding model editable

* Embedding regeneration when models change

package.json

@@ -1,6 +1,6 @@
 {
   "description": "Pile",
-  "version": "0.9.7",
+  "version": "0.9.8",
   "keywords": [
     "pile"
   ],
@@ -96,6 +96,7 @@
     "@radix-ui/react-dialog": "^1.0.4",
     "@radix-ui/react-dropdown-menu": "^2.0.5",
     "@radix-ui/react-switch": "^1.0.3",
+    "@radix-ui/react-tabs": "^1.1.0",
     "@tanstack/react-virtual": "^3.0.1",
     "@tiptap/extension-character-count": "^2.1.12",
     "@tiptap/extension-link": "^2.1.12",
@@ -109,19 +110,20 @@
     "dotenv": "^16.3.1",
     "electron-debug": "^3.2.0",
     "electron-log": "^4.4.8",
+    "electron-settings": "^4.0.4",
     "electron-updater": "^5.3.0",
     "framer-motion": "^11.2.4",
     "gray-matter": "^4.0.3",
     "lunr": "^2.3.9",
     "luxon": "^3.3.0",
     "openai": "^4.44.0",
-    "react": "^19.0.0-beta-26f2496093-20240514",
-    "react-dom": "^19.0.0-beta-26f2496093-20240514",
+    "react": "19.0.0-rc-14a4699f-20240725",
+    "react-dom": "19.0.0-rc-14a4699f-20240725",
     "react-markdown": "^9.0.1",
     "react-router": "^6.23.1",
     "react-router-dom": "^6.23.1",
     "react-textarea-autosize": "^8.5.3",
-    "react-virtuoso": "^4.7.10"
+    "react-virtuoso": "^4.7.13"
   },
   "devDependencies": {
     "@adobe/css-tools": "^4.3.2",

release/app/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pile",
-  "version": "0.9.7",
+  "version": "0.9.8",
   "description": "Pile: Everyday journal and thought companion.",
   "license": "MIT",
   "author": {

src/main/handlers/index.ts

@@ -11,6 +11,11 @@ ipcMain.handle('index-get', (event) => {
   return index;
 });
 
+ipcMain.handle('index-regenerate-embeddings', (event) => {
+  const index = pileIndex.regenerateEmbeddings();
+  return index;
+});
+
 ipcMain.handle('index-add', (event, filePath) => {
   const index = pileIndex.add(filePath);
   return index;

src/main/handlers/store.ts

@@ -0,0 +1,10 @@
+import { ipcMain } from 'electron';
+import settings from 'electron-settings';
+
+ipcMain.handle('electron-store-get', async (event, key) => {
+  return await settings.get(key);
+});
+
+ipcMain.handle('electron-store-set', async (event, key, value) => {
+  await settings.set(key, value);
+});

src/main/ipc.ts

@@ -5,3 +5,4 @@ import './handlers/tags';
 import './handlers/highlights';
 import './handlers/index';
 import './handlers/links';
+import './handlers/store';

src/main/preload.ts

@@ -66,6 +66,9 @@ const electronHandler = {
   isMac: process.platform === 'darwin',
   isWindows: process.platform === 'win32',
   pathSeparator: path.sep,
+  settingsGet: (key: string) => ipcRenderer.invoke('electron-store-get', key),
+  settingsSet: (key: string, value: string) =>
+    ipcRenderer.invoke('electron-store-set', key, value),
 };
 
 contextBridge.exposeInMainWorld('electron', electronHandler);

src/main/utils/pileEmbeddings.js

@@ -4,12 +4,13 @@ const path = require('path');
 const keytar = require('keytar');
 const { walk } = require('../util');
 const matter = require('gray-matter');
+const settings = require('electron-settings');
 
 // Todo: Cache the norms alongside embeddings at some point
 // to avoid recomputing them for every query
 function cosineSimilarity(embedding, queryEmbedding) {
-  if (embedding.length !== queryEmbedding.length) {
-    throw new Error('Vectors have different dimensions');
+  if (embedding?.length !== queryEmbedding?.length) {
+    return 0;
   }
 
   let dotProduct = 0;
@@ -35,7 +36,7 @@ function cosineSimilarity(embedding, queryEmbedding) {
 class PileEmbeddings {
   constructor() {
     this.pilePath = null;
-    this.fileName = 'embeddings.json';
+    this.fileName = `embeddings.json`;
     this.apiKey = null;
     this.embeddings = new Map();
   }
@@ -96,13 +97,13 @@ class PileEmbeddings {
 
   async addDocument(entryPath, metadata) {
     try {
-      //  we only index parent documents
+      // we only index parent documents
       // the replies are concatenated to the contents
       if (metadata.isReply) return;
+
       let fullPath = path.join(this.pilePath, entryPath);
       let fileContent = fs.readFileSync(fullPath, 'utf8');
       let { content } = matter(fileContent);
-
       content =
         'Entry on ' + metadata.createdAt + '\n\n' + content + '\n\nReplies:\n';
 
@@ -114,34 +115,75 @@ class PileEmbeddings {
         content += '\n' + replyContent;
       }
 
-      const embedding = await this.generateEmbedding(content);
-      this.embeddings.set(entryPath, embedding);
+      try {
+        const embedding = await this.generateEmbedding(content);
+        this.embeddings.set(entryPath, embedding);
+        console.log('🧮 Embeddings created for thread: ', entryPath);
+      } catch (embeddingError) {
+        console.warn(
+          `Failed to generate embedding for thread: ${entryPath}`,
+          embeddingError
+        );
+        // Skip this document and continue with the next one
+        return;
+      }
+
       this.saveEmbeddings();
-      console.log('🧮 Embeddings created for threads: ', entryPath);
     } catch (error) {
       console.error('Failed to process thread for vector index.', error);
-      return;
     }
   }
 
+  // todo: based on which ai is configured this should either
+  // use ollama or openai
   async generateEmbedding(document) {
-    const url = 'https://api.openai.com/v1/embeddings';
-    const headers = {
-      Authorization: `Bearer ${this.apiKey}`,
-      'Content-Type': 'application/json',
-    };
-    const data = {
-      model: 'text-embedding-3-small',
-      input: document,
-    };
+    const pileAIProvider = await settings.get('pileAIProvider');
+    const embeddingModel = await settings.get('embeddingModel');
+    const isOllama = pileAIProvider === 'ollama';
+
+    if (isOllama) {
+      const url = 'http://127.0.0.1:11434/api/embed';
+      const data = {
+        model: 'mxbai-embed-large',
+        input: document,
+      };
+      try {
+        const response = await axios.post(url, data);
+        const embeddings = response.data.embeddings;
+        return response.data.embeddings[0];
+      } catch (error) {
+        console.error('Error generating embedding with Ollama:', error);
+        return null;
+      }
+    } else {
+      const url = 'https://api.openai.com/v1/embeddings';
+      const headers = {
+        Authorization: `Bearer ${this.apiKey}`,
+        'Content-Type': 'application/json',
+      };
+      const data = {
+        model: 'text-embedding-3-small',
+        input: document,
+      };
+
+      try {
+        const response = await axios.post(url, data, { headers });
+        return response.data.data[0].embedding;
+      } catch (error) {
+        console.error('Error generating embedding with OpenAI:', error);
+        return null;
+      }
+    }
+  }
 
-    try {
-      const response = await axios.post(url, data, { headers });
-      return response.data.data[0].embedding;
-    } catch (error) {
-      console.error('Error generating embedding:', error);
-      return null;
+  async regenerateEmbeddings(index) {
+    console.log('🧮 Regenerating embeddings for index:', index.size);
+    this.embeddings.clear();
+    for (let [entryPath, metadata] of index) {
+      await this.addDocument(entryPath, metadata);
     }
+    this.saveEmbeddings();
+    console.log('✅ Embeddings regeneration complete');
   }
 
   async search(query, topN = 50) {

src/main/utils/pileIndex.js

@@ -175,6 +175,12 @@ class PileIndex {
     }
   }
 
+  regenerateEmbeddings() {
+    pileEmbeddings.regenerateEmbeddings(this.index);
+    this.save();
+    return;
+  }
+
   update(relativeFilePath, data) {
     this.index.set(relativeFilePath, data);
     pileSearchIndex.initialize(this.pilePath, this.index);