mirror of
https://github.com/MindWorkAI/AI-Studio.git
synced 2026-10-06 08:09:40 +00:00
current main is now successfully merged into branch
This commit is contained in:
1 parent
00b6c98f3b
commit
ea87d79ca9
32 files changed
+607
-478
No files matched your search
@@ -1,6 +1,6 @@
|
||||
namespace AIStudio.Tools.Databases;
|
||||
|
||||
public abstract class EmbeddingStore(string name, string path)
|
||||
public abstract class DatabaseClient(string name, string path)
|
||||
{
|
||||
public string Name => name;
|
||||
|
||||
@@ -12,7 +12,7 @@ public abstract class EmbeddingStore(string name, string path)
|
||||
|
||||
private string Path => path;
|
||||
|
||||
private ILogger<EmbeddingStore>? logger;
|
||||
private ILogger<DatabaseClient>? logger;
|
||||
|
||||
public abstract IAsyncEnumerable<(string Label, string Value)> GetDisplayInfo();
|
||||
|
||||
@@ -50,7 +50,7 @@ public abstract class EmbeddingStore(string name, string path)
|
||||
return $"{convertedSize:0.##} {suffixes[suffixIndex]}";
|
||||
}
|
||||
|
||||
public void SetLogger(ILogger<EmbeddingStore> logService)
|
||||
public void SetLogger(ILogger<DatabaseClient> logService)
|
||||
{
|
||||
this.logger = logService;
|
||||
}
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
using AIStudio.Tools.Rust;
|
||||
using AIStudio.Tools.Databases.VectorStore;
|
||||
using AIStudio.Tools.Rust;
|
||||
|
||||
namespace AIStudio.Tools.Databases;
|
||||
|
||||
@@ -6,7 +7,7 @@ public sealed partial class DatabaseClientProvider
|
||||
{
|
||||
private async Task<DatabaseClient> CreateQdrantClientAsync(CancellationToken cancellationToken)
|
||||
{
|
||||
var qdrantInfo = await this.rustService.GetQdrantInfo(cancellationToken);
|
||||
var qdrantInfo = await rustService.GetQdrantInfo(cancellationToken);
|
||||
if (qdrantInfo.Status is QdrantStatus.STARTING)
|
||||
{
|
||||
return this.CreateNoDatabaseClient(
|
||||
|
||||
@@ -1,15 +1,14 @@
|
||||
using AIStudio.Tools.Databases.Qdrant;
|
||||
using AIStudio.Tools.Rust;
|
||||
using AIStudio.Tools.Services;
|
||||
using AIStudio.Tools.Databases.VectorStore;
|
||||
|
||||
namespace AIStudio.Tools.Databases;
|
||||
|
||||
public sealed class EmbeddingStoreProvider(RustService rustService, ILoggerFactory loggerFactory) : IDisposable
|
||||
public sealed partial class DatabaseClientProvider(RustService rustService, ILoggerFactory loggerFactory) : IDisposable
|
||||
{
|
||||
private readonly Dictionary<DatabaseRole, EmbeddingStore> clients = new();
|
||||
private readonly Dictionary<DatabaseRole, DatabaseClient> clients = new();
|
||||
private readonly Dictionary<DatabaseRole, SemaphoreSlim> locks = new();
|
||||
private readonly Lock locksLock = new();
|
||||
private readonly ILogger<EmbeddingStoreProvider> logger = loggerFactory.CreateLogger<EmbeddingStoreProvider>();
|
||||
private readonly ILogger<DatabaseClientProvider> logger = loggerFactory.CreateLogger<DatabaseClientProvider>();
|
||||
private readonly ILogger<DatabaseClient> databaseClientLogger = loggerFactory.CreateLogger<DatabaseClient>();
|
||||
|
||||
public async Task<DatabaseClient> GetClientAsync(DatabaseRole databaseRole, CancellationToken cancellationToken = default)
|
||||
@@ -45,6 +44,18 @@ public sealed class EmbeddingStoreProvider(RustService rustService, ILoggerFacto
|
||||
}
|
||||
}
|
||||
|
||||
public async Task<IVectorStoreClient> GetVectorStoreAsync(CancellationToken cancellationToken = default)
|
||||
{
|
||||
var client = await this.GetClientAsync(DatabaseRole.VECTOR_STORE, cancellationToken);
|
||||
if (client is IVectorStoreClient vectorStore)
|
||||
return vectorStore;
|
||||
|
||||
return new NoVectorStoreClient(
|
||||
client.Name,
|
||||
"The configured database client does not support vector store operations.",
|
||||
client.Status);
|
||||
}
|
||||
|
||||
private DatabaseClient CacheIfAvailable(DatabaseRole databaseRole, DatabaseClient client)
|
||||
{
|
||||
if (!client.IsAvailable)
|
||||
@@ -84,79 +95,6 @@ public sealed class EmbeddingStoreProvider(RustService rustService, ILoggerFacto
|
||||
_ => new NoDatabaseClient(databaseRole.ToString(), "The requested database role is not supported.")
|
||||
};
|
||||
|
||||
private async Task<DatabaseClient> CreateQdrantClientAsync(CancellationToken cancellationToken)
|
||||
{
|
||||
var qdrantInfo = await rustService.GetQdrantInfo(cancellationToken);
|
||||
if (qdrantInfo.Status is QdrantStatus.STARTING)
|
||||
{
|
||||
return this.CreateNoDatabaseClient(
|
||||
"Qdrant",
|
||||
"Qdrant is starting. Details will appear shortly.",
|
||||
DatabaseClientStatus.STARTING);
|
||||
}
|
||||
|
||||
if (!qdrantInfo.IsAvailable || qdrantInfo.Status is QdrantStatus.UNAVAILABLE)
|
||||
{
|
||||
var reason = qdrantInfo.UnavailableReason ?? "unknown";
|
||||
this.logger.LogWarning("Qdrant is not available. Starting without vector database. Reason: '{Reason}'.", reason);
|
||||
return this.CreateNoDatabaseClient("Qdrant", qdrantInfo.UnavailableReason, DatabaseClientStatus.UNAVAILABLE);
|
||||
}
|
||||
|
||||
if (!HasValidQdrantConnectionInfo(qdrantInfo, out var invalidReason))
|
||||
return this.CreateNoDatabaseClient("Qdrant", invalidReason, DatabaseClientStatus.UNAVAILABLE);
|
||||
|
||||
var client = new QdrantClientImplementation("Qdrant", qdrantInfo.Path, qdrantInfo.PortHttp, qdrantInfo.PortGrpc, qdrantInfo.Fingerprint, qdrantInfo.ApiToken);
|
||||
client.SetLogger(this.databaseClientLogger);
|
||||
|
||||
try
|
||||
{
|
||||
await client.CheckAvailabilityAsync();
|
||||
return client;
|
||||
}
|
||||
catch (Exception e)
|
||||
{
|
||||
client.Dispose();
|
||||
this.logger.LogWarning(e, "Qdrant reported as available by Rust, but the health check failed.");
|
||||
return this.CreateNoDatabaseClient("Qdrant", e.Message, DatabaseClientStatus.STARTING);
|
||||
}
|
||||
}
|
||||
|
||||
private static bool HasValidQdrantConnectionInfo(QdrantInfo qdrantInfo, out string invalidReason)
|
||||
{
|
||||
if (qdrantInfo.Path == string.Empty)
|
||||
{
|
||||
invalidReason = "Failed to get the Qdrant path from Rust.";
|
||||
return false;
|
||||
}
|
||||
|
||||
if (qdrantInfo.PortHttp == 0)
|
||||
{
|
||||
invalidReason = "Failed to get the Qdrant HTTP port from Rust.";
|
||||
return false;
|
||||
}
|
||||
|
||||
if (qdrantInfo.PortGrpc == 0)
|
||||
{
|
||||
invalidReason = "Failed to get the Qdrant gRPC port from Rust.";
|
||||
return false;
|
||||
}
|
||||
|
||||
if (qdrantInfo.Fingerprint == string.Empty)
|
||||
{
|
||||
invalidReason = "Failed to get the Qdrant fingerprint from Rust.";
|
||||
return false;
|
||||
}
|
||||
|
||||
if (qdrantInfo.ApiToken == string.Empty)
|
||||
{
|
||||
invalidReason = "Failed to get the Qdrant API token from Rust.";
|
||||
return false;
|
||||
}
|
||||
|
||||
invalidReason = string.Empty;
|
||||
return true;
|
||||
}
|
||||
|
||||
private NoDatabaseClient CreateNoDatabaseClient(string name, string? unavailableReason, DatabaseClientStatus status)
|
||||
{
|
||||
var client = new NoDatabaseClient(name, unavailableReason, status);
|
||||
|
||||
@@ -1,33 +0,0 @@
|
||||
using AIStudio.Tools.Databases.Qdrant;
|
||||
|
||||
namespace AIStudio.Tools.Databases;
|
||||
|
||||
public class EmbeddingStoreFactory
|
||||
{
|
||||
public static EmbeddingStore Create(EmbeddingStoreConfiguration configuration) => configuration.Kind switch
|
||||
{
|
||||
EmbeddingStoreKind.NONE => new NoEmbeddingStore(configuration.Name, configuration.UnavailableReason ?? "unknown"),
|
||||
_ when configuration.Location is null => new NoEmbeddingStore(configuration.Name, $"No location specified for {configuration.Name}"),
|
||||
EmbeddingStoreKind.QDRANT_REMOTE when configuration.Location is RemoteLocation location=> new QdrantClientImplementation(configuration.Name, location.Path, location.HttpPort, location.GrpcPort, location.Fingerprint, location.ApiToken),
|
||||
_ => throw new ArgumentException("Invalid configuration for " + configuration.Name, nameof(configuration)),
|
||||
};
|
||||
}
|
||||
|
||||
public enum EmbeddingStoreKind
|
||||
{
|
||||
NONE,
|
||||
QDRANT_EMBED,
|
||||
QDRANT_REMOTE,
|
||||
}
|
||||
|
||||
public abstract record EmbeddingStoreLocation;
|
||||
|
||||
public sealed record EmbeddedLocation(string Path) : EmbeddingStoreLocation;
|
||||
|
||||
public sealed record RemoteLocation(string Path, int? HttpPort, int? GrpcPort, string? Fingerprint, string? ApiToken) : EmbeddingStoreLocation;
|
||||
|
||||
public sealed record EmbeddingStoreConfiguration(
|
||||
EmbeddingStoreKind Kind,
|
||||
string Name,
|
||||
EmbeddingStoreLocation? Location,
|
||||
string? UnavailableReason);
|
||||
@@ -2,9 +2,9 @@ using AIStudio.Tools.PluginSystem;
|
||||
|
||||
namespace AIStudio.Tools.Databases;
|
||||
|
||||
public sealed class NoEmbeddingStore(string name, string? unavailableReason, DatabaseClientStatus status = DatabaseClientStatus.UNAVAILABLE) : EmbeddingStore(name, string.Empty)
|
||||
public sealed class NoDatabaseClient(string name, string? unavailableReason, DatabaseClientStatus status = DatabaseClientStatus.UNAVAILABLE) : DatabaseClient(name, string.Empty)
|
||||
{
|
||||
private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(NoEmbeddingStore).Namespace, nameof(NoEmbeddingStore));
|
||||
private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(NoDatabaseClient).Namespace, nameof(NoDatabaseClient));
|
||||
|
||||
public override DatabaseClientStatus Status => status;
|
||||
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
namespace AIStudio.Tools.Databases;
|
||||
|
||||
public sealed record EmbeddingStoragePoint(
|
||||
public sealed record VectorStoragePoint(
|
||||
string PointId,
|
||||
IReadOnlyList<float> Vector,
|
||||
string DataSourceId,
|
||||
@@ -0,0 +1,20 @@
|
||||
namespace AIStudio.Tools.Databases.VectorStore;
|
||||
|
||||
public interface IVectorStoreClient
|
||||
{
|
||||
string Name { get; }
|
||||
|
||||
DatabaseClientStatus Status { get; }
|
||||
|
||||
bool IsAvailable { get; }
|
||||
|
||||
IAsyncEnumerable<(string Label, string Value)> GetDisplayInfo();
|
||||
|
||||
Task EnsureVectorStoreExists(string storeName, int vectorSize, CancellationToken token);
|
||||
|
||||
Task InsertEmbedding(string storeName, IReadOnlyList<VectorStoragePoint> points, CancellationToken token);
|
||||
|
||||
Task DeleteEmbeddingByFile(string storeName, string filePath, CancellationToken token);
|
||||
|
||||
Task DeleteVectorStore(string storeName, CancellationToken token);
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
using AIStudio.Tools.PluginSystem;
|
||||
|
||||
namespace AIStudio.Tools.Databases.VectorStore;
|
||||
|
||||
public sealed class NoVectorStoreClient(string name, string? unavailableReason, DatabaseClientStatus status = DatabaseClientStatus.UNAVAILABLE) : IVectorStoreClient
|
||||
{
|
||||
private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(NoVectorStoreClient).Namespace, nameof(NoVectorStoreClient));
|
||||
|
||||
public string Name => name;
|
||||
|
||||
public DatabaseClientStatus Status => status;
|
||||
|
||||
public bool IsAvailable => false;
|
||||
|
||||
public async IAsyncEnumerable<(string Label, string Value)> GetDisplayInfo()
|
||||
{
|
||||
yield return (TB("Status"), TB("Unavailable"));
|
||||
|
||||
if (!string.IsNullOrWhiteSpace(unavailableReason))
|
||||
yield return (TB("Reason"), unavailableReason);
|
||||
|
||||
await Task.CompletedTask;
|
||||
}
|
||||
|
||||
public Task EnsureVectorStoreExists(string storeName, int vectorSize, CancellationToken token) =>
|
||||
Task.FromException(this.CreateUnavailableException());
|
||||
|
||||
public Task InsertEmbedding(string storeName, IReadOnlyList<VectorStoragePoint> points, CancellationToken token) =>
|
||||
Task.FromException(this.CreateUnavailableException());
|
||||
|
||||
public Task DeleteEmbeddingByFile(string storeName, string filePath, CancellationToken token) =>
|
||||
Task.FromException(this.CreateUnavailableException());
|
||||
|
||||
public Task DeleteVectorStore(string storeName, CancellationToken token) =>
|
||||
Task.FromException(this.CreateUnavailableException());
|
||||
|
||||
private InvalidOperationException CreateUnavailableException() =>
|
||||
new(unavailableReason ?? "The vector store is not available.");
|
||||
}
|
||||
+6
-8
@@ -4,9 +4,9 @@ using Grpc.Core;
|
||||
using AIStudio.Tools.PluginSystem;
|
||||
using static Qdrant.Client.Grpc.Conditions;
|
||||
|
||||
namespace AIStudio.Tools.Databases.Qdrant;
|
||||
namespace AIStudio.Tools.Databases.VectorStore;
|
||||
|
||||
public class QdrantClientImplementation : EmbeddingStore
|
||||
public class QdrantClientImplementation : DatabaseClient, IVectorStoreClient
|
||||
{
|
||||
private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(QdrantClientImplementation).Namespace, nameof(QdrantClientImplementation));
|
||||
|
||||
@@ -71,7 +71,7 @@ public class QdrantClientImplementation : EmbeddingStore
|
||||
yield return (TB("Number of collections"), await this.GetCollectionsAmount());
|
||||
}
|
||||
|
||||
public override async Task EnsureEmbeddingStoreExists(string collectionName, int vectorSize, CancellationToken token)
|
||||
public async Task EnsureVectorStoreExists(string collectionName, int vectorSize, CancellationToken token)
|
||||
{
|
||||
var exists = await this.GrpcClient.CollectionExistsAsync(collectionName, token);
|
||||
if (exists)
|
||||
@@ -87,7 +87,7 @@ public class QdrantClientImplementation : EmbeddingStore
|
||||
cancellationToken: token);
|
||||
}
|
||||
|
||||
public override Task InsertEmbedding(string collectionName, IReadOnlyList<EmbeddingStoragePoint> points, CancellationToken token)
|
||||
public Task InsertEmbedding(string collectionName, IReadOnlyList<VectorStoragePoint> points, CancellationToken token)
|
||||
{
|
||||
var qdrantPoints = points.Select(point => new PointStruct
|
||||
{
|
||||
@@ -112,7 +112,7 @@ public class QdrantClientImplementation : EmbeddingStore
|
||||
return this.GrpcClient.UpsertAsync(collectionName, qdrantPoints, true, null, null, token);
|
||||
}
|
||||
|
||||
public override async Task DeleteEmbeddingByFile(string collectionName, string filePath, CancellationToken token)
|
||||
public async Task DeleteEmbeddingByFile(string collectionName, string filePath, CancellationToken token)
|
||||
{
|
||||
try
|
||||
{
|
||||
@@ -120,11 +120,10 @@ public class QdrantClientImplementation : EmbeddingStore
|
||||
}
|
||||
catch (RpcException exception) when (exception.StatusCode is StatusCode.NotFound)
|
||||
{
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
public override async Task DeleteEmbeddingStore(string collectionName, CancellationToken token)
|
||||
public async Task DeleteVectorStore(string collectionName, CancellationToken token)
|
||||
{
|
||||
var exists = await this.GrpcClient.CollectionExistsAsync(collectionName, token);
|
||||
if (!exists)
|
||||
@@ -136,7 +135,6 @@ public class QdrantClientImplementation : EmbeddingStore
|
||||
}
|
||||
catch (RpcException exception) when (exception.StatusCode is StatusCode.NotFound)
|
||||
{
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,6 +0,0 @@
|
||||
namespace AIStudio.Tools.Metadata;
|
||||
|
||||
public class MetaDataDatabasesAttribute(string databaseVersion) : Attribute
|
||||
{
|
||||
public string DatabaseVersion => databaseVersion;
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
namespace AIStudio.Tools.Metadata;
|
||||
|
||||
public class MetaDataVectorStoreAttribute(string vectorStoreVersion) : Attribute
|
||||
{
|
||||
public string VectorStoreVersion => vectorStoreVersion;
|
||||
}
|
||||
@@ -438,7 +438,7 @@ public sealed record PluginConfigurationObject
|
||||
return string.Empty;
|
||||
}
|
||||
|
||||
return storeResult.Message;
|
||||
return storeResult.StoredPath;
|
||||
}
|
||||
|
||||
private static string? ResolvePluginTokenizerPath(string configuredTokenizerPath, string pluginPath)
|
||||
|
||||
@@ -1,3 +1,8 @@
|
||||
namespace AIStudio.Tools.Rust;
|
||||
|
||||
public readonly record struct TokenizerResponse(bool Success, int TokenCount, string Message);
|
||||
public readonly record struct TokenizerResponse(
|
||||
bool Success,
|
||||
int TokenCount,
|
||||
string Message,
|
||||
TokenizerStatus Status = TokenizerStatus.UNAVAILABLE,
|
||||
string StoredPath = "");
|
||||
@@ -0,0 +1,8 @@
|
||||
namespace AIStudio.Tools.Rust;
|
||||
|
||||
public enum TokenizerStatus
|
||||
{
|
||||
UNAVAILABLE,
|
||||
RUNNING,
|
||||
AVAILABLE,
|
||||
}
|
||||
@@ -23,7 +23,7 @@ public sealed partial class DataSourceEmbeddingService
|
||||
|
||||
var currentChunk = new StringBuilder();
|
||||
|
||||
await foreach (var segment in this.rustService.StreamArbitraryFileData(filePath, token: token))
|
||||
await foreach (var segment in rustService.StreamArbitraryFileData(filePath, token: token))
|
||||
{
|
||||
var normalized = NormalizeChunkSegment(segment);
|
||||
if (string.IsNullOrWhiteSpace(normalized))
|
||||
@@ -118,7 +118,7 @@ public sealed partial class DataSourceEmbeddingService
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
this.logger.LogWarning(exception, "Cannot access directory '{DirectoryPath}' while indexing.", currentPath);
|
||||
logger.LogWarning(exception, "Cannot access directory '{DirectoryPath}' while indexing.", currentPath);
|
||||
result.FailedFiles++;
|
||||
result.LastError = $"The directory '{currentPath}' could not be accessed.";
|
||||
continue;
|
||||
@@ -135,7 +135,7 @@ public sealed partial class DataSourceEmbeddingService
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
this.logger.LogWarning(exception, "Cannot inspect file '{FilePath}' while indexing.", filePath);
|
||||
logger.LogWarning(exception, "Cannot inspect file '{FilePath}' while indexing.", filePath);
|
||||
result.FailedFiles++;
|
||||
result.LastError = $"The file '{filePath}' could not be inspected.";
|
||||
continue;
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
using System.Text.Json;
|
||||
|
||||
using AIStudio.Settings;
|
||||
using AIStudio.Tools.Databases.VectorStore;
|
||||
|
||||
namespace AIStudio.Tools.Services;
|
||||
|
||||
@@ -70,13 +71,13 @@ public sealed partial class DataSourceEmbeddingService
|
||||
await File.WriteAllTextAsync(statePath, json, token);
|
||||
}
|
||||
|
||||
private async Task ResetPersistedStateAsync(string dataSourceId)
|
||||
private async Task ResetPersistedStateAsync(string dataSourceId, IVectorStoreClient? vectorStore, CancellationToken token)
|
||||
{
|
||||
await this.EnsureStateLoadedAsync(CancellationToken.None);
|
||||
await this.EnsureStateLoadedAsync(token);
|
||||
this.manifests.Remove(dataSourceId);
|
||||
await this.DeleteCollectionAsync(this.GetCollectionName(dataSourceId));
|
||||
await this.SaveStateAsync(CancellationToken.None);
|
||||
this.logger.LogInformation("Reset persisted embedding state for data source '{DataSourceId}'.", dataSourceId);
|
||||
await this.DeleteCollectionAsync(this.GetCollectionName(dataSourceId), vectorStore, token);
|
||||
await this.SaveStateAsync(token);
|
||||
logger.LogInformation("Reset persisted embedding state for data source '{DataSourceId}'.", dataSourceId);
|
||||
}
|
||||
|
||||
private string GetStatePath()
|
||||
|
||||
@@ -14,13 +14,13 @@ public sealed partial class DataSourceEmbeddingService
|
||||
|
||||
private void RefreshWatchers()
|
||||
{
|
||||
if (!this.settingsManager.ConfigurationData.DataSourceIndexing.AutomaticRefresh)
|
||||
if (!settingsManager.ConfigurationData.DataSourceIndexing.AutomaticRefresh)
|
||||
{
|
||||
this.RemoveAllWatchers();
|
||||
return;
|
||||
}
|
||||
|
||||
var supportedSources = this.settingsManager.ConfigurationData.DataSources
|
||||
var supportedSources = settingsManager.ConfigurationData.DataSources
|
||||
.Where(this.IsSupportedInternalDataSource)
|
||||
.ToDictionary(source => source.Id, StringComparer.OrdinalIgnoreCase);
|
||||
|
||||
@@ -33,7 +33,7 @@ public sealed partial class DataSourceEmbeddingService
|
||||
|
||||
private void EnsureWatcher(IDataSource dataSource)
|
||||
{
|
||||
if (!this.settingsManager.ConfigurationData.DataSourceIndexing.AutomaticRefresh)
|
||||
if (!settingsManager.ConfigurationData.DataSourceIndexing.AutomaticRefresh)
|
||||
return;
|
||||
|
||||
var configuration = GetWatchConfiguration(dataSource);
|
||||
@@ -73,7 +73,7 @@ public sealed partial class DataSourceEmbeddingService
|
||||
watcher.Renamed += (_, _) => this.OnWatchedDataSourceChanged(dataSourceId);
|
||||
watcher.Error += (_, args) =>
|
||||
{
|
||||
this.logger.LogWarning(args.GetException(), "The file watcher for data source '{DataSourceId}' failed. Recreating it.", dataSourceId);
|
||||
logger.LogWarning(args.GetException(), "The file watcher for data source '{DataSourceId}' failed. Recreating it.", dataSourceId);
|
||||
this.RemoveWatcher(dataSourceId);
|
||||
this.EnsureWatcher(dataSourceId);
|
||||
this.OnWatchedDataSourceChanged(dataSourceId);
|
||||
@@ -83,7 +83,7 @@ public sealed partial class DataSourceEmbeddingService
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
this.logger.LogWarning(exception, "Failed to create file watcher for data source '{DataSourceId}' at '{RootPath}'.", dataSourceId, configuration.RootPath);
|
||||
logger.LogWarning(exception, "Failed to create file watcher for data source '{DataSourceId}' at '{RootPath}'.", dataSourceId, configuration.RootPath);
|
||||
return null;
|
||||
}
|
||||
}
|
||||
@@ -114,10 +114,10 @@ public sealed partial class DataSourceEmbeddingService
|
||||
|
||||
private void OnWatchedDataSourceChanged(string dataSourceId)
|
||||
{
|
||||
if (!this.settingsManager.ConfigurationData.DataSourceIndexing.AutomaticRefresh)
|
||||
if (!settingsManager.ConfigurationData.DataSourceIndexing.AutomaticRefresh)
|
||||
return;
|
||||
|
||||
this.logger.LogDebug("Detected file system change for data source '{DataSourceId}'. Scheduling a debounced embedding run.", dataSourceId);
|
||||
logger.LogDebug("Detected file system change for data source '{DataSourceId}'. Scheduling a debounced embedding run.", dataSourceId);
|
||||
var debounceToken = new CancellationTokenSource();
|
||||
|
||||
lock (this.watcherDebounceLock)
|
||||
@@ -136,12 +136,12 @@ public sealed partial class DataSourceEmbeddingService
|
||||
if (!this.TryCompletePendingWatcherRefresh(dataSourceId, debounceToken))
|
||||
return;
|
||||
|
||||
var dataSource = this.settingsManager.ConfigurationData.DataSources
|
||||
var dataSource = settingsManager.ConfigurationData.DataSources
|
||||
.FirstOrDefault(source => source.Id.Equals(dataSourceId, StringComparison.OrdinalIgnoreCase));
|
||||
|
||||
if (dataSource is not null)
|
||||
{
|
||||
this.logger.LogInformation("Queueing data source '{DataSourceName}' ({DataSourceId}) after file system changes settled.", dataSource.Name, dataSource.Id);
|
||||
logger.LogInformation("Queueing data source '{DataSourceName}' ({DataSourceId}) after file system changes settled.", dataSource.Name, dataSource.Id);
|
||||
await this.QueueDataSourceAsync(dataSource);
|
||||
}
|
||||
}
|
||||
@@ -150,7 +150,7 @@ public sealed partial class DataSourceEmbeddingService
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
this.logger.LogWarning(exception, "Failed to queue watched data source '{DataSourceId}' after a file system change.", dataSourceId);
|
||||
logger.LogWarning(exception, "Failed to queue watched data source '{DataSourceId}' after a file system change.", dataSourceId);
|
||||
}
|
||||
finally
|
||||
{
|
||||
@@ -161,7 +161,7 @@ public sealed partial class DataSourceEmbeddingService
|
||||
|
||||
private void EnsureWatcher(string dataSourceId)
|
||||
{
|
||||
var dataSource = this.settingsManager.ConfigurationData.DataSources
|
||||
var dataSource = settingsManager.ConfigurationData.DataSources
|
||||
.FirstOrDefault(source => source.Id.Equals(dataSourceId, StringComparison.OrdinalIgnoreCase));
|
||||
|
||||
if (dataSource is not null)
|
||||
|
||||
@@ -6,22 +6,19 @@ using AIStudio.Provider;
|
||||
using AIStudio.Settings;
|
||||
using AIStudio.Settings.DataModel;
|
||||
using AIStudio.Tools.Databases;
|
||||
using AIStudio.Tools.Databases.VectorStore;
|
||||
using AIStudio.Tools.PluginSystem;
|
||||
using AIStudio.Tools.Rust;
|
||||
|
||||
namespace AIStudio.Tools.Services;
|
||||
|
||||
public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
public sealed partial class DataSourceEmbeddingService(SettingsManager settingsManager, RustService rustService, DatabaseClientProvider databaseClientProvider, ILogger<DataSourceEmbeddingService> logger)
|
||||
: BackgroundService
|
||||
{
|
||||
private const int MAX_CHUNK_LENGTH = 3_200;
|
||||
private const int MIN_CHUNK_LENGTH = 800;
|
||||
private const int CHUNK_OVERLAP_LENGTH = 320;
|
||||
private const int EMBEDDING_BATCH_SIZE = 16;
|
||||
|
||||
private readonly SettingsManager settingsManager;
|
||||
private readonly RustService rustService;
|
||||
private readonly EmbeddingStore embeddingStore;
|
||||
private readonly ILogger<DataSourceEmbeddingService> logger;
|
||||
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);
|
||||
@@ -32,14 +29,6 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
private Dictionary<string, DataSourceEmbeddingManifest> manifests = new(StringComparer.OrdinalIgnoreCase);
|
||||
private bool stateLoaded;
|
||||
|
||||
public DataSourceEmbeddingService(SettingsManager settingsManager, RustService rustService, EmbeddingStore embeddingStore, ILogger<DataSourceEmbeddingService> logger)
|
||||
{
|
||||
this.settingsManager = settingsManager;
|
||||
this.rustService = rustService;
|
||||
this.embeddingStore = embeddingStore;
|
||||
this.logger = logger;
|
||||
}
|
||||
|
||||
public IReadOnlyList<DataSourceEmbeddingStatus> GetStatuses()
|
||||
{
|
||||
return this.statuses.Values
|
||||
@@ -78,7 +67,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
{
|
||||
this.RefreshWatchers();
|
||||
|
||||
var tasks = this.settingsManager.ConfigurationData.DataSources
|
||||
var tasks = settingsManager.ConfigurationData.DataSources
|
||||
.Where(this.IsSupportedInternalDataSource)
|
||||
.Select(this.QueueDataSourceAsync);
|
||||
|
||||
@@ -87,7 +76,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
|
||||
public Task QueueAllInternalDataSourcesIfAutomaticRefreshAsync()
|
||||
{
|
||||
if (!this.settingsManager.ConfigurationData.DataSourceIndexing.AutomaticRefresh)
|
||||
if (!settingsManager.ConfigurationData.DataSourceIndexing.AutomaticRefresh)
|
||||
{
|
||||
this.RefreshWatchers();
|
||||
return Task.CompletedTask;
|
||||
@@ -106,16 +95,16 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
if (!this.IsSupportedInternalDataSource(dataSource))
|
||||
return;
|
||||
|
||||
this.logger.LogInformation("Queueing data source '{DataSourceName}' ({DataSourceId}) for background embeddings.", dataSource.Name, dataSource.Id);
|
||||
logger.LogInformation("Queueing data source '{DataSourceName}' ({DataSourceId}) for background embeddings.", dataSource.Name, dataSource.Id);
|
||||
this.RefreshWatchers();
|
||||
this.logger.LogDebug("Adding watcher for data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id);
|
||||
logger.LogDebug("Adding watcher for data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id);
|
||||
|
||||
if (!this.statuses.TryGetValue(dataSource.Id, out var currentStatus) || currentStatus.State is not DataSourceEmbeddingState.RUNNING)
|
||||
this.UpsertStatus(this.CreateStatus(dataSource, DataSourceEmbeddingState.QUEUED, currentStatus?.TotalFiles ?? 0, currentStatus?.IndexedFiles ?? 0, currentStatus?.FailedFiles ?? 0));
|
||||
this.logger.LogDebug("Upserting status for data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id);
|
||||
logger.LogDebug("Upserting status for data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id);
|
||||
if (this.queuedIds.TryAdd(dataSource.Id, 0))
|
||||
await this.queue.Writer.WriteAsync(dataSource.Id);
|
||||
this.logger.LogDebug("Queued data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id);
|
||||
logger.LogDebug("Queued data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id);
|
||||
}
|
||||
|
||||
public async Task RemoveDataSourceAsync(IDataSource dataSource)
|
||||
@@ -125,7 +114,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
|
||||
this.RemoveWatcher(dataSource.Id);
|
||||
this.statuses.TryRemove(dataSource.Id, out _);
|
||||
await this.ResetPersistedStateAsync(dataSource.Id);
|
||||
await this.ResetPersistedStateAsync(dataSource.Id, null, CancellationToken.None);
|
||||
this.PublishStatusChanged();
|
||||
}
|
||||
|
||||
@@ -138,7 +127,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
var dataSourceId = await this.queue.Reader.ReadAsync(stoppingToken);
|
||||
this.queuedIds.TryRemove(dataSourceId, out _);
|
||||
|
||||
var dataSource = this.settingsManager.ConfigurationData.DataSources
|
||||
var dataSource = settingsManager.ConfigurationData.DataSources
|
||||
.FirstOrDefault(source => source.Id.Equals(dataSourceId, StringComparison.OrdinalIgnoreCase));
|
||||
|
||||
if (dataSource is null || !this.IsSupportedInternalDataSource(dataSource))
|
||||
@@ -154,7 +143,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
this.logger.LogError(exception, "Background embedding failed for data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id);
|
||||
logger.LogError(exception, "Background embedding failed for data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id);
|
||||
this.UpsertStatus(this.GetFallbackStatus(dataSource, exception.Message));
|
||||
}
|
||||
}
|
||||
@@ -170,15 +159,17 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
private async Task ProcessDataSourceAsync(IDataSource dataSource, CancellationToken token)
|
||||
{
|
||||
await this.EnsureStateLoadedAsync(token);
|
||||
this.logger.LogInformation("Starting background embeddings for data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id);
|
||||
logger.LogInformation("Starting background embeddings for data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id);
|
||||
|
||||
if (!this.embeddingStore.IsAvailable)
|
||||
var vectorStore = await databaseClientProvider.GetVectorStoreAsync(token);
|
||||
|
||||
if (!vectorStore.IsAvailable)
|
||||
{
|
||||
this.logger.LogWarning(
|
||||
logger.LogWarning(
|
||||
"Skipping background embeddings for data source '{DataSourceName}' ({DataSourceId}) because the database client '{DatabaseName}' is unavailable.",
|
||||
dataSource.Name,
|
||||
dataSource.Id,
|
||||
this.embeddingStore.Name);
|
||||
vectorStore.Name);
|
||||
this.UpsertStatus(this.GetFallbackStatus(dataSource, "The vector database is not available."));
|
||||
return;
|
||||
}
|
||||
@@ -189,7 +180,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
return;
|
||||
}
|
||||
|
||||
this.logger.LogInformation(
|
||||
logger.LogInformation(
|
||||
"Using embedding provider '{EmbeddingProviderId}' with model '{EmbeddingModelId}' for data source '{DataSourceName}' ({DataSourceId}).",
|
||||
embeddingProvider.Id,
|
||||
embeddingProvider.Model.Id,
|
||||
@@ -197,12 +188,12 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
dataSource.Id);
|
||||
|
||||
var collectionName = this.GetCollectionName(dataSource.Id);
|
||||
var manifest = await this.EnsureCompatibleManifestAsync(dataSource, embeddingProvider, collectionName, token);
|
||||
var manifest = await this.EnsureCompatibleManifestAsync(dataSource, embeddingProvider, collectionName, vectorStore, token);
|
||||
var inputFiles = this.GetInputFiles(dataSource);
|
||||
var indexedFiles = inputFiles.Files;
|
||||
var totalFiles = indexedFiles.Count + inputFiles.FailedFiles;
|
||||
|
||||
this.logger.LogInformation(
|
||||
logger.LogInformation(
|
||||
"Prepared data source '{DataSourceName}' ({DataSourceId}) for embedding. AccessibleFiles={AccessibleFiles}, FailedFiles={FailedFiles}, Collection='{CollectionName}'.",
|
||||
dataSource.Name,
|
||||
dataSource.Id,
|
||||
@@ -210,7 +201,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
inputFiles.FailedFiles,
|
||||
collectionName);
|
||||
|
||||
await this.RemoveMissingFileEmbeddingsAsync(dataSource, collectionName, manifest, indexedFiles, token);
|
||||
await this.RemoveMissingFileEmbeddingsAsync(vectorStore, dataSource, collectionName, manifest, indexedFiles, token);
|
||||
await this.SaveStateAsync(token);
|
||||
|
||||
this.UpsertStatus(this.CreateStatus(
|
||||
@@ -235,7 +226,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
if (manifest.Files.TryGetValue(file.FullName, out var existingRecord) &&
|
||||
string.Equals(existingRecord.Fingerprint, fingerprint, StringComparison.Ordinal))
|
||||
{
|
||||
this.logger.LogDebug(
|
||||
logger.LogDebug(
|
||||
"Skipping unchanged file '{FilePath}' for data source '{DataSourceName}' ({DataSourceId}).",
|
||||
file.FullName,
|
||||
dataSource.Name,
|
||||
@@ -249,7 +240,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
|
||||
try
|
||||
{
|
||||
this.logger.LogInformation(
|
||||
logger.LogInformation(
|
||||
"Embedding file '{FilePath}' for data source '{DataSourceName}' ({DataSourceId}). Progress={CompletedFiles}/{TotalFiles}.",
|
||||
file.FullName,
|
||||
dataSource.Name,
|
||||
@@ -257,7 +248,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
skippedFiles + completedFiles + 1,
|
||||
totalFiles);
|
||||
var startedAtUtc = DateTime.UtcNow;
|
||||
var chunkCount = await this.IndexOneFileAsync(dataSource, file, fingerprint, embeddingProvider, provider, manifest, token);
|
||||
var chunkCount = await this.IndexOneFileAsync(vectorStore, dataSource, file, fingerprint, embeddingProvider, provider, manifest, token);
|
||||
manifest.Files[file.FullName] = new EmbeddedFileRecord(
|
||||
fingerprint,
|
||||
file.Length,
|
||||
@@ -266,7 +257,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
chunkCount);
|
||||
await this.SaveStateAsync(token);
|
||||
completedFiles++;
|
||||
this.logger.LogInformation(
|
||||
logger.LogInformation(
|
||||
"Embedded file '{FilePath}' for data source '{DataSourceName}' ({DataSourceId}) successfully. Chunks={ChunkCount}, DurationMs={DurationMs}.",
|
||||
file.FullName,
|
||||
dataSource.Name,
|
||||
@@ -279,16 +270,16 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
failedFiles++;
|
||||
lastError = exception.Message;
|
||||
manifest.Files.Remove(file.FullName);
|
||||
await this.DeleteFilePointsAsync(collectionName, file.FullName, token);
|
||||
await this.DeleteFilePointsAsync(vectorStore, collectionName, file.FullName, token);
|
||||
await this.SaveStateAsync(token);
|
||||
|
||||
this.logger.LogWarning(exception, "Failed to embed file '{FilePath}' for data source '{DataSourceName}'.", file.FullName, dataSource.Name);
|
||||
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));
|
||||
}
|
||||
}
|
||||
|
||||
this.UpsertStatus(this.CreateCompletedStatus(dataSource, totalFiles, skippedFiles + completedFiles, failedFiles, lastError));
|
||||
this.logger.LogInformation(
|
||||
logger.LogInformation(
|
||||
"Finished background embeddings for data source '{DataSourceName}' ({DataSourceId}). Indexed={IndexedFiles}, Failed={FailedFiles}, Total={TotalFiles}.",
|
||||
dataSource.Name,
|
||||
dataSource.Id,
|
||||
@@ -298,6 +289,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
}
|
||||
|
||||
private async Task<int> IndexOneFileAsync(
|
||||
IVectorStoreClient vectorStore,
|
||||
IDataSource dataSource,
|
||||
FileInfo file,
|
||||
string fingerprint,
|
||||
@@ -307,11 +299,11 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
CancellationToken token)
|
||||
{
|
||||
var collectionName = this.GetCollectionName(dataSource.Id);
|
||||
this.logger.LogDebug(
|
||||
logger.LogDebug(
|
||||
"Resetting stored embeddings for file '{FilePath}' in collection '{CollectionName}' before re-indexing.",
|
||||
file.FullName,
|
||||
collectionName);
|
||||
await this.DeleteFilePointsAsync(collectionName, file.FullName, token);
|
||||
await this.DeleteFilePointsAsync(vectorStore, collectionName, file.FullName, token);
|
||||
|
||||
var batch = new List<(string Text, int ChunkIndex)>(EMBEDDING_BATCH_SIZE);
|
||||
var totalChunkCount = 0;
|
||||
@@ -322,16 +314,16 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
totalChunkCount++;
|
||||
|
||||
if (batch.Count >= EMBEDDING_BATCH_SIZE)
|
||||
await this.FlushBatchAsync(dataSource, file, fingerprint, embeddingProvider, provider, manifest, collectionName, batch, token);
|
||||
await this.FlushBatchAsync(vectorStore, dataSource, file, fingerprint, embeddingProvider, provider, manifest, collectionName, batch, token);
|
||||
}
|
||||
|
||||
if (batch.Count > 0)
|
||||
await this.FlushBatchAsync(dataSource, file, fingerprint, embeddingProvider, provider, manifest, collectionName, batch, token);
|
||||
await this.FlushBatchAsync(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.");
|
||||
|
||||
this.logger.LogDebug(
|
||||
logger.LogDebug(
|
||||
"Generated {ChunkCount} chunks for file '{FilePath}' in data source '{DataSourceName}' ({DataSourceId}).",
|
||||
totalChunkCount,
|
||||
file.FullName,
|
||||
@@ -342,6 +334,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
}
|
||||
|
||||
private async Task FlushBatchAsync(
|
||||
IVectorStoreClient vectorStore,
|
||||
IDataSource dataSource,
|
||||
FileInfo file,
|
||||
string fingerprint,
|
||||
@@ -352,7 +345,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
List<(string Text, int ChunkIndex)> batch,
|
||||
CancellationToken token)
|
||||
{
|
||||
this.logger.LogDebug(
|
||||
logger.LogDebug(
|
||||
"Requesting embeddings for batch of {ChunkCount} chunks from file '{FilePath}' in data source '{DataSourceName}' ({DataSourceId}).",
|
||||
batch.Count,
|
||||
file.FullName,
|
||||
@@ -360,7 +353,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
dataSource.Id);
|
||||
|
||||
var texts = batch.Select(item => item.Text).ToList();
|
||||
var vectors = await provider.EmbedTextAsync(embeddingProvider.Model, this.settingsManager, token, texts);
|
||||
var vectors = await provider.EmbedTextAsync(embeddingProvider.Model, settingsManager, token, texts);
|
||||
if (vectors.Count != batch.Count)
|
||||
throw new InvalidOperationException($"The embedding provider returned {vectors.Count} vectors for {batch.Count} text chunks.");
|
||||
|
||||
@@ -374,9 +367,9 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
if (manifest.VectorSize == 0)
|
||||
{
|
||||
manifest.VectorSize = vectorSize;
|
||||
await this.EnsureCollectionExistsAsync(collectionName, vectorSize, token);
|
||||
await this.EnsureCollectionExistsAsync(vectorStore, collectionName, vectorSize, token);
|
||||
await this.SaveStateAsync(token);
|
||||
this.logger.LogInformation(
|
||||
logger.LogInformation(
|
||||
"Created embedding collection '{CollectionName}' with vector size {VectorSize} for data source '{DataSourceName}' ({DataSourceId}).",
|
||||
collectionName,
|
||||
vectorSize,
|
||||
@@ -385,6 +378,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
}
|
||||
|
||||
await this.UpsertPointsAsync(
|
||||
vectorStore,
|
||||
collectionName,
|
||||
dataSource,
|
||||
file,
|
||||
@@ -394,7 +388,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
this.TryGetRelativePath(dataSource, file),
|
||||
token);
|
||||
|
||||
this.logger.LogDebug(
|
||||
logger.LogDebug(
|
||||
"Stored {ChunkCount} embedded chunks for file '{FilePath}' in collection '{CollectionName}'.",
|
||||
batch.Count,
|
||||
file.FullName,
|
||||
@@ -403,12 +397,13 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
batch.Clear();
|
||||
}
|
||||
|
||||
private async Task EnsureCollectionExistsAsync(string collectionName, int vectorSize, CancellationToken token)
|
||||
private async Task EnsureCollectionExistsAsync(IVectorStoreClient vectorStore, string collectionName, int vectorSize, CancellationToken token)
|
||||
{
|
||||
await this.embeddingStore.EnsureEmbeddingStoreExists(collectionName, vectorSize, token);
|
||||
await vectorStore.EnsureVectorStoreExists(collectionName, vectorSize, token);
|
||||
}
|
||||
|
||||
private async Task UpsertPointsAsync(
|
||||
IVectorStoreClient vectorStore,
|
||||
string collectionName,
|
||||
IDataSource dataSource,
|
||||
FileInfo file,
|
||||
@@ -419,7 +414,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
CancellationToken token)
|
||||
{
|
||||
var embeddedAtUtc = DateTime.UtcNow;
|
||||
var points = batch.Select((item, index) => new EmbeddingStoragePoint(
|
||||
var points = batch.Select((item, index) => new VectorStoragePoint(
|
||||
this.CreatePointId(dataSource.Id, fingerprint, item.ChunkIndex),
|
||||
vectors[index],
|
||||
dataSource.Id,
|
||||
@@ -434,24 +429,31 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
file.LastWriteTimeUtc,
|
||||
embeddedAtUtc)).ToList();
|
||||
|
||||
await this.embeddingStore.InsertEmbedding(collectionName, points, token);
|
||||
await vectorStore.InsertEmbedding(collectionName, points, token);
|
||||
}
|
||||
|
||||
private async Task DeleteFilePointsAsync(string collectionName, string filePath, CancellationToken token)
|
||||
private async Task DeleteFilePointsAsync(IVectorStoreClient vectorStore, string collectionName, string filePath, CancellationToken token)
|
||||
{
|
||||
await this.embeddingStore.DeleteEmbeddingByFile(collectionName, filePath, token);
|
||||
await vectorStore.DeleteEmbeddingByFile(collectionName, filePath, token);
|
||||
}
|
||||
|
||||
private async Task DeleteCollectionAsync(string collectionName)
|
||||
private async Task DeleteCollectionAsync(string collectionName, IVectorStoreClient? vectorStore, CancellationToken token)
|
||||
{
|
||||
await this.embeddingStore.DeleteEmbeddingStore(collectionName, CancellationToken.None);
|
||||
vectorStore ??= await databaseClientProvider.GetVectorStoreAsync(token);
|
||||
if (!vectorStore.IsAvailable)
|
||||
{
|
||||
logger.LogWarning("Could not delete embedding collection '{CollectionName}' because the vector store '{VectorStoreName}' is unavailable.", collectionName, vectorStore.Name);
|
||||
return;
|
||||
}
|
||||
|
||||
await vectorStore.DeleteVectorStore(collectionName, token);
|
||||
}
|
||||
|
||||
private async Task WaitForInitialSettingsAndBootstrapAsync(CancellationToken token)
|
||||
{
|
||||
while (!token.IsCancellationRequested)
|
||||
{
|
||||
if (this.settingsManager.HasCompletedInitialSettingsLoad
|
||||
if (settingsManager.HasCompletedInitialSettingsLoad
|
||||
&& !string.IsNullOrWhiteSpace(SettingsManager.ConfigDirectory)
|
||||
&& !string.IsNullOrWhiteSpace(SettingsManager.DataDirectory))
|
||||
{
|
||||
@@ -463,7 +465,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
|
||||
token.ThrowIfCancellationRequested();
|
||||
|
||||
this.logger.LogInformation("Embedding background service is ready. Checking whether automatic data source refresh is enabled.");
|
||||
logger.LogInformation("Embedding background service is ready. Checking whether automatic data source refresh is enabled.");
|
||||
await this.QueueAllInternalDataSourcesIfAutomaticRefreshAsync();
|
||||
}
|
||||
|
||||
@@ -474,26 +476,26 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
|
||||
private bool TryResolveEmbeddingProvider(IDataSource dataSource, [NotNullWhen(true)] out EmbeddingProvider? embeddingProvider)
|
||||
{
|
||||
embeddingProvider = this.settingsManager.ConfigurationData.EmbeddingProviders.FirstOrDefault(provider =>
|
||||
embeddingProvider = settingsManager.ConfigurationData.EmbeddingProviders.FirstOrDefault(provider =>
|
||||
dataSource is IInternalDataSource internalDataSource &&
|
||||
provider.Id.Equals(internalDataSource.EmbeddingId, StringComparison.OrdinalIgnoreCase));
|
||||
|
||||
return embeddingProvider != default && embeddingProvider.UsedLLMProvider is not LLMProviders.NONE;
|
||||
}
|
||||
|
||||
private async Task<DataSourceEmbeddingManifest> EnsureCompatibleManifestAsync(IDataSource dataSource, EmbeddingProvider embeddingProvider, string collectionName, CancellationToken token)
|
||||
private async Task<DataSourceEmbeddingManifest> EnsureCompatibleManifestAsync(IDataSource dataSource, EmbeddingProvider embeddingProvider, string collectionName, IVectorStoreClient vectorStore, CancellationToken token)
|
||||
{
|
||||
var embeddingSignature = this.BuildEmbeddingSignature(embeddingProvider);
|
||||
var manifest = await this.GetManifestAsync(dataSource.Id, token);
|
||||
|
||||
if (!string.Equals(manifest.EmbeddingSignature, embeddingSignature, StringComparison.Ordinal))
|
||||
{
|
||||
this.logger.LogInformation(
|
||||
logger.LogInformation(
|
||||
"Embedding configuration changed for data source '{DataSourceName}' ({DataSourceId}). Resetting persisted state and collection '{CollectionName}'.",
|
||||
dataSource.Name,
|
||||
dataSource.Id,
|
||||
collectionName);
|
||||
await this.ResetPersistedStateAsync(dataSource.Id);
|
||||
await this.ResetPersistedStateAsync(dataSource.Id, vectorStore, token);
|
||||
manifest = await this.GetManifestAsync(dataSource.Id, token);
|
||||
}
|
||||
|
||||
@@ -509,6 +511,7 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
}
|
||||
|
||||
private async Task RemoveMissingFileEmbeddingsAsync(
|
||||
IVectorStoreClient vectorStore,
|
||||
IDataSource dataSource,
|
||||
string collectionName,
|
||||
DataSourceEmbeddingManifest manifest,
|
||||
@@ -521,9 +524,9 @@ public sealed partial class DataSourceEmbeddingService : BackgroundService
|
||||
|
||||
foreach (var removedFilePath in manifest.Files.Keys.Except(existingPaths, StringComparer.OrdinalIgnoreCase).ToList())
|
||||
{
|
||||
await this.DeleteFilePointsAsync(collectionName, removedFilePath, token);
|
||||
await this.DeleteFilePointsAsync(vectorStore, collectionName, removedFilePath, token);
|
||||
manifest.Files.Remove(removedFilePath);
|
||||
this.logger.LogInformation(
|
||||
logger.LogInformation(
|
||||
"Removed stale embeddings for deleted file '{FilePath}' from data source '{DataSourceName}' ({DataSourceId}).",
|
||||
removedFilePath,
|
||||
dataSource.Name,
|
||||
|
||||
@@ -9,6 +9,31 @@ public sealed partial class RustService
|
||||
private string currentTokenizerPath = string.Empty;
|
||||
private bool hasInitializedTokenizer;
|
||||
|
||||
private static TokenizerResponse CreateUnavailableTokenizerResponse(string message) => new(
|
||||
false,
|
||||
0,
|
||||
message,
|
||||
TokenizerStatus.UNAVAILABLE,
|
||||
string.Empty);
|
||||
|
||||
public async Task<TokenizerResponse> GetTokenizerInfo(CancellationToken cancellationToken = default)
|
||||
{
|
||||
try
|
||||
{
|
||||
return await this.http.GetFromJsonAsync<TokenizerResponse>("/system/tokenizer/info", this.jsonRustSerializerOptions, cancellationToken);
|
||||
}
|
||||
catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
|
||||
{
|
||||
this.logger?.LogWarning("Fetching tokenizer info from Rust service was cancelled by caller.");
|
||||
return CreateUnavailableTokenizerResponse("Operation cancelled by caller.");
|
||||
}
|
||||
catch (Exception e)
|
||||
{
|
||||
this.logger?.LogError(e, "Error while fetching tokenizer info from Rust service.");
|
||||
return CreateUnavailableTokenizerResponse(e.Message);
|
||||
}
|
||||
}
|
||||
|
||||
public async Task<TokenizerResponse> ValidateTokenizer(string filePath)
|
||||
{
|
||||
var result = await this.http.PostAsJsonAsync("/tokenizer/validate", new {
|
||||
@@ -18,15 +43,12 @@ public sealed partial class RustService
|
||||
if (!result.IsSuccessStatusCode)
|
||||
{
|
||||
this.logger!.LogError($"Failed to validate the tokenizer '{result.StatusCode}'");
|
||||
return new TokenizerResponse
|
||||
{
|
||||
Success = false,
|
||||
Message = "An error occured while sending the path to the Rust framework for validation: "+result.StatusCode,
|
||||
TokenCount = 0
|
||||
};
|
||||
return CreateUnavailableTokenizerResponse("An error occured while sending the path to the Rust framework for validation: "+result.StatusCode);
|
||||
}
|
||||
|
||||
return await result.Content.ReadFromJsonAsync<TokenizerResponse>(this.jsonRustSerializerOptions);
|
||||
var response = await result.Content.ReadFromJsonAsync<TokenizerResponse>(this.jsonRustSerializerOptions);
|
||||
|
||||
return response;
|
||||
}
|
||||
|
||||
public async Task<TokenizerResponse> StoreTokenizer(string modelId, string filePath)
|
||||
@@ -40,11 +62,7 @@ public sealed partial class RustService
|
||||
if (!result.IsSuccessStatusCode)
|
||||
{
|
||||
this.logger!.LogError($"Failed to store the tokenizer '{result.StatusCode}'");
|
||||
return new TokenizerResponse{
|
||||
Success = false,
|
||||
Message = "An error occured while sending the path to the Rust framework for storing: "+result.StatusCode,
|
||||
TokenCount = 0
|
||||
};
|
||||
return CreateUnavailableTokenizerResponse("An error occured while sending the path to the Rust framework for storing: "+result.StatusCode);
|
||||
}
|
||||
|
||||
return await result.Content.ReadFromJsonAsync<TokenizerResponse>(this.jsonRustSerializerOptions);
|
||||
@@ -60,11 +78,7 @@ public sealed partial class RustService
|
||||
if (!result.IsSuccessStatusCode)
|
||||
{
|
||||
this.logger!.LogError($"Failed to delete the tokenizer '{result.StatusCode}'");
|
||||
return new TokenizerResponse{
|
||||
Success = false,
|
||||
Message = "An error occured while sending the tokenizer delete request to the Rust framework: "+result.StatusCode,
|
||||
TokenCount = 0
|
||||
};
|
||||
return CreateUnavailableTokenizerResponse("An error occured while sending the tokenizer delete request to the Rust framework: "+result.StatusCode);
|
||||
}
|
||||
|
||||
return await result.Content.ReadFromJsonAsync<TokenizerResponse>(this.jsonRustSerializerOptions);
|
||||
@@ -79,14 +93,15 @@ public sealed partial class RustService
|
||||
if (!result.IsSuccessStatusCode)
|
||||
{
|
||||
this.logger!.LogError($"Failed to get the token count '{result.StatusCode}'");
|
||||
return new TokenizerResponse{
|
||||
Success = false,
|
||||
Message = "Error while getting token count from Rust service: "+result.StatusCode,
|
||||
TokenCount = 0
|
||||
};
|
||||
this.hasInitializedTokenizer = false;
|
||||
return CreateUnavailableTokenizerResponse("Error while getting token count from Rust service: "+result.StatusCode);
|
||||
}
|
||||
|
||||
return await result.Content.ReadFromJsonAsync<TokenizerResponse>(this.jsonRustSerializerOptions);
|
||||
var response = await result.Content.ReadFromJsonAsync<TokenizerResponse>(this.jsonRustSerializerOptions);
|
||||
if (response is not { Status: TokenizerStatus.AVAILABLE })
|
||||
this.hasInitializedTokenizer = false;
|
||||
|
||||
return response;
|
||||
}
|
||||
|
||||
public async Task<TokenizerResponse?> SetTokenizer(string providerName, string path)
|
||||
@@ -99,14 +114,15 @@ public sealed partial class RustService
|
||||
if (!result.IsSuccessStatusCode)
|
||||
{
|
||||
this.logger!.LogError($"Failed to set the tokenizer '{result.StatusCode}'");
|
||||
return new TokenizerResponse{
|
||||
Success = false,
|
||||
Message = "An error occured while sending the path to the Rust framework for setting a tokenizer: "+result.StatusCode,
|
||||
TokenCount = 0
|
||||
};
|
||||
this.hasInitializedTokenizer = false;
|
||||
return CreateUnavailableTokenizerResponse("An error occured while sending the path to the Rust framework for setting a tokenizer: "+result.StatusCode);
|
||||
}
|
||||
|
||||
return await result.Content.ReadFromJsonAsync<TokenizerResponse>(this.jsonRustSerializerOptions);
|
||||
var response = await result.Content.ReadFromJsonAsync<TokenizerResponse>(this.jsonRustSerializerOptions);
|
||||
if (response is not { Success: true, Status: TokenizerStatus.AVAILABLE })
|
||||
this.hasInitializedTokenizer = false;
|
||||
|
||||
return response;
|
||||
}
|
||||
|
||||
public async Task<TokenizerResponse?> EnsureTokenizer(string providerName, string path)
|
||||
@@ -115,14 +131,19 @@ public sealed partial class RustService
|
||||
try
|
||||
{
|
||||
if (this.hasInitializedTokenizer && this.currentTokenizerPath == path)
|
||||
return new TokenizerResponse(true, 0, "Success");
|
||||
return new TokenizerResponse(true, 0, string.Empty, TokenizerStatus.AVAILABLE);
|
||||
|
||||
var response = await this.SetTokenizer(providerName, path);
|
||||
if (response is { Success: true })
|
||||
if (response is { Success: true, Status: TokenizerStatus.AVAILABLE })
|
||||
{
|
||||
this.currentTokenizerPath = path;
|
||||
this.hasInitializedTokenizer = true;
|
||||
}
|
||||
else
|
||||
{
|
||||
this.currentTokenizerPath = string.Empty;
|
||||
this.hasInitializedTokenizer = false;
|
||||
}
|
||||
|
||||
return response;
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user