Skip to content

Commit 4cae8ed

Browse files
Copilotdavidfowl
Copilot
and
davidfowl
committed
Implement adaptive OpenAI health check based on endpoint configuration
Co-authored-by: davidfowl <[email protected]>
1 parent b5f8fda commit 4cae8ed

File tree

3 files changed

+146
-39
lines changed

3 files changed

+146
-39
lines changed

‎src/Aspire.Hosting.OpenAI/OpenAIExtensions.cs‎

Lines changed: 12 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -39,16 +39,21 @@ You can obtain an API key from the [OpenAI API Keys page](https://platform.opena
3939

4040
defaultApiKeyParameter.WithParentRelationship(resource);
4141

42-
// Register the health check
42+
// Register the adaptive health check
4343
var healthCheckKey = $"{name}_check";
4444

45-
builder.AddStatusPageCheck(
46-
healthCheckKey,
47-
statusJsonUrl: "https://status.openai.com/api/v2/status.json",
48-
httpClientName: "OpenAIHealthCheck",
49-
timeout: TimeSpan.FromSeconds(5),
45+
// Ensure IHttpClientFactory is available by registering HTTP client services
46+
builder.Services.AddHttpClient();
47+
48+
builder.Services.AddHealthChecks().Add(new HealthCheckRegistration(
49+
name: healthCheckKey,
50+
factory: sp =>
51+
{
52+
var httpFactory = sp.GetRequiredService<IHttpClientFactory>();
53+
return new OpenAIHealthCheck(httpFactory, resource);
54+
},
5055
failureStatus: HealthStatus.Unhealthy,
51-
tags: ["openai", "healthcheck"]);
56+
tags: ["openai", "healthcheck"]));
5257

5358
return builder.AddResource(resource)
5459
.WithInitialState(new()
@@ -135,16 +140,6 @@ public static IResourceBuilder<OpenAIResource> WithEndpoint(this IResourceBuilde
135140
ArgumentException.ThrowIfNullOrEmpty(endpoint);
136141

137142
builder.Resource.Endpoint = endpoint;
138-
139-
// Remove the StatusPage health check annotation since it's only relevant for the default OpenAI endpoint
140-
var healthCheckKey = $"{builder.Resource.Name}_check";
141-
var healthCheckAnnotation = builder.Resource.Annotations.OfType<HealthCheckAnnotation>()
142-
.FirstOrDefault(a => a.Key == healthCheckKey);
143-
if (healthCheckAnnotation is not null)
144-
{
145-
builder.Resource.Annotations.Remove(healthCheckAnnotation);
146-
}
147-
148143
return builder;
149144
}
150145

‎src/Aspire.Hosting.OpenAI/OpenAIHealthCheck.cs‎

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
// Licensed to the .NET Foundation under one or more agreements.
2+
// The .NET Foundation licenses this file to you under the MIT license.
3+
4+
using System.Text.Json;
5+
using Microsoft.Extensions.Diagnostics.HealthChecks;
6+
7+
namespace Aspire.Hosting.OpenAI;
8+
9+
/// <summary>
10+
/// An adaptive health check for OpenAI resources that changes behavior based on configuration.
11+
/// </summary>
12+
/// <param name="httpClientFactory">The HTTP client factory.</param>
13+
/// <param name="resource">The OpenAI resource.</param>
14+
internal sealed class OpenAIHealthCheck(IHttpClientFactory httpClientFactory, OpenAIResource resource) : IHealthCheck
15+
{
16+
private const string DefaultEndpoint = "https://api.openai.com/v1";
17+
private HealthCheckResult? _result;
18+
19+
/// <summary>
20+
/// Checks the health of the OpenAI resource.
21+
/// </summary>
22+
public async Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context, CancellationToken cancellationToken = default)
23+
{
24+
if (_result is not null)
25+
{
26+
return _result.Value;
27+
}
28+
29+
try
30+
{
31+
// Case 1: Default endpoint - use StatusPageHealthCheck
32+
if (resource.Endpoint == DefaultEndpoint)
33+
{
34+
return await CheckStatusPageAsync(cancellationToken).ConfigureAwait(false);
35+
}
36+
37+
// Case 2: Custom endpoint without model health check - return healthy
38+
// We can't check the endpoint without a model, so we just return healthy
39+
// The model-level health check will do the actual verification if WithHealthCheck is called
40+
_result = HealthCheckResult.Healthy("Custom OpenAI endpoint configured");
41+
return _result.Value;
42+
}
43+
catch (Exception ex)
44+
{
45+
_result = HealthCheckResult.Unhealthy($"Failed to check OpenAI resource: {ex.Message}", ex);
46+
return _result.Value;
47+
}
48+
}
49+
50+
private async Task<HealthCheckResult> CheckStatusPageAsync(CancellationToken cancellationToken)
51+
{
52+
var client = httpClientFactory.CreateClient("OpenAIHealthCheck");
53+
var statusEndpoint = new Uri("https://status.openai.com/api/v2/status.json");
54+
var timeout = TimeSpan.FromSeconds(5);
55+
56+
using var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
57+
cts.CancelAfter(timeout);
58+
59+
using var req = new HttpRequestMessage(HttpMethod.Get, statusEndpoint);
60+
req.Headers.Accept.ParseAdd("application/json");
61+
62+
HttpResponseMessage resp;
63+
try
64+
{
65+
resp = await client.SendAsync(req, HttpCompletionOption.ResponseHeadersRead, cts.Token).ConfigureAwait(false);
66+
}
67+
catch (OperationCanceledException oce) when (!cancellationToken.IsCancellationRequested)
68+
{
69+
_result = HealthCheckResult.Unhealthy($"StatusPage request timed out after {timeout.TotalSeconds:0.#}s.", oce);
70+
return _result.Value;
71+
}
72+
catch (Exception ex)
73+
{
74+
_result = HealthCheckResult.Unhealthy("StatusPage request failed.", ex);
75+
return _result.Value;
76+
}
77+
78+
if (!resp.IsSuccessStatusCode)
79+
{
80+
_result = HealthCheckResult.Unhealthy($"StatusPage returned {(int)resp.StatusCode} {resp.ReasonPhrase}.");
81+
return _result.Value;
82+
}
83+
84+
try
85+
{
86+
using var stream = await resp.Content.ReadAsStreamAsync(cts.Token).ConfigureAwait(false);
87+
using var doc = await JsonDocument.ParseAsync(stream, cancellationToken: cts.Token).ConfigureAwait(false);
88+
89+
if (!doc.RootElement.TryGetProperty("status", out var statusEl))
90+
{
91+
_result = HealthCheckResult.Unhealthy("Missing 'status' object in StatusPage response.");
92+
return _result.Value;
93+
}
94+
95+
var indicator = statusEl.TryGetProperty("indicator", out var indEl) && indEl.ValueKind == JsonValueKind.String
96+
? indEl.GetString() ?? string.Empty
97+
: string.Empty;
98+
99+
var description = statusEl.TryGetProperty("description", out var descEl) && descEl.ValueKind == JsonValueKind.String
100+
? descEl.GetString() ?? string.Empty
101+
: string.Empty;
102+
103+
var data = new Dictionary<string, object>
104+
{
105+
["indicator"] = indicator,
106+
["description"] = description,
107+
["endpoint"] = statusEndpoint.ToString()
108+
};
109+
110+
_result = indicator switch
111+
{
112+
"none" => HealthCheckResult.Healthy(description.Length > 0 ? description : "All systems operational."),
113+
"minor" => HealthCheckResult.Degraded(description.Length > 0 ? description : "Minor service issues."),
114+
"major" => HealthCheckResult.Unhealthy(description.Length > 0 ? description : "Major service outage."),
115+
"critical" => HealthCheckResult.Unhealthy(description.Length > 0 ? description : "Critical service outage."),
116+
_ => HealthCheckResult.Unhealthy($"Unknown indicator '{indicator}'", data: data)
117+
};
118+
119+
return _result.Value;
120+
}
121+
catch (JsonException jex)
122+
{
123+
_result = HealthCheckResult.Unhealthy("Failed to parse StatusPage JSON.", jex);
124+
return _result.Value;
125+
}
126+
}
127+
}

‎tests/Aspire.Hosting.OpenAI.Tests/OpenAIExtensionTests.cs‎

Lines changed: 7 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -368,48 +368,33 @@ public void AddOpenAIModelWorksWithDifferentModels(string modelName)
368368
}
369369

