MindWorkAI/AI-Studio
MindWorkAI/AI-Studio
1
0
Fork 0
mirror of https://github.com/MindWorkAI/AI-Studio.git synced 2026-08-12 12:52:12 +00:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

run optimise on qdrant_edge store

Browse Source
This commit is contained in:
PaulKoudelka 2026-08-04 15:59:13 +02:00
parent 2a3a82820c
commit 23416b21fc
6 changed files with 126 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
app/MindWork AI Studio/Tools/Databases/VectorStore/NoVectorStoreClient.cs
View File

@ -34,6 +34,9 @@ public sealed class NoVectorStoreClient(string name, string? unavailableReason,
public override Task DeleteEmbeddingByFile(string storeName, string filePath, CancellationToken token) =>
Task.FromException(this.CreateUnavailableException());
public override Task OptimizeVectorStore(string storeName, CancellationToken token) =>
Task.FromException(this.CreateUnavailableException());
public override Task DeleteVectorStore(string storeName, CancellationToken token) =>
Task.FromException(this.CreateUnavailableException());

6
app/MindWork AI Studio/Tools/Databases/VectorStore/QdrantEdgeClientImplementation.cs
View File

@ -17,6 +17,7 @@ public sealed class QdrantEdgeClientImplementation(
private const string INSERT_PATH = "/system/qdrant-edge/insert";
private const string SEARCH_PATH = "/system/qdrant-edge/search";
private const string DELETE_FILE_PATH = "/system/qdrant-edge/delete-file";
private const string OPTIMIZE_PATH = "/system/qdrant-edge/optimize";
private const string DELETE_STORE_PATH = "/system/qdrant-edge/delete-store";
private readonly string path = path;
@ -102,6 +103,9 @@ public sealed class QdrantEdgeClientImplementation(
public override Task DeleteEmbeddingByFile(string storeName, string filePath, CancellationToken token) =>
rustService.ExecuteDatabaseOperation(DATABASE_NAME, DELETE_FILE_PATH, new DeleteEmbeddingByFileRequest(storeName, filePath), token);
public override Task OptimizeVectorStore(string storeName, CancellationToken token) =>
rustService.ExecuteDatabaseOperation(DATABASE_NAME, OPTIMIZE_PATH, new OptimizeVectorStoreRequest(storeName), token);
public override Task DeleteVectorStore(string storeName, CancellationToken token) =>
rustService.ExecuteDatabaseOperation(DATABASE_NAME, DELETE_STORE_PATH, new DeleteVectorStoreRequest(storeName), token);
@ -124,6 +128,8 @@ public sealed class QdrantEdgeClientImplementation(
private sealed record SearchEmbeddingRequest(string StoreName, IReadOnlyList<float> Vector, int MaxMatches);
private sealed record DeleteEmbeddingByFileRequest(string StoreName, string FilePath);
private sealed record OptimizeVectorStoreRequest(string StoreName);
private sealed record DeleteVectorStoreRequest(string StoreName);
// ReSharper restore NotAccessedPositionalProperty.Local

2
app/MindWork AI Studio/Tools/Databases/VectorStore/VectorStoreClient.cs
View File

@ -10,5 +10,7 @@ public abstract class VectorStoreClient(string name, string path): DatabaseClien
public abstract Task DeleteEmbeddingByFile(string storeName, string filePath, CancellationToken token);
public abstract Task OptimizeVectorStore(string storeName, CancellationToken token);
public abstract Task DeleteVectorStore(string storeName, CancellationToken token);
}

93
app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs
View File

@ -16,6 +16,8 @@ namespace AIStudio.Tools.Services;
public sealed partial class DataSourceEmbeddingService(SettingsManager settingsManager, RustService rustService, DatabaseClientProvider databaseClientProvider, ILogger<DataSourceEmbeddingService> logger)
: BackgroundService
{
private const int VECTOR_STORE_OPTIMIZATION_CHUNK_THRESHOLD = 100_000;
private readonly Channel<DataSourceEmbeddingQueueItem> queue = Channel.CreateUnbounded<DataSourceEmbeddingQueueItem>();
private readonly ConcurrentDictionary<string, byte> queuedIds = new(StringComparer.OrdinalIgnoreCase);
private readonly ConcurrentDictionary<string, byte> runningIds = new(StringComparer.OrdinalIgnoreCase);
@ -48,6 +50,33 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
private sealed record DataSourceRunControl(CancellationTokenSource TokenSource, TaskCompletionSource<object?> Completion);
private sealed class VectorStoreOptimizationTracker
{
public long StoredChunksSinceLastOptimization { get; private set; }
public bool HasPendingChanges { get; private set; }
public void MarkChanged()
{
this.HasPendingChanges = true;
}
public void RecordStoredChunks(int chunkCount)
{
if (chunkCount <= 0)
return;
this.HasPendingChanges = true;
this.StoredChunksSinceLastOptimization += chunkCount;
}
public void Reset()
{
this.StoredChunksSinceLastOptimization = 0;
this.HasPendingChanges = false;
}
}
public IReadOnlyList<DataSourceEmbeddingStatus> GetStatuses()
{
return this.statuses.Values
@ -431,6 +460,9 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
var metadataSnapshot = this.BuildDataSourceMetadataSnapshot(dataSource, indexedFiles);
var removedMissingFiles = await this.RemoveMissingFileEmbeddingsAsync(vectorStore, embeddingState, dataSource, collectionName, manifest, indexedFiles, token);
var optimizationTracker = new VectorStoreOptimizationTracker();
if (removedMissingFiles > 0)
optimizationTracker.MarkChanged();
token.ThrowIfCancellationRequested();
logger.LogInformation(
@ -451,6 +483,14 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
dataSource.Id,
refreshMode);
await this.OptimizeCollectionIfNeededAsync(
optimizationTracker,
vectorStore,
collectionName,
dataSource,
"data source finished after removing missing files",
token);
token.ThrowIfCancellationRequested();
await embeddingState.UpdateDataSourceHashAsync(dataSource.Id, metadataSnapshot.SourceHash, token);
this.UpsertStatus(this.CreateCompletedStatus(dataSource, totalFiles, indexedFiles.Count, inputFiles.FailedFiles, inputFiles.LastError, inputFiles.Failures));
@ -511,7 +551,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
skippedFiles + completedFiles + 1,
totalFiles);
var startedAtUtc = DateTime.UtcNow;
var chunkCount = await this.IndexOneFileAsync(embeddingState, vectorStore, dataSource, file, fingerprint, embeddingProvider, provider, manifest, token);
var chunkCount = await this.IndexOneFileAsync(embeddingState, vectorStore, dataSource, file, fingerprint, embeddingProvider, provider, manifest, optimizationTracker, token);
token.ThrowIfCancellationRequested();
var embeddedAtUtc = DateTime.UtcNow;
var record = new EmbeddedFileRecord(
@ -550,6 +590,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
failureDetails.Add(new DataSourceEmbeddingFailure(file.FullName, exception.Message));
manifest.Files.Remove(file.FullName);
await this.DeleteFilePointsAsync(vectorStore, collectionName, file.FullName, token);
optimizationTracker.MarkChanged();
await embeddingState.DeleteFileAsync(dataSource.Id, file.FullName, token);
logger.LogWarning(exception, "Failed to embed file '{FilePath}' for data source '{DataSourceName}'.", file.FullName, dataSource.Name);
@ -558,6 +599,15 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
}
manifest.SourceHash = metadataSnapshot.SourceHash;
token.ThrowIfCancellationRequested();
await this.OptimizeCollectionIfNeededAsync(
optimizationTracker,
vectorStore,
collectionName,
dataSource,
"data source embedding run finished",
token);
token.ThrowIfCancellationRequested();
await embeddingState.UpdateDataSourceHashAsync(dataSource.Id, metadataSnapshot.SourceHash, token);
token.ThrowIfCancellationRequested();
@ -587,6 +637,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
EmbeddingProvider embeddingProvider,
IProvider provider,
DataSourceEmbeddingManifest manifest,
VectorStoreOptimizationTracker optimizationTracker,
CancellationToken token)
{
var collectionName = this.GetCollectionName(dataSource.Name, dataSource.Id);
@ -595,6 +646,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
file.FullName,
collectionName);
await this.DeleteFilePointsAsync(vectorStore, collectionName, file.FullName, token);
optimizationTracker.MarkChanged();
await embeddingState.DeleteFileAsync(dataSource.Id, file.FullName, token);
var parentFile = this.CreateEmbeddingStateFile(dataSource, file, fingerprint, 0, DateTime.UtcNow);
@ -610,11 +662,11 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
totalChunkCount++;
if (batch.Count >= embeddingBatchSize)
await this.FlushBatchAsync(embeddingState, vectorStore, dataSource, file, fingerprint, parentFile, embeddingProvider, provider, manifest, collectionName, batch, token);
await this.FlushBatchAsync(embeddingState, vectorStore, dataSource, file, fingerprint, parentFile, embeddingProvider, provider, manifest, optimizationTracker, collectionName, batch, token);
}
if (batch.Count > 0)
await this.FlushBatchAsync(embeddingState, vectorStore, dataSource, file, fingerprint, parentFile, embeddingProvider, provider, manifest, collectionName, batch, token);
await this.FlushBatchAsync(embeddingState, vectorStore, dataSource, file, fingerprint, parentFile, embeddingProvider, provider, manifest, optimizationTracker, collectionName, batch, token);
if (totalChunkCount == 0)
throw new InvalidOperationException($"The file '{file.Name}' did not yield any text chunks.");
@ -639,6 +691,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
EmbeddingProvider embeddingProvider,
IProvider provider,
DataSourceEmbeddingManifest manifest,
VectorStoreOptimizationTracker optimizationTracker,
string collectionName,
List<EmbeddingChunkDraft> batch,
CancellationToken token)
@ -709,6 +762,16 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
this.CreateEmbeddingStateChunks(parentFile, batch, embeddedAtUtc),
token);
optimizationTracker.RecordStoredChunks(batch.Count);
if (optimizationTracker.StoredChunksSinceLastOptimization >= VECTOR_STORE_OPTIMIZATION_CHUNK_THRESHOLD)
await this.OptimizeCollectionIfNeededAsync(
optimizationTracker,
vectorStore,
collectionName,
dataSource,
"stored chunk threshold reached",
token);
logger.LogDebug(
"Stored {ChunkCount} embedded chunks for file '{FilePath}' in collection '{CollectionName}'.",
batch.Count,
@ -766,6 +829,30 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
await vectorStore.DeleteEmbeddingByFile(collectionName, filePath, token);
}
private async Task OptimizeCollectionIfNeededAsync(
VectorStoreOptimizationTracker optimizationTracker,
VectorStoreClient vectorStore,
string collectionName,
IDataSource dataSource,
string reason,
CancellationToken token)
{
if (!optimizationTracker.HasPendingChanges)
return;
logger.LogInformation(
"Optimizing embedding collection '{CollectionName}' for data source '{DataSourceName}' ({DataSourceId}). Reason='{Reason}', StoredChunksSinceLastOptimization={StoredChunksSinceLastOptimization}, ChunkThreshold={ChunkThreshold}.",
collectionName,
dataSource.Name,
dataSource.Id,
reason,
optimizationTracker.StoredChunksSinceLastOptimization,
VECTOR_STORE_OPTIMIZATION_CHUNK_THRESHOLD);
await vectorStore.OptimizeVectorStore(collectionName, token);
optimizationTracker.Reset();
}
private async Task DeleteCollectionAsync(string collectionName, VectorStoreClient? vectorStore, CancellationToken token)
{
vectorStore ??= await databaseClientProvider.GetVectorStoreAsync(token);

24
runtime/src/qdrant_edge_database.rs
View File

@ -111,6 +111,11 @@ pub struct DeleteQdrantEdgeEmbeddingByFileRequest {
pub file_path: String,
}
#[derive(Deserialize)]
pub struct OptimizeQdrantEdgeStoreRequest {
pub store_name: String,
}
#[derive(Deserialize)]
pub struct DeleteQdrantEdgeStoreRequest {
pub store_name: String,
@ -305,6 +310,19 @@ impl QdrantEdgeDatabase {
Ok(())
}
fn optimize_store(&mut self, store_name: &str) -> QdrantEdgeResult<()> {
let Some(shard) = self.get_existing_store(store_name)? else {
return Ok(());
};
let optimized = shard.optimize()?;
if optimized {
info!(Source = "Qdrant Edge"; "Optimized vector store '{}'.", store_name);
}
shard.flush();
Ok(())
}
fn delete_store(&mut self, store_name: &str) -> QdrantEdgeResult<()> {
self.shards.remove(store_name);
@ -378,6 +396,12 @@ pub async fn delete_qdrant_edge_embedding_by_file(_token: APIToken, Json(request
})
}
pub async fn optimize_qdrant_edge_store(_token: APIToken, Json(request): Json<OptimizeQdrantEdgeStoreRequest>) -> Json<QdrantEdgeOperationResponse> {
execute_qdrant_edge_operation(|database| {
database.optimize_store(&request.store_name)
})
}
pub async fn delete_qdrant_edge_store(_token: APIToken, Json(request): Json<DeleteQdrantEdgeStoreRequest>) -> Json<QdrantEdgeOperationResponse> {
execute_qdrant_edge_operation(|database| {
database.delete_store(&request.store_name)

1
runtime/src/runtime_api.rs
View File

@ -38,6 +38,7 @@ pub fn start_runtime_api() {
.route("/system/qdrant-edge/insert", post(crate::qdrant_edge_database::insert_qdrant_edge_embedding))
.route("/system/qdrant-edge/search", post(crate::qdrant_edge_database::search_qdrant_edge_embeddings))
.route("/system/qdrant-edge/delete-file", post(crate::qdrant_edge_database::delete_qdrant_edge_embedding_by_file))
.route("/system/qdrant-edge/optimize", post(crate::qdrant_edge_database::optimize_qdrant_edge_store))
.route("/system/qdrant-edge/delete-store", post(crate::qdrant_edge_database::delete_qdrant_edge_store))
.route("/clipboard/set", post(crate::clipboard::set_clipboard))
.route("/events", get(crate::app_window::get_event_stream))