106 lines
4.5 KiB
C#
106 lines
4.5 KiB
C#
using Wiki.Core.Model;
|
|
using Wiki.Core.Ports;
|
|
|
|
namespace Wiki.Contracts;
|
|
|
|
/// <summary>Every <see cref="IVectorIndex"/> adapter derives from this. Vectors are unit length; scores are cosine.</summary>
|
|
public abstract class VectorIndexContract
|
|
{
|
|
/// <summary>A fresh, empty index. Called again after <see cref="ReopenAsync"/> semantics are needed.</summary>
|
|
protected abstract Task<IVectorIndex> CreateIndexAsync();
|
|
|
|
/// <summary>The same index reopened from persisted state, or the same instance for adapters that keep no files.</summary>
|
|
protected abstract Task<IVectorIndex> ReopenAsync(IVectorIndex index);
|
|
|
|
private const string Model = "test-model";
|
|
|
|
private static float[] Unit(params float[] v)
|
|
{
|
|
var n = MathF.Sqrt(v.Sum(x => x * x));
|
|
return v.Select(x => x / n).ToArray();
|
|
}
|
|
|
|
[Fact]
|
|
public async Task EmptyIndexDescribesEmptyAndQueriesNothing()
|
|
{
|
|
var index = await CreateIndexAsync();
|
|
var state = await index.DescribeAsync(Model, CancellationToken.None);
|
|
Assert.Equal(0, state.Count);
|
|
Assert.Empty(await index.QueryAsync(Model, Unit(1, 0, 0), 5, CancellationToken.None));
|
|
}
|
|
|
|
[Fact]
|
|
public async Task QueryRanksByCosineAndHonoursK()
|
|
{
|
|
var index = await CreateIndexAsync();
|
|
await index.UpsertAsync(
|
|
[
|
|
new Embedding(Model, "a", "h1", Unit(1, 0, 0)),
|
|
new Embedding(Model, "b", "h2", Unit(0.9f, 0.1f, 0)),
|
|
new Embedding(Model, "c", "h3", Unit(0, 1, 0)),
|
|
new Embedding(Model, "d", "h4", Unit(0, 0, 1)),
|
|
], CancellationToken.None);
|
|
|
|
var hits = await index.QueryAsync(Model, Unit(1, 0, 0), 2, CancellationToken.None);
|
|
Assert.Equal(["a", "b"], hits.Select(h => h.Slug));
|
|
Assert.True(hits[0].Score > hits[1].Score);
|
|
Assert.InRange(hits[0].Score, 0.999f, 1.001f);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task UpsertReplacesAndDescribeTracksHashes()
|
|
{
|
|
var index = await CreateIndexAsync();
|
|
await index.UpsertAsync([new Embedding(Model, "a", "h1", Unit(1, 0, 0))], CancellationToken.None);
|
|
await index.UpsertAsync([new Embedding(Model, "a", "h2", Unit(0, 1, 0))], CancellationToken.None);
|
|
var state = await index.DescribeAsync(Model, CancellationToken.None);
|
|
Assert.Equal(1, state.Count);
|
|
Assert.Equal("h2", state.ContentHashBySlug["a"]);
|
|
var hits = await index.QueryAsync(Model, Unit(0, 1, 0), 1, CancellationToken.None);
|
|
Assert.Equal("a", hits[0].Slug);
|
|
Assert.InRange(hits[0].Score, 0.999f, 1.001f);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task RemoveDropsSlugs()
|
|
{
|
|
var index = await CreateIndexAsync();
|
|
await index.UpsertAsync([new Embedding(Model, "a", "h1", Unit(1, 0, 0)), new Embedding(Model, "b", "h2", Unit(0, 1, 0))], CancellationToken.None);
|
|
await index.RemoveAsync(Model, ["a", "missing"], CancellationToken.None);
|
|
var state = await index.DescribeAsync(Model, CancellationToken.None);
|
|
Assert.Equal(["b"], state.ContentHashBySlug.Keys);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task ModelsAreIsolatedNamespaces()
|
|
{
|
|
var index = await CreateIndexAsync();
|
|
await index.UpsertAsync([new Embedding("m1", "a", "h1", Unit(1, 0, 0))], CancellationToken.None);
|
|
await index.UpsertAsync([new Embedding("m2", "b", "h2", Unit(1, 0, 0))], CancellationToken.None);
|
|
Assert.Equal(["a"], (await index.QueryAsync("m1", Unit(1, 0, 0), 10, CancellationToken.None)).Select(h => h.Slug));
|
|
Assert.Equal(["b"], (await index.QueryAsync("m2", Unit(1, 0, 0), 10, CancellationToken.None)).Select(h => h.Slug));
|
|
Assert.Equal(1, (await index.DescribeAsync("m1", CancellationToken.None)).Count);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task SurvivesFlushAndReopen()
|
|
{
|
|
var index = await CreateIndexAsync();
|
|
await index.UpsertAsync([new Embedding(Model, "a", "h1", Unit(1, 0, 0)), new Embedding(Model, "b", "h2", Unit(0, 1, 0))], CancellationToken.None);
|
|
await index.FlushAsync(CancellationToken.None);
|
|
var reopened = await ReopenAsync(index);
|
|
var state = await reopened.DescribeAsync(Model, CancellationToken.None);
|
|
Assert.Equal(2, state.Count);
|
|
var hits = await reopened.QueryAsync(Model, Unit(0, 1, 0), 1, CancellationToken.None);
|
|
Assert.Equal("b", hits[0].Slug);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task HealthReportsOk()
|
|
{
|
|
var index = await CreateIndexAsync();
|
|
var health = await index.CheckAsync(CancellationToken.None);
|
|
Assert.True(health.Ok, health.Detail);
|
|
}
|
|
}
|