Add embedding health check execution and update JSON schema with endpoint/health sub-objects

Co-authored-by: JerryNixon <1749983+JerryNixon@users.noreply.github.com>

Author: Copilot

schemas/dab.draft.schema.json

@@ -686,6 +686,59 @@
               "default": 30000,
               "minimum": 1,
               "maximum": 300000
+            },
+            "endpoint": {
+              "type": "object",
+              "description": "REST endpoint configuration for the embedding service.",
+              "additionalProperties": false,
+              "properties": {
+                "enabled": {
+                  "type": "boolean",
+                  "description": "Whether the /embed REST endpoint is enabled. Defaults to false.",
+                  "default": false
+                },
+                "path": {
+                  "type": "string",
+                  "description": "The endpoint path. Defaults to '/embed'.",
+                  "default": "/embed"
+                },
+                "roles": {
+                  "type": "array",
+                  "description": "The roles allowed to access the embedding endpoint. In development mode, defaults to ['anonymous'].",
+                  "items": {
+                    "type": "string"
+                  }
+                }
+              }
+            },
+            "health": {
+              "type": "object",
+              "description": "Health check configuration for the embedding service.",
+              "additionalProperties": false,
+              "properties": {
+                "enabled": {
+                  "type": "boolean",
+                  "description": "Whether health checks are enabled for embeddings. Defaults to true.",
+                  "default": true
+                },
+                "threshold-ms": {
+                  "type": "integer",
+                  "description": "The maximum response time in milliseconds to be considered healthy.",
+                  "default": 5000,
+                  "minimum": 1,
+                  "maximum": 300000
+                },
+                "test-text": {
+                  "type": "string",
+                  "description": "The text to use for health check validation.",
+                  "default": "health check"
+                },
+                "expected-dimensions": {
+                  "type": "integer",
+                  "description": "The expected number of dimensions in the embedding result. If specified, dimension validation is performed.",
+                  "minimum": 1
+                }
+              }
             }
           },
           "required": ["provider", "base-url", "api-key"],

src/Service/HealthCheck/HealthCheckHelper.cs

@@ -10,7 +10,9 @@
 using System.Threading.Tasks;
 using Azure.DataApiBuilder.Config.HealthCheck;
 using Azure.DataApiBuilder.Config.ObjectModel;
+using Azure.DataApiBuilder.Config.ObjectModel.Embeddings;
 using Azure.DataApiBuilder.Core.Authorization;
+using Azure.DataApiBuilder.Core.Services.Embeddings;
 using Azure.DataApiBuilder.Product;
 using Microsoft.AspNetCore.Http;
 using Microsoft.Extensions.Logging;
@@ -27,20 +29,24 @@ public class HealthCheckHelper
         // Dependencies
         private ILogger<HealthCheckHelper> _logger;
         private HttpUtilities _httpUtility;
+        private IEmbeddingService? _embeddingService;
         private string _incomingRoleHeader = string.Empty;
         private string _incomingRoleToken = string.Empty;
 
         private const string TIME_EXCEEDED_ERROR_MESSAGE = "The threshold for executing the request has exceeded.";
+        private const string DIMENSIONS_MISMATCH_ERROR_MESSAGE = "The embedding dimensions do not match the expected dimensions.";
 
         /// <summary>
         /// Constructor to inject the logger and HttpUtility class.
         /// </summary>
         /// <param name="logger">Logger to track the log statements.</param>
         /// <param name="httpUtility">HttpUtility to call methods from the internal class.</param>
-        public HealthCheckHelper(ILogger<HealthCheckHelper> logger, HttpUtilities httpUtility)
+        /// <param name="embeddingService">Optional embedding service for embedding health checks.</param>
+        public HealthCheckHelper(ILogger<HealthCheckHelper> logger, HttpUtilities httpUtility, IEmbeddingService? embeddingService = null)
         {
             _logger = logger;
             _httpUtility = httpUtility;
+            _embeddingService = embeddingService;
         }
 
         /// <summary>
@@ -159,6 +165,7 @@ private async Task UpdateHealthCheckDetailsAsync(ComprehensiveHealthCheckReport
             comprehensiveHealthCheckReport.Checks = new List<HealthCheckResultEntry>();
             await UpdateDataSourceHealthCheckResultsAsync(comprehensiveHealthCheckReport, runtimeConfig);
             await UpdateEntityHealthCheckResultsAsync(comprehensiveHealthCheckReport, runtimeConfig);
+            await UpdateEmbeddingsHealthCheckResultsAsync(comprehensiveHealthCheckReport, runtimeConfig);
         }
 
         // Updates the DataSource Health Check Results in the response.
@@ -351,5 +358,108 @@ private async Task PopulateEntityHealthAsync(ComprehensiveHealthCheckReport comp
 
             return (HealthCheckConstants.ERROR_RESPONSE_TIME_MS, errorMessage);
         }