370370
[Fact]
371-
public void WithEndpointRemovesStatusPageHealthCheckAnnotation()
371+
public void HealthCheckIsAlwaysPresent()
372372
{
373373
using var builder = TestDistributedApplicationBuilder.Create();
374374
builder.Configuration["Parameters:openai-openai-apikey"] = "test-api-key";
375375

376376
var parent = builder.AddOpenAI("openai");
377377

378-
// Verify that the StatusPage health check annotation is added by AddOpenAI
378+
// Verify that the health check annotation is always added by AddOpenAI
379379
var healthCheckAnnotations = parent.Resource.Annotations.OfType<HealthCheckAnnotation>().ToList();
380380
Assert.Single(healthCheckAnnotations);
381381
Assert.Equal("openai_check", healthCheckAnnotations[0].Key);
382-
383-
// Call WithEndpoint with a custom endpoint
384-
parent.WithEndpoint("http://localhost:12434/engines/v1");
385-
386-
// Verify that the StatusPage health check annotation is removed
387-
healthCheckAnnotations = parent.Resource.Annotations.OfType<HealthCheckAnnotation>().ToList();
388-
Assert.Empty(healthCheckAnnotations);
389382
}
390383

391384
[Fact]
392-
public void WithEndpointDoesNotAffectModelHealthCheck()
385+
public void HealthCheckRemainsAfterWithEndpoint()
393386
{
394387
using var builder = TestDistributedApplicationBuilder.Create();
395388
builder.Configuration["Parameters:openai-openai-apikey"] = "test-api-key";
396389

397390
var parent = builder.AddOpenAI("openai");
398-
var model = parent.AddModel("chat", "gpt-4o-mini");
399391

400392
// Call WithEndpoint with a custom endpoint
401393
parent.WithEndpoint("http://localhost:12434/engines/v1");
402394

403-
// Add a model health check
404-
model.WithHealthCheck();
405-
406-
// Verify that the model health check annotation is still present
407-
var modelHealthCheckAnnotations = model.Resource.Annotations.OfType<HealthCheckAnnotation>().ToList();
408-
Assert.Single(modelHealthCheckAnnotations);
409-
Assert.Equal("chat_check", modelHealthCheckAnnotations[0].Key);
410-
411-
// Verify that the parent's StatusPage health check annotation is removed
412-
var parentHealthCheckAnnotations = parent.Resource.Annotations.OfType<HealthCheckAnnotation>().ToList();
413-
Assert.Empty(parentHealthCheckAnnotations);
395+
// Verify that the health check annotation is still present (adaptive health check)
396+
var healthCheckAnnotations = parent.Resource.Annotations.OfType<HealthCheckAnnotation>().ToList();
397+
Assert.Single(healthCheckAnnotations);
398+
Assert.Equal("openai_check", healthCheckAnnotations[0].Key);
414399
}
415400
}

0 commit comments

Comments
 (0)