-
Notifications
You must be signed in to change notification settings - Fork 8
feat: Update embedding logic to bulk #1037
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| @@ -1,51 +1,122 @@ | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| 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> | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| /// 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<float>.</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 a single batch call 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. Embed all chunks in one batch API call (Azure OpenAI supports up to 2048 inputs). | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| /// 3. Batch-upsert all chunks into staging. | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| /// 4. Atomically swap staging → live via three SQL RENAMEs in a single transaction. | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| /// PostgreSQL ALTER TABLE acquires AccessExclusiveLock automatically; no explicit | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| /// LOCK TABLE is needed. The transaction ensures no reader sees an intermediate state. | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| /// 5. Drop the old live backup table. | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
Comment on lines
+37
to
+40
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| /// 4. Atomically swap staging → live via three SQL RENAMEs in a single transaction. | |
| /// PostgreSQL ALTER TABLE acquires AccessExclusiveLock automatically; no explicit | |
| /// LOCK TABLE is needed. The transaction ensures no reader sees an intermediate state. | |
| /// 5. Drop the old live backup table. | |
| /// 4. 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. | |
| /// 5. Drop the old live backup table with DROP TABLE. |
Copilot
AI
Apr 26, 2026
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The embedding generator call batches all chunks into a single request, but Azure OpenAI embeddings have an input-count limit (noted in the comment as 2048). If the book produces more than that many chunks, this will fail at runtime. Consider chunking texts into batches (<= provider limit) and merging the returned vectors back into chunkList (or throw a clear exception when the limit is exceeded).
| // ── Step 2: Batch-embed all chunks in a single API call ─────────────────────── | |
| // IEmbeddingGenerator.GenerateAsync natively accepts IEnumerable<string>. | |
| // The single-string overload used previously is a convenience extension method | |
| // that wraps one item and calls this same method. | |
| var chunkList = bookContents.ToList(); | |
| var texts = chunkList.Select(c => c.ChunkText).ToList(); | |
| GeneratedEmbeddings<Embedding<float>> embeddings = | |
| await embeddingGenerator.GenerateAsync(texts, cancellationToken: cancellationToken); | |
| if (embeddings.Count != chunkList.Count) | |
| throw new InvalidOperationException( | |
| $"Embedding count mismatch: expected {chunkList.Count}, got {embeddings.Count}."); | |
| for (int i = 0; i < chunkList.Count; i++) | |
| { | |
| chunkList[i].TextEmbedding = embeddings[i].Vector; | |
| // ── Step 2: Batch-embed all chunks in provider-safe API calls ───────────────── | |
| // Azure OpenAI embeddings impose an input-count limit per request. | |
| // Process the texts in batches and merge the returned vectors back into the | |
| // original chunk list to preserve ordering. | |
| var chunkList = bookContents.ToList(); | |
| var texts = chunkList.Select(c => c.ChunkText).ToList(); | |
| const int maxEmbeddingBatchSize = 2048; | |
| for (int batchStart = 0; batchStart < texts.Count; batchStart += maxEmbeddingBatchSize) | |
| { | |
| int batchSize = Math.Min(maxEmbeddingBatchSize, texts.Count - batchStart); | |
| List<string> batchTexts = texts.GetRange(batchStart, batchSize); | |
| GeneratedEmbeddings<Embedding<float>> embeddings = | |
| await embeddingGenerator.GenerateAsync(batchTexts, cancellationToken: cancellationToken); | |
| if (embeddings.Count != batchSize) | |
| throw new InvalidOperationException( | |
| $"Embedding count mismatch for batch starting at index {batchStart}: expected {batchSize}, got {embeddings.Count}."); | |
| for (int i = 0; i < batchSize; i++) | |
| { | |
| chunkList[batchStart + i].TextEmbedding = embeddings[i].Vector; | |
| } |
Copilot
AI
Apr 26, 2026
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
GenerateBookContentEmbeddingsAndUploadToVectorStore now contains significant new behavior (batch embedding + staging/swap workflow) but there are no unit/integration tests covering it in the Chat test suite. Adding tests that at least verify batching semantics (e.g., embedding generator called with N inputs and staging collection UpsertAsync invoked once with embeddings assigned) would help prevent regressions; the swap step can be isolated behind an interface to make it testable without a live database.
Copilot
AI
Apr 26, 2026
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
These SQL statements interpolate collectionName/derived names directly into identifier-quoted SQL. If collectionName can be influenced outside trusted code, this becomes identifier-injection (quotes can be escaped/broken). Consider restricting collectionName to a safe identifier regex (e.g., letters/digits/underscore) before composing SQL, and use Npgsql's identifier-quoting helper to build the final identifiers consistently.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Local variable name
candidates_listuses snake_case, which is inconsistent with the surrounding camelCase naming in this method (e.g.,vectorSearchOptions,searchVector). Rename tocandidatesList(or similar) for consistency and readability.