Skip to content

Commit 65d3d0e

Browse files
Use SHA384 and make key more explicit (#6237)
* Minimal change in AIJsonUtilities from SHA256 to SHA384 * Update GetCacheKey API to have explicit messages/options params The point of this is to make it easier and more reliable for people overriding the logic to obtain specific values from messages/options to include in cache key without having to search through an opaque list of `object` values.
1 parent 9759023 commit 65d3d0e

File tree

5 files changed

+24
-27
lines changed

5 files changed

+24
-27
lines changed

src/Libraries/Microsoft.Extensions.AI.Abstractions/Utilities/AIJsonUtilities.cs

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -90,7 +90,7 @@ public static string HashDataToString(ReadOnlySpan<object?> values, JsonSerializ
9090

9191
// For cases where the hash may be used as a cache key, we rely on collision resistance for security purposes.
9292
// If a collision occurs, we'd serve a cached LLM response for a potentially unrelated prompt, leading to information
93-
// disclosure. Use of SHA256 is an implementation detail and can be easily swapped in the future if needed, albeit
93+
// disclosure. Use of SHA384 is an implementation detail and can be easily swapped in the future if needed, albeit
9494
// invalidating any existing cache entries.
9595
#if NET
9696
IncrementalHashStream? stream = IncrementalHashStream.ThreadStaticInstance;
@@ -107,7 +107,7 @@ public static string HashDataToString(ReadOnlySpan<object?> values, JsonSerializ
107107
stream = new();
108108
}
109109

110-
Span<byte> hashData = stackalloc byte[SHA256.HashSizeInBytes];
110+
Span<byte> hashData = stackalloc byte[SHA384.HashSizeInBytes];
111111
try
112112
{
113113
foreach (object? value in values)
@@ -133,8 +133,8 @@ public static string HashDataToString(ReadOnlySpan<object?> values, JsonSerializ
133133
JsonSerializer.Serialize(stream, value, jti);
134134
}
135135

136-
using var sha256 = SHA256.Create();
137-
var hashData = sha256.ComputeHash(stream.GetBuffer(), 0, (int)stream.Length);
136+
using var hashAlgorithm = SHA384.Create();
137+
var hashData = hashAlgorithm.ComputeHash(stream.GetBuffer(), 0, (int)stream.Length);
138138

139139
return ConvertToHexString(hashData);
140140

@@ -185,7 +185,7 @@ private sealed class IncrementalHashStream : Stream
185185
public static IncrementalHashStream? ThreadStaticInstance;
186186

187187
/// <summary>The <see cref="IncrementalHash"/> used by this instance.</summary>
188-
private readonly IncrementalHash _hash = IncrementalHash.CreateHash(HashAlgorithmName.SHA256);
188+
private readonly IncrementalHash _hash = IncrementalHash.CreateHash(HashAlgorithmName.SHA384);
189189

190190
/// <summary>Gets the current hash and resets.</summary>
191191
public void GetHashAndReset(Span<byte> bytes) => _hash.GetHashAndReset(bytes);

src/Libraries/Microsoft.Extensions.AI.Evaluation.Reporting/CSharp/ResponseCachingChatClient.cs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,6 @@ protected override async Task WriteCacheStreamingAsync(
125125
}
126126
}
127127

128-
protected override string GetCacheKey(params ReadOnlySpan<object?> values)
129-
=> base.GetCacheKey([.. values, .. _cachingKeys]);
128+
protected override string GetCacheKey(IEnumerable<ChatMessage> messages, ChatOptions? options, params ReadOnlySpan<object?> additionalValues)
129+
=> base.GetCacheKey(messages, options, [.. additionalValues, .. _cachingKeys]);
130130
}

src/Libraries/Microsoft.Extensions.AI/ChatCompletion/CachingChatClient.cs

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ public override async Task<ChatResponse> GetResponseAsync(
5353
// We're only storing the final result, not the in-flight task, so that we can avoid caching failures
5454
// or having problems when one of the callers cancels but others don't. This has the drawback that
5555
// concurrent callers might trigger duplicate requests, but that's acceptable.
56-
var cacheKey = GetCacheKey(_boxedFalse, messages, options);
56+
var cacheKey = GetCacheKey(messages, options, _boxedFalse);
5757

5858
if (await ReadCacheAsync(cacheKey, cancellationToken).ConfigureAwait(false) is not { } result)
5959
{
@@ -76,7 +76,7 @@ public override async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseA
7676
// we make a streaming request, yielding those results, but then convert those into a non-streaming
7777
// result and cache it. When we get a cache hit, we yield the non-streaming result as a streaming one.
7878

79-
var cacheKey = GetCacheKey(_boxedTrue, messages, options);
79+
var cacheKey = GetCacheKey(messages, options, _boxedTrue);
8080
if (await ReadCacheAsync(cacheKey, cancellationToken).ConfigureAwait(false) is { } chatResponse)
8181
{
8282
// Yield all of the cached items.
@@ -101,7 +101,7 @@ public override async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseA
101101
}
102102
else
103103
{
104-
var cacheKey = GetCacheKey(_boxedTrue, messages, options);
104+
var cacheKey = GetCacheKey(messages, options, _boxedTrue);
105105
if (await ReadCacheStreamingAsync(cacheKey, cancellationToken).ConfigureAwait(false) is { } existingChunks)
106106
{
107107
// Yield all of the cached items.
@@ -129,9 +129,11 @@ public override async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseA
129129
}
130130

131131
/// <summary>Computes a cache key for the specified values.</summary>
132-
/// <param name="values">The values to inform the key.</param>
132+
/// <param name="messages">The messages to inform the key.</param>
133+
/// <param name="options">The <see cref="ChatOptions"/> to inform the key.</param>
134+
/// <param name="additionalValues">Any other values to inform the key.</param>
133135
/// <returns>The computed key.</returns>
134-
protected abstract string GetCacheKey(params ReadOnlySpan<object?> values);
136+
protected abstract string GetCacheKey(IEnumerable<ChatMessage> messages, ChatOptions? options, params ReadOnlySpan<object?> additionalValues);
135137

136138
/// <summary>
137139
/// Returns a previously cached <see cref="ChatResponse"/>, if available.

src/Libraries/Microsoft.Extensions.AI/ChatCompletion/DistributedCachingChatClient.cs

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -97,21 +97,24 @@ protected override async Task WriteCacheStreamingAsync(string key, IReadOnlyList
9797
}
9898

9999
/// <summary>Computes a cache key for the specified values.</summary>
100-
/// <param name="values">The values to inform the key.</param>
100+
/// <param name="messages">The messages to inform the key.</param>
101+
/// <param name="options">The <see cref="ChatOptions"/> to inform the key.</param>
102+
/// <param name="additionalValues">Any other values to inform the key.</param>
101103
/// <returns>The computed key.</returns>
102104
/// <remarks>
103105
/// <para>
104-
/// The <paramref name="values"/> are serialized to JSON using <see cref="JsonSerializerOptions"/> in order to compute the key.
106+
/// The <paramref name="messages"/>, <paramref name="options"/>, and <paramref name="additionalValues"/> are serialized to JSON using <see cref="JsonSerializerOptions"/>
107+
/// in order to compute the key.
105108
/// </para>
106109
/// <para>
107110
/// The generated cache key is not guaranteed to be stable across releases of the library.
108111
/// </para>
109112
/// </remarks>
110-
protected override string GetCacheKey(params ReadOnlySpan<object?> values)
113+
protected override string GetCacheKey(IEnumerable<ChatMessage> messages, ChatOptions? options, params ReadOnlySpan<object?> additionalValues)
111114
{
112115
// Bump the cache version to invalidate existing caches if the serialization format changes in a breaking way.
113116
const int CacheVersion = 1;
114117

115-
return AIJsonUtilities.HashDataToString([CacheVersion, .. values], _jsonSerializerOptions);
118+
return AIJsonUtilities.HashDataToString([CacheVersion, messages, options, .. additionalValues], _jsonSerializerOptions);
116119
}
117120
}

test/Libraries/Microsoft.Extensions.AI.Tests/ChatCompletion/DistributedCachingChatClientTest.cs

Lines changed: 3 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -801,18 +801,10 @@ private static async Task AssertResponsesEqualAsync(IReadOnlyList<ChatResponseUp
801801
private sealed class CachingChatClientWithCustomKey(IChatClient innerClient, IDistributedCache storage)
802802
: DistributedCachingChatClient(innerClient, storage)
803803
{
804-
protected override string GetCacheKey(params ReadOnlySpan<object?> values)
804+
protected override string GetCacheKey(IEnumerable<ChatMessage> messages, ChatOptions? options, params ReadOnlySpan<object?> additionalValues)
805805
{
806-
var baseKey = base.GetCacheKey(values);
807-
foreach (var value in values)
808-
{
809-
if (value is ChatOptions options)
810-
{
811-
return baseKey + options.AdditionalProperties?["someKey"]?.ToString();
812-
}
813-
}
814-
815-
return baseKey;
806+
var baseKey = base.GetCacheKey(messages, options, additionalValues);
807+
return baseKey + options?.AdditionalProperties?["someKey"]?.ToString();
816808
}
817809
}
818810

0 commit comments

Comments
 (0)