ocean-brain 0.8.2 → 0.11.0

package/server/dist/features/search/search-manager.js

@@ -0,0 +1,431 @@
+import models from "../../models.js";
+import { subscribeServerEvents } from "../../modules/server-events.js";
+import { paths } from "../../paths.js";
+import {
+  createEmbeddingApiKeyFingerprint,
+  FileEmbeddingApiKeyStore
+} from "./embedding-api-key-store.js";
+import {
+  createOpenAiCompatibleEmbeddingClient,
+  listOpenAiCompatibleEmbeddingModels
+} from "./embedding-client.js";
+import { subscribeSemanticSearchNoteChanges } from "./note-change.js";
+import {
+  buildNoteEmbeddingChunks,
+  NOTE_EMBEDDING_TEXT_SCHEMA_VERSION
+} from "./note-chunking.js";
+import { SemanticSearchConfigStore } from "./search-config.js";
+import {
+  buildSemanticSearchIndex,
+  searchSemanticIndex,
+  updateSemanticSearchNotes
+} from "./semantic-indexer.js";
+import {
+  SqliteSemanticVectorIndex
+} from "./sqlite-vector-index.js";
+const errorMessage = (error) => error instanceof Error ? error.message : "Unknown semantic search error";
+const NOTE_SYNC_QUIET_PERIOD_MS = 1e4;
+const NOTE_SYNC_MAX_WAIT_MS = 6e4;
+const NOTE_SYNC_POLL_INTERVAL_MS = 5e3;
+const NOTE_SYNC_RECONCILIATION_INTERVAL_MS = 5 * 6e4;
+const NOTE_SYNC_BATCH_SIZE = 20;
+const NOTE_SYNC_RETRY_BASE_MS = 15e3;
+const NOTE_SYNC_RETRY_MAX_MS = 5 * 6e4;
+const profileMatchesConfig = (profile, config) => {
+  return Boolean(
+    profile && profile.baseUrl === config.baseUrl && profile.model === config.model && profile.textSchemaVersion === NOTE_EMBEDDING_TEXT_SCHEMA_VERSION
+  );
+};
+class SemanticSearchManager {
+  constructor(dependencies) {
+    this.dependencies = dependencies;
+  }
+  dependencies;
+  activeReindex = null;
+  activeNoteSync = null;
+  activeReconciliation = null;
+  noteSyncPollTimer = null;
+  noteSyncReconciliationTimer = null;
+  progress = null;
+  lastError = null;
+  lastQueueError = null;
+  resolveApiKey(input) {
+    return input?.provided ? input.apiKey : this.dependencies.apiKeyStore.get();
+  }
+  createClient(config, apiKeyInput) {
+    const factory = this.dependencies.createEmbeddingClient ?? createOpenAiCompatibleEmbeddingClient;
+    return factory({
+      ...config,
+      apiKey: this.resolveApiKey(apiKeyInput)
+    });
+  }
+  getAuthFingerprint(apiKeyInput) {
+    return createEmbeddingApiKeyFingerprint(this.resolveApiKey(apiKeyInput));
+  }
+  now() {
+    return this.dependencies.now?.() ?? Date.now();
+  }
+  async getStatus() {
+    const [config, indexStatus, noteSyncStatus] = await Promise.all([
+      this.dependencies.configStore.get(),
+      this.dependencies.vectorIndex.getStatus(),
+      this.dependencies.vectorIndex.getNoteSyncQueueStatus()
+    ]);
+    const connectionValidated = await this.dependencies.configStore.isConnectionValidated(
+      config,
+      this.getAuthFingerprint()
+    );
+    const profileMatches = profileMatchesConfig(indexStatus.profile, config);
+    const indexReady = indexStatus.ready && profileMatches;
+    const available = config.enabled && connectionValidated && indexReady;
+    const needsReindex = config.enabled && connectionValidated && !indexReady;
+    let phase;
+    if (this.activeReindex) {
+      phase = "indexing";
+    } else if (this.lastError) {
+      phase = "error";
+    } else if (!config.enabled) {
+      phase = "disabled";
+    } else if (!connectionValidated) {
+      phase = "needs-connection";
+    } else if (available) {
+      phase = "ready";
+    } else {
+      phase = "needs-index";
+    }
+    return {
+      config,
+      connectionValidated,
+      apiKeyConfigured: Boolean(this.dependencies.apiKeyStore.get()),
+      phase,
+      available,
+      needsReindex,
+      noteCount: indexStatus.noteCount,
+      chunkCount: indexStatus.chunkCount,
+      indexedAt: indexStatus.indexedAt,
+      dimensions: indexStatus.profile?.dimensions ?? null,
+      pendingNoteCount: noteSyncStatus.pendingNoteCount,
+      lastSyncedAt: noteSyncStatus.lastSyncedAt ?? indexStatus.indexedAt,
+      syncError: noteSyncStatus.error ?? this.lastQueueError,
+      progress: this.progress,
+      error: this.lastError
+    };
+  }
+  async saveConfig(config, apiKeyInput) {
+    if (this.activeReindex) {
+      throw new Error("Semantic search settings cannot change while indexing is running.");
+    }
+    const nextApiKey = this.resolveApiKey(apiKeyInput);
+    const apiKeyChanged = Boolean(apiKeyInput?.provided && nextApiKey !== this.dependencies.apiKeyStore.get());
+    this.lastError = null;
+    this.progress = null;
+    if (apiKeyChanged) {
+      this.dependencies.apiKeyStore.set(nextApiKey);
+    }
+    await this.dependencies.configStore.set(config);
+    return this.getStatus();
+  }
+  async validateConnection(config, apiKeyInput) {
+    const client = this.createClient({ ...config, enabled: true }, apiKeyInput);
+    const [embedding] = await client.embedDocuments(["Ocean Brain embedding connection test"]);
+    await this.dependencies.configStore.markConnectionValidated(config, this.getAuthFingerprint(apiKeyInput));
+    return embedding;
+  }
+  async testConnection(configInput, apiKeyInput) {
+    const config = configInput ?? await this.dependencies.configStore.get();
+    const embedding = await this.validateConnection(config, apiKeyInput);
+    return {
+      ok: true,
+      dimensions: embedding.length,
+      model: config.model
+    };
+  }
+  async listModels(baseUrl, apiKeyInput) {
+    return listOpenAiCompatibleEmbeddingModels(baseUrl, {
+      apiKey: this.resolveApiKey(apiKeyInput)
+    });
+  }
+  async startReindex() {
+    if (this.activeReindex) {
+      return { started: false, status: await this.getStatus() };
+    }
+    if (this.activeNoteSync) {
+      await this.activeNoteSync;
+    }
+    const config = await this.dependencies.configStore.get();
+    if (!config.enabled) {
+      throw new Error("Enable semantic search before building its index.");
+    }
+    const currentStatus = await this.getStatus();
+    if (!currentStatus.needsReindex) {
+      return { started: false, status: currentStatus };
+    }
+    this.lastError = null;
+    this.progress = { processedChunks: 0, totalChunks: 0 };
+    const run = async () => {
+      const notes = await this.dependencies.listNotes();
+      await buildSemanticSearchIndex({
+        notes,
+        embeddingClient: this.createClient(config),
+        vectorIndex: this.dependencies.vectorIndex,
+        model: config.model,
+        baseUrl: config.baseUrl,
+        queryInstruction: config.queryInstruction,
+        onProgress: (progress) => {
+          this.progress = progress;
+        }
+      });
+      await this.dependencies.vectorIndex.recordNoteSyncSuccess(this.now());
+    };
+    this.activeReindex = run().then(() => {
+      this.progress = null;
+      this.lastError = null;
+    }).catch((error) => {
+      this.progress = null;
+      this.lastError = errorMessage(error);
+    }).finally(() => {
+      this.activeReindex = null;
+      void this.processNextNoteSyncBatch().catch((error) => {
+        this.lastQueueError = errorMessage(error);
+      });
+    });
+    return { started: true, status: await this.getStatus() };
+  }
+  async waitForActiveReindex() {
+    await this.activeReindex;
+    return this.getStatus();
+  }
+  getRetryAt(entries) {
+    const attemptCount = Math.max(0, ...entries.map((entry) => entry.attemptCount));
+    const retryDelay = Math.min(NOTE_SYNC_RETRY_BASE_MS * 2 ** attemptCount, NOTE_SYNC_RETRY_MAX_MS);
+    return this.now() + retryDelay;
+  }
+  processNextNoteSyncBatch(force = false) {
+    if (this.activeNoteSync) {
+      return this.activeNoteSync;
+    }
+    const run = async () => {
+      if (this.activeReindex) {
+        if (!force) {
+          return { failed: false, processed: false };
+        }
+        await this.activeReindex;
+      }
+      const status = await this.getStatus();
+      if (!status.config.enabled || !status.available) {
+        return { failed: false, processed: false };
+      }
+      const entries = await this.dependencies.vectorIndex.listPendingNoteSyncs({
+        now: this.now(),
+        quietPeriodMs: NOTE_SYNC_QUIET_PERIOD_MS,
+        maxWaitMs: NOTE_SYNC_MAX_WAIT_MS,
+        limit: NOTE_SYNC_BATCH_SIZE,
+        force
+      });
+      if (entries.length === 0) {
+        return { failed: false, processed: false };
+      }
+      try {
+        const noteIds = entries.map((entry) => entry.noteId);
+        const notes = await this.dependencies.findNotes(noteIds);
+        const existingNoteIds = new Set(notes.map((note) => note.id));
+        const removedNoteIds = noteIds.filter((noteId) => !existingNoteIds.has(noteId));
+        await updateSemanticSearchNotes({
+          notes,
+          removedNoteIds,
+          embeddingClient: this.createClient(status.config),
+          vectorIndex: this.dependencies.vectorIndex
+        });
+        await this.dependencies.vectorIndex.completeNoteSyncs(entries, this.now());
+        this.lastQueueError = null;
+        return { failed: false, processed: true };
+      } catch (error) {
+        const message = errorMessage(error);
+        await this.dependencies.vectorIndex.failNoteSyncs(entries, message, this.getRetryAt(entries));
+        this.lastQueueError = message;
+        return { failed: true, processed: true };
+      }
+    };
+    const activeNoteSync = run().finally(() => {
+      if (this.activeNoteSync === activeNoteSync) {
+        this.activeNoteSync = null;
+      }
+    });
+    this.activeNoteSync = activeNoteSync;
+    return activeNoteSync;
+  }
+  async scheduleNoteSync(noteId) {
+    if (!Number.isInteger(noteId) || noteId <= 0) {
+      return;
+    }
+    try {
+      await this.dependencies.vectorIndex.enqueueNoteSync(noteId, this.now());
+      this.lastQueueError = null;
+    } catch (error) {
+      this.lastQueueError = errorMessage(error);
+      throw error;
+    }
+  }
+  async waitForPendingNoteSync() {
+    while ((await this.dependencies.vectorIndex.getNoteSyncQueueStatus()).pendingNoteCount > 0) {
+      const result = await this.processNextNoteSyncBatch(true);
+      if (!result.processed || result.failed) {
+        break;
+      }
+    }
+    return this.getStatus();
+  }
+  runNoteSyncReconciliation() {
+    if (this.activeReconciliation) {
+      return this.activeReconciliation;
+    }
+    const run = async () => {
+      if (this.activeReindex) {
+        return;
+      }
+      const status = await this.getStatus();
+      if (!status.available) {
+        return;
+      }
+      const [notes, indexedSourceHashes] = await Promise.all([
+        this.dependencies.listNotes(),
+        this.dependencies.vectorIndex.getAllNoteSourceHashes()
+      ]);
+      const currentSourceHashes = /* @__PURE__ */ new Map();
+      for (const note of notes) {
+        const sourceHash = buildNoteEmbeddingChunks(note)[0]?.sourceHash;
+        if (sourceHash) {
+          currentSourceHashes.set(note.id, sourceHash);
+        }
+      }
+      const changedNoteIds = /* @__PURE__ */ new Set();
+      for (const [noteId, sourceHash] of currentSourceHashes) {
+        if (indexedSourceHashes.get(noteId) !== sourceHash) {
+          changedNoteIds.add(noteId);
+        }
+      }
+      for (const noteId of indexedSourceHashes.keys()) {
+        if (!currentSourceHashes.has(noteId)) {
+          changedNoteIds.add(noteId);
+        }
+      }
+      const reconciledAt = this.now();
+      for (const noteId of changedNoteIds) {
+        await this.dependencies.vectorIndex.enqueueNoteSync(noteId, reconciledAt);
+      }
+      await this.dependencies.vectorIndex.recordNoteSyncReconciliation(reconciledAt);
+      this.lastQueueError = null;
+    };
+    const activeReconciliation = run().catch((error) => {
+      this.lastQueueError = errorMessage(error);
+    }).finally(() => {
+      if (this.activeReconciliation === activeReconciliation) {
+        this.activeReconciliation = null;
+      }
+    });
+    this.activeReconciliation = activeReconciliation;
+    return activeReconciliation;
+  }
+  startBackgroundSync() {
+    if (this.noteSyncPollTimer || this.noteSyncReconciliationTimer) {
+      return;
+    }
+    const poll = () => {
+      void this.processNextNoteSyncBatch().catch((error) => {
+        this.lastQueueError = errorMessage(error);
+      });
+    };
+    const reconcile = () => {
+      void this.runNoteSyncReconciliation();
+    };
+    poll();
+    reconcile();
+    this.noteSyncPollTimer = setInterval(poll, NOTE_SYNC_POLL_INTERVAL_MS);
+    this.noteSyncPollTimer.unref?.();
+    this.noteSyncReconciliationTimer = setInterval(reconcile, NOTE_SYNC_RECONCILIATION_INTERVAL_MS);
+    this.noteSyncReconciliationTimer.unref?.();
+  }
+  stopBackgroundSync() {
+    if (this.noteSyncPollTimer) {
+      clearInterval(this.noteSyncPollTimer);
+      this.noteSyncPollTimer = null;
+    }
+    if (this.noteSyncReconciliationTimer) {
+      clearInterval(this.noteSyncReconciliationTimer);
+      this.noteSyncReconciliationTimer = null;
+    }
+  }
+  async trySearch(query, limit) {
+    const status = await this.getStatus();
+    if (!status.available) {
+      return {
+        available: false,
+        used: false,
+        matches: [],
+        error: status.error
+      };
+    }
+    try {
+      const matches = await searchSemanticIndex(
+        query,
+        limit,
+        this.createClient(status.config),
+        this.dependencies.vectorIndex
+      );
+      return {
+        available: true,
+        used: true,
+        matches,
+        error: null
+      };
+    } catch (error) {
+      return {
+        available: true,
+        used: false,
+        matches: [],
+        error: errorMessage(error)
+      };
+    }
+  }
+}
+let defaultSemanticSearchManager = null;
+const getDefaultSemanticSearchManager = () => {
+  if (!defaultSemanticSearchManager) {
+    defaultSemanticSearchManager = new SemanticSearchManager({
+      configStore: new SemanticSearchConfigStore({
+        findUnique: (args) => models.cache.findUnique(args),
+        upsert: (args) => models.cache.upsert(args)
+      }),
+      vectorIndex: new SqliteSemanticVectorIndex(paths.searchIndex),
+      listNotes: () => models.note.findMany({
+        orderBy: { id: "asc" },
+        select: {
+          id: true,
+          title: true,
+          content: true
+        }
+      }),
+      findNotes: (noteIds) => models.note.findMany({
+        where: { id: { in: noteIds } },
+        select: {
+          id: true,
+          title: true,
+          content: true
+        }
+      }),
+      apiKeyStore: new FileEmbeddingApiKeyStore(paths.embeddingApiKey)
+    });
+    subscribeServerEvents((event) => {
+      void defaultSemanticSearchManager?.scheduleNoteSync(Number(event.noteId)).catch(() => void 0);
+    });
+    subscribeSemanticSearchNoteChanges((noteId) => {
+      void defaultSemanticSearchManager?.scheduleNoteSync(noteId).catch(() => void 0);
+    });
+    defaultSemanticSearchManager.startBackgroundSync();
+  }
+  return defaultSemanticSearchManager;
+};
+export {
+  SemanticSearchManager,
+  getDefaultSemanticSearchManager
+};

