.Net: Improve Milvus filtering (#14498)

### Motivation and Context

<!-- Thank you for your contribution to the semantic-kernel repo!
Please help reviewers and future users, providing the following
information:
  1. Why is this change required?
  2. What problem does it solve?
  3. What scenario does it contribute to?
  4. If it fixes an open issue, please link to the issue here.
-->

### Description

Improve Milvus filtering to allow more id types to be supported

### Contribution Checklist

<!-- Before submitting this PR, please make sure: -->

- [ ] The code builds clean without any errors or warnings
- [ ] The PR follows the [SK Contribution
Guidelines](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md)
and the [pre-submission formatting
script](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md#development-scripts)
raises no violations
- [ ] All unit tests pass, and I have added new tests where possible
- [ ] I didn't break anyone 😄
This commit is contained in:
westey
2026-09-30 18:11:20 +00:00
committed by GitHub
parent 082700e24b
commit 58ad1a3780
9 changed files with 534 additions and 67 deletions
+1
View File
@@ -26,5 +26,6 @@
<Folder Name="/test/" />
<Folder Name="/test/VectorData/">
<Project Path="test/VectorData/Chroma.UnitTests/Chroma.UnitTests.csproj" />
<Project Path="test/VectorData/Milvus.UnitTests/Milvus.UnitTests.csproj" />
</Folder>
</Solution>
@@ -28,6 +28,7 @@
<ItemGroup>
<ProjectReference Include="..\..\SemanticKernel.Core\SemanticKernel.Core.csproj" />
<InternalsVisibleTo Include="SemanticKernel.Connectors.Milvus.UnitTests" />
</ItemGroup>
</Project>
@@ -247,6 +247,7 @@ public class MilvusMemoryStore : IMemoryStore, IDisposable
/// <inheritdoc />
public async Task<string> UpsertAsync(string collectionName, MemoryRecord record, CancellationToken cancellationToken = default)
{
await this.EnsureSupportedServerAsync(cancellationToken).ConfigureAwait(false);
MilvusCollection collection = this.Client.GetCollection(collectionName);
var metadata = record.Metadata;
@@ -276,8 +277,7 @@ public class MilvusMemoryStore : IMemoryStore, IDisposable
IEnumerable<MemoryRecord> records,
[EnumeratorCancellation] CancellationToken cancellationToken = default)
{
StringBuilder idString = new();
await this.EnsureSupportedServerAsync(cancellationToken).ConfigureAwait(false);
List<bool> isReferenceData = [];
List<string> externalSourceNameData = [];
List<string> idData = [];
@@ -292,13 +292,6 @@ public class MilvusMemoryStore : IMemoryStore, IDisposable
{
var metadata = record.Metadata;
if (idString.Length > 0)
{
idString.Append(',');
}
idString.Append('"').Append(metadata.Id).Append('"');
isReferenceData.Add(metadata.IsReference);
externalSourceNameData.Add(metadata.ExternalSourceName);
idData.Add(record.Metadata.Id);
@@ -356,22 +349,11 @@ public class MilvusMemoryStore : IMemoryStore, IDisposable
bool withEmbeddings = false,
[EnumeratorCancellation] CancellationToken cancellationToken = default)
{
StringBuilder idString = new();
foreach (string key in keys)
{
if (idString.Length > 0)
{
idString.Append(',');
}
idString.Append('"').Append(key).Append('"');
}
await this.EnsureSupportedServerAsync(cancellationToken).ConfigureAwait(false);
IReadOnlyList<FieldData> fields = await this.Client
.GetCollection(collectionName)
.QueryAsync(
$"{IdFieldName} in [{idString}]",
BuildIdFilter(keys),
withEmbeddings ? this._queryParametersWithEmbedding : this._queryParametersWithoutEmbedding,
cancellationToken: cancellationToken)
.ConfigureAwait(false);
@@ -386,36 +368,16 @@ public class MilvusMemoryStore : IMemoryStore, IDisposable
/// <inheritdoc />
public Task RemoveAsync(string collectionName, string key, CancellationToken cancellationToken = default)
=> this.Client.GetCollection(collectionName)
.DeleteAsync($@"{IdFieldName} in [""{key}""]", cancellationToken: cancellationToken);
=> this.RemoveBatchAsync(collectionName, [key], cancellationToken);
/// <inheritdoc />
public Task RemoveBatchAsync(string collectionName, IEnumerable<string> keys, CancellationToken cancellationToken = default)
public async Task RemoveBatchAsync(string collectionName, IEnumerable<string> keys, CancellationToken cancellationToken = default)
{
StringBuilder idString = new();
idString.Append(IdFieldName).Append(" in [");
bool first = true;
foreach (string id in keys)
{
if (first)
{
first = false;
}
else
{
idString.Append(',');
}
idString.Append('"').Append(id).Append('"');
}
idString.Append(']');
return this.Client
await this.EnsureSupportedServerAsync(cancellationToken).ConfigureAwait(false);
await this.Client
.GetCollection(collectionName)
.DeleteAsync(idString.ToString(), cancellationToken: cancellationToken);
.DeleteAsync(BuildIdFilter(keys), cancellationToken: cancellationToken)
.ConfigureAwait(false);
}
/// <inheritdoc />
@@ -443,6 +405,7 @@ public class MilvusMemoryStore : IMemoryStore, IDisposable
bool withEmbeddings = false,
[EnumeratorCancellation] CancellationToken cancellationToken = default)
{
await this.EnsureSupportedServerAsync(cancellationToken).ConfigureAwait(false);
MilvusCollection collection = this.Client.GetCollection(collectionName);
SearchResults results = await collection
@@ -459,23 +422,8 @@ public class MilvusMemoryStore : IMemoryStore, IDisposable
Dictionary<string, ReadOnlyMemory<float>>? embeddingMap = null;
if (withEmbeddings)
{
StringBuilder filter = new();
filter.Append(IdFieldName).Append(" in [");
for (int rowNum = 0; rowNum < ids.Count; rowNum++)
{
if (rowNum > 0)
{
filter.Append(',');
}
filter.Append('"').Append(ids[rowNum]).Append('"');
}
filter.Append(']');
IReadOnlyList<FieldData> fieldData = await collection.QueryAsync(
filter.ToString(),
BuildIdFilter(ids),
new() { OutputFields = { EmbeddingFieldName } },
cancellationToken: cancellationToken)
.ConfigureAwait(false);
@@ -502,6 +450,83 @@ public class MilvusMemoryStore : IMemoryStore, IDisposable
}
}
private async Task EnsureSupportedServerAsync(CancellationToken cancellationToken)
{
string version = await this.Client.GetVersionAsync(cancellationToken).ConfigureAwait(false);
ValidateServerVersion(version);
}
internal static void ValidateServerVersion(string version)
{
string numericVersion = version.StartsWith("v", StringComparison.Ordinal) ? version.Substring(1) : version;
// The official 2.3.0 release image reports "v2.3.0-dev".
if (numericVersion.EndsWith("-dev", StringComparison.Ordinal))
{
numericVersion = numericVersion.Substring(0, numericVersion.Length - 4);
}
// Older servers accept escaped strings but retain the escapes as part of the ID.
if (!Version.TryParse(numericVersion, out Version? parsedVersion) ||
parsedVersion < new Version(2, 3, 0))
{
throw new NotSupportedException($"Milvus 2.3.0 or newer is required to handle record IDs safely. The server reported version '{version}'.");
}
}
internal static string BuildIdFilter(IEnumerable<string> keys)
{
StringBuilder filter = new();
filter.Append(IdFieldName).Append(" in [");
bool first = true;
foreach (string key in keys)
{
if (!first)
{
filter.Append(',');
}
first = false;
filter.Append('"');
foreach (char character in key)
{
switch (character)
{
case '\\':
filter.Append(@"\\");
break;
case '"':
filter.Append("\\\"");
break;
case '\r':
filter.Append(@"\r");
break;
case '\n':
filter.Append(@"\n");
break;
case '\t':
filter.Append(@"\t");
break;
default:
if (char.IsControl(character))
{
filter.Append(@"\u").Append(((int)character).ToString("x4", CultureInfo.InvariantCulture));
}
else
{
filter.Append(character);
}
break;
}
}
filter.Append('"');
}
return filter.Append(']').ToString();
}
private MemoryRecord ReadMemoryRecord(IReadOnlyList<FieldData> data, int rowNum, ReadOnlyMemory<float>? externalEmbedding = null)
{
bool isReference = false;
+5 -3
View File
@@ -2,14 +2,16 @@
This is an implementation of the Semantic Kernel Memory Store abstraction for the [Milvus vector database](https://milvus.io).
**Note:** Currently, only Milvus v2.2 is supported. v2.3 is coming soon, older versions are untested.
**Server requirement:** Record operations require Milvus v2.3.0 or newer. The connector checks the server version before saving, reading, removing, or searching records and throws `NotSupportedException` for older or unrecognized versions. Collection management and direct calls through `Client` are not covered by this check.
Record IDs are treated as literal values, including quotes, backslashes, brackets, and line breaks. IDs are not renamed or normalized. Milvus v2.2.14 does not decode escaped IDs correctly and could otherwise target a different record. Upgrade the server before using this connector with existing collections.
## Quickstart using a standalone Milvus installation
1. Download the Milvus docker-compose.yml:
```bash
wget https://github.com/milvus-io/milvus/releases/download/v2.2.14/milvus-standalone-docker-compose.yml -O docker-compose.yml
wget https://github.com/milvus-io/milvus/releases/download/v2.3.0/milvus-standalone-docker-compose.yml -O docker-compose.yml
```
2. Start Milvus:
@@ -30,4 +32,4 @@ SemanticTextMemory textMemory = new(memoryStore, embeddingGenerator);
var memoryPlugin = kernel.ImportPluginFromObject(new TextMemoryPlugin(textMemory));
```
More information on setting up Milvus can be found [here](https://milvus.io/docs/v2.2.x/install_standalone-docker.md). The `MilvusMemoryStore` constructor provides additional configuration options, such as the vector size, the similarity metric type, etc.
More information on setting up Milvus can be found [here](https://milvus.io/docs/v2.3.x/install_standalone-docker.md). The `MilvusMemoryStore` constructor provides additional configuration options, such as the vector size, the similarity metric type, etc.
@@ -0,0 +1,24 @@
<Project Sdk="Microsoft.NET.Sdk">
<PropertyGroup>
<AssemblyName>SemanticKernel.Connectors.Milvus.UnitTests</AssemblyName>
<RootNamespace>SemanticKernel.Connectors.Milvus.UnitTests</RootNamespace>
<TargetFramework>net10.0</TargetFramework>
<IsTestProject>true</IsTestProject>
<IsPackable>false</IsPackable>
<NoWarn>$(NoWarn);SKEXP0001;SKEXP0020</NoWarn>
</PropertyGroup>
<ItemGroup>
<PackageReference Include="Microsoft.NET.Test.Sdk" />
<PackageReference Include="xunit" />
<PackageReference Include="xunit.runner.visualstudio">
<IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
<PrivateAssets>all</PrivateAssets>
</PackageReference>
<PackageReference Include="Testcontainers.Milvus" />
</ItemGroup>
<ItemGroup>
<ProjectReference Include="..\..\..\src\VectorData\Milvus\Milvus.csproj" />
</ItemGroup>
</Project>
@@ -0,0 +1,84 @@
// Copyright (c) Microsoft. All rights reserved.
using System;
using System.Collections.Generic;
using Microsoft.SemanticKernel.Connectors.Milvus;
using Xunit;
namespace SemanticKernel.Connectors.Milvus.UnitTests;
public sealed class MilvusFilterTests
{
[Theory]
[InlineData("v2.3.0")]
[InlineData("v2.3.0-dev")]
[InlineData("2.3.0")]
[InlineData("v2.6.1")]
[InlineData("3.0.0")]
public void AcceptsSupportedServerVersions(string version)
{
MilvusMemoryStore.ValidateServerVersion(version);
}
[Theory]
[InlineData("v2.2.14")]
[InlineData("v2.2.14-dev")]
[InlineData("v2.2.16")]
[InlineData("v2.3.0-beta")]
[InlineData("unknown")]
[InlineData("")]
public void RejectsOlderOrUnrecognizedServerVersions(string version)
{
Assert.Throws<NotSupportedException>(() => MilvusMemoryStore.ValidateServerVersion(version));
}
[Theory]
[InlineData("", "")]
[InlineData("ordinary-id", "ordinary-id")]
[InlineData("[],'", "[],'")]
[InlineData("\"", "\\\"")]
[InlineData("\\", "\\\\")]
[InlineData("\\\"", "\\\\\\\"")]
[InlineData("trailing\\", "trailing\\\\")]
[InlineData("\r\n\t", "\\r\\n\\t")]
[InlineData("\\r\\n\\t", "\\\\r\\\\n\\\\t")]
[InlineData("\0\b\f\v\u001f\u007f\u0085", "\\u0000\\u0008\\u000c\\u000b\\u001f\\u007f\\u0085")]
[InlineData("\u00e9\u4e2d\ud83d\ude00", "\u00e9\u4e2d\ud83d\ude00")]
[InlineData("a\"] || true || id in [\"", "a\\\"] || true || id in [\\\"")]
[InlineData("z\"] && text like \"%secret%\" || id in [\"", "z\\\"] && text like \\\"%secret%\\\" || id in [\\\"")]
[InlineData("x\" || id != \"", "x\\\" || id != \\\"")]
[InlineData("a\",\"victim", "a\\\",\\\"victim")]
public void EncodesEachIdAsOneString(string key, string expected)
{
Assert.Equal("id in [\"" + expected + "\"]", MilvusMemoryStore.BuildIdFilter([key]));
}
[Fact]
public void PreservesEmptyLists()
{
Assert.Equal("id in []", MilvusMemoryStore.BuildIdFilter([]));
}
[Fact]
public void PreservesOrderDuplicatesAndEmptyIds()
{
Assert.Equal("id in [\"a\",\"\",\"a\",\"b\"]", MilvusMemoryStore.BuildIdFilter(["a", "", "a", "b"]));
}
[Fact]
public void EnumeratesKeysOnlyOnce()
{
bool enumerated = false;
IEnumerable<string> Keys()
{
Assert.False(enumerated);
enumerated = true;
yield return "first";
yield return "second";
}
Assert.Equal("id in [\"first\",\"second\"]", MilvusMemoryStore.BuildIdFilter(Keys()));
Assert.True(enumerated);
}
}
@@ -0,0 +1,76 @@
// Copyright (c) Microsoft. All rights reserved.
using System;
using System.Threading.Tasks;
using DotNet.Testcontainers.Builders;
using DotNet.Testcontainers.Containers;
using DotNet.Testcontainers.Networks;
using Testcontainers.Milvus;
using Xunit;
namespace SemanticKernel.Connectors.Milvus.UnitTests;
public class MilvusFixture : IAsyncLifetime
{
private readonly INetwork _network = new NetworkBuilder().Build();
private readonly IContainer _storage;
public MilvusFixture() : this("milvusdb/milvus:v2.3.0")
{
}
protected MilvusFixture(string image)
{
string password = Guid.NewGuid().ToString("N");
this._storage = new ContainerBuilder("milvusdb/minio:RELEASE.2024-12-18T13-15-44Z")
.WithNetwork(this._network)
.WithNetworkAliases("minio")
.WithEnvironment("MINIO_ROOT_USER", "test-user")
.WithEnvironment("MINIO_ROOT_PASSWORD", password)
.WithCommand("server", "/data")
.WithPortBinding(9000, true)
.WithWaitStrategy(Wait.ForUnixContainer().UntilHttpRequestIsSucceeded(
request => request.ForPort(9000).ForPath("/minio/health/live")))
.Build();
// Milvus 2.2 needs object storage; the module's local-storage default requires a newer server.
this.Container = new MilvusBuilder(image)
.WithNetwork(this._network)
.WithEnvironment("COMMON_STORAGETYPE", "minio")
.WithEnvironment("MINIO_ADDRESS", "minio:9000")
.WithEnvironment("MINIO_ACCESS_KEY_ID", "test-user")
.WithEnvironment("MINIO_SECRET_ACCESS_KEY", password)
.WithWaitStrategy(Wait.ForUnixContainer().UntilHttpRequestIsSucceeded(
request => request.ForPort(MilvusBuilder.MilvusManagementPort).ForPath("/healthz"),
strategy => strategy.WithTimeout(TimeSpan.FromMinutes(2))))
.Build();
}
public MilvusContainer Container { get; }
public async Task InitializeAsync()
{
await this._network.CreateAsync();
await this._storage.StartAsync();
await this.Container.StartAsync();
}
public async Task DisposeAsync()
{
try
{
await this.Container.DisposeAsync();
}
finally
{
try
{
await this._storage.DisposeAsync();
}
finally
{
await this._network.DisposeAsync();
}
}
}
}
@@ -0,0 +1,50 @@
// Copyright (c) Microsoft. All rights reserved.
using System;
using System.Linq;
using System.Threading.Tasks;
using Microsoft.SemanticKernel.Connectors.Milvus;
using Microsoft.SemanticKernel.Memory;
using Testcontainers.Milvus;
using Xunit;
namespace SemanticKernel.Connectors.Milvus.UnitTests;
public sealed class MilvusLegacyFixture : MilvusFixture
{
public MilvusLegacyFixture() : base("milvusdb/milvus:v2.2.14")
{
}
}
[Trait("Category", "MilvusIntegration")]
public sealed class MilvusLegacyServerTests : IAsyncLifetime
{
private readonly MilvusLegacyFixture _fixture = new();
public Task InitializeAsync() => this._fixture.InitializeAsync();
public Task DisposeAsync() => this._fixture.DisposeAsync();
// For manual testing with Docker only. Remove Skip locally to enable this test.
[Fact(Skip = "For manual testing only; requires Docker.")]
public async Task RejectsOldServerBeforeReadingOrChangingRecordsAsync()
{
using MilvusMemoryStore store = new(
this._fixture.Container.Hostname,
port: this._fixture.Container.GetMappedPublicPort(MilvusBuilder.MilvusGrpcPort));
const string CollectionName = "does_not_exist";
const string Key = "a\",\"victim";
MemoryRecord record = MemoryRecord.LocalRecord(Key, "synthetic", null, new float[] { 1, 0 });
await Assert.ThrowsAsync<NotSupportedException>(() => store.GetAsync(CollectionName, Key));
await Assert.ThrowsAsync<NotSupportedException>(() => store.GetBatchAsync(CollectionName, [Key]).ToListAsync().AsTask());
await Assert.ThrowsAsync<NotSupportedException>(() => store.RemoveAsync(CollectionName, Key));
await Assert.ThrowsAsync<NotSupportedException>(() => store.RemoveBatchAsync(CollectionName, [Key]));
await Assert.ThrowsAsync<NotSupportedException>(() => store.UpsertAsync(CollectionName, record));
await Assert.ThrowsAsync<NotSupportedException>(() => store.UpsertBatchAsync(CollectionName, [record]).ToListAsync().AsTask());
await Assert.ThrowsAsync<NotSupportedException>(() => store.GetNearestMatchAsync(CollectionName, new float[] { 1, 0 }, withEmbedding: true));
await Assert.ThrowsAsync<NotSupportedException>(() => store.GetNearestMatchesAsync(
CollectionName, new float[] { 1, 0 }, limit: 1, withEmbeddings: true).ToListAsync().AsTask());
}
}
@@ -0,0 +1,204 @@
// Copyright (c) Microsoft. All rights reserved.
using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading.Tasks;
using Microsoft.SemanticKernel.Connectors.Milvus;
using Microsoft.SemanticKernel.Memory;
using Milvus.Client;
using Testcontainers.Milvus;
using Xunit;
namespace SemanticKernel.Connectors.Milvus.UnitTests;
[Trait("Category", "MilvusIntegration")]
public sealed class MilvusMemoryStoreTests : IAsyncLifetime
{
// For manual testing with Docker only. Set to null locally to enable these tests.
private const string? SkipReason = "For manual testing only; requires Docker.";
private readonly MilvusFixture _fixture = new();
private readonly string _collectionName = "id_regression_" + Guid.NewGuid().ToString("N");
private MilvusMemoryStore _store = null!;
public static TheoryData<string> UnusualIds => new()
{
"ordinary",
"quotes\"and'single",
"brackets[],commas",
"backslash\\and\\\"quote\\",
"line\r\nbreak\ttab",
"literal\\r\\n\\t",
"controls\0\b\f\v\u001f\u007f\u0085",
"\u00e9\u4e2d\ud83d\ude00",
"a\"] || true || id in [\"",
"z\"] && text like \"%secret%\" || id in [\"",
"x\" || id != \"",
"a\",\"victim"
};
public async Task InitializeAsync()
{
await this._fixture.InitializeAsync();
this._store = new MilvusMemoryStore(
this._fixture.Container.Hostname,
port: this._fixture.Container.GetMappedPublicPort(MilvusBuilder.MilvusGrpcPort),
vectorSize: 2,
consistencyLevel: ConsistencyLevel.Strong);
// Create the fixture directly so these tests only depend on the ID operations under test.
MilvusCollection collection = await this._store.Client.CreateCollectionAsync(
this._collectionName,
new CollectionSchema
{
Fields =
{
FieldSchema.CreateVarchar("id", maxLength: 1024, isPrimaryKey: true, autoId: false),
FieldSchema.CreateFloatVector("embedding", 2)
},
EnableDynamicFields = true
},
ConsistencyLevel.Strong);
await collection.CreateIndexAsync("embedding", IndexType.Flat, SimilarityMetricType.Ip);
await collection.LoadAsync();
await collection.WaitForCollectionLoadAsync(
waitingInterval: TimeSpan.FromMilliseconds(100), timeout: TimeSpan.FromMinutes(1));
}
public async Task DisposeAsync()
{
try
{
if (this._store is not null)
{
await this._store.DeleteCollectionAsync(this._collectionName);
}
}
finally
{
this._store?.Dispose();
await this._fixture.DisposeAsync();
}
}
[Theory(Skip = SkipReason)]
[MemberData(nameof(UnusualIds))]
public async Task ReadsAndDeletesOnlyTheExactIdsAsync(string key)
{
MemoryRecord record = CreateRecord(key, [1, 0]);
MemoryRecord victim = CreateRecord("victim", [0, 1]);
await this.InsertRecordsAsync(record, victim);
foreach (bool withEmbeddings in new[] { false, true })
{
MemoryRecord? single = await this._store.GetAsync(this._collectionName, key, withEmbeddings);
AssertRecord(record, single, withEmbeddings);
List<MemoryRecord> batch = await this._store.GetBatchAsync(this._collectionName, [key], withEmbeddings).ToListAsync();
AssertRecord(record, Assert.Single(batch), withEmbeddings);
}
await this._store.RemoveAsync(this._collectionName, key);
Assert.Null(await this._store.GetAsync(this._collectionName, key));
AssertRecord(victim, await this._store.GetAsync(this._collectionName, "victim"), false);
await this.InsertRecordsAsync(record, CreateRecord("second", [0.5f, 0.5f]));
List<MemoryRecord> restored = await this._store.GetBatchAsync(
this._collectionName, [key, "second"], withEmbeddings: true).ToListAsync();
Assert.Equal(2, restored.Count);
AssertRecord(record, Assert.Single(restored, item => item.Metadata.Id == key), true);
await this._store.RemoveBatchAsync(this._collectionName, [key, "second"]);
Assert.Empty(await this._store.GetBatchAsync(this._collectionName, [key, "second"]).ToListAsync());
AssertRecord(victim, await this._store.GetAsync(this._collectionName, "victim"), false);
}
[Theory(Skip = SkipReason)]
[MemberData(nameof(UnusualIds))]
public async Task SearchUsesStoredIdsLiterallyAsync(string key)
{
MemoryRecord record = CreateRecord(key, [1, 0]);
MemoryRecord victim = CreateRecord("victim", [0, 1]);
await this.InsertRecordsAsync(record, victim);
await this._store.Client.GetCollection(this._collectionName).FlushAsync();
foreach (bool withEmbeddings in new[] { false, true })
{
List<(MemoryRecord, double)> matches = await this._store.GetNearestMatchesAsync(
this._collectionName, new float[] { 1, 0 }, limit: 2, withEmbeddings: withEmbeddings).ToListAsync();
Assert.Equal(2, matches.Count);
AssertRecord(record, matches[0].Item1, withEmbeddings);
AssertRecord(victim, matches[1].Item1, withEmbeddings);
(MemoryRecord, double)? nearest = await this._store.GetNearestMatchAsync(
this._collectionName, new float[] { 1, 0 }, withEmbedding: withEmbeddings);
Assert.NotNull(nearest);
AssertRecord(record, nearest.Value.Item1, withEmbeddings);
}
}
private static MemoryRecord CreateRecord(string key, float[] embedding)
=> MemoryRecord.LocalRecord(key, "text for " + key, description: null, embedding, key: key);
[Theory(Skip = SkipReason)]
[InlineData("quotes\"", "quotes\\\"")]
[InlineData("line\nbreak", "line\\nbreak")]
[InlineData("slash\\", "slash\\\\")]
public async Task DoesNotConfuseAnIdWithItsEscapedSpellingAsync(string key, string otherKey)
{
MemoryRecord expected = CreateRecord(key, [1, 0]);
MemoryRecord other = CreateRecord(otherKey, [0, 1]);
await this.InsertRecordsAsync(expected, other);
AssertRecord(expected, await this._store.GetAsync(this._collectionName, key), false);
await this._store.RemoveAsync(this._collectionName, key);
Assert.Null(await this._store.GetAsync(this._collectionName, key));
AssertRecord(other, await this._store.GetAsync(this._collectionName, otherKey), false);
}
[Fact(Skip = SkipReason)]
public async Task MissingHostileIdsDoNotReadOrDeleteOtherRecordsAsync()
{
MemoryRecord victim = CreateRecord("victim", [1, 0]);
await this.InsertRecordsAsync(victim);
string[] keys = ["a\",\"victim", "a\"] || true || id in [\""];
foreach (string key in keys)
{
Assert.Null(await this._store.GetAsync(this._collectionName, key));
await this._store.RemoveAsync(this._collectionName, key);
}
Assert.Empty(await this._store.GetBatchAsync(this._collectionName, keys).ToListAsync());
await this._store.RemoveBatchAsync(this._collectionName, keys);
AssertRecord(victim, await this._store.GetAsync(this._collectionName, "victim"), false);
}
private async Task InsertRecordsAsync(params MemoryRecord[] records)
{
// Seed typed records independently of the legacy client's Upsert response handling.
FieldData[] fields =
[
FieldData.Create("id", records.Select(record => record.Metadata.Id).ToArray()),
FieldData.CreateFloatVector("embedding", records.Select(record => record.Embedding).ToArray()),
FieldData.Create("text", records.Select(record => record.Metadata.Text).ToArray(), isDynamic: true),
FieldData.Create("is_reference", records.Select(record => record.Metadata.IsReference).ToArray(), isDynamic: true),
FieldData.Create("external_source_name", records.Select(record => record.Metadata.ExternalSourceName).ToArray(), isDynamic: true),
FieldData.Create("description", records.Select(record => record.Metadata.Description).ToArray(), isDynamic: true),
FieldData.Create("additional_metadata", records.Select(record => record.Metadata.AdditionalMetadata).ToArray(), isDynamic: true),
FieldData.Create("key", records.Select(record => record.Key).ToArray(), isDynamic: true),
FieldData.Create("timestamp", records.Select(_ => string.Empty).ToArray(), isDynamic: true)
];
await this._store.Client.GetCollection(this._collectionName).InsertAsync(fields);
}
private static void AssertRecord(MemoryRecord expected, MemoryRecord? actual, bool withEmbedding)
{
Assert.NotNull(actual);
Assert.Equal(expected.Metadata.Id, actual.Metadata.Id);
Assert.Equal(expected.Metadata.Text, actual.Metadata.Text);
Assert.Equal(withEmbedding ? expected.Embedding.ToArray() : [], actual.Embedding.ToArray());
}
}