Skip to content

Commit 4e0c866

Browse files
adamsitnikadamsitnik
committed
Sqlite
1 parent bbb2842 commit 4e0c866

File tree

7 files changed

+308
-213
lines changed

7 files changed

+308
-213
lines changed

‎dotnet/src/Connectors/Connectors.Memory.SqliteVec/SqliteCollectionOptions.cs

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,21 @@ public sealed class SqliteCollectionOptions
1212
{
1313
internal static readonly SqliteCollectionOptions Default = new();
1414

15+
/// <summary>
16+
/// Initializes a new instance of the <see cref="SqliteCollectionOptions"/> class.
17+
/// </summary>
18+
public SqliteCollectionOptions()
19+
{
20+
}
21+
22+
internal SqliteCollectionOptions(SqliteCollectionOptions? source, IEmbeddingGenerator embeddingGenerator)
23+
{
24+
this.VectorStoreRecordDefinition = source?.VectorStoreRecordDefinition;
25+
this.VectorVirtualTableName = source?.VectorVirtualTableName;
26+
this.VectorSearchExtensionName = source?.VectorSearchExtensionName;
27+
this.EmbeddingGenerator = embeddingGenerator;
28+
}
29+
1530
/// <summary>
1631
/// Gets or sets an optional record definition that defines the schema of the record type.
1732
/// </summary>

‎dotnet/src/Connectors/Connectors.Memory.SqliteVec/SqliteServiceCollectionExtensions.cs

Lines changed: 193 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -1,84 +1,227 @@
11
// Copyright (c) Microsoft. All rights reserved.
22

3-
using Microsoft.Data.Sqlite;
3+
using System;
44
using Microsoft.Extensions.AI;
5-
using Microsoft.Extensions.DependencyInjection;
65
using Microsoft.Extensions.VectorData;
6+
using Microsoft.SemanticKernel;
77
using Microsoft.SemanticKernel.Connectors.SqliteVec;
88

9-
namespace Microsoft.SemanticKernel;
9+
namespace Microsoft.Extensions.DependencyInjection;
1010

1111
/// <summary>
1212
/// Extension methods to register SQLite <see cref="VectorStore"/> instances on an <see cref="IServiceCollection"/>.
1313
/// </summary>
1414
public static class SqliteServiceCollectionExtensions
1515
{
1616
/// <summary>
17-
/// Register a SQLite <see cref="VectorStore"/> with the specified service ID.
18-
/// <see cref="SqliteConnection"/> instance will be initialized, connection will be opened and vector search extension with be loaded.
17+
/// Registers a <see cref="SqliteVectorStore"/> as <see cref="VectorStore"/>, with the specified connection string and service lifetime.
1918
/// </summary>
20-
/// <param name="services">The <see cref="IServiceCollection"/> to register the <see cref="VectorStore"/> on.</param>
21-
/// <param name="connectionString">Connection string for <see cref="SqliteConnection"/>.</param>
22-
/// <param name="options">Optional options to further configure the <see cref="VectorStore"/>.</param>
23-
/// <param name="serviceId">An optional service id to use as the service key.</param>
24-
/// <returns>Service collection.</returns>
19+
/// <inheritdoc cref="AddVectorStore"/>
2520
public static IServiceCollection AddSqliteVectorStore(
2621
this IServiceCollection services,
27-
string connectionString,
28-
SqliteVectorStoreOptions? options = default,
29-
string? serviceId = default)
30-
=> services.AddKeyedSingleton<VectorStore>(
31-
serviceId,
32-
(sp, _) => new SqliteVectorStore(connectionString, options ?? sp.GetService<SqliteVectorStoreOptions>() ?? new() { EmbeddingGenerator = sp.GetService<IEmbeddingGenerator>() }));
22+
Func<IServiceProvider, string> connectionStringProvider,
23+
Func<IServiceProvider, SqliteVectorStoreOptions>? optionsProvider = null,
24+
ServiceLifetime lifetime = ServiceLifetime.Singleton)
25+
=> AddVectorStore(services, serviceKey: null, connectionStringProvider, optionsProvider, lifetime);
26+
27+
/// <inheritdoc cref="AddVectorStore"/>
28+
public static IServiceCollection AddKeyedSqliteVectorStore(
29+
this IServiceCollection services,
30+
object serviceKey,
31+
Func<IServiceProvider, string> connectionStringProvider,
32+
Func<IServiceProvider, SqliteVectorStoreOptions>? optionsProvider = null,
33+
ServiceLifetime lifetime = ServiceLifetime.Singleton)
34+
{
35+
Verify.NotNull(serviceKey);
36+
37+
return AddVectorStore(services, serviceKey, connectionStringProvider, optionsProvider, lifetime);
38+
}
39+
40+
/// <summary>
41+
/// Registers a keyed <see cref="SqliteVectorStore"/> as <see cref="VectorStore"/>, with the specified connection string and service lifetime.
42+
/// </summary>
43+
/// <param name="services">The <see cref="IServiceCollection"/> to register the <see cref="VectorStore"/> on.</param>
44+
/// <param name="serviceKey">The key with which to associate the vector store.</param>
45+
/// <param name="connectionStringProvider">The connection string provider.</param>
46+
/// <param name="optionsProvider">Options provider to further configure the vector store.</param>
47+
/// <param name="lifetime">The service lifetime for the store. Defaults to <see cref="ServiceLifetime.Singleton"/>.</param>
48+
/// <returns>The service collection.</returns>
49+
private static IServiceCollection AddVectorStore(
50+
IServiceCollection services,
51+
object? serviceKey,
52+
Func<IServiceProvider, string> connectionStringProvider,
53+
Func<IServiceProvider, SqliteVectorStoreOptions>? optionsProvider,
54+
ServiceLifetime lifetime)
55+
{
56+
Verify.NotNull(services);
57+
Verify.NotNull(connectionStringProvider);
58+
59+
services.Add(new ServiceDescriptor(typeof(SqliteVectorStore), serviceKey, (sp, _) =>
60+
{
61+
var connectionString = connectionStringProvider(sp);
62+
var options = GetStoreOptions(sp, optionsProvider);
63+
return new SqliteVectorStore(connectionString, options);
64+
}, lifetime));
65+
66+
services.Add(new ServiceDescriptor(typeof(VectorStore), serviceKey,
67+
static (sp, key) => sp.GetRequiredKeyedService<SqliteVectorStore>(key), lifetime));
68+
69+
return services;
70+
}
71+
72+
/// <summary>
73+
/// Registers a <see cref="SqliteCollection{TKey, TRecord}"/> as <see cref="VectorStoreCollection{TKey, TRecord}"/>, with the specified connection string and service lifetime.
74+
/// </summary>
75+
/// <inheritdoc cref="AddCollection{TKey, TRecord}(IServiceCollection, object?, string, Func{IServiceProvider, string}, Func{IServiceProvider, SqliteCollectionOptions}?, ServiceLifetime)"/>
76+
public static IServiceCollection AddSqliteCollection<TKey, TRecord>(
77+
this IServiceCollection services,
78+
string collectionName,
79+
Func<IServiceProvider, string> connectionStringProvider,
80+
Func<IServiceProvider, SqliteCollectionOptions>? optionsProvider = null,
81+
ServiceLifetime lifetime = ServiceLifetime.Singleton)
82+
where TKey : notnull
83+
where TRecord : class
84+
=> AddCollection<TKey, TRecord>(services, serviceKey: null, collectionName, connectionStringProvider, optionsProvider, lifetime);
85+
86+
/// <inheritdoc cref="AddCollection{TKey, TRecord}(IServiceCollection, object?, string, Func{IServiceProvider, string}, Func{IServiceProvider, SqliteCollectionOptions}?, ServiceLifetime)"/>
87+
public static IServiceCollection AddKeyedSqliteCollection<TKey, TRecord>(
88+
this IServiceCollection services,
89+
object serviceKey,
90+
string collectionName,
91+
Func<IServiceProvider, string> connectionStringProvider,
92+
Func<IServiceProvider, SqliteCollectionOptions>? optionsProvider = null,
93+
ServiceLifetime lifetime = ServiceLifetime.Singleton)
94+
where TKey : notnull
95+
where TRecord : class
96+
{
97+
Verify.NotNull(serviceKey);
98+
99+
return AddCollection<TKey, TRecord>(services, serviceKey, collectionName, connectionStringProvider, optionsProvider, lifetime);
100+
}
33101

34102
/// <summary>
35-
/// Register a SQLite <see cref="VectorStoreCollection{TKey, TRecord}"/> and <see cref="IVectorSearchable{TRecord}"/> with the specified service ID.
36-
/// <see cref="SqliteConnection"/> instance will be initialized, connection will be opened and vector search extension with be loaded.
103+
/// Registers a keyed <see cref="SqliteCollection{TKey, TRecord}"/> as <see cref="VectorStoreCollection{TKey, TRecord}"/>, with the specified connection string and service lifetime.
37104
/// </summary>
38-
/// <typeparam name="TKey">The type of the key.</typeparam>
39-
/// <typeparam name="TRecord">The type of the record.</typeparam>
40105
/// <param name="services">The <see cref="IServiceCollection"/> to register the <see cref="VectorStoreCollection{TKey, TRecord}"/> on.</param>
106+
/// <param name="serviceKey">The key with which to associate the collection.</param>
41107
/// <param name="collectionName">The name of the collection.</param>
42-
/// <param name="connectionString">Connection string for <see cref="SqliteConnection"/>.</param>
43-
/// <param name="options">Optional options to further configure the <see cref="VectorStoreCollection{TKey, TRecord}"/>.</param>
44-
/// <param name="serviceId">An optional service id to use as the service key.</param>
45-
/// <returns>Service collection.</returns>
46-
public static IServiceCollection AddSqliteVectorStoreRecordCollection<TKey, TRecord>(
108+
/// <param name="connectionStringProvider">The connection string provider.</param>
109+
/// <param name="optionsProvider">Options provider to further configure the collection.</param>
110+
/// <param name="lifetime">The service lifetime for the store. Defaults to <see cref="ServiceLifetime.Singleton"/>.</param>
111+
/// <returns>The service collection.</returns>
112+
private static IServiceCollection AddCollection<TKey, TRecord>(
47113
this IServiceCollection services,
114+
object? serviceKey,
48115
string collectionName,
49-
string connectionString,
50-
SqliteCollectionOptions? options = default,
51-
string? serviceId = default)
116+
Func<IServiceProvider, string> connectionStringProvider,
117+
Func<IServiceProvider, SqliteCollectionOptions?>? optionsProvider = null,
118+
ServiceLifetime lifetime = ServiceLifetime.Singleton)
52119
where TKey : notnull
53120
where TRecord : class
54121
{
55-
services.AddKeyedSingleton<VectorStoreCollection<TKey, TRecord>>(
56-
serviceId,
57-
(sp, _) => (
58-
new SqliteCollection<TKey, TRecord>(
59-
connectionString,
60-
collectionName,
61-
options ?? sp.GetService<SqliteCollectionOptions>() ?? new()
62-
{
63-
EmbeddingGenerator = sp.GetService<IEmbeddingGenerator>()
64-
})
65-
as VectorStoreCollection<TKey, TRecord>)!);
66-
67-
AddVectorizedSearch<TKey, TRecord>(services, serviceId);
122+
Verify.NotNull(services);
123+
Verify.NotNullOrWhiteSpace(collectionName);
124+
Verify.NotNull(connectionStringProvider);
125+
126+
services.Add(new ServiceDescriptor(typeof(SqliteCollection<TKey, TRecord>), serviceKey, (sp, _) =>
127+
{
128+
var connectionString = connectionStringProvider(sp);
129+
var options = GetCollectionOptions(sp, optionsProvider);
130+
return new SqliteCollection<TKey, TRecord>(connectionString, collectionName, options);
131+
}, lifetime));
132+
133+
services.Add(new ServiceDescriptor(typeof(VectorStoreCollection<TKey, TRecord>), serviceKey,
134+
static (sp, key) => sp.GetRequiredKeyedService<SqliteCollection<TKey, TRecord>>(key), lifetime));
135+
136+
services.Add(new ServiceDescriptor(typeof(IVectorSearchable<TRecord>), serviceKey,
137+
static (sp, key) => sp.GetRequiredKeyedService<SqliteCollection<TKey, TRecord>>(key), lifetime));
138+
139+
// Once HybridSearch supports get implemented by SqliteCollection
140+
// we need to add IKeywordHybridSearchable abstraction here as well.
68141

69142
return services;
70143
}
71144

72145
/// <summary>
73-
/// Also register the <see cref="VectorStoreCollection{TKey, TRecord}"/> with the given <paramref name="serviceId"/> as a <see cref="IVectorSearchable{TRecord}"/>.
146+
/// Registers a <see cref="SqliteCollection{TKey, TRecord}"/> as <see cref="VectorStoreCollection{TKey, TRecord}"/>, with the specified connection string and service lifetime.
74147
/// </summary>
75-
/// <typeparam name="TKey">The type of the key.</typeparam>
76-
/// <typeparam name="TRecord">The type of the data model that the collection should contain.</typeparam>
77-
/// <param name="services">The service collection to register on.</param>
78-
/// <param name="serviceId">The service id that the registrations should use.</param>
79-
private static void AddVectorizedSearch<TKey, TRecord>(IServiceCollection services, string? serviceId) where TRecord : class
148+
/// <inheritdoc cref="AddCollection{TKey, TRecord}(IServiceCollection, object?, string, string, SqliteCollectionOptions?, ServiceLifetime)"/>/>
149+
public static IServiceCollection AddSqliteCollection<TKey, TRecord>(
150+
this IServiceCollection services,
151+
string collectionName,
152+
string connectionString,
153+
SqliteCollectionOptions? options = null,
154+
ServiceLifetime lifetime = ServiceLifetime.Singleton)
155+
where TKey : notnull
156+
where TRecord : class
157+
=> AddCollection<TKey, TRecord>(services, serviceKey: null, collectionName, connectionString, options, lifetime);
158+
159+
/// <inheritdoc cref="AddCollection{TKey, TRecord}(IServiceCollection, object?, string, string, SqliteCollectionOptions?, ServiceLifetime)"/>/>
160+
public static IServiceCollection AddKeyedSqliteCollection<TKey, TRecord>(
161+
this IServiceCollection services,
162+
object serviceKey,
163+
string collectionName,
164+
string connectionString,
165+
SqliteCollectionOptions? options = null,
166+
ServiceLifetime lifetime = ServiceLifetime.Singleton)
80167
where TKey : notnull
81-
=> services.AddKeyedSingleton<IVectorSearchable<TRecord>>(
82-
serviceId,
83-
(sp, _) => sp.GetRequiredKeyedService<VectorStoreCollection<TKey, TRecord>>(serviceId));
168+
where TRecord : class
169+
{
170+
Verify.NotNull(serviceKey);
171+
172+
return AddCollection<TKey, TRecord>(services, serviceKey, collectionName, connectionString, options, lifetime);
173+
}
174+
175+
/// <summary>
176+
/// Registers a keyed <see cref="SqliteCollection{TKey, TRecord}"/> as <see cref="VectorStoreCollection{TKey, TRecord}"/>, with the specified connection string and service lifetime.
177+
/// </summary>
178+
/// <param name="services">The <see cref="IServiceCollection"/> to register the <see cref="VectorStoreCollection{TKey, TRecord}"/> on.</param>
179+
/// <param name="serviceKey">The key with which to associate the collection.</param>
180+
/// <param name="collectionName">The name of the collection.</param>
181+
/// <param name="connectionString">The connection string.</param>
182+
/// <param name="options">Options to further configure the collection.</param>
183+
/// <param name="lifetime">The service lifetime for the store. Defaults to <see cref="ServiceLifetime.Singleton"/>.</param>
184+
/// <returns>The service collection.</returns>
185+
private static IServiceCollection AddCollection<TKey, TRecord>(
186+
this IServiceCollection services,
187+
object? serviceKey,
188+
string collectionName,
189+
string connectionString,
190+
SqliteCollectionOptions? options = null,
191+
ServiceLifetime lifetime = ServiceLifetime.Singleton)
192+
where TKey : notnull
193+
where TRecord : class
194+
{
195+
Verify.NotNullOrWhiteSpace(connectionString);
196+
197+
return AddCollection<TKey, TRecord>(services, serviceKey, collectionName, _ => connectionString, _ => options, lifetime);
198+
}
199+
200+
private static SqliteVectorStoreOptions? GetStoreOptions(IServiceProvider sp, Func<IServiceProvider, SqliteVectorStoreOptions?>? optionsProvider)
201+
{
202+
var options = optionsProvider?.Invoke(sp);
203+
if (options?.EmbeddingGenerator is not null)
204+
{
205+
return options; // The user has provided everything, there is nothing to change.
206+
}
207+
208+
var embeddingGenerator = sp.GetService<IEmbeddingGenerator>();
209+
return embeddingGenerator is null
210+
? options // There is nothing to change.
211+
: new(options, embeddingGenerator); // Create a brand new copy in order to avoid modifying the original options.
212+
}
213+
214+
private static SqliteCollectionOptions? GetCollectionOptions(IServiceProvider sp, Func<IServiceProvider, SqliteCollectionOptions?>? optionsProvider)
215+
{
216+
var options = optionsProvider?.Invoke(sp);
217+
if (options?.EmbeddingGenerator is not null)
218+
{
219+
return options; // The user has provided everything, there is nothing to change.
220+
}
221+
222+
var embeddingGenerator = sp.GetService<IEmbeddingGenerator>();
223+
return embeddingGenerator is null
224+
? options // There is nothing to change.
225+
: new(options, embeddingGenerator); // Create a brand new copy in order to avoid modifying the original options.
226+
}
84227
}

‎dotnet/src/Connectors/Connectors.Memory.SqliteVec/SqliteVectorStoreOptions.cs

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,19 @@ public sealed class SqliteVectorStoreOptions
1111
{
1212
internal static readonly SqliteVectorStoreOptions Default = new();
1313

14+
/// <summary>
15+
/// Initializes a new instance of the <see cref="SqliteVectorStoreOptions"/> class.
16+
/// </summary>
17+
public SqliteVectorStoreOptions()
18+
{
19+
}
20+
21+
internal SqliteVectorStoreOptions(SqliteVectorStoreOptions? source, IEmbeddingGenerator embeddingGenerator)
22+
{
23+
this.VectorVirtualTableName = source?.VectorVirtualTableName;
24+
this.EmbeddingGenerator = embeddingGenerator;
25+
}
26+
1427
/// <summary>
1528
/// Custom virtual table name to store vectors.
1629
/// </summary>

‎dotnet/src/Connectors/Connectors.SqliteVec.UnitTests/SqliteServiceCollectionExtensionsTests.cs

Lines changed: 0 additions & 79 deletions
This file was deleted.

0 commit comments

Comments
 (0)