synchronize method

Future<({int embeddedChunks, int removedChunks, int writtenChunks})> synchronize(KnowledgeSnapshot input)

Embeds missing/changed inputs before one atomic replacement. Failed loads or inference leave the prior store/index usable. Other bundles are retained.

Implementation

Future<({int embeddedChunks, int removedChunks, int writtenChunks})>
synchronize(KnowledgeSnapshot input) async {
  if (_synchronizing || _searches > 0) {
    throw StateError(
      'Await active searches/synchronization before updating.',
    );
  }
  if (_snapshot != null && _snapshot!.bundleId != input.bundleId) {
    throw ArgumentError('Create a separate index for a different bundle ID.');
  }
  _synchronizing = true;
  try {
    final snapshot = countTokens == null
        ? input
        : await input.fitInputs(
            countTokens: countTokens!,
            maxTokens: maxTokens!,
            includeContext: includeContext,
          );
    final previous = {
      for (final chunk in await store.getAllChunks()) chunk.id: chunk,
    };
    final ids = snapshot.chunks.map((chunk) => chunk.id).toSet();
    final removed = previous.values
        .where(
          (chunk) =>
              (chunk.metadata['okf'] as Map?)?['bundleId'] ==
                  snapshot.bundleId &&
              !ids.contains(chunk.id),
        )
        .map((chunk) => chunk.id)
        .toSet();
    final model = embeddingModelName;
    final cached = embedder == null
        ? <String>{}
        : (await store.getEmbeddingsForChunks(
            ids,
            source: embedder!.sourceName,
            modelName: model!,
          )).map((item) => item.chunkId).toSet();
    final pending = <({Chunk chunk, String text})>[];
    final writes = <Chunk>[];
    for (final chunk in snapshot.chunks) {
      final inputHashes = <String, Object?>{
        ...?previous[chunk.id]?.metadata['okfEmbeddingInputs']
            as Map<String, Object?>?,
      };
      if (embedder != null) {
        final text = snapshot.textFor(chunk, includeContext: includeContext);
        final hash = sha256.convert(utf8.encode(text)).toString();
        if (!cached.contains(chunk.id) || inputHashes[model] != hash) {
          pending.add((chunk: chunk, text: text));
        }
        inputHashes[model!] = hash;
      }
      final updated = chunk.copyWith(
        metadata: {
          ...chunk.metadata,
          if (inputHashes.isNotEmpty) 'okfEmbeddingInputs': inputHashes,
        },
      );
      if (previous[chunk.id] != updated) writes.add(updated);
    }
    final vectors = <Embedding>[];
    for (var offset = 0; offset < pending.length; offset += 32) {
      final batch = pending.skip(offset).take(32).toList();
      final encoded = await embedder!.generateEmbeddings(
        batch.map((item) => item.text).toList(),
      );
      if (encoded.length != batch.length ||
          encoded.any((vector) => vector.length != embedder!.dimension)) {
        throw StateError('Embedder returned an invalid batch shape.');
      }
      for (var i = 0; i < batch.length; i++) {
        vectors.add(
          Embedding(
            chunkId: batch[i].chunk.id,
            source: embedder!.sourceName,
            modelName: model!,
            vector: encoded[i],
          ),
        );
      }
    }
    // Construct the new read snapshot before committing the store update.
    final lexical = BM25LexicalIndex.fromChunks(
      snapshot.chunks.map(
        (chunk) => chunk.copyWith(
          id: chunk.id,
          content: snapshot.textFor(chunk, includeContext: includeContext),
        ),
      ),
    );
    if (writes.isNotEmpty || vectors.isNotEmpty || removed.isNotEmpty) {
      await store.replaceChunks(
        chunks: writes,
        embeddings: vectors,
        removeChunkIds: removed,
      );
    }
    _snapshot = snapshot;
    _lexical = lexical;
    return (
      embeddedChunks: pending.length,
      removedChunks: removed.length,
      writtenChunks: writes.length,
    );
  } finally {
    _synchronizing = false;
  }
}