feat: Update embedding logic to bulk (#1037)

## Description

Describe your changes here.

Fixes #Issue_Number (if available)

### Ensure that your pull request has followed all the steps below:
- [ ] Code compilation
- [ ] Created tests which fail without the change (if possible)
- [ ] All tests passing
- [ ] Extended the README / documentation, if necessary

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: BenjaminMichaelis

EssentialCSharp.Chat.Shared/Models/BookContentChunk.cs

@@ -32,11 +32,24 @@ public sealed class BookContentChunk
     public string ChunkText { get; set; } = string.Empty;
 
     /// <summary>
-    /// Chapter number extracted from filename (e.g., "Chapter01.md" -> 1)
+    /// Chapter number extracted from filename (e.g., "Chapter01.md" -> 1).
+    /// Null for files that do not follow the ChapterNN naming pattern.
     /// </summary>
     [VectorStoreData]
     public int? ChapterNumber { get; set; }
 
+    /// <summary>
+    /// Zero-based ordinal of this chunk within its source file.
+    /// Together with FileName, forms the basis for the deterministic Id.
+    /// </summary>
+    /// <remarks>
+    /// This column was added as part of the bulk-embedding refactor. Existing vector-store
+    /// collections created before this change must be rebuilt (via the staging-swap upload
+    /// command) before reads against the live table will succeed with this schema.
+    /// </remarks>
+    [VectorStoreData]
+    public int ChunkIndex { get; set; }
+
     /// <summary>
     /// SHA256 hash of the chunk content for change detection
     /// </summary>

EssentialCSharp.Chat.Shared/Services/AISearchService.cs

@@ -1,4 +1,4 @@
-using System.Diagnostics;
+using System.Diagnostics;
 using EssentialCSharp.Chat.Common.Models;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.VectorData;
@@ -35,11 +35,29 @@ public async Task<IReadOnlyList<VectorSearchResult<BookContentChunk>>> ExecuteVe
         {
             try
             {
-                var results = new List<VectorSearchResult<BookContentChunk>>();
-                await foreach (var result in collection.SearchAsync(searchVector, options: vectorSearchOptions, top: top, cancellationToken: cancellationToken))
+                // Fetch more candidates than needed so we can deduplicate by heading.
+                // Multiple chunks from the same section share the same Heading; without dedup
+                // all top-N results could come from one long section, reducing context diversity.
+                int candidates = top * 3;
+
+                var candidatesList = new List<VectorSearchResult<BookContentChunk>>();
+                await foreach (var result in collection.SearchAsync(searchVector, options: vectorSearchOptions, top: candidates, cancellationToken: cancellationToken))
                 {
-                    results.Add(result);
+                    candidatesList.Add(result);
                 }
+
+                // Keep only the highest-scoring chunk per unique heading, then take the globally
+                // top-N by score. GroupBy on a materialized list preserves insertion (score desc)
+                // order, but we make the ordering explicit via OrderByDescending so the result
+                // is correct regardless of provider sort guarantees.
+                // MaxBy on a non-empty IGrouping never returns null; ! asserts this invariant.
+                var results = candidatesList
+                    .GroupBy(r => r.Record.Heading)
+                    .Select(g => g.MaxBy(r => r.Score)!)
+                    .OrderByDescending(r => r.Score)
+                    .Take(top)
+                    .ToList();
+
                 return results;
             }
             catch (PostgresException ex) when (ex.SqlState == "28000" && attempt == 0)

EssentialCSharp.Chat.Shared/Services/ChunkingResultExtensions.cs

@@ -1,27 +1,35 @@
 using System.Security.Cryptography;
 using System.Text;
-using System.Linq;
 using EssentialCSharp.Chat.Common.Models;
 
 namespace EssentialCSharp.Chat.Common.Services;
 
 public static partial class ChunkingResultExtensions
 {
+    /// <summary>
+    /// Converts a <see cref="FileChunkingResult"/> into a list of <see cref="BookContentChunk"/> records
+    /// ready for embedding and vector store upload.
+    /// </summary>
+    /// <remarks>
+    /// <see cref="BookContentChunk.ChapterNumber"/> is set to null for files that do not match
+    /// the <c>ChapterNN</c> naming pattern (e.g. appendix or non-chapter markdown files).
+    /// </remarks>
     public static List<BookContentChunk> ToBookContentChunks(this FileChunkingResult result)
     {
         int? chapterNumber = ExtractChapterNumber(result.FileName);
 
         var chunks = result.Chunks
-            .Select(chunkText =>
+            .Select((markdownChunk, index) =>
             {
-                var contentHash = ComputeSha256Hash(chunkText);
+                var contentHash = ComputeSha256Hash(markdownChunk.ChunkText);
                 return new BookContentChunk
                 {
-                    Id = Guid.NewGuid().ToString(),
+                    Id = $"{result.FileName}_{index}",
                     FileName = result.FileName,
-                    Heading = ExtractHeading(chunkText),
-                    ChunkText = chunkText,
+                    Heading = markdownChunk.Heading,
+                    ChunkText = markdownChunk.ChunkText,
                     ChapterNumber = chapterNumber,
+                    ChunkIndex = index,
                     ContentHash = contentHash
                 };
             })
@@ -30,25 +38,13 @@ public static List<BookContentChunk> ToBookContentChunks(this FileChunkingResult
         return chunks;
     }
 
-    private static string ExtractHeading(string chunkText)
+    private static int? ExtractChapterNumber(string fileName)
     {
-        // get characters until the first " - " or newline
-        var firstLine = chunkText.Split(["\r\n", "\r", "\n"], StringSplitOptions.None)[0];
-        var headingParts = firstLine.Split([" - "], StringSplitOptions.None);
-        return headingParts.Length > 0 ? headingParts[0].Trim() : string.Empty;
-    }
-
-    private static int ExtractChapterNumber(string fileName)
-    {
-        // Example: "Chapter01.md" -> 1
-        // Regex: Chapter(?<ChapterNumber>[0-9]{2})
+        // Example: "Chapter01.md" -> 1; non-chapter files return null.
         var match = ChapterNumberRegex().Match(fileName);
         if (match.Success && int.TryParse(match.Groups["ChapterNumber"].Value, out int chapterNumber))
-
-        {
             return chapterNumber;
-        }
-        throw new InvalidOperationException($"File name '{fileName}' does not contain a valid chapter number in the expected format.");
+        return null;
     }
 
     private static string ComputeSha256Hash(string text)

EssentialCSharp.Chat.Shared/Services/EmbeddingService.cs

@@ -1,51 +1,170 @@
+using System.Text.RegularExpressions;
 using EssentialCSharp.Chat.Common.Models;
 using Microsoft.Extensions.AI;
 using Microsoft.Extensions.VectorData;
+using Npgsql;
 
 namespace EssentialCSharp.Chat.Common.Services;
 
 /// <summary>
-/// Service for generating embeddings for markdown chunks using Azure OpenAI
+/// Service for generating embeddings for markdown chunks using Azure OpenAI and uploading
+/// them to a PostgreSQL vector store via a staging-then-swap pattern to avoid downtime.
 /// </summary>
-public class EmbeddingService(VectorStore vectorStore, IEmbeddingGenerator<string, Embedding<float>> embeddingGenerator)
+public class EmbeddingService(
+    VectorStore vectorStore,
+    IEmbeddingGenerator<string, Embedding<float>> embeddingGenerator,
+    NpgsqlDataSource? dataSource = null)
 {
     public static string CollectionName { get; } = "markdown_chunks";
 
+    /// <summary>
+    /// Maximum number of inputs per Azure OpenAI embedding batch call.
+    /// </summary>
+    private const int EmbeddingBatchSize = 2048;
+
+    // Only allow simple identifiers: letters, digits, and underscores, starting with a letter or underscore.
+    private static readonly Regex _safeIdentifierRegex = new(@"^[a-zA-Z_][a-zA-Z0-9_]*$", RegexOptions.Compiled);
+
     /// <summary>
     /// Generate an embedding for the given text.
     /// </summary>
-    /// <param name="text">The text to generate an embedding for.</param>
-    /// <param name="cancellationToken">The cancellation token.</param>
-    /// <returns>A search vector as ReadOnlyMemory&lt;float&gt;.</returns>
     public async Task<ReadOnlyMemory<float>> GenerateEmbeddingAsync(string text, CancellationToken cancellationToken = default)
     {
         var embedding = await embeddingGenerator.GenerateAsync(text, cancellationToken: cancellationToken);
         return embedding.Vector;
     }
 
     /// <summary>
-    /// Generate an embedding for each text paragraph and upload it to the specified collection.
+    /// Generate embeddings for all chunks in batches and upload them to the vector store
+    /// using a staging-then-atomic-swap pattern so the live collection stays queryable
+    /// throughout the rebuild.
+    ///
+    /// Steps:
+    ///   1. Create a staging collection ({collectionName}_staging).
+    ///   2. For each batch of <see cref="EmbeddingBatchSize"/> chunks: embed the batch
+    ///      and immediately upsert it into staging, keeping peak memory bounded.
+    ///   3. Atomically swap tables in a single transaction using two SQL RENAME operations
+    ///      (live → old, staging → live). PostgreSQL ALTER TABLE acquires
+    ///      AccessExclusiveLock automatically; no explicit LOCK TABLE is needed. The
+    ///      transaction ensures no reader sees an intermediate state.
+    ///   4. Drop the old live backup table with DROP TABLE.
+    ///
+    /// If an error occurs before the swap, only the staging table is affected — the live
+    /// collection is untouched.
     /// </summary>
-    /// <param name="collectionName">The name of the collection to upload the text paragraphs to.</param>
-    /// <returns>An async task.</returns>
-    public async Task GenerateBookContentEmbeddingsAndUploadToVectorStore(IEnumerable<BookContentChunk> bookContents, CancellationToken cancellationToken, string? collectionName = null)
+    public async Task GenerateBookContentEmbeddingsAndUploadToVectorStore(
+        IEnumerable<BookContentChunk> bookContents,
+        CancellationToken cancellationToken,
+        string? collectionName = null)
     {
         collectionName ??= CollectionName;
 
-        var collection = vectorStore.GetCollection<string, BookContentChunk>(collectionName);
-        await collection.EnsureCollectionDeletedAsync(cancellationToken);
-        await collection.EnsureCollectionExistsAsync(cancellationToken);
+        if (dataSource is null)
+            throw new InvalidOperationException(
+                $"{nameof(NpgsqlDataSource)} must be provided to upload embeddings. Ensure it is registered in DI.");
 
-        int uploadedCount = 0;
+        if (!_safeIdentifierRegex.IsMatch(collectionName))
+            throw new ArgumentException(
+                $"Collection name '{collectionName}' contains unsafe characters. Use only letters, digits, and underscores.",
+                nameof(collectionName));
 
-        foreach (var chunk in bookContents)
+        string stagingName = $"{collectionName}_staging";
+        string oldName = $"{collectionName}_old";
+
+        // ── Step 1: Prepare staging collection ────────────────────────────────────────
+        var staging = vectorStore.GetCollection<string, BookContentChunk>(stagingName);
+        await staging.EnsureCollectionDeletedAsync(cancellationToken);
+        await staging.EnsureCollectionExistsAsync(cancellationToken);
+
+        // ── Step 2 & 3: Batch-embed and immediately upsert each batch ─────────────────
+        // Azure OpenAI supports at most EmbeddingBatchSize inputs per GenerateAsync call.
+        // bookContents is streamed in fixed-size batches without full upfront materialization,
+        // keeping peak memory bounded to one batch of chunk objects and their embeddings at a time.
+        // The staging-swap (Step 3) is safe because it only runs after all batches have
+        // been successfully upserted.
+        var buffer = new List<BookContentChunk>(EmbeddingBatchSize);
+        int totalCount = 0;
+
+        async Task EmbedAndUpsertBatchAsync()
         {
-            cancellationToken.ThrowIfCancellationRequested();
-            chunk.TextEmbedding = await GenerateEmbeddingAsync(chunk.ChunkText, cancellationToken);
-            await collection.UpsertAsync(chunk, cancellationToken);
-            Console.WriteLine($"Uploaded chunk '{chunk.Id}' to collection '{collectionName}' for file '{chunk.FileName}' with heading '{chunk.Heading}'.");
-            uploadedCount++;
+            var batchEmbeddings = await embeddingGenerator.GenerateAsync(
+                buffer.Select(c => c.ChunkText), cancellationToken: cancellationToken);
+
+            if (batchEmbeddings.Count != buffer.Count)
+                throw new InvalidOperationException(
+                    $"Embedding count mismatch: expected {buffer.Count}, got {batchEmbeddings.Count}.");
+
+            for (int i = 0; i < buffer.Count; i++)
+                buffer[i].TextEmbedding = batchEmbeddings[i].Vector;
+
+            await staging.UpsertAsync(buffer, cancellationToken);
+            totalCount += buffer.Count;
+            buffer.Clear();
+        }
+
+        try
+        {
+            foreach (var chunk in bookContents)
+            {
+                buffer.Add(chunk);
+                if (buffer.Count == EmbeddingBatchSize)
+                    await EmbedAndUpsertBatchAsync();
+            }
+
+            if (buffer.Count > 0)
+                await EmbedAndUpsertBatchAsync();
+
+            Console.WriteLine($"Uploaded {totalCount} chunks to staging collection '{stagingName}'.");
         }
-        Console.WriteLine($"Successfully generated embeddings and uploaded {uploadedCount} chunks to collection '{collectionName}'.");
+        catch
+        {
+            // Best-effort cleanup: drop the partially-populated staging table so the
+            // next run starts clean. Do not let this secondary failure mask the original.
+            try
+            {
+                await staging.EnsureCollectionDeletedAsync(cancellationToken);
+            }
+            catch (Exception cleanupEx) when (cleanupEx is not OperationCanceledException)
+            {
+                Console.Error.WriteLine($"Warning: failed to clean up staging collection '{stagingName}' after upsert failure: {cleanupEx.Message}");
+            }
+            throw;
+        }
+
+        // ── Step 3: Atomic swap — staging → live ──────────────────────────────────────
+        // Two ALTER TABLE RENAME operations in one transaction (live → old, staging → live).
+        // Each RENAME auto-acquires AccessExclusiveLock on its table; the transaction
+        // guarantees both renames are visible atomically to other sessions.
+        await using var conn = await dataSource.OpenConnectionAsync(cancellationToken);
+        await using var tx = await conn.BeginTransactionAsync(cancellationToken);
+
+        await using (var cmd = conn.CreateCommand())
+        {
+            cmd.Transaction = tx;
+
+            // Drop any leftover backup from a previous run
+            cmd.CommandText = $"DROP TABLE IF EXISTS \"{oldName}\"";
+            await cmd.ExecuteNonQueryAsync(cancellationToken);
+
+            // Rename live → old. IF EXISTS is a no-op on first run when no live table exists.
+            cmd.CommandText = $"ALTER TABLE IF EXISTS \"{collectionName}\" RENAME TO \"{oldName}\"";
+            await cmd.ExecuteNonQueryAsync(cancellationToken);
+
+            // Rename staging → live
+            cmd.CommandText = $"ALTER TABLE \"{stagingName}\" RENAME TO \"{collectionName}\"";
+            await cmd.ExecuteNonQueryAsync(cancellationToken);
+        }
+
+        await tx.CommitAsync(cancellationToken);
+        Console.WriteLine($"Swapped '{stagingName}' → '{collectionName}' atomically.");
+
+        // ── Step 4: Drop the old backup ───────────────────────────────────────────────
+        await using (var cmd = conn.CreateCommand())
+        {
+            cmd.CommandText = $"DROP TABLE IF EXISTS \"{oldName}\"";
+            await cmd.ExecuteNonQueryAsync(cancellationToken);
+        }
+
+        Console.WriteLine($"Successfully generated embeddings and uploaded {totalCount} chunks to collection '{collectionName}'.");
     }
 }

EssentialCSharp.Chat.Shared/Services/FileChunkingResult.cs

@@ -1,5 +1,12 @@
 namespace EssentialCSharp.Chat.Common.Services;
 
+/// <summary>
+/// A single chunk from a markdown file, paired with the section heading it belongs to.
+/// </summary>
+/// <param name="Heading">Full breadcrumb heading for the section (e.g. "Chapter: 1: Intro: Summary").</param>
+/// <param name="ChunkText">The raw chunk text, including the "Heading - " prefix prepended by TextChunker.</param>
+public record MarkdownChunk(string Heading, string ChunkText);
+
 /// <summary>
 /// Data structure to hold chunking results for a single file
 /// </summary>
@@ -9,6 +16,6 @@ public class FileChunkingResult
     public string FilePath { get; set; } = string.Empty;
     public int OriginalCharCount { get; set; }
     public int ChunkCount { get; set; }
-    public List<string> Chunks { get; set; } = [];
+    public List<MarkdownChunk> Chunks { get; set; } = [];
     public int TotalChunkCharacters { get; set; }
 }