+
+        /// <summary>
+        /// Updates the Embeddings Health Check Results in the response.
+        /// Executes a test embedding and validates response time and optionally dimensions.
+        /// </summary>
+        private async Task UpdateEmbeddingsHealthCheckResultsAsync(ComprehensiveHealthCheckReport comprehensiveHealthCheckReport, RuntimeConfig runtimeConfig)
+        {
+            EmbeddingsOptions? embeddingsOptions = runtimeConfig?.Runtime?.Embeddings;
+            EmbeddingsHealthCheckConfig? healthConfig = embeddingsOptions?.Health;
+
+            // Only run health check if embeddings is enabled, health check is enabled, and embedding service is available
+            if (embeddingsOptions is null || !embeddingsOptions.Enabled || healthConfig is null || !healthConfig.Enabled || _embeddingService is null)
+            {
+                return;
+            }
+
+            if (comprehensiveHealthCheckReport.Checks is null)
+            {
+                comprehensiveHealthCheckReport.Checks = new List<HealthCheckResultEntry>();
+            }
+
+            string testText = healthConfig.TestText;
+            int thresholdMs = healthConfig.ThresholdMs;
+            int? expectedDimensions = healthConfig.ExpectedDimensions;
+
+            try
+            {
+                Stopwatch stopwatch = new();
+                stopwatch.Start();
+                EmbeddingResult result = await _embeddingService.TryEmbedAsync(testText);
+                stopwatch.Stop();
+
+                int responseTimeMs = (int)stopwatch.ElapsedMilliseconds;
+                bool isResponseTimeWithinThreshold = responseTimeMs <= thresholdMs;
+                bool isDimensionsValid = true;
+                string? errorMessage = null;
+
+                if (!result.Success)
+                {
+                    errorMessage = result.ErrorMessage ?? "Embedding request failed.";
+                    comprehensiveHealthCheckReport.Checks.Add(new HealthCheckResultEntry
+                    {
+                        Name = "embeddings",
+                        ResponseTimeData = new ResponseTimeData
+                        {
+                            ResponseTimeMs = HealthCheckConstants.ERROR_RESPONSE_TIME_MS,
+                            ThresholdMs = thresholdMs
+                        },
+                        Exception = errorMessage,
+                        Tags = new List<string> { HealthCheckConstants.EMBEDDING },
+                        Status = HealthStatus.Unhealthy
+                    });
+                    return;
+                }
+
+                // Validate dimensions if expected dimensions is specified
+                if (expectedDimensions.HasValue && result.Embedding is not null)
+                {
+                    isDimensionsValid = result.Embedding.Length == expectedDimensions.Value;
+                    if (!isDimensionsValid)
+                    {
+                        errorMessage = $"{DIMENSIONS_MISMATCH_ERROR_MESSAGE} Expected: {expectedDimensions.Value}, Actual: {result.Embedding.Length}";
+                    }
+                }
+
+                // Check response time threshold
+                if (!isResponseTimeWithinThreshold)
+                {
+                    errorMessage = errorMessage is null ? TIME_EXCEEDED_ERROR_MESSAGE : $"{errorMessage} {TIME_EXCEEDED_ERROR_MESSAGE}";
+                }
+
+                bool isHealthy = isResponseTimeWithinThreshold && isDimensionsValid;
+
+                comprehensiveHealthCheckReport.Checks.Add(new HealthCheckResultEntry
+                {
+                    Name = "embeddings",
+                    ResponseTimeData = new ResponseTimeData
+                    {
+                        ResponseTimeMs = responseTimeMs,
+                        ThresholdMs = thresholdMs
+                    },
+                    Exception = errorMessage,
+                    Tags = new List<string> { HealthCheckConstants.EMBEDDING },
+                    Status = isHealthy ? HealthStatus.Healthy : HealthStatus.Unhealthy
+                });
+            }
+            catch (Exception ex)
+            {
+                _logger.LogError(ex, "Error executing embeddings health check.");
+                comprehensiveHealthCheckReport.Checks.Add(new HealthCheckResultEntry
+                {
+                    Name = "embeddings",
+                    ResponseTimeData = new ResponseTimeData
+                    {
+                        ResponseTimeMs = HealthCheckConstants.ERROR_RESPONSE_TIME_MS,
+                        ThresholdMs = thresholdMs
+                    },
+                    Exception = ex.Message,
+                    Tags = new List<string> { HealthCheckConstants.EMBEDDING },
+                    Status = HealthStatus.Unhealthy
+                });
+            }
+        }
     }
 }

src/Service/Startup.cs

@@ -262,7 +262,13 @@ public void ConfigureServices(IServiceCollection services)
             services.AddSingleton<GQLFilterParser>();
             services.AddSingleton<RequestValidator>();
             services.AddSingleton<RestService>();
-            services.AddSingleton<HealthCheckHelper>();
+            services.AddSingleton<HealthCheckHelper>(sp =>
+            {
+                ILogger<HealthCheckHelper> logger = sp.GetRequiredService<ILogger<HealthCheckHelper>>();
+                HttpUtilities httpUtility = sp.GetRequiredService<HttpUtilities>();
+                IEmbeddingService? embeddingService = sp.GetService<IEmbeddingService>();
+                return new HealthCheckHelper(logger, httpUtility, embeddingService);
+            });
             services.AddSingleton<HttpUtilities>();
             services.AddSingleton<BasicHealthReportResponseWriter>();
             services.AddSingleton<ComprehensiveHealthReportResponseWriter>();