Skip to content

Commit 2f5f5db

Browse files
fix: stream bookContents in batches without ToList(); document ChunkIndex schema migration
- EmbeddingService: replace bookContents.ToList() + Chunk() with a streaming buffer pattern (List<BookContentChunk>(EmbeddingBatchSize)). Each batch is filled from the IEnumerable, embedded, and upserted before moving on, so only one batch of chunks + vectors lives in memory at a time. Track total count with a running int. - BookContentChunk: add <remarks> to ChunkIndex documenting that it is a new column requiring a collection rebuild on existing deployments.
1 parent c044f16 commit 2f5f5db

2 files changed

Lines changed: 37 additions & 19 deletions

File tree

EssentialCSharp.Chat.Shared/Models/BookContentChunk.cs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,11 @@ public sealed class BookContentChunk
4242
/// Zero-based ordinal of this chunk within its source file.
4343
/// Together with FileName, forms the basis for the deterministic Id.
4444
/// </summary>
45+
/// <remarks>
46+
/// This column was added as part of the bulk-embedding refactor. Existing vector-store
47+
/// collections created before this change must be rebuilt (via the staging-swap upload
48+
/// command) before reads against the live table will succeed with this schema.
49+
/// </remarks>
4550
[VectorStoreData]
4651
public int ChunkIndex { get; set; }
4752

EssentialCSharp.Chat.Shared/Services/EmbeddingService.cs

Lines changed: 32 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -78,30 +78,43 @@ public async Task GenerateBookContentEmbeddingsAndUploadToVectorStore(
7878

7979
// ── Step 2 & 3: Batch-embed and immediately upsert each batch ─────────────────
8080
// Azure OpenAI supports at most EmbeddingBatchSize inputs per GenerateAsync call.
81-
// All chunks are materialized upfront (bookContents.ToList()) so Chunk() can
82-
// slice them by index. Embedding vectors are generated and upserted one batch at
83-
// a time, keeping peak vector memory bounded to a single batch rather than the
84-
// full dataset. The staging-swap (Step 4) is safe because it only runs after all
85-
// batches have been successfully upserted.
86-
var chunkList = bookContents.ToList();
87-
try
81+
// bookContents is streamed in fixed-size batches without full upfront materialization,
82+
// keeping peak memory bounded to one batch of chunk objects and their embeddings at a time.
83+
// The staging-swap (Step 3) is safe because it only runs after all batches have
84+
// been successfully upserted.
85+
var buffer = new List<BookContentChunk>(EmbeddingBatchSize);
86+
int totalCount = 0;
87+
88+
async Task EmbedAndUpsertBatchAsync()
8889
{
89-
foreach (var batch in chunkList.Chunk(EmbeddingBatchSize))
90-
{
91-
var batchEmbeddings = await embeddingGenerator.GenerateAsync(
92-
batch.Select(c => c.ChunkText), cancellationToken: cancellationToken);
90+
var batchEmbeddings = await embeddingGenerator.GenerateAsync(
91+
buffer.Select(c => c.ChunkText), cancellationToken: cancellationToken);
9392

94-
if (batchEmbeddings.Count != batch.Length)
95-
throw new InvalidOperationException(
96-
$"Embedding count mismatch: expected {batch.Length} for batch, got {batchEmbeddings.Count}.");
93+
if (batchEmbeddings.Count != buffer.Count)
94+
throw new InvalidOperationException(
95+
$"Embedding count mismatch: expected {buffer.Count}, got {batchEmbeddings.Count}.");
9796

98-
for (int i = 0; i < batch.Length; i++)
99-
batch[i].TextEmbedding = batchEmbeddings[i].Vector;
97+
for (int i = 0; i < buffer.Count; i++)
98+
buffer[i].TextEmbedding = batchEmbeddings[i].Vector;
99+
100+
await staging.UpsertAsync(buffer, cancellationToken);
101+
totalCount += buffer.Count;
102+
buffer.Clear();
103+
}
100104

101-
await staging.UpsertAsync(batch, cancellationToken);
105+
try
106+
{
107+
foreach (var chunk in bookContents)
108+
{
109+
buffer.Add(chunk);
110+
if (buffer.Count == EmbeddingBatchSize)
111+
await EmbedAndUpsertBatchAsync();
102112
}
103113

104-
Console.WriteLine($"Uploaded {chunkList.Count} chunks to staging collection '{stagingName}'.");
114+
if (buffer.Count > 0)
115+
await EmbedAndUpsertBatchAsync();
116+
117+
Console.WriteLine($"Uploaded {totalCount} chunks to staging collection '{stagingName}'.");
105118
}
106119
catch
107120
{
@@ -152,6 +165,6 @@ public async Task GenerateBookContentEmbeddingsAndUploadToVectorStore(
152165
await cmd.ExecuteNonQueryAsync(cancellationToken);
153166
}
154167

155-
Console.WriteLine($"Successfully generated embeddings and uploaded {chunkList.Count} chunks to collection '{collectionName}'.");
168+
Console.WriteLine($"Successfully generated embeddings and uploaded {totalCount} chunks to collection '{collectionName}'.");
156169
}
157170
}

0 commit comments

Comments
 (0)