added sql integration instead of json files

This commit is contained in:
PaulKoudelka 2026-07-28 15:25:10 +02:00
parent 2e06af13e0
commit c59e33e1b1
22 changed files with 700 additions and 156 deletions

View File

@ -2896,11 +2896,12 @@ UI_TEXT_CONTENT["AISTUDIO::COMPONENTS::SETTINGS::SETTINGSPANELEMBEDDINGS::T32678
-- Close
UI_TEXT_CONTENT["AISTUDIO::COMPONENTS::SETTINGS::SETTINGSPANELEMBEDDINGS::T3448155331"] = "Close"
-- Couldn't delete the embedding provider '{0}'. The issue: {1}. We can ignore this issue and delete the embedding provider anyway. Do you want to ignore it and delete this embedding provider?
UI_TEXT_CONTENT["AISTUDIO::COMPONENTS::SETTINGS::SETTINGSPANELEMBEDDINGS::T3703173892"] = "Couldn't delete the embedding provider '{0}'. The issue: {1}. We can ignore this issue and delete the embedding provider anyway. Do you want to ignore it and delete this embedding provider?"
-- This embedding provider is trusted by your organization for data source security checks. Local data can be sent to it without security warnings.
UI_TEXT_CONTENT["AISTUDIO::COMPONENTS::SETTINGS::SETTINGSPANELEMBEDDINGS::T3459188215"] = "This embedding provider is trusted by your organization for data source security checks. Local data can be sent to it without security warnings."
-- Couldn't delete the embedding provider '{0}'. The issue: {1}. We can ignore this issue and delete the embedding provider anyway. Do you want to ignore it and delete this embedding provider?
UI_TEXT_CONTENT["AISTUDIO::COMPONENTS::SETTINGS::SETTINGSPANELEMBEDDINGS::T3703173892"] = "Couldn't delete the embedding provider '{0}'. The issue: {1}. We can ignore this issue and delete the embedding provider anyway. Do you want to ignore it and delete this embedding provider?"
-- Actions
UI_TEXT_CONTENT["AISTUDIO::COMPONENTS::SETTINGS::SETTINGSPANELEMBEDDINGS::T3865031940"] = "Actions"
@ -4956,9 +4957,6 @@ UI_TEXT_CONTENT["AISTUDIO::DIALOGS::SETTINGS::SETTINGSDIALOGCODING::T585868261"]
-- You might configure different data sources. A data source can include one file, all files in a directory, or data from your company. Later, you can incorporate these data sources as needed when the AI requires this data to complete a certain task.
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::SETTINGS::SETTINGSDIALOGDATASOURCES::T1084943026"] = "You might configure different data sources. A data source can include one file, all files in a directory, or data from your company. Later, you can incorporate these data sources as needed when the AI requires this data to complete a certain task."
-- Are you sure you want to delete the data source '{0}' of type {1}?
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::SETTINGS::SETTINGSDIALOGDATASOURCES::T1096979935"] = "Are you sure you want to delete the data source '{0}' of type {1}?"
-- Automatic local data source refresh
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::SETTINGS::SETTINGSDIALOGDATASOURCES::T1208397349"] = "Automatic local data source refresh"
@ -5019,6 +5017,9 @@ UI_TEXT_CONTENT["AISTUDIO::DIALOGS::SETTINGS::SETTINGSDIALOGDATASOURCES::T303146
-- Edit
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::SETTINGS::SETTINGSDIALOGDATASOURCES::T3267849393"] = "Edit"
-- Are you sure you want to delete the data source '{0}' of type '{1}'?
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::SETTINGS::SETTINGSDIALOGDATASOURCES::T3337072977"] = "Are you sure you want to delete the data source '{0}' of type '{1}'?"
-- Add Data Source
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::SETTINGS::SETTINGSDIALOGDATASOURCES::T3387511033"] = "Add Data Source"
@ -5916,11 +5917,12 @@ UI_TEXT_CONTENT["AISTUDIO::LAYOUT::MAINLAYOUT::T1614176092"] = "Assistants"
-- Update
UI_TEXT_CONTENT["AISTUDIO::LAYOUT::MAINLAYOUT::T1847791252"] = "Update"
-- Data sync
UI_TEXT_CONTENT["AISTUDIO::LAYOUT::MAINLAYOUT::T1903948824"] = "Data sync"
-- Check for updates
UI_TEXT_CONTENT["AISTUDIO::LAYOUT::MAINLAYOUT::T1890416390"] = "Check for updates"
-- Data sync
UI_TEXT_CONTENT["AISTUDIO::LAYOUT::MAINLAYOUT::T1903948824"] = "Data sync"
-- Your settings were created by a newer AI Studio version. Changes in this session will not be saved. Please install or start the latest available update.
UI_TEXT_CONTENT["AISTUDIO::LAYOUT::MAINLAYOUT::T1988273622"] = "Your settings were created by a newer AI Studio version. Changes in this session will not be saved. Please install or start the latest available update."
@ -5969,11 +5971,12 @@ UI_TEXT_CONTENT["AISTUDIO::LAYOUT::MAINLAYOUT::T671981715"] = "Some embeddings f
-- Embeddings are running: {0} of {1} files are indexed.
UI_TEXT_CONTENT["AISTUDIO::LAYOUT::MAINLAYOUT::T714077986"] = "Embeddings are running: {0} of {1} files are indexed."
-- Embeddings
UI_TEXT_CONTENT["AISTUDIO::LAYOUT::MAINLAYOUT::T951463987"] = "Embeddings"
-- AI Studio does not recognize your settings-format version. Changes in this session will not be saved to avoid overwriting your settings. Please check for updates or contact support.
UI_TEXT_CONTENT["AISTUDIO::LAYOUT::MAINLAYOUT::T915412625"] = "AI Studio does not recognize your settings-format version. Changes in this session will not be saved to avoid overwriting your settings. Please check for updates or contact support."
-- Embeddings
UI_TEXT_CONTENT["AISTUDIO::LAYOUT::MAINLAYOUT::T951463987"] = "Embeddings"
-- Get coding and debugging support from an LLM.
UI_TEXT_CONTENT["AISTUDIO::PAGES::ASSISTANTS::T1243850917"] = "Get coding and debugging support from an LLM."
@ -6358,6 +6361,9 @@ UI_TEXT_CONTENT["AISTUDIO::PAGES::INFORMATION::T2173617769"] = "This library is
-- For the secure communication between the user interface and the runtime, we need to create certificates. This Rust library is great for this purpose.
UI_TEXT_CONTENT["AISTUDIO::PAGES::INFORMATION::T2174764529"] = "For the secure communication between the user interface and the runtime, we need to create certificates. This Rust library is great for this purpose."
-- SQLite stores local RAG indexing metadata and the file fingerprints used to decide whether local files need to be indexed again, without requiring a separate database server or a system SQLite installation.
UI_TEXT_CONTENT["AISTUDIO::PAGES::INFORMATION::T224267373"] = "SQLite stores local RAG indexing metadata and the file fingerprints used to decide whether local files need to be indexed again, without requiring a separate database server or a system SQLite installation."
-- OK
UI_TEXT_CONTENT["AISTUDIO::PAGES::INFORMATION::T2246359087"] = "OK"
@ -7312,6 +7318,33 @@ UI_TEXT_CONTENT["AISTUDIO::TOOLS::CONFIDENCESCHEMESEXTENSIONS::T3893997203"] = "
-- Trust all LLM providers
UI_TEXT_CONTENT["AISTUDIO::TOOLS::CONFIDENCESCHEMESEXTENSIONS::T4107860491"] = "Trust all LLM providers"
-- Reason
UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::EMBEDDINGSTATE::NOEMBEDDINGSTATECLIENT::T1093747001"] = "Reason"
-- Starting
UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::EMBEDDINGSTATE::NOEMBEDDINGSTATECLIENT::T1233211769"] = "Starting"
-- Unavailable
UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::EMBEDDINGSTATE::NOEMBEDDINGSTATECLIENT::T3662391977"] = "Unavailable"
-- Status
UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::EMBEDDINGSTATE::NOEMBEDDINGSTATECLIENT::T6222351"] = "Status"
-- Database path
UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::EMBEDDINGSTATE::SQLITEEMBEDDINGSTATECLIENTIMPLEMENTATION::T1100578143"] = "Database path"
-- Storage size
UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::EMBEDDINGSTATE::SQLITEEMBEDDINGSTATECLIENTIMPLEMENTATION::T1230141403"] = "Storage size"
-- Indexed files
UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::EMBEDDINGSTATE::SQLITEEMBEDDINGSTATECLIENTIMPLEMENTATION::T2235289713"] = "Indexed files"
-- Indexed data sources
UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::EMBEDDINGSTATE::SQLITEEMBEDDINGSTATECLIENTIMPLEMENTATION::T3524534748"] = "Indexed data sources"
-- Reported version
UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::EMBEDDINGSTATE::SQLITEEMBEDDINGSTATECLIENTIMPLEMENTATION::T3556099842"] = "Reported version"
-- Reason
UI_TEXT_CONTENT["AISTUDIO::TOOLS::DATABASES::NODATABASECLIENT::T1093747001"] = "Reason"

View File

@ -161,4 +161,4 @@
}
</MudButton>
</DialogActions>
</MudDialog>
</MudDialog>

View File

@ -341,6 +341,11 @@ public partial class EmbeddingProviderDialog : MSGComponentBase, ISecretId
}
}
/// <summary>
/// Stores a new tokenizer or deletes the existing one, based on the specified tokenizer path.
/// If the path is null or empty, any existing tokenizer is removed.
/// Otherwise, the tokenizer is stored at the specified path.
/// </summary>
private Task<TokenizerResponse> StoreOrDeleteTokenizerAsync()
{
var tokenizerId = TokenizerModelId.ForEmbeddingProviderId(this.DataId);

View File

@ -370,6 +370,11 @@ public partial class ProviderDialog : MSGComponentBase, ISecretId
}
}
/// <summary>
/// Stores a new tokenizer or deletes the existing one, based on the specified tokenizer path.
/// If the path is null or empty, any existing tokenizer is removed.
/// Otherwise, the tokenizer is stored at the specified path.
/// </summary>
private Task<TokenizerResponse> StoreOrDeleteTokenizerAsync()
{
var tokenizerId = TokenizerModelId.ForProviderId(this.DataId);

View File

@ -294,7 +294,7 @@ public partial class SettingsDialogDataSources : SettingsDialogBase
var dialogParameters = new DialogParameters<ConfirmDialog>
{
{ x => x.Message, string.Format(T("Are you sure you want to delete the data source '{0}' of type {1}?"), dataSource.Name, dataSource.Type.GetDisplayName()) },
{ x => x.Message, string.Format(T("Are you sure you want to delete the data source '{0}' of type '{1}'?"), dataSource.Name, dataSource.Type.GetDisplayName()) },
};
var dialogReference = await this.DialogService.ShowAsync<ConfirmDialog>(T("Delete Data Source"), dialogParameters, DialogOptions.FULLSCREEN);

View File

@ -104,4 +104,4 @@
</MudLayout>
</MudPaper>
<MudThemeProvider @ref="@this.themeProvider" Theme="@this.ColorTheme" IsDarkMode="@this.useDarkMode" />
<MudThemeProvider @ref="@this.themeProvider" Theme="@this.ColorTheme" IsDarkMode="@this.useDarkMode" />

View File

@ -559,4 +559,4 @@ public partial class MainLayout : LayoutComponentBase, IMessageBusReceiver, ILan
}
#endregion
}
}

View File

@ -290,6 +290,7 @@
}
<ThirdPartyComponent Name="Qdrant Edge" Developer="Andrey Vasnetsov, Tim Visée, Arnaud Gourlay, Luis Cossío, Ivan Pleshkov, Roman Titov, xzfc, JojiiOfficial & Open Source Community" LicenseName="Apache-2.0" LicenseUrl="https://github.com/qdrant/qdrant/blob/master/LICENSE" RepositoryUrl="https://github.com/qdrant/qdrant" UseCase="@T("Qdrant Edge is an embedded vector database and vector similarity search engine. We use it to realize local RAG—retrieval-augmented generation—within AI Studio. Thanks for the effort and great work that has been and is being put into Qdrant.")"/>
<ThirdPartyComponent Name="SQLite" Developer="SQLite developers & Open Source Community" LicenseName="Public Domain" LicenseUrl="https://www.sqlite.org/copyright.html" RepositoryUrl="https://www.sqlite.org/src/" UseCase="@T("SQLite stores local RAG indexing metadata and the file fingerprints used to decide whether local files need to be indexed again, without requiring a separate database server or a system SQLite installation.")"/>
<ThirdPartyComponent Name="axum" Developer="David Pedersen, Jonas Platte, tottoto, David Mládek, Yann Simon, Tobias Bieniek, Open Source Community & Tokio Project" LicenseName="MIT" LicenseUrl="https://github.com/tokio-rs/axum/blob/main/LICENSE" RepositoryUrl="https://github.com/tokio-rs/axum" UseCase="@T("Axum is used to provide the small internal service that connects the Rust runtime with the app's user interface. This lets both parts of AI Studio exchange information while the app is running.")"/>
<ThirdPartyComponent Name="axum-server" Developer="Eray Karatay, Adi Salimgereyev, daxpedda & Open Source Community" LicenseName="MIT" LicenseUrl="https://github.com/programatik29/axum-server/blob/master/LICENSE" RepositoryUrl="https://github.com/programatik29/axum-server" UseCase="@T("Axum server runs the internal axum service over a secure local connection. This helps AI Studio protect the communication between the Rust runtime and the user interface.")"/>
<ThirdPartyComponent Name="Rustls" Developer="Joe Birr-Pixton, Dirkjan Ochtman, Daniel McCarney, Brian Smith, Jacob Hoffman-Andrews, Jorge Aparicio & Open Source Community" LicenseName="MIT" LicenseUrl="https://github.com/rustls/rustls/blob/main/LICENSE-MIT" RepositoryUrl="https://github.com/rustls/rustls" UseCase="@T("Rustls helps secure the internal connection between the app's user interface and the Rust runtime. This protects the local communication that AI Studio needs while it is running.")"/>

View File

@ -116,4 +116,4 @@ public interface IProvider
/// <param name="token">>The cancellation token.</param>
/// <returns>>The list of transcription models.</returns>
public Task<ModelLoadResult> GetTranscriptionModels(string? apiKeyProvisional = null, CancellationToken token = default);
}
}

View File

@ -225,7 +225,7 @@ public static class LLMProvidersExtensions
/// <returns>The provider instance.</returns>
public static IProvider CreateProvider(this AIStudio.Settings.Provider providerSettings)
{
return providerSettings.UsedLLMProvider.CreateProvider(providerSettings.InstanceName, providerSettings.Host, providerSettings.Hostname, providerSettings.Model, providerSettings.HFInferenceProvider, providerSettings.Id, providerSettings.AdditionalJsonApiParameters, providerSettings.TokenizerPath, providerSettings.IsEnterpriseConfiguration);
return providerSettings.UsedLLMProvider.CreateProvider(providerSettings.InstanceName, providerSettings.Host, providerSettings.Hostname, providerSettings.Model, providerSettings.HFInferenceProvider, providerSettings.Id, providerSettings.AdditionalJsonApiParameters, tokenizerPath: providerSettings.TokenizerPath, providerSettings.IsEnterpriseConfiguration);
}
/// <summary>
@ -235,7 +235,7 @@ public static class LLMProvidersExtensions
/// <returns>The provider instance.</returns>
public static IProvider CreateProvider(this EmbeddingProvider embeddingProviderSettings)
{
return embeddingProviderSettings.UsedLLMProvider.CreateProvider(embeddingProviderSettings.Name, embeddingProviderSettings.Host, embeddingProviderSettings.Hostname, embeddingProviderSettings.Model, HFInferenceProvider.NONE, configuredProviderId: embeddingProviderSettings.Id, embeddingProviderSettings.TokenizerPath, isEnterpriseConfiguration: embeddingProviderSettings.IsEnterpriseConfiguration);
return embeddingProviderSettings.UsedLLMProvider.CreateProvider(embeddingProviderSettings.Name, embeddingProviderSettings.Host, embeddingProviderSettings.Hostname, embeddingProviderSettings.Model, HFInferenceProvider.NONE, configuredProviderId: embeddingProviderSettings.Id, tokenizerPath: embeddingProviderSettings.TokenizerPath, isEnterpriseConfiguration: embeddingProviderSettings.IsEnterpriseConfiguration);
}
/// <summary>
@ -245,10 +245,10 @@ public static class LLMProvidersExtensions
/// <returns>The provider instance.</returns>
public static IProvider CreateProvider(this TranscriptionProvider transcriptionProviderSettings)
{
return transcriptionProviderSettings.UsedLLMProvider.CreateProvider(transcriptionProviderSettings.Name, transcriptionProviderSettings.Host, transcriptionProviderSettings.Hostname, transcriptionProviderSettings.Model, HFInferenceProvider.NONE, configuredProviderId: transcriptionProviderSettings.Id, string.Empty, isEnterpriseConfiguration: transcriptionProviderSettings.IsEnterpriseConfiguration);
return transcriptionProviderSettings.UsedLLMProvider.CreateProvider(transcriptionProviderSettings.Name, transcriptionProviderSettings.Host, transcriptionProviderSettings.Hostname, transcriptionProviderSettings.Model, HFInferenceProvider.NONE, configuredProviderId: transcriptionProviderSettings.Id, tokenizerPath: string.Empty, isEnterpriseConfiguration: transcriptionProviderSettings.IsEnterpriseConfiguration);
}
private static IProvider CreateProvider(this LLMProviders provider, string instanceName, Host host, string hostname, Model model, HFInferenceProvider inferenceProvider, string configuredProviderId = "", string tokenizerPath = "", string expertProviderApiParameter = "", bool isEnterpriseConfiguration = false)
private static IProvider CreateProvider(this LLMProviders provider, string instanceName, Host host, string hostname, Model model, HFInferenceProvider inferenceProvider, string configuredProviderId = "", string expertProviderApiParameter = "", string tokenizerPath = "", bool isEnterpriseConfiguration = false)
{
try
{

View File

@ -52,4 +52,4 @@ public class NoProvider : IProvider
public IReadOnlyCollection<Capability> GetModelCapabilities(Model model) => [ Capability.NONE ];
#endregion
}
}

View File

@ -1,5 +1,6 @@
using AIStudio.Tools.Services;
using AIStudio.Tools.Databases.EmbeddingState;
using AIStudio.Tools.Databases.VectorStore;
using AIStudio.Tools.Services;
namespace AIStudio.Tools.Databases;
@ -56,6 +57,18 @@ public sealed partial class DatabaseClientProvider(RustService rustService, ILog
client.Status);
}
public async Task<EmbeddingStateClient> GetEmbeddingStateAsync(CancellationToken cancellationToken = default)
{
var client = await this.GetClientAsync(DatabaseRole.EMBEDDING_STATE, cancellationToken);
if (client is EmbeddingStateClient embeddingState)
return embeddingState;
return new NoEmbeddingStateClient(
client.Name,
"The configured database client does not support embedding state operations.",
client.Status);
}
private DatabaseClient CacheIfAvailable(DatabaseRole databaseRole, DatabaseClient client)
{
if (!client.IsAvailable)
@ -92,6 +105,7 @@ public sealed partial class DatabaseClientProvider(RustService rustService, ILog
private async Task<DatabaseClient> CreateClientAsync(DatabaseRole databaseRole, CancellationToken cancellationToken) => databaseRole switch
{
DatabaseRole.VECTOR_STORE => await QdrantEdgeClientImplementation.CreateAsync(rustService, this.logger, this.databaseClientLogger, cancellationToken),
DatabaseRole.EMBEDDING_STATE => await SqliteEmbeddingStateClientImplementation.CreateAsync(this.logger, this.databaseClientLogger, cancellationToken),
_ => new NoDatabaseClient(databaseRole.ToString(), "The requested database role is not supported.")
};

View File

@ -3,4 +3,5 @@ namespace AIStudio.Tools.Databases;
public enum DatabaseRole
{
VECTOR_STORE,
EMBEDDING_STATE,
}

View File

@ -0,0 +1,35 @@
using AIStudio.Tools.Services;
namespace AIStudio.Tools.Databases.EmbeddingState;
public abstract class EmbeddingStateClient(string name, string path) : DatabaseClient(name, path)
{
public abstract Task<DataSourceEmbeddingManifest> GetManifestAsync(string dataSourceId, CancellationToken token);
public abstract Task UpsertDataSourceAsync(
string dataSourceId,
string dataSourceName,
string dataSourceType,
string embeddingProviderId,
string embeddingSignature,
int vectorSize,
CancellationToken token);
public abstract Task UpdateVectorSizeAsync(string dataSourceId, int vectorSize, CancellationToken token);
public abstract Task UpsertFileAsync(string dataSourceId, EmbeddingStateFile file, CancellationToken token);
public abstract Task DeleteFileAsync(string dataSourceId, string filePath, CancellationToken token);
public abstract Task DeleteDataSourceAsync(string dataSourceId, CancellationToken token);
}
public sealed record EmbeddingStateFile(
string FilePath,
string FileName,
string RelativePath,
string Fingerprint,
long FileSize,
DateTime LastWriteUtc,
DateTime EmbeddedAtUtc,
int ChunkCount);

View File

@ -0,0 +1,49 @@
using AIStudio.Tools.PluginSystem;
using AIStudio.Tools.Services;
namespace AIStudio.Tools.Databases.EmbeddingState;
public sealed class NoEmbeddingStateClient(string name, string? unavailableReason, DatabaseClientStatus status = DatabaseClientStatus.UNAVAILABLE) : EmbeddingStateClient(name, string.Empty)
{
private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(NoEmbeddingStateClient).Namespace, nameof(NoEmbeddingStateClient));
public override DatabaseClientStatus Status => status;
public override async IAsyncEnumerable<(string Label, string Value)> GetDisplayInfo()
{
yield return (TB("Status"), status switch
{
DatabaseClientStatus.STARTING => TB("Starting"),
_ => TB("Unavailable")
});
if (!string.IsNullOrWhiteSpace(unavailableReason))
yield return (TB("Reason"), unavailableReason);
await Task.CompletedTask;
}
public override Task<DataSourceEmbeddingManifest> GetManifestAsync(string dataSourceId, CancellationToken token) =>
Task.FromResult(new DataSourceEmbeddingManifest());
public override Task UpsertDataSourceAsync(
string dataSourceId,
string dataSourceName,
string dataSourceType,
string embeddingProviderId,
string embeddingSignature,
int vectorSize,
CancellationToken token) => Task.CompletedTask;
public override Task UpdateVectorSizeAsync(string dataSourceId, int vectorSize, CancellationToken token) => Task.CompletedTask;
public override Task UpsertFileAsync(string dataSourceId, EmbeddingStateFile file, CancellationToken token) => Task.CompletedTask;
public override Task DeleteFileAsync(string dataSourceId, string filePath, CancellationToken token) => Task.CompletedTask;
public override Task DeleteDataSourceAsync(string dataSourceId, CancellationToken token) => Task.CompletedTask;
public override void Dispose()
{
}
}

View File

@ -0,0 +1,345 @@
using System.Globalization;
using AIStudio.Settings;
using AIStudio.Tools.PluginSystem;
using AIStudio.Tools.Services;
using Microsoft.Data.Sqlite;
namespace AIStudio.Tools.Databases.EmbeddingState;
public sealed class SqliteEmbeddingStateClientImplementation(
string name,
string databasePath,
string basePath,
string version) : EmbeddingStateClient(name, basePath)
{
private const string DATABASE_NAME = "SQLite";
private const string DATABASE_FILENAME = "rag-embedding-state.sqlite3";
private readonly string databasePath = databasePath;
private readonly string connectionString = new SqliteConnectionStringBuilder
{
DataSource = databasePath,
Mode = SqliteOpenMode.ReadWriteCreate,
Cache = SqliteCacheMode.Shared,
ForeignKeys = true,
DefaultTimeout = 30,
}.ToString();
private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(SqliteEmbeddingStateClientImplementation).Namespace, nameof(SqliteEmbeddingStateClientImplementation));
public override string CacheKey => $"{this.Name}:{this.databasePath}:{version}";
public static async Task<DatabaseClient> CreateAsync(
ILogger logger,
ILogger<DatabaseClient> databaseClientLogger,
CancellationToken cancellationToken)
{
if (string.IsNullOrWhiteSpace(SettingsManager.DataDirectory))
return CreateNoEmbeddingStateClient(DATABASE_NAME, "The application data directory is not available yet.", DatabaseClientStatus.STARTING, databaseClientLogger);
try
{
SQLitePCL.Batteries_V2.Init();
var basePath = Path.Combine(SettingsManager.DataDirectory, "databases", "sqlite");
Directory.CreateDirectory(basePath);
var databasePath = Path.Combine(basePath, DATABASE_FILENAME);
var client = new SqliteEmbeddingStateClientImplementation(DATABASE_NAME, databasePath, basePath, string.Empty);
await client.InitializeAsync(cancellationToken);
var version = await client.GetSqliteVersionAsync(cancellationToken);
client = new SqliteEmbeddingStateClientImplementation(DATABASE_NAME, databasePath, basePath, version);
client.SetLogger(databaseClientLogger);
return client;
}
catch (Exception exception)
{
logger.LogWarning(exception, "{DatabaseName} is not available. Indexed file fingerprints are disabled.", DATABASE_NAME);
return CreateNoEmbeddingStateClient(DATABASE_NAME, exception.Message, DatabaseClientStatus.UNAVAILABLE, databaseClientLogger);
}
}
public override async IAsyncEnumerable<(string Label, string Value)> GetDisplayInfo()
{
yield return (TB("Reported version"), version);
yield return (TB("Database path"), this.databasePath);
yield return (TB("Storage size"), this.GetStorageSize());
yield return (TB("Indexed data sources"), (await this.CountAsync("data_sources", CancellationToken.None)).ToString(CultureInfo.InvariantCulture));
yield return (TB("Indexed files"), (await this.CountAsync("embedded_files", CancellationToken.None)).ToString(CultureInfo.InvariantCulture));
}
public override async Task<DataSourceEmbeddingManifest> GetManifestAsync(string dataSourceId, CancellationToken token)
{
await using var connection = await this.OpenConnectionAsync(token);
var manifest = new DataSourceEmbeddingManifest();
await using (var command = CreateCommand(connection, """
SELECT embedding_provider_id, embedding_signature, vector_size
FROM data_sources
WHERE data_source_id = $dataSourceId
"""))
{
command.Parameters.AddWithValue("$dataSourceId", dataSourceId);
await using var reader = await command.ExecuteReaderAsync(token);
if (await reader.ReadAsync(token))
{
manifest.EmbeddingProviderId = reader.GetString(0);
manifest.EmbeddingSignature = reader.GetString(1);
manifest.VectorSize = reader.GetInt32(2);
}
}
await using (var command = CreateCommand(connection, """
SELECT file_path, fingerprint, file_size, last_write_utc, embedded_at_utc, chunk_count
FROM embedded_files
WHERE data_source_id = $dataSourceId
"""))
{
command.Parameters.AddWithValue("$dataSourceId", dataSourceId);
await using var reader = await command.ExecuteReaderAsync(token);
while (await reader.ReadAsync(token))
{
manifest.Files[reader.GetString(0)] = new EmbeddedFileRecord(
reader.GetString(1),
reader.GetInt64(2),
ParseUtc(reader.GetString(3)),
ParseUtc(reader.GetString(4)),
reader.GetInt32(5));
}
}
return manifest;
}
public override async Task UpsertDataSourceAsync(
string dataSourceId,
string dataSourceName,
string dataSourceType,
string embeddingProviderId,
string embeddingSignature,
int vectorSize,
CancellationToken token)
{
await using var connection = await this.OpenConnectionAsync(token);
await ExecuteNonQueryAsync(connection, """
INSERT INTO data_sources (
data_source_id,
data_source_name,
data_source_type,
embedding_provider_id,
embedding_signature,
vector_size,
updated_at_utc)
VALUES (
$dataSourceId,
$dataSourceName,
$dataSourceType,
$embeddingProviderId,
$embeddingSignature,
$vectorSize,
$updatedAtUtc)
ON CONFLICT(data_source_id) DO UPDATE SET
data_source_name = excluded.data_source_name,
data_source_type = excluded.data_source_type,
embedding_provider_id = excluded.embedding_provider_id,
embedding_signature = excluded.embedding_signature,
vector_size = excluded.vector_size,
updated_at_utc = excluded.updated_at_utc
""", token,
("$dataSourceId", dataSourceId),
("$dataSourceName", dataSourceName),
("$dataSourceType", dataSourceType),
("$embeddingProviderId", embeddingProviderId),
("$embeddingSignature", embeddingSignature),
("$vectorSize", vectorSize),
("$updatedAtUtc", ToUtcText(DateTime.UtcNow)));
}
public override async Task UpdateVectorSizeAsync(string dataSourceId, int vectorSize, CancellationToken token)
{
await using var connection = await this.OpenConnectionAsync(token);
await ExecuteNonQueryAsync(connection, """
UPDATE data_sources
SET vector_size = $vectorSize,
updated_at_utc = $updatedAtUtc
WHERE data_source_id = $dataSourceId
""", token,
("$dataSourceId", dataSourceId),
("$vectorSize", vectorSize),
("$updatedAtUtc", ToUtcText(DateTime.UtcNow)));
}
public override async Task UpsertFileAsync(string dataSourceId, EmbeddingStateFile file, CancellationToken token)
{
await using var connection = await this.OpenConnectionAsync(token);
await ExecuteNonQueryAsync(connection, """
INSERT INTO embedded_files (
data_source_id,
file_path,
file_name,
relative_path,
fingerprint,
file_size,
last_write_utc,
embedded_at_utc,
chunk_count)
VALUES (
$dataSourceId,
$filePath,
$fileName,
$relativePath,
$fingerprint,
$fileSize,
$lastWriteUtc,
$embeddedAtUtc,
$chunkCount)
ON CONFLICT(data_source_id, file_path) DO UPDATE SET
file_name = excluded.file_name,
relative_path = excluded.relative_path,
fingerprint = excluded.fingerprint,
file_size = excluded.file_size,
last_write_utc = excluded.last_write_utc,
embedded_at_utc = excluded.embedded_at_utc,
chunk_count = excluded.chunk_count
""", token,
("$dataSourceId", dataSourceId),
("$filePath", file.FilePath),
("$fileName", file.FileName),
("$relativePath", file.RelativePath),
("$fingerprint", file.Fingerprint),
("$fileSize", file.FileSize),
("$lastWriteUtc", ToUtcText(file.LastWriteUtc)),
("$embeddedAtUtc", ToUtcText(file.EmbeddedAtUtc)),
("$chunkCount", file.ChunkCount));
}
public override async Task DeleteFileAsync(string dataSourceId, string filePath, CancellationToken token)
{
await using var connection = await this.OpenConnectionAsync(token);
await ExecuteNonQueryAsync(connection, """
DELETE FROM embedded_files
WHERE data_source_id = $dataSourceId
AND file_path = $filePath
""", token,
("$dataSourceId", dataSourceId),
("$filePath", filePath));
}
public override async Task DeleteDataSourceAsync(string dataSourceId, CancellationToken token)
{
await using var connection = await this.OpenConnectionAsync(token);
await ExecuteNonQueryAsync(connection, """
DELETE FROM data_sources
WHERE data_source_id = $dataSourceId
""", token,
("$dataSourceId", dataSourceId));
}
public override void Dispose()
{
}
private async Task InitializeAsync(CancellationToken token)
{
await using var connection = await this.OpenConnectionAsync(token);
await ExecuteNonQueryAsync(connection, """
PRAGMA journal_mode=WAL;
PRAGMA foreign_keys=ON;
CREATE TABLE IF NOT EXISTS data_sources (
data_source_id TEXT PRIMARY KEY,
data_source_name TEXT NOT NULL,
data_source_type TEXT NOT NULL,
embedding_provider_id TEXT NOT NULL,
embedding_signature TEXT NOT NULL,
vector_size INTEGER NOT NULL DEFAULT 0,
updated_at_utc TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS embedded_files (
data_source_id TEXT NOT NULL,
file_path TEXT COLLATE NOCASE NOT NULL,
file_name TEXT NOT NULL,
relative_path TEXT NOT NULL,
fingerprint TEXT NOT NULL,
file_size INTEGER NOT NULL,
last_write_utc TEXT NOT NULL,
embedded_at_utc TEXT NOT NULL,
chunk_count INTEGER NOT NULL,
PRIMARY KEY (data_source_id, file_path),
FOREIGN KEY (data_source_id)
REFERENCES data_sources(data_source_id)
ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS idx_embedded_files_data_source
ON embedded_files(data_source_id);
""", token);
}
private async Task<string> GetSqliteVersionAsync(CancellationToken token)
{
await using var connection = await this.OpenConnectionAsync(token);
await using var command = CreateCommand(connection, "SELECT sqlite_version()");
var versionObject = await command.ExecuteScalarAsync(token);
return Convert.ToString(versionObject, CultureInfo.InvariantCulture) ?? string.Empty;
}
private async Task<long> CountAsync(string tableName, CancellationToken token)
{
await using var connection = await this.OpenConnectionAsync(token);
await using var command = CreateCommand(connection, $"SELECT COUNT(*) FROM {tableName}");
var countObject = await command.ExecuteScalarAsync(token);
return Convert.ToInt64(countObject, CultureInfo.InvariantCulture);
}
private async Task<SqliteConnection> OpenConnectionAsync(CancellationToken token)
{
var connection = new SqliteConnection(this.connectionString);
await connection.OpenAsync(token);
return connection;
}
private static SqliteCommand CreateCommand(SqliteConnection connection, string commandText)
{
var command = connection.CreateCommand();
command.CommandText = commandText;
return command;
}
private static async Task ExecuteNonQueryAsync(
SqliteConnection connection,
string commandText,
CancellationToken token,
params (string Name, object? Value)[] parameters)
{
await using var command = CreateCommand(connection, commandText);
foreach (var (name, value) in parameters)
command.Parameters.AddWithValue(name, value ?? DBNull.Value);
await command.ExecuteNonQueryAsync(token);
}
private static string ToUtcText(DateTime dateTime)
{
var utc = dateTime.Kind is DateTimeKind.Utc ? dateTime : dateTime.ToUniversalTime();
return utc.ToString("O", CultureInfo.InvariantCulture);
}
private static DateTime ParseUtc(string value)
{
return DateTime.TryParse(value, CultureInfo.InvariantCulture, DateTimeStyles.RoundtripKind, out var dateTime)
? dateTime.ToUniversalTime()
: DateTime.UnixEpoch;
}
private static NoEmbeddingStateClient CreateNoEmbeddingStateClient(string name, string? unavailableReason, DatabaseClientStatus status, ILogger<DatabaseClient> databaseClientLogger)
{
var client = new NoEmbeddingStateClient(name, unavailableReason, status);
client.SetLogger(databaseClientLogger);
return client;
}
}

View File

@ -17,7 +17,8 @@ public sealed class NoDatabaseClient(string name, string? unavailableReason, Dat
await Task.CompletedTask;
}
public override void Dispose()
{
}
}
}

View File

@ -171,10 +171,11 @@ public sealed record PluginConfigurationObject
public static async Task<bool> SyncManagedTokenizersAsync(Guid configPluginId, string pluginPath)
{
var wasConfigurationChanged = false;
var localSettingsManager = SettingsManagerAccess;
for (var i = 0; i < SETTINGS_MANAGER.ConfigurationData.Providers.Count; i++)
for (var i = 0; i < localSettingsManager.ConfigurationData.Providers.Count; i++)
{
var provider = SETTINGS_MANAGER.ConfigurationData.Providers[i];
var provider = localSettingsManager.ConfigurationData.Providers[i];
if (!provider.IsEnterpriseConfiguration || provider.EnterpriseConfigurationPluginId != configPluginId)
continue;
@ -182,13 +183,13 @@ public sealed record PluginConfigurationObject
if (syncedProvider == provider)
continue;
SETTINGS_MANAGER.ConfigurationData.Providers[i] = syncedProvider;
localSettingsManager.ConfigurationData.Providers[i] = syncedProvider;
wasConfigurationChanged = true;
}
for (var i = 0; i < SETTINGS_MANAGER.ConfigurationData.EmbeddingProviders.Count; i++)
for (var i = 0; i < localSettingsManager.ConfigurationData.EmbeddingProviders.Count; i++)
{
var provider = SETTINGS_MANAGER.ConfigurationData.EmbeddingProviders[i];
var provider = localSettingsManager.ConfigurationData.EmbeddingProviders[i];
if (!provider.IsEnterpriseConfiguration || provider.EnterpriseConfigurationPluginId != configPluginId)
continue;
@ -196,7 +197,7 @@ public sealed record PluginConfigurationObject
if (syncedProvider == provider)
continue;
SETTINGS_MANAGER.ConfigurationData.EmbeddingProviders[i] = syncedProvider;
localSettingsManager.ConfigurationData.EmbeddingProviders[i] = syncedProvider;
wasConfigurationChanged = true;
}
@ -345,13 +346,13 @@ public sealed record PluginConfigurationObject
{
if (item is Settings.Provider provider)
{
var deleteTokenizerResult = await RUST_SERVICE.DeleteTokenizer(TokenizerModelId.ForProvider(provider));
var deleteTokenizerResult = await RustService.DeleteTokenizer(TokenizerModelId.ForProvider(provider));
if (!deleteTokenizerResult.Success)
LOG.LogWarning("Failed to delete tokenizer for removed enterprise provider '{ProviderName}': {Issue}", provider.InstanceName, deleteTokenizerResult.Message);
}
else if (item is EmbeddingProvider embeddingProvider)
{
var deleteTokenizerResult = await RUST_SERVICE.DeleteTokenizer(TokenizerModelId.ForEmbeddingProvider(embeddingProvider));
var deleteTokenizerResult = await RustService.DeleteTokenizer(TokenizerModelId.ForEmbeddingProvider(embeddingProvider));
if (!deleteTokenizerResult.Success)
LOG.LogWarning("Failed to delete tokenizer for removed enterprise embedding provider '{ProviderName}': {Issue}", embeddingProvider.Name, deleteTokenizerResult.Message);
}
@ -406,7 +407,7 @@ public sealed record PluginConfigurationObject
{
if (string.IsNullOrWhiteSpace(configuredTokenizerPath))
{
var deleteResult = await RUST_SERVICE.DeleteTokenizer(modelId);
var deleteResult = await RustService.DeleteTokenizer(modelId);
if (!deleteResult.Success)
LOG.LogWarning("Failed to delete tokenizer for {LogName}: {Issue}", logName, deleteResult.Message);
@ -416,7 +417,7 @@ public sealed record PluginConfigurationObject
var resolvedPath = ResolvePluginTokenizerPath(configuredTokenizerPath, pluginPath);
if (resolvedPath is null)
{
var deleteResult = await RUST_SERVICE.DeleteTokenizer(modelId);
var deleteResult = await RustService.DeleteTokenizer(modelId);
if (!deleteResult.Success)
LOG.LogWarning("Failed to delete tokenizer after invalid path for {LogName}: {Issue}", logName, deleteResult.Message);
@ -424,10 +425,10 @@ public sealed record PluginConfigurationObject
return string.Empty;
}
var validateResult = await RUST_SERVICE.ValidateTokenizer(resolvedPath);
var validateResult = await RustService.ValidateTokenizer(resolvedPath);
if (!validateResult.Success)
{
var deleteResult = await RUST_SERVICE.DeleteTokenizer(modelId);
var deleteResult = await RustService.DeleteTokenizer(modelId);
if (!deleteResult.Success)
LOG.LogWarning("Failed to delete tokenizer after validation failure for {LogName}: {Issue}", logName, deleteResult.Message);
@ -435,7 +436,7 @@ public sealed record PluginConfigurationObject
return string.Empty;
}
var storeResult = await RUST_SERVICE.StoreTokenizer(modelId, resolvedPath);
var storeResult = await RustService.StoreTokenizer(modelId, resolvedPath);
if (!storeResult.Success)
{
LOG.LogWarning("Failed to store tokenizer for {LogName}. Path='{TokenizerPath}', issue='{Issue}'", logName, resolvedPath, storeResult.Message);

View File

@ -62,11 +62,6 @@ public sealed class FileEnumerationResult
public string LastError { get; set; } = string.Empty;
}
public sealed class PersistedEmbeddingState
{
public Dictionary<string, DataSourceEmbeddingManifest> DataSources { get; init; } = new(StringComparer.OrdinalIgnoreCase);
}
public sealed class DataSourceEmbeddingManifest
{
public string EmbeddingProviderId { get; set; } = string.Empty;

View File

@ -10,9 +10,20 @@ namespace AIStudio.Tools.Services;
public sealed partial class DataSourceEmbeddingService
{
private static readonly string[] ADDITIONAL_RAG_FILE_EXTENSIONS = ["csv", "tsv", "ods", "xlsm", "xlsb", "xla", "xlam"];
private const string OFFICE_LOCK_FILE_PREFIX = "~$";
private static readonly string[] RAG_DELIMITED_TABLE_FILE_EXTENSIONS = ["csv", "tsv"];
private static readonly string[] RAG_SPREADSHEET_FILE_EXTENSIONS = ["ods", "xlsm", "xlsb"];
private static readonly string[] RAG_SPREADSHEET_ADD_IN_FILE_EXTENSIONS = ["xla", "xlam"];
private static readonly string[] SKIPPED_RAG_FILE_EXTENSIONS = ["lnk"];
private enum RagFileIndexingDecision
{
INDEXABLE,
EXCLUDED,
UNSUPPORTED,
}
private async IAsyncEnumerable<string> StreamEmbeddingChunksAsync(string filePath, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token)
{
if (this.IsImageFilePath(filePath))
@ -67,14 +78,21 @@ public sealed partial class DataSourceEmbeddingService
switch (dataSource)
{
case DataSourceLocalFile localFile when File.Exists(localFile.FilePath):
if (this.IsSupportedRagFilePath(localFile.FilePath))
var file = new FileInfo(localFile.FilePath);
switch (this.GetRagFileIndexingDecision(file))
{
result.Files.Add(new FileInfo(localFile.FilePath));
}
else
{
result.FailedFiles = 1;
result.LastError = $"The selected file '{localFile.FilePath}' is not supported for background embeddings.";
case RagFileIndexingDecision.INDEXABLE:
result.Files.Add(file);
break;
case RagFileIndexingDecision.EXCLUDED:
logger.LogDebug("Skipping excluded file '{FilePath}' while indexing.", file.FullName);
break;
default:
result.FailedFiles = 1;
result.LastError = $"The selected file '{localFile.FilePath}' is not supported for background embeddings.";
break;
}
return result;
@ -141,14 +159,25 @@ public sealed partial class DataSourceEmbeddingService
continue;
}
if (!this.IsSupportedRagFilePath(fileInfo.FullName))
continue;
switch (this.GetRagFileIndexingDecision(fileInfo))
{
case RagFileIndexingDecision.INDEXABLE:
result.Files.Add(fileInfo);
break;
result.Files.Add(fileInfo);
case RagFileIndexingDecision.EXCLUDED:
logger.LogDebug("Skipping excluded file '{FilePath}' while indexing.", fileInfo.FullName);
break;
}
}
foreach (var subDirectory in subDirectories)
{
if (this.IsSkippedRagDirectory(subDirectory))
continue;
pendingDirectories.Push(subDirectory);
}
}
}
@ -174,11 +203,59 @@ public sealed partial class DataSourceEmbeddingService
private bool IsSupportedRagFilePath(string filePath)
{
var extension = Path.GetExtension(filePath).TrimStart('.');
if (SKIPPED_RAG_FILE_EXTENSIONS.Contains(extension, StringComparer.OrdinalIgnoreCase))
return false;
return FileTypes.IsAllowedPath(filePath, FileTypes.DOCUMENT, FileTypes.IMAGE)
|| ADDITIONAL_RAG_FILE_EXTENSIONS.Contains(extension, StringComparer.OrdinalIgnoreCase);
|| RAG_DELIMITED_TABLE_FILE_EXTENSIONS.Contains(extension, StringComparer.OrdinalIgnoreCase)
|| RAG_SPREADSHEET_FILE_EXTENSIONS.Contains(extension, StringComparer.OrdinalIgnoreCase)
|| RAG_SPREADSHEET_ADD_IN_FILE_EXTENSIONS.Contains(extension, StringComparer.OrdinalIgnoreCase);
}
private RagFileIndexingDecision GetRagFileIndexingDecision(FileInfo file)
{
if (this.IsSkippedRagFile(file))
return RagFileIndexingDecision.EXCLUDED;
return this.IsSupportedRagFilePath(file.FullName)
? RagFileIndexingDecision.INDEXABLE
: RagFileIndexingDecision.UNSUPPORTED;
}
private bool IsSkippedRagFile(FileInfo file)
{
var extension = file.Extension.TrimStart('.');
if (SKIPPED_RAG_FILE_EXTENSIONS.Contains(extension, StringComparer.OrdinalIgnoreCase))
return true;
if (file.Name.StartsWith(OFFICE_LOCK_FILE_PREFIX, StringComparison.Ordinal))
return true;
try
{
return file.Attributes.HasFlag(FileAttributes.ReparsePoint)
|| file.Attributes.HasFlag(FileAttributes.Offline)
|| file.Attributes.HasFlag(FileAttributes.Temporary)
|| file.Attributes.HasFlag(FileAttributes.System);
}
catch (Exception exception)
{
logger.LogWarning(exception, "Cannot inspect file '{FilePath}' while indexing.", file.FullName);
return true;
}
}
private bool IsSkippedRagDirectory(string path)
{
try
{
var directory = new DirectoryInfo(path);
return directory.Attributes.HasFlag(FileAttributes.ReparsePoint)
|| directory.Attributes.HasFlag(FileAttributes.Offline)
|| directory.Attributes.HasFlag(FileAttributes.System);
}
catch (Exception exception)
{
logger.LogWarning(exception, "Cannot inspect directory '{DirectoryPath}' while indexing.", path);
return true;
}
}
private string BuildImageIndexText(string filePath)
@ -212,9 +289,17 @@ public sealed partial class DataSourceEmbeddingService
embeddingProvider.TokenizerPath);
}
private string BuildFingerprint(FileInfo file)
private async Task<string> BuildFingerprintAsync(FileInfo file, CancellationToken token)
{
var fingerprintSource = $"{file.FullName}|{file.Length}|{file.LastWriteTimeUtc.Ticks}";
await using var stream = new FileStream(
file.FullName,
FileMode.Open,
FileAccess.Read,
FileShare.ReadWrite | FileShare.Delete,
1024 * 128,
FileOptions.Asynchronous | FileOptions.SequentialScan);
var contentHash = await SHA256.HashDataAsync(stream, token);
var fingerprintSource = $"{file.FullName}|{Convert.ToHexString(contentHash)}";
var bytes = SHA256.HashData(Encoding.UTF8.GetBytes(fingerprintSource));
return Convert.ToHexString(bytes);
}

View File

@ -1,90 +1,27 @@
using System.Text.Json;
using AIStudio.Settings;
using AIStudio.Tools.Databases.EmbeddingState;
using AIStudio.Tools.Databases.VectorStore;
namespace AIStudio.Tools.Services;
public sealed partial class DataSourceEmbeddingService
{
private const string STATE_FILENAME = "rag-embedding-state.json";
private readonly JsonSerializerOptions jsonOptions = new()
private async Task ResetPersistedStateAsync(
string dataSourceName,
string dataSourceId,
VectorStoreClient? vectorStore,
EmbeddingStateClient? embeddingState,
CancellationToken token)
{
WriteIndented = true,
};
await this.DeleteCollectionAsync(this.GetCollectionName(dataSourceName, dataSourceId), vectorStore, token);
private async Task EnsureStateLoadedAsync(CancellationToken token)
{
if (this.stateLoaded)
embeddingState ??= await databaseClientProvider.GetEmbeddingStateAsync(token);
if (!embeddingState.IsAvailable)
{
logger.LogWarning("Could not delete SQLite embedding state for data source '{DataSourceId}' because the embedding state database '{DatabaseName}' is unavailable.", dataSourceId, embeddingState.Name);
return;
await this.stateLock.WaitAsync(token);
try
{
if (this.stateLoaded)
return;
var statePath = this.GetStatePath();
if (!string.IsNullOrWhiteSpace(statePath) && File.Exists(statePath))
{
var json = await File.ReadAllTextAsync(statePath, token);
var persistedState = JsonSerializer.Deserialize<PersistedEmbeddingState>(json, this.jsonOptions);
this.manifests = persistedState?.DataSources ?? new Dictionary<string, DataSourceEmbeddingManifest>(StringComparer.OrdinalIgnoreCase);
}
this.stateLoaded = true;
}
finally
{
this.stateLock.Release();
}
}
private async Task<DataSourceEmbeddingManifest> GetManifestAsync(string dataSourceId, CancellationToken token)
{
await this.EnsureStateLoadedAsync(token);
if (this.manifests.TryGetValue(dataSourceId, out var manifest))
return manifest;
manifest = new DataSourceEmbeddingManifest();
this.manifests[dataSourceId] = manifest;
return manifest;
}
private async Task SaveStateAsync(CancellationToken token)
{
var statePath = this.GetStatePath();
if (string.IsNullOrWhiteSpace(statePath))
return;
var directory = Path.GetDirectoryName(statePath);
if (!string.IsNullOrWhiteSpace(directory))
Directory.CreateDirectory(directory);
var persistedState = new PersistedEmbeddingState
{
DataSources = this.manifests
};
var json = JsonSerializer.Serialize(persistedState, this.jsonOptions);
await File.WriteAllTextAsync(statePath, json, token);
}
private async Task ResetPersistedStateAsync(string dataSourcename, string dataSourceId, VectorStoreClient? vectorStore, CancellationToken token)
{
await this.EnsureStateLoadedAsync(token);
this.manifests.Remove(dataSourceId);
await this.DeleteCollectionAsync(this.GetCollectionName(dataSourcename, dataSourceId), vectorStore, token);
await this.SaveStateAsync(token);
await embeddingState.DeleteDataSourceAsync(dataSourceId, token);
logger.LogInformation("Reset persisted embedding state for data source '{DataSourceId}'.", dataSourceId);
}
private string GetStatePath()
{
if (string.IsNullOrWhiteSpace(SettingsManager.ConfigDirectory))
return string.Empty;
return Path.Combine(SettingsManager.ConfigDirectory, STATE_FILENAME);
}
}

View File

@ -6,6 +6,7 @@ using AIStudio.Provider;
using AIStudio.Settings;
using AIStudio.Settings.DataModel;
using AIStudio.Tools.Databases;
using AIStudio.Tools.Databases.EmbeddingState;
using AIStudio.Tools.Databases.VectorStore;
using AIStudio.Tools.PluginSystem;
@ -22,13 +23,9 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
private readonly Channel<string> queue = Channel.CreateUnbounded<string>();
private readonly ConcurrentDictionary<string, byte> queuedIds = new(StringComparer.OrdinalIgnoreCase);
private readonly ConcurrentDictionary<string, DataSourceEmbeddingStatus> statuses = new(StringComparer.OrdinalIgnoreCase);
private readonly SemaphoreSlim stateLock = new(1, 1);
private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(DataSourceEmbeddingService).Namespace, nameof(DataSourceEmbeddingService));
private Dictionary<string, DataSourceEmbeddingManifest> manifests = new(StringComparer.OrdinalIgnoreCase);
private bool stateLoaded;
public IReadOnlyList<DataSourceEmbeddingStatus> GetStatuses()
{
return this.statuses.Values
@ -114,7 +111,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
this.RemoveWatcher(dataSource.Id);
this.statuses.TryRemove(dataSource.Id, out _);
await this.ResetPersistedStateAsync(dataSource.Name, dataSource.Id, null, CancellationToken.None);
await this.ResetPersistedStateAsync(dataSource.Name, dataSource.Id, null, null, CancellationToken.None);
this.PublishStatusChanged();
}
@ -152,16 +149,15 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
public override void Dispose()
{
this.DisposeWatchers();
this.stateLock.Dispose();
base.Dispose();
}
private async Task ProcessDataSourceAsync(IDataSource dataSource, CancellationToken token)
{
await this.EnsureStateLoadedAsync(token);
logger.LogInformation("Starting background embeddings for data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id);
var vectorStore = await databaseClientProvider.GetVectorStoreAsync(token);
var embeddingState = await databaseClientProvider.GetEmbeddingStateAsync(token);
if (!vectorStore.IsAvailable)
{
@ -174,6 +170,17 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
return;
}
if (!embeddingState.IsAvailable)
{
logger.LogWarning(
"Skipping background embeddings for data source '{DataSourceName}' ({DataSourceId}) because the database client '{DatabaseName}' is unavailable.",
dataSource.Name,
dataSource.Id,
embeddingState.Name);
this.UpsertStatus(this.GetFallbackStatus(dataSource, "The SQLite embedding state database is not available."));
return;
}
if (!this.TryResolveEmbeddingProvider(dataSource, out var embeddingProvider))
{
this.UpsertStatus(this.GetFallbackStatus(dataSource, "The selected embedding provider is not available."));
@ -189,7 +196,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
dataSource.Id);
var collectionName = this.GetCollectionName(dataSource.Name, dataSource.Id);
var manifest = await this.EnsureCompatibleManifestAsync(dataSource, embeddingProvider, collectionName, vectorStore, token);
var manifest = await this.EnsureCompatibleManifestAsync(dataSource, embeddingProvider, collectionName, vectorStore, embeddingState, token);
var inputFiles = this.GetInputFiles(dataSource);
var indexedFiles = inputFiles.Files;
var totalFiles = indexedFiles.Count + inputFiles.FailedFiles;
@ -202,8 +209,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
inputFiles.FailedFiles,
collectionName);
await this.RemoveMissingFileEmbeddingsAsync(vectorStore, dataSource, collectionName, manifest, indexedFiles, token);
await this.SaveStateAsync(token);
await this.RemoveMissingFileEmbeddingsAsync(vectorStore, embeddingState, dataSource, collectionName, manifest, indexedFiles, token);
this.UpsertStatus(this.CreateStatus(
dataSource,
@ -223,7 +229,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
{
token.ThrowIfCancellationRequested();
var fingerprint = this.BuildFingerprint(file);
var fingerprint = await this.BuildFingerprintAsync(file, token);
if (manifest.Files.TryGetValue(file.FullName, out var existingRecord) &&
string.Equals(existingRecord.Fingerprint, fingerprint, StringComparison.Ordinal))
{
@ -249,14 +255,27 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
skippedFiles + completedFiles + 1,
totalFiles);
var startedAtUtc = DateTime.UtcNow;
var chunkCount = await this.IndexOneFileAsync(vectorStore, dataSource, file, fingerprint, embeddingProvider, provider, manifest, token);
manifest.Files[file.FullName] = new EmbeddedFileRecord(
var chunkCount = await this.IndexOneFileAsync(embeddingState, vectorStore, dataSource, file, fingerprint, embeddingProvider, provider, manifest, token);
var embeddedAtUtc = DateTime.UtcNow;
var record = new EmbeddedFileRecord(
fingerprint,
file.Length,
file.LastWriteTimeUtc,
DateTime.UtcNow,
embeddedAtUtc,
chunkCount);
await this.SaveStateAsync(token);
await embeddingState.UpsertFileAsync(
dataSource.Id,
new EmbeddingStateFile(
file.FullName,
file.Name,
this.TryGetRelativePath(dataSource, file),
fingerprint,
file.Length,
file.LastWriteTimeUtc,
embeddedAtUtc,
chunkCount),
token);
manifest.Files[file.FullName] = record;
completedFiles++;
logger.LogInformation(
"Embedded file '{FilePath}' for data source '{DataSourceName}' ({DataSourceId}) successfully. Chunks={ChunkCount}, DurationMs={DurationMs}.",
@ -272,7 +291,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
lastError = exception.Message;
manifest.Files.Remove(file.FullName);
await this.DeleteFilePointsAsync(vectorStore, collectionName, file.FullName, token);
await this.SaveStateAsync(token);
await embeddingState.DeleteFileAsync(dataSource.Id, file.FullName, token);
logger.LogWarning(exception, "Failed to embed file '{FilePath}' for data source '{DataSourceName}'.", file.FullName, dataSource.Name);
this.UpsertStatus(this.CreateStatus(dataSource, DataSourceEmbeddingState.RUNNING, totalFiles, skippedFiles + completedFiles, failedFiles, file.Name, exception.Message));
@ -290,6 +309,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
}
private async Task<int> IndexOneFileAsync(
EmbeddingStateClient embeddingState,
VectorStoreClient vectorStore,
IDataSource dataSource,
FileInfo file,
@ -315,11 +335,11 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
totalChunkCount++;
if (batch.Count >= EMBEDDING_BATCH_SIZE)
await this.FlushBatchAsync(vectorStore, dataSource, file, fingerprint, embeddingProvider, provider, manifest, collectionName, batch, token);
await this.FlushBatchAsync(embeddingState, vectorStore, dataSource, file, fingerprint, embeddingProvider, provider, manifest, collectionName, batch, token);
}
if (batch.Count > 0)
await this.FlushBatchAsync(vectorStore, dataSource, file, fingerprint, embeddingProvider, provider, manifest, collectionName, batch, token);
await this.FlushBatchAsync(embeddingState, vectorStore, dataSource, file, fingerprint, embeddingProvider, provider, manifest, collectionName, batch, token);
if (totalChunkCount == 0)
throw new InvalidOperationException($"The file '{file.Name}' did not yield any text chunks.");
@ -335,6 +355,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
}
private async Task FlushBatchAsync(
EmbeddingStateClient embeddingState,
VectorStoreClient vectorStore,
IDataSource dataSource,
FileInfo file,
@ -369,7 +390,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
{
manifest.VectorSize = vectorSize;
await this.EnsureCollectionExistsAsync(vectorStore, collectionName, vectorSize, token);
await this.SaveStateAsync(token);
await embeddingState.UpdateVectorSizeAsync(dataSource.Id, vectorSize, token);
logger.LogInformation(
"Created embedding collection '{CollectionName}' with vector size {VectorSize} for data source '{DataSourceName}' ({DataSourceId}).",
collectionName,
@ -484,10 +505,16 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
return embeddingProvider != default && embeddingProvider.UsedLLMProvider is not LLMProviders.NONE;
}
private async Task<DataSourceEmbeddingManifest> EnsureCompatibleManifestAsync(IDataSource dataSource, EmbeddingProvider embeddingProvider, string collectionName, VectorStoreClient vectorStore, CancellationToken token)
private async Task<DataSourceEmbeddingManifest> EnsureCompatibleManifestAsync(
IDataSource dataSource,
EmbeddingProvider embeddingProvider,
string collectionName,
VectorStoreClient vectorStore,
EmbeddingStateClient embeddingState,
CancellationToken token)
{
var embeddingSignature = this.BuildEmbeddingSignature(embeddingProvider);
var manifest = await this.GetManifestAsync(dataSource.Id, token);
var manifest = await embeddingState.GetManifestAsync(dataSource.Id, token);
if (!string.Equals(manifest.EmbeddingSignature, embeddingSignature, StringComparison.Ordinal))
{
@ -496,8 +523,8 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
dataSource.Name,
dataSource.Id,
collectionName);
await this.ResetPersistedStateAsync(dataSource.Name, dataSource.Id, vectorStore, token);
manifest = await this.GetManifestAsync(dataSource.Id, token);
await this.ResetPersistedStateAsync(dataSource.Name, dataSource.Id, vectorStore, embeddingState, token);
manifest = await embeddingState.GetManifestAsync(dataSource.Id, token);
}
if (!string.Equals(manifest.EmbeddingProviderId, embeddingProvider.Id, StringComparison.OrdinalIgnoreCase) ||
@ -505,14 +532,23 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
{
manifest.EmbeddingProviderId = embeddingProvider.Id;
manifest.EmbeddingSignature = embeddingSignature;
await this.SaveStateAsync(token);
}
await embeddingState.UpsertDataSourceAsync(
dataSource.Id,
dataSource.Name,
dataSource.Type.ToString(),
manifest.EmbeddingProviderId,
manifest.EmbeddingSignature,
manifest.VectorSize,
token);
return manifest;
}
private async Task RemoveMissingFileEmbeddingsAsync(
VectorStoreClient vectorStore,
EmbeddingStateClient embeddingState,
IDataSource dataSource,
string collectionName,
DataSourceEmbeddingManifest manifest,
@ -526,6 +562,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
foreach (var removedFilePath in manifest.Files.Keys.Except(existingPaths, StringComparer.OrdinalIgnoreCase).ToList())
{
await this.DeleteFilePointsAsync(vectorStore, collectionName, removedFilePath, token);
await embeddingState.DeleteFileAsync(dataSource.Id, removedFilePath, token);
manifest.Files.Remove(removedFilePath);
logger.LogInformation(
"Removed stale embeddings for deleted file '{FilePath}' from data source '{DataSourceName}' ({DataSourceId}).",