package/server/dist/features/search/semantic-indexer.js

@@ -0,0 +1,138 @@
+import {
+  buildNoteEmbeddingChunks,
+  NOTE_EMBEDDING_TEXT_SCHEMA_VERSION
+} from "./note-chunking.js";
+const DEFAULT_EMBEDDING_BATCH_SIZE = 16;
+const validateBatchSize = (batchSize) => {
+  if (!Number.isInteger(batchSize) || batchSize <= 0 || batchSize > 256) {
+    throw new Error("Embedding batch size must be an integer between 1 and 256.");
+  }
+};
+const embedChunks = async (chunks, embeddingClient, batchSize, onProgress) => {
+  const indexedChunks = [];
+  let dimensions = null;
+  for (let start = 0; start < chunks.length; start += batchSize) {
+    const batch = chunks.slice(start, start + batchSize);
+    const embeddings = await embeddingClient.embedDocuments(batch.map((chunk) => chunk.text));
+    if (embeddings.length !== batch.length) {
+      throw new Error(`Embedding provider returned ${embeddings.length} vectors for ${batch.length} chunks.`);
+    }
+    embeddings.forEach((embedding, index) => {
+      dimensions ??= embedding.length;
+      if (embedding.length !== dimensions) {
+        throw new Error("Embedding provider changed vector dimensions while building the index.");
+      }
+      indexedChunks.push({
+        ...batch[index],
+        embedding
+      });
+    });
+    onProgress?.(indexedChunks.length);
+  }
+  return { indexedChunks, dimensions };
+};
+const buildSemanticSearchIndex = async ({
+  notes,
+  embeddingClient,
+  vectorIndex,
+  model,
+  baseUrl,
+  queryInstruction,
+  batchSize = DEFAULT_EMBEDDING_BATCH_SIZE,
+  onProgress
+}) => {
+  validateBatchSize(batchSize);
+  const chunks = notes.flatMap((note) => buildNoteEmbeddingChunks(note));
+  const totalChunks = chunks.length;
+  onProgress?.({ processedChunks: 0, totalChunks });
+  if (totalChunks === 0) {
+    await vectorIndex.clear();
+    return vectorIndex.getStatus();
+  }
+  const { indexedChunks, dimensions } = await embedChunks(
+    chunks,
+    embeddingClient,
+    batchSize,
+    (processedChunks) => onProgress?.({ processedChunks, totalChunks })
+  );
+  if (!dimensions) {
+    throw new Error("Embedding provider returned no vector dimensions.");
+  }
+  return vectorIndex.replaceAll(
+    {
+      model,
+      baseUrl,
+      dimensions,
+      queryInstruction,
+      textSchemaVersion: NOTE_EMBEDDING_TEXT_SCHEMA_VERSION
+    },
+    indexedChunks
+  );
+};
+const updateSemanticSearchNote = async ({
+  note,
+  embeddingClient,
+  vectorIndex,
+  batchSize = DEFAULT_EMBEDDING_BATCH_SIZE
+}) => {
+  validateBatchSize(batchSize);
+  const chunks = buildNoteEmbeddingChunks(note);
+  const sourceHash = chunks[0]?.sourceHash ?? null;
+  const indexedSourceHash = await vectorIndex.getNoteSourceHash(note.id);
+  if (sourceHash === indexedSourceHash) {
+    return vectorIndex.getStatus();
+  }
+  if (chunks.length === 0) {
+    return vectorIndex.removeNote(note.id);
+  }
+  const { indexedChunks } = await embedChunks(chunks, embeddingClient, batchSize);
+  return vectorIndex.replaceNote(note.id, indexedChunks);
+};
+const updateSemanticSearchNotes = async ({
+  notes,
+  removedNoteIds,
+  embeddingClient,
+  vectorIndex,
+  batchSize = DEFAULT_EMBEDDING_BATCH_SIZE
+}) => {
+  validateBatchSize(batchSize);
+  const changedNotes = [];
+  for (const note of notes) {
+    const chunks = buildNoteEmbeddingChunks(note);
+    const sourceHash = chunks[0]?.sourceHash ?? null;
+    const indexedSourceHash = await vectorIndex.getNoteSourceHash(note.id);
+    if (sourceHash !== indexedSourceHash) {
+      changedNotes.push({ note, chunks });
+    }
+  }
+  const chunksToEmbed = changedNotes.flatMap(({ chunks }) => chunks);
+  const { indexedChunks } = await embedChunks(chunksToEmbed, embeddingClient, batchSize);
+  for (const { note, chunks } of changedNotes) {
+    if (chunks.length === 0) {
+      await vectorIndex.removeNote(note.id);
+      continue;
+    }
+    await vectorIndex.replaceNote(
+      note.id,
+      indexedChunks.filter((chunk) => chunk.noteId === note.id)
+    );
+  }
+  for (const noteId of new Set(removedNoteIds)) {
+    await vectorIndex.removeNote(noteId);
+  }
+  return vectorIndex.getStatus();
+};
+const searchSemanticIndex = async (query, limit, embeddingClient, vectorIndex) => {
+  const status = await vectorIndex.getStatus();
+  if (!status.ready) {
+    return [];
+  }
+  const queryEmbedding = await embeddingClient.embedQuery(query);
+  return vectorIndex.search(queryEmbedding, limit);
+};
+export {
+  buildSemanticSearchIndex,
+  searchSemanticIndex,
+  updateSemanticSearchNote,
+  updateSemanticSearchNotes
+};