Skip to content

Commit 9ac4029

Browse files
floberndflobernd
authored
Merge pull request #6 from elastic/builder-extensions
Add `VectorStoreRecordCollection` targeted `KernelBuilder` and `ServiceCollection` extensions
2 parents 3c5317e + 0c26b26 commit 9ac4029

File tree

2 files changed

+142
-10
lines changed

2 files changed

+142
-10
lines changed

‎Elastic.SemanticKernel.Connectors.Elasticsearch/ElasticsearchKernelBuilderExtensions.cs

Lines changed: 50 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -31,16 +31,62 @@ public static IKernelBuilder AddElasticsearchVectorStore(this IKernelBuilder bui
3131
}
3232

3333
/// <summary>
34-
/// Register an Elasticsearch <see cref="IVectorStore"/> with the specified service ID and where <see cref="ElasticsearchClient"/> is constructed using the provided settings.
34+
/// Register an Elasticsearch <see cref="IVectorStore"/> with the specified service ID and where <see cref="ElasticsearchClient"/> is constructed using the provided client settings.
3535
/// </summary>
3636
/// <param name="builder">The builder to register the <see cref="IVectorStore"/> on.</param>
37-
/// <param name="settings">The Elasticsearch client settings.</param>
37+
/// <param name="clientSettings">The Elasticsearch client settings.</param>
3838
/// <param name="options">Optional options to further configure the <see cref="IVectorStore"/>.</param>
3939
/// <param name="serviceId">An optional service id to use as the service key.</param>
4040
/// <returns>The kernel builder.</returns>
41-
public static IKernelBuilder AddElasticsearchVectorStore(this IKernelBuilder builder, IElasticsearchClientSettings settings, ElasticsearchVectorStoreOptions? options = default, string? serviceId = default)
41+
public static IKernelBuilder AddElasticsearchVectorStore(this IKernelBuilder builder, IElasticsearchClientSettings clientSettings, ElasticsearchVectorStoreOptions? options = default, string? serviceId = default)
4242
{
43-
builder.Services.AddElasticsearchVectorStore(settings, options, serviceId);
43+
builder.Services.AddElasticsearchVectorStore(clientSettings, options, serviceId);
44+
return builder;
45+
}
46+
47+
/// <summary>
48+
/// Register an Elasticsearch <see cref="IVectorStoreRecordCollection{TKey, TRecord}"/> and <see cref="IVectorizedSearch{TRecord}"/> with the specified service ID
49+
/// and where the <see cref="ElasticsearchClient"/> is retrieved from the dependency injection container.
50+
/// </summary>
51+
/// <typeparam name="TKey">The type of the key.</typeparam>
52+
/// <typeparam name="TRecord">The type of the record.</typeparam>
53+
/// <param name="builder">The builder to register the <see cref="IVectorStoreRecordCollection{TKey, TRecord}"/> on.</param>
54+
/// <param name="collectionName">The name of the collection.</param>
55+
/// <param name="options">Optional options to further configure the <see cref="IVectorStoreRecordCollection{TKey, TRecord}"/>.</param>
56+
/// <param name="serviceId">An optional service id to use as the service key.</param>
57+
/// <returns>The kernel builder.</returns>
58+
public static IKernelBuilder AddElasticsearchVectorStoreRecordCollection<TKey, TRecord>(
59+
this IKernelBuilder builder,
60+
string collectionName,
61+
ElasticsearchVectorStoreRecordCollectionOptions<TRecord>? options = default,
62+
string? serviceId = default)
63+
where TKey : notnull
64+
{
65+
builder.Services.AddElasticsearchVectorStoreRecordCollection<TKey, TRecord>(collectionName, options, serviceId);
66+
return builder;
67+
}
68+
69+
/// <summary>
70+
/// Register an Elasticsearch <see cref="IVectorStoreRecordCollection{TKey, TRecord}"/> and <see cref="IVectorizedSearch{TRecord}"/> with the specified service ID
71+
/// and where the <see cref="ElasticsearchClient"/> is constructed using the provided client settings.
72+
/// </summary>
73+
/// <typeparam name="TKey">The type of the key.</typeparam>
74+
/// <typeparam name="TRecord">The type of the record.</typeparam>
75+
/// <param name="builder">The builder to register the <see cref="IVectorStoreRecordCollection{TKey, TRecord}"/> on.</param>
76+
/// <param name="collectionName">The name of the collection.</param>
77+
/// <param name="clientSettings">The Elasticsearch client settings.</param>
78+
/// <param name="options">Optional options to further configure the <see cref="IVectorStoreRecordCollection{TKey, TRecord}"/>.</param>
79+
/// <param name="serviceId">An optional service id to use as the service key.</param>
80+
/// <returns>The kernel builder.</returns>
81+
public static IKernelBuilder AddElasticsearchVectorStoreRecordCollection<TKey, TRecord>(
82+
this IKernelBuilder builder,
83+
string collectionName,
84+
IElasticsearchClientSettings clientSettings,
85+
ElasticsearchVectorStoreRecordCollectionOptions<TRecord>? options = default,
86+
string? serviceId = default)
87+
where TKey : notnull
88+
{
89+
builder.Services.AddElasticsearchVectorStoreRecordCollection<TKey, TRecord>(collectionName, clientSettings, options, serviceId);
4490
return builder;
4591
}
4692
}

‎Elastic.SemanticKernel.Connectors.Elasticsearch/ElasticsearchServiceCollectionExtensions.cs

Lines changed: 92 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ public static IServiceCollection AddElasticsearchVectorStore(this IServiceCollec
2929
// cannot make assumptions about how ElasticsearchClient is being managed.
3030
services.AddKeyedTransient<IVectorStore>(
3131
serviceId,
32-
(sp, obj) =>
32+
(sp, _) =>
3333
{
3434
var elasticsearchClient = sp.GetRequiredService<ElasticsearchClient>();
3535
var selectedOptions = options ?? sp.GetService<ElasticsearchVectorStoreOptions>();
@@ -43,20 +43,20 @@ public static IServiceCollection AddElasticsearchVectorStore(this IServiceCollec
4343
}
4444

4545
/// <summary>
46-
/// Register an Elasticsearch <see cref="IVectorStore"/> with the specified service ID and where <see cref="ElasticsearchClient"/> is constructed using the provided settings.
46+
/// Register an Elasticsearch <see cref="IVectorStore"/> with the specified service ID and where <see cref="ElasticsearchClient"/> is constructed using the provided client settings.
4747
/// </summary>
4848
/// <param name="services">The <see cref="IServiceCollection"/> to register the <see cref="IVectorStore"/> on.</param>
49-
/// <param name="settings">The Elasticsearch client settings.</param>
49+
/// <param name="clientSettings">The Elasticsearch client settings.</param>
5050
/// <param name="options">Optional options to further configure the <see cref="IVectorStore"/>.</param>
5151
/// <param name="serviceId">An optional service id to use as the service key.</param>
5252
/// <returns>The service collection.</returns>
53-
public static IServiceCollection AddElasticsearchVectorStore(this IServiceCollection services, IElasticsearchClientSettings settings, ElasticsearchVectorStoreOptions? options = default, string? serviceId = default)
53+
public static IServiceCollection AddElasticsearchVectorStore(this IServiceCollection services, IElasticsearchClientSettings clientSettings, ElasticsearchVectorStoreOptions? options = default, string? serviceId = default)
5454
{
5555
services.AddKeyedSingleton<IVectorStore>(
5656
serviceId,
57-
(sp, obj) =>
57+
(sp, _) =>
5858
{
59-
var elasticsearchClient = new ElasticsearchClient(settings);
59+
var elasticsearchClient = new ElasticsearchClient(clientSettings);
6060
var selectedOptions = options ?? sp.GetService<ElasticsearchVectorStoreOptions>();
6161

6262
return new ElasticsearchVectorStore(
@@ -66,4 +66,90 @@ public static IServiceCollection AddElasticsearchVectorStore(this IServiceCollec
6666

6767
return services;
6868
}
69+
70+
/// <summary>
71+
/// Register an Elasticsearch <see cref="IVectorStoreRecordCollection{TKey, TRecord}"/> and <see cref="IVectorizedSearch{TRecord}"/> with the specified service ID
72+
/// and where the <see cref="ElasticsearchClient"/> is retrieved from the dependency injection container.
73+
/// </summary>
74+
/// <typeparam name="TKey">The type of the key.</typeparam>
75+
/// <typeparam name="TRecord">The type of the record.</typeparam>
76+
/// <param name="services">The <see cref="IServiceCollection"/> to register the <see cref="IVectorStoreRecordCollection{TKey, TRecord}"/> on.</param>
77+
/// <param name="collectionName">The name of the collection.</param>
78+
/// <param name="options">Optional options to further configure the <see cref="IVectorStoreRecordCollection{TKey, TRecord}"/>.</param>
79+
/// <param name="serviceId">An optional service id to use as the service key.</param>
80+
/// <returns>Service collection.</returns>
81+
public static IServiceCollection AddElasticsearchVectorStoreRecordCollection<TKey, TRecord>(
82+
this IServiceCollection services,
83+
string collectionName,
84+
ElasticsearchVectorStoreRecordCollectionOptions<TRecord>? options = default,
85+
string? serviceId = default)
86+
where TKey : notnull
87+
{
88+
services.AddKeyedTransient<IVectorStoreRecordCollection<TKey, TRecord>>(
89+
serviceId,
90+
(sp, _) =>
91+
{
92+
var elasticsearchClient = sp.GetRequiredService<ElasticsearchClient>();
93+
var selectedOptions = options ?? sp.GetService<ElasticsearchVectorStoreRecordCollectionOptions<TRecord>>();
94+
95+
return (new ElasticsearchVectorStoreRecordCollection<TRecord>(elasticsearchClient, collectionName, selectedOptions) as IVectorStoreRecordCollection<TKey, TRecord>)!;
96+
});
97+
98+
AddVectorizedSearch<TKey, TRecord>(services, serviceId);
99+
100+
return services;
101+
}
102+
103+
/// <summary>
104+
/// Register an Elasticsearch <see cref="IVectorStoreRecordCollection{TKey, TRecord}"/> and <see cref="IVectorizedSearch{TRecord}"/> with the specified service ID
105+
/// and where the <see cref="ElasticsearchClient"/> is constructed using the provided client settings.
106+
/// </summary>
107+
/// <typeparam name="TKey">The type of the key.</typeparam>
108+
/// <typeparam name="TRecord">The type of the record.</typeparam>
109+
/// <param name="services">The <see cref="IServiceCollection"/> to register the <see cref="IVectorStoreRecordCollection{TKey, TRecord}"/> on.</param>
110+
/// <param name="collectionName">The name of the collection.</param>
111+
/// <param name="clientSettings">The Elasticsearch client settings.</param>
112+
/// <param name="options">Optional options to further configure the <see cref="IVectorStoreRecordCollection{TKey, TRecord}"/>.</param>
113+
/// <param name="serviceId">An optional service id to use as the service key.</param>
114+
/// <returns>Service collection.</returns>
115+
public static IServiceCollection AddElasticsearchVectorStoreRecordCollection<TKey, TRecord>(
116+
this IServiceCollection services,
117+
string collectionName,
118+
IElasticsearchClientSettings clientSettings,
119+
ElasticsearchVectorStoreRecordCollectionOptions<TRecord>? options = default,
120+
string? serviceId = default)
121+
where TKey : notnull
122+
{
123+
services.AddKeyedSingleton<IVectorStoreRecordCollection<TKey, TRecord>>(
124+
serviceId,
125+
(sp, _) =>
126+
{
127+
var elasticsearchClient = new ElasticsearchClient(clientSettings);
128+
var selectedOptions = options ?? sp.GetService<ElasticsearchVectorStoreRecordCollectionOptions<TRecord>>();
129+
130+
return (new ElasticsearchVectorStoreRecordCollection<TRecord>(elasticsearchClient, collectionName, selectedOptions) as IVectorStoreRecordCollection<TKey, TRecord>)!;
131+
});
132+
133+
AddVectorizedSearch<TKey, TRecord>(services, serviceId);
134+
135+
return services;
136+
}
137+
138+
/// <summary>
139+
/// Also register the <see cref="IVectorStoreRecordCollection{TKey, TRecord}"/> with the given <paramref name="serviceId"/> as a <see cref="IVectorizedSearch{TRecord}"/>.
140+
/// </summary>
141+
/// <typeparam name="TKey">The type of the key.</typeparam>
142+
/// <typeparam name="TRecord">The type of the data model that the collection should contain.</typeparam>
143+
/// <param name="services">The service collection to register on.</param>
144+
/// <param name="serviceId">The service id that the registrations should use.</param>
145+
private static void AddVectorizedSearch<TKey, TRecord>(IServiceCollection services, string? serviceId)
146+
where TKey : notnull
147+
{
148+
services.AddKeyedTransient<IVectorizedSearch<TRecord>>(
149+
serviceId,
150+
(sp, _) =>
151+
{
152+
return sp.GetRequiredKeyedService<IVectorStoreRecordCollection<TKey, TRecord>>(serviceId);
153+
});
154+
}
69155
}

0 commit comments

Comments
 (0)