// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
using Microsoft.Extensions.AI;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.VectorData;
using VectorData.ConformanceTests.Support;
using Xunit;
namespace VectorData.ConformanceTests;
#pragma warning disable CA2000 // Don't actually need to dispose FakeEmbeddingGenerator
#pragma warning disable CS8605 // Unboxing a possibly null value.
public abstract class EmbeddingGenerationTests<TKey>(
EmbeddingGenerationTests<TKey>.StringVectorFixture stringVectorFixture,
EmbeddingGenerationTests<TKey>.RomOfFloatVectorFixture romOfFloatVectorFixture)
where TKey : notnull
{
#region Search
[Fact]
public virtual async Task SearchAsync_with_property_generator()
{
// Property level: embedding generators are defined at all levels. The property generator should take precedence.
var collection = GetCollection<Record>(storeGenerator: true, collectionGenerator: true, propertyGenerator: true);
var result = await collection.SearchAsync("[1, 1, 0]", top: 1).SingleAsync();
Assert.Equal("Property ([1, 1, 3])", result.Record.Text);
}
[Fact]
public virtual async Task SearchAsync_with_property_generator_dynamic()
{
// Property level: embedding generators are defined at all levels. The property generator should take precedence.
var collection = GetDynamicCollection(storeGenerator: true, collectionGenerator: true, propertyGenerator: true);
var result = await collection.SearchAsync("[1, 1, 0]", top: 1).SingleAsync();
Assert.Equal("Property ([1, 1, 3])", result.Record[nameof(Record.Text)]);
}
[Fact]
public virtual async Task SearchAsync_with_collection_generator()
{
// Collection level: embedding generators are defined at the collection and store level - the collection generator should take precedence.
var collection = GetCollection<Record>(storeGenerator: true, collectionGenerator: true, propertyGenerator: false);
var result = await collection.SearchAsync("[1, 1, 0]", top: 1).SingleAsync();
Assert.Equal("Collection ([1, 1, 2])", result.Record.Text);
}
[Fact]
public virtual async Task SearchAsync_with_store_generator()
{
// Store level: an embedding generator is defined at the store level only.
var collection = GetCollection<Record>(storeGenerator: true, collectionGenerator: false, propertyGenerator: false);
var result = await collection.SearchAsync("[1, 1, 0]", top: 1).SingleAsync();
Assert.Equal("Store ([1, 1, 1])", result.Record.Text);
}
[Fact]
public virtual async Task SearchAsync_with_store_dependency_injection()
{
foreach (var registrationDelegate in stringVectorFixture.DependencyInjectionStoreRegistrationDelegates)
{
IServiceCollection serviceCollection = new ServiceCollection();
serviceCollection.AddSingleton<IEmbeddingGenerator>(new FakeEmbeddingGenerator(replaceLast: 1));
registrationDelegate(serviceCollection);
await using var serviceProvider = serviceCollection.BuildServiceProvider();
var vectorStore = serviceProvider.GetRequiredService<VectorStore>();
var collection = vectorStore.GetCollection<TKey, Record>(stringVectorFixture.CollectionName, stringVectorFixture.CreateRecordDefinition());
var result = await collection.SearchAsync("[1, 1, 0]", top: 1).SingleAsync();
Assert.Equal("Store ([1, 1, 1])", result.Record.Text);
}
}
[Fact]
public virtual async Task SearchAsync_with_collection_dependency_injection()
{
foreach (var registrationDelegate in stringVectorFixture.DependencyInjectionCollectionRegistrationDelegates)
{
IServiceCollection serviceCollection = new ServiceCollection();
serviceCollection.AddSingleton<IEmbeddingGenerator>(new FakeEmbeddingGenerator(replaceLast: 1));
registrationDelegate(serviceCollection);
await using var serviceProvider = serviceCollection.BuildServiceProvider();
var collection = serviceProvider.GetRequiredService<VectorStoreCollection<TKey, RecordWithAttributes>>();
var result = await collection.SearchAsync("[1, 1, 0]", top: 1).SingleAsync();
Assert.Equal("Store ([1, 1, 1])", result.Record.Text);
}
}
[Fact]
public virtual async Task SearchAsync_with_custom_input_type()
{
var recordDefinition = new VectorStoreCollectionDefinition
{
Properties = stringVectorFixture.CreateRecordDefinition().Properties
.Select(p => p is VectorStoreVectorProperty vectorProperty
? new VectorStoreVectorProperty<Customer>(nameof(Record.Embedding), dimensions: 3)
{
DistanceFunction = stringVectorFixture.DefaultDistanceFunction,
IndexKind = stringVectorFixture.DefaultIndexKind
}
: p)
.ToList()
};
var collection = stringVectorFixture.GetCollection<RecordWithCustomerVectorProperty>(
stringVectorFixture.CreateVectorStore(new FakeCustomerEmbeddingGenerator([1, 1, 1])),
stringVectorFixture.CollectionName,
recordDefinition);
var result = await collection.SearchAsync(new Customer(), top: 1).SingleAsync();
Assert.Equal("Store ([1, 1, 1])", result.Record.Text);
}
[Fact]
public virtual async Task SearchAsync_with_search_only_embedding_generation()
{
var collection = romOfFloatVectorFixture.GetCollection<RecordWithRomOfFloatVectorProperty>(
romOfFloatVectorFixture.CreateVectorStore(new FakeEmbeddingGenerator()),
romOfFloatVectorFixture.CollectionName,
romOfFloatVectorFixture.CreateRecordDefinition());
var result = await collection.SearchAsync("[1, 1, 1]", top: 1).SingleAsync();
Assert.Equal("Store ([1, 1, 1])", result.Record.Text);
}
[Fact]
public virtual async Task SearchAsync_string_without_generator_throws()
{
// The database doesn't support embedding generation, and no client-side generator has been configured at any level,
// so SearchAsync should throw for string.
var collection = stringVectorFixture.GetCollection<RawRecord>(stringVectorFixture.TestStore.DefaultVectorStore, stringVectorFixture.CollectionName + "withoutgenerator");
var exception = await Assert.ThrowsAsync<NotSupportedException>(() => collection.SearchAsync("foo", top: 1).ToListAsync().AsTask());
Assert.StartsWith(
"A value of type 'string' was passed to 'SearchAsync', but that isn't a supported vector type by your provider and no embedding generator was configured. The supported vector types are:",
exception.Message);
}
public class RawRecord
{
[VectorStoreKey]
public TKey Key { get; set; } = default!;
[VectorStoreVector(dimensions: 3)]
public ReadOnlyMemory<float> Embedding { get; set; }
}
[Fact]
public virtual async Task SearchAsync_with_incompatible_generator_throws()
{
var collection = GetCollection<Record>(storeGenerator: true, collectionGenerator: true, propertyGenerator: true);
// We have a generator configured for string, not int.
var exception = await Assert.ThrowsAsync<InvalidOperationException>(() => collection.SearchAsync(8, top: 1).ToListAsync().AsTask());
Assert.Equal($"An input of type 'int' was provided, but an incompatible embedding generator of type '{nameof(FakeEmbeddingGenerator)}' was configured.", exception.Message);
}
#endregion Search
#region Upsert
[Fact]
public virtual async Task UpsertAsync()
{
var counter = stringVectorFixture.GenerateNextCounter();
var record = new Record
{
Key = stringVectorFixture.GenerateNextKey<TKey>(),
Embedding = "[100, 1, 0]",
Counter = counter,
Text = nameof(UpsertAsync)
};
// Property level: embedding generators are defined at all levels. The property generator should take precedence.
var collection = GetCollection<Record>(storeGenerator: true, collectionGenerator: true, propertyGenerator: true);
await collection.UpsertAsync(record);
await stringVectorFixture.TestStore.WaitForDataAsync(collection, 1, filter: r => r.Counter == counter);
var result = await collection.SearchAsync(new ReadOnlyMemory<float>([100, 1, 3]), top: 1).SingleAsync();
Assert.Equal(counter, result.Record.Counter);
}
[Fact]
public virtual async Task UpsertAsync_dynamic()
{
var counter = stringVectorFixture.GenerateNextCounter();
var record = new Dictionary<string, object?>
{
[nameof(Record.Key)] = stringVectorFixture.GenerateNextKey<TKey>(),
[nameof(Record.Embedding)] = "[200, 1, 0]",
[nameof(Record.Counter)] = counter,
[nameof(Record.Text)] = nameof(UpsertAsync_dynamic)
};
// Property level: embedding generators are defined at all levels. The property generator should take precedence.
var collection = GetDynamicCollection(storeGenerator: true, collectionGenerator: true, propertyGenerator: true);
await collection.UpsertAsync(record);
await stringVectorFixture.TestStore.WaitForDataAsync(collection, 1, filter: r => (int)r[nameof(Record.Counter)] == counter);
var result = await collection.SearchAsync(new ReadOnlyMemory<float>([200, 1, 3]), top: 1).SingleAsync();
Assert.Equal(counter, result.Record[nameof(Record.Counter)]);
}
[Fact]
public virtual async Task UpsertAsync_batch()
{
var (counter1, counter2) = (stringVectorFixture.GenerateNextCounter(), stringVectorFixture.GenerateNextCounter());
Record[] records =
[
new()
{
Key = stringVectorFixture.GenerateNextKey<TKey>(),
Embedding = "[300, 1, 0]",
Counter = counter1,
Text = nameof(UpsertAsync_batch) + "1"
},
new()
{
Key = stringVectorFixture.GenerateNextKey<TKey>(),
Embedding = "[400, 1, 0]",
Counter = counter2,
Text = nameof(UpsertAsync_batch) + "2"
}
];
var collection = GetCollection<Record>(storeGenerator: true, collectionGenerator: true, propertyGenerator: true);
await collection.UpsertAsync(records);
await stringVectorFixture.TestStore.WaitForDataAsync(collection, 2, filter: r => r.Counter == counter1 || r.Counter == counter2);
var result = await collection.SearchAsync(new ReadOnlyMemory<float>([300, 1, 3]), top: 1).SingleAsync();
Assert.Equal(counter1, result.Record.Counter);
result = await collection.SearchAsync(new ReadOnlyMemory<float>([400, 1, 3]), top: 1).SingleAsync();
Assert.Equal(counter2, result.Record.Counter);
}
[Fact]
public virtual async Task UpsertAsync_batch_dynamic()
{
var (counter1, counter2) = (stringVectorFixture.GenerateNextCounter(), stringVectorFixture.GenerateNextCounter());
Dictionary<string, object?>[] records =
[
new()
{
[nameof(Record.Key)] = stringVectorFixture.GenerateNextKey<TKey>(),
[nameof(Record.Embedding)] = "[500, 1, 0]",
[nameof(Record.Counter)] = counter1,
[nameof(Record.Text)] = nameof(UpsertAsync_batch_dynamic) + "1"
},
new()
{
[nameof(Record.Key)] = stringVectorFixture.GenerateNextKey<TKey>(),
[nameof(Record.Embedding)] = "[600, 1, 0]",
[nameof(Record.Counter)] = counter2,
[nameof(Record.Text)] = nameof(UpsertAsync_batch_dynamic) + "2"
}
];
var collection = GetDynamicCollection(storeGenerator: true, collectionGenerator: true, propertyGenerator: true);
await collection.UpsertAsync(records);
await stringVectorFixture.TestStore.WaitForDataAsync(collection, 2, filter: r => (int)r[nameof(Record.Counter)] == counter1 || (int)r[nameof(Record.Counter)] == counter2);
var result = await collection.SearchAsync(new ReadOnlyMemory<float>([500, 1, 3]), top: 1).SingleAsync();
Assert.Equal(counter1, result.Record[nameof(Record.Counter)]);
result = await collection.SearchAsync(new ReadOnlyMemory<float>([600, 1, 3]), top: 1).SingleAsync();
Assert.Equal(counter2, result.Record[nameof(Record.Counter)]);
}
#endregion Upsert
#region IncludeVectors
[Fact]
public virtual async Task SearchAsync_with_IncludeVectors_throws()
{
var collection = GetCollection<Record>(storeGenerator: true, collectionGenerator: true, propertyGenerator: true);
var exception = await Assert.ThrowsAsync<NotSupportedException>(() => collection.SearchAsync("[1, 0, 0]", top: 1, new() { IncludeVectors = true }).ToListAsync().AsTask());
Assert.Equal("When an embedding generator is configured, `Include Vectors` cannot be enabled.", exception.Message);
}
[Fact]
public virtual async Task GetAsync_with_IncludeVectors_throws()
{
var collection = GetCollection<Record>(storeGenerator: true, collectionGenerator: true, propertyGenerator: true);
var exception = await Assert.ThrowsAsync<NotSupportedException>(() => collection.GetAsync(stringVectorFixture.TestData[0].Key, new() { IncludeVectors = true }));
Assert.Equal("When an embedding generator is configured, `Include Vectors` cannot be enabled.", exception.Message);
}
[Fact]
public virtual async Task GetAsync_enumerable_with_IncludeVectors_throws()
{
var collection = GetCollection<Record>(storeGenerator: true, collectionGenerator: true, propertyGenerator: true);
var exception = await Assert.ThrowsAsync<NotSupportedException>(() =>
collection.GetAsync(
[stringVectorFixture.TestData[0].Key, stringVectorFixture.TestData[1].Key],
new() { IncludeVectors = true })
.ToListAsync().AsTask());
Assert.Equal("When an embedding generator is configured, `Include Vectors` cannot be enabled.", exception.Message);
}
#endregion IncludeVectors
#region Support
public class Record : TestRecord<TKey>
{
public string? Embedding { get; set; }
public int Counter { get; set; }
public string? Text { get; set; }
}
public class RecordWithAttributes
{
[VectorStoreKey]
public TKey Key { get; set; } = default!;
[VectorStoreVector(dimensions: 3)]
public string? Embedding { get; set; }
[VectorStoreData(IsIndexed = true)]
public int Counter { get; set; }
[VectorStoreData]
public string? Text { get; set; }
}
public class RecordWithCustomerVectorProperty
{
public TKey Key { get; set; } = default!;
public Customer? Embedding { get; set; }
public int Counter { get; set; }
public string? Text { get; set; }
}
public class RecordWithRomOfFloatVectorProperty : TestRecord<TKey>
{
public ReadOnlyMemory<float> Embedding { get; set; }
public int Counter { get; set; }
public string? Text { get; set; }
}
public class Customer
{
public string? FirstName { get; set; }
public string? LastName { get; set; }
}
private VectorStoreCollection<TKey, TRecord> GetCollection<TRecord>(
bool storeGenerator = false,
bool collectionGenerator = false,
bool propertyGenerator = false)
where TRecord : class
{
var (vectorStore, recordDefinition) = GetStoreAndRecordDefinition(storeGenerator, collectionGenerator, propertyGenerator);
return stringVectorFixture.GetCollection<TRecord>(vectorStore, stringVectorFixture.CollectionName, recordDefinition);
}
private VectorStoreCollection<object, Dictionary<string, object?>> GetDynamicCollection(
bool storeGenerator = false,
bool collectionGenerator = false,
bool propertyGenerator = false)
{
var (vectorStore, recordDefinition) = GetStoreAndRecordDefinition(storeGenerator, collectionGenerator, propertyGenerator);
return stringVectorFixture.GetDynamicCollection(vectorStore, stringVectorFixture.CollectionName, recordDefinition);
}
private (VectorStore vectorStore, VectorStoreCollectionDefinition definition) GetStoreAndRecordDefinition(
bool storeGenerator = false,
bool collectionGenerator = false,
bool propertyGenerator = false)
{
var recordDefinition = stringVectorFixture.CreateRecordDefinition();
if (propertyGenerator)
{
foreach (var vectorProperty in recordDefinition.Properties.OfType<VectorStoreVectorProperty>())
{
vectorProperty.EmbeddingGenerator = new FakeEmbeddingGenerator(replaceLast: 3);
}
}
if (collectionGenerator)
{
recordDefinition.EmbeddingGenerator = new FakeEmbeddingGenerator(replaceLast: 2);
}
var vectorStore = stringVectorFixture.CreateVectorStore(storeGenerator ? new FakeEmbeddingGenerator(replaceLast: 1) : null);
return (vectorStore, recordDefinition);
}
public abstract class StringVectorFixture : VectorStoreCollectionFixture<TKey, Record>
{
private int _counter;
protected override string CollectionNameBase => nameof(EmbeddingGenerationTests<int>);
public override VectorStoreCollectionDefinition CreateRecordDefinition()
=> new()
{
Properties =
[
new VectorStoreKeyProperty(nameof(Record.Key), typeof(TKey)),
new VectorStoreVectorProperty(nameof(Record.Embedding), typeof(string), dimensions: 3)
{
DistanceFunction = DefaultDistanceFunction,
IndexKind = DefaultIndexKind
},
new VectorStoreDataProperty(nameof(Record.Counter), typeof(int)) { IsIndexed = true },
new VectorStoreDataProperty(nameof(Record.Text), typeof(string))
],
EmbeddingGenerator = new FakeEmbeddingGenerator()
};
protected override List<Record> BuildTestData() =>
[
new()
{
Key = GenerateNextKey<TKey>(),
Embedding = "[1, 1, 1]",
Counter = GenerateNextCounter(),
Text = "Store ([1, 1, 1])"
},
new()
{
Key = GenerateNextKey<TKey>(),
Embedding = "[1, 1, 2]",
Counter = GenerateNextCounter(),
Text = "Collection ([1, 1, 2])"
},
new()
{
Key = GenerateNextKey<TKey>(),
Embedding = "[1, 1, 3]",
Counter = GenerateNextCounter(),
Text = "Property ([1, 1, 3])"
}
];
public virtual VectorStoreCollection<TKey, TRecord> GetCollection<TRecord>(
VectorStore vectorStore,
string collectionName,
VectorStoreCollectionDefinition? recordDefinition = null)
where TRecord : class
=> vectorStore.GetCollection<TKey, TRecord>(collectionName, recordDefinition);
public virtual VectorStoreCollection<object, Dictionary<string, object?>> GetDynamicCollection(
VectorStore vectorStore,
string collectionName,
VectorStoreCollectionDefinition recordDefinition)
=> vectorStore.GetDynamicCollection(collectionName, recordDefinition);
public abstract VectorStore CreateVectorStore(IEmbeddingGenerator? embeddingGenerator = null);
public abstract Func<IServiceCollection, IServiceCollection>[] DependencyInjectionStoreRegistrationDelegates { get; }
public abstract Func<IServiceCollection, IServiceCollection>[] DependencyInjectionCollectionRegistrationDelegates { get; }
public virtual int GenerateNextCounter()
=> Interlocked.Increment(ref _counter);
}
public abstract class RomOfFloatVectorFixture : VectorStoreCollectionFixture<TKey, RecordWithRomOfFloatVectorProperty>
{
private int _counter;
protected override string CollectionNameBase => "SearchOnlyEmbeddingGenerationTests";
public override VectorStoreCollectionDefinition CreateRecordDefinition()
=> new()
{
Properties =
[
new VectorStoreKeyProperty(nameof(RecordWithRomOfFloatVectorProperty.Key), typeof(TKey)),
new VectorStoreVectorProperty(nameof(RecordWithRomOfFloatVectorProperty.Embedding), typeof(ReadOnlyMemory<float>), dimensions: 3)
{
DistanceFunction = DefaultDistanceFunction,
IndexKind = DefaultIndexKind
},
new VectorStoreDataProperty(nameof(RecordWithRomOfFloatVectorProperty.Counter), typeof(int)) { IsIndexed = true },
new VectorStoreDataProperty(nameof(RecordWithRomOfFloatVectorProperty.Text), typeof(string))
],
EmbeddingGenerator = new FakeEmbeddingGenerator()
};
protected override List<RecordWithRomOfFloatVectorProperty> BuildTestData() =>
[
new()
{
Key = GenerateNextKey<TKey>(),
Embedding = new ReadOnlyMemory<float>([1, 1, 1]),
Counter = GenerateNextCounter(),
Text = "Store ([1, 1, 1])",
},
new()
{
Key = GenerateNextKey<TKey>(),
Embedding = new ReadOnlyMemory<float>([1, 1, 2]),
Counter = GenerateNextCounter(),
Text = "Collection ([1, 1, 2])"
},
new()
{
Key = GenerateNextKey<TKey>(),
Embedding = new ReadOnlyMemory<float>([1, 1, 3]),
Counter = GenerateNextCounter(),
Text = "Property ([1, 1, 3])"
}
];
public virtual VectorStoreCollection<TKey, TRecord> GetCollection<TRecord>(
VectorStore vectorStore,
string collectionName,
VectorStoreCollectionDefinition? recordDefinition = null)
where TRecord : class
=> vectorStore.GetCollection<TKey, TRecord>(collectionName, recordDefinition);
public virtual VectorStoreCollection<object, Dictionary<string, object?>> GetDynamicCollection(
VectorStore vectorStore,
string collectionName,
VectorStoreCollectionDefinition recordDefinition)
=> vectorStore.GetDynamicCollection(collectionName, recordDefinition);
public abstract VectorStore CreateVectorStore(IEmbeddingGenerator? embeddingGenerator = null);
public abstract Func<IServiceCollection, IServiceCollection>[] DependencyInjectionStoreRegistrationDelegates { get; }
public abstract Func<IServiceCollection, IServiceCollection>[] DependencyInjectionCollectionRegistrationDelegates { get; }
public virtual int GenerateNextCounter()
=> Interlocked.Increment(ref _counter);
}
protected sealed class FakeEmbeddingGenerator(int? replaceLast = null) : IEmbeddingGenerator<string, Embedding<float>>
{
public Task<GeneratedEmbeddings<Embedding<float>>> GenerateAsync(
IEnumerable<string> values,
EmbeddingGenerationOptions? options = null,
CancellationToken cancellationToken = default)
{
var results = new GeneratedEmbeddings<Embedding<float>>();
foreach (var value in values)
{
var vector = value.TrimStart('[').TrimEnd(']').Split(',').Select(s => float.Parse(s.Trim())).ToArray();
if (replaceLast is not null)
{
vector[vector.Length - 1] = replaceLast.Value;
}
results.Add(new Embedding<float>(vector));
}
return Task.FromResult(results);
}
public object? GetService(Type serviceType, object? serviceKey = null)
=> null;
public void Dispose()
{
}
}
private sealed class FakeCustomerEmbeddingGenerator(float[] embedding) : IEmbeddingGenerator<Customer, Embedding<float>>
{
public Task<GeneratedEmbeddings<Embedding<float>>> GenerateAsync(IEnumerable<Customer> values, EmbeddingGenerationOptions? options = null, CancellationToken cancellationToken = default)
=> Task.FromResult(new GeneratedEmbeddings<Embedding<float>> { new(embedding) });
public object? GetService(Type serviceType, object? serviceKey = null)
=> null;
public void Dispose()
{
}
}
#endregion Support
}