diff --git a/app/MindWork AI Studio/Tools/Databases/VectorStore/NoVectorStoreClient.cs b/app/MindWork AI Studio/Tools/Databases/VectorStore/NoVectorStoreClient.cs index 150bf6e3..fcf0baac 100644 --- a/app/MindWork AI Studio/Tools/Databases/VectorStore/NoVectorStoreClient.cs +++ b/app/MindWork AI Studio/Tools/Databases/VectorStore/NoVectorStoreClient.cs @@ -22,8 +22,8 @@ public sealed class NoVectorStoreClient(string name, string? unavailableReason, await Task.CompletedTask; } - public override Task EnsureVectorStoreExists(string storeName, string dataSourceName, int vectorSize, CancellationToken token) => - Task.FromException(this.CreateUnavailableException()); + public override Task EnsureVectorStoreExists(string storeName, string dataSourceName, int vectorSize, CancellationToken token) => + Task.FromException(this.CreateUnavailableException()); public override Task InsertEmbedding(string storeName, IReadOnlyList points, CancellationToken token) => Task.FromException(this.CreateUnavailableException()); diff --git a/app/MindWork AI Studio/Tools/Databases/VectorStore/QdrantEdgeClientImplementation.cs b/app/MindWork AI Studio/Tools/Databases/VectorStore/QdrantEdgeClientImplementation.cs index fc04cd32..a72ef8cd 100644 --- a/app/MindWork AI Studio/Tools/Databases/VectorStore/QdrantEdgeClientImplementation.cs +++ b/app/MindWork AI Studio/Tools/Databases/VectorStore/QdrantEdgeClientImplementation.cs @@ -82,8 +82,9 @@ public sealed class QdrantEdgeClientImplementation( yield return (TB("Number of vector stores"), displayStoresCount.ToString()); } - public override Task EnsureVectorStoreExists(string storeName, string dataSourceName, int vectorSize, CancellationToken token) => - rustService.ExecuteDatabaseOperation(DATABASE_NAME, ENSURE_PATH, new EnsureVectorStoreRequest(storeName, dataSourceName, vectorSize), token); + public override async Task EnsureVectorStoreExists(string storeName, string dataSourceName, int vectorSize, CancellationToken token) => + await rustService.ExecuteDatabaseQuery( DATABASE_NAME, ENSURE_PATH, + new EnsureVectorStoreRequest(storeName, dataSourceName, vectorSize), token) ?? throw new InvalidOperationException("The vector store ensure response was empty."); public override Task InsertEmbedding(string storeName, IReadOnlyList points, CancellationToken token) => rustService.ExecuteDatabaseOperation(DATABASE_NAME, INSERT_PATH, new InsertEmbeddingRequest(storeName, points), token); diff --git a/app/MindWork AI Studio/Tools/Databases/VectorStore/VectorStoreClient.cs b/app/MindWork AI Studio/Tools/Databases/VectorStore/VectorStoreClient.cs index a9b6017c..1fef76ab 100644 --- a/app/MindWork AI Studio/Tools/Databases/VectorStore/VectorStoreClient.cs +++ b/app/MindWork AI Studio/Tools/Databases/VectorStore/VectorStoreClient.cs @@ -2,7 +2,7 @@ public abstract class VectorStoreClient(string name, string path): DatabaseClient(name, path) { - public abstract Task EnsureVectorStoreExists(string storeName, string dataSourceName, int vectorSize, CancellationToken token); + public abstract Task EnsureVectorStoreExists(string storeName, string dataSourceName, int vectorSize, CancellationToken token); public abstract Task InsertEmbedding(string storeName, IReadOnlyList points, CancellationToken token); @@ -14,3 +14,5 @@ public abstract class VectorStoreClient(string name, string path): DatabaseClien public abstract Task DeleteVectorStore(string storeName, CancellationToken token); } + +public sealed record VectorStoreEnsureResult(bool Created); diff --git a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs b/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs index 94a71317..244dcf9e 100644 --- a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs +++ b/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs @@ -428,6 +428,22 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM return; } + var collectionName = this.GetCollectionName(dataSource.Id); + var persistedManifest = await embeddingState.GetManifestAsync(dataSource.Id, token); + if (persistedManifest.VectorSize > 0) + { + var ensureResult = await this.EnsureCollectionExistsAsync(vectorStore, collectionName, dataSource.Name, persistedManifest.VectorSize, token); + if (ensureResult.Created) + { + logger.LogWarning( + "Vector store '{CollectionName}' for data source '{DataSourceName}' ({DataSourceId}) was missing although persisted embedding state exists. Resetting the stale state so all vectors are rebuilt.", + collectionName, + dataSource.Name, + dataSource.Id); + await this.ResetPersistedStateAsync(dataSource.Id, vectorStore, embeddingState, token); + } + } + if (!this.TryResolveEmbeddingProvider(dataSource, out var embeddingProvider)) { token.ThrowIfCancellationRequested(); @@ -460,11 +476,8 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM dataSource.Name, dataSource.Id); - var collectionName = this.GetCollectionName(dataSource.Id); var manifest = await this.EnsureCompatibleManifestAsync(dataSource, embeddingProvider, collectionName, vectorStore, embeddingState, token); token.ThrowIfCancellationRequested(); - if (manifest.VectorSize > 0) - await this.EnsureCollectionExistsAsync(vectorStore, collectionName, dataSource.Name, manifest.VectorSize, token); var inputFiles = this.GetInputFiles(dataSource); var indexedFiles = inputFiles.Files; @@ -770,7 +783,20 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM if (manifest.VectorSize == 0) { token.ThrowIfCancellationRequested(); - await this.EnsureCollectionExistsAsync(vectorStore, collectionName, dataSource.Name, vectorSize, token); + var ensureResult = await this.EnsureCollectionExistsAsync(vectorStore, collectionName, dataSource.Name, vectorSize, token); + if (!ensureResult.Created) + { + logger.LogWarning( + "Vector store '{CollectionName}' exists for data source '{DataSourceName}' ({DataSourceId}) although no persisted embedding state exists. Replacing the orphaned store before indexing.", + collectionName, + dataSource.Name, + dataSource.Id); + await vectorStore.DeleteVectorStore(collectionName, token); + ensureResult = await this.EnsureCollectionExistsAsync(vectorStore, collectionName, dataSource.Name, vectorSize, token); + if (!ensureResult.Created) + throw new InvalidOperationException($"Vector store '{collectionName}' could not be recreated cleanly."); + } + await embeddingState.UpdateVectorSizeAsync(dataSource.Id, vectorSize, token); manifest.VectorSize = vectorSize; logger.LogInformation( @@ -819,9 +845,9 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM batch.Clear(); } - private async Task EnsureCollectionExistsAsync(VectorStoreClient vectorStore, string collectionName, string dataSourceName, int vectorSize, CancellationToken token) + private async Task EnsureCollectionExistsAsync(VectorStoreClient vectorStore, string collectionName, string dataSourceName, int vectorSize, CancellationToken token) { - await vectorStore.EnsureVectorStoreExists(collectionName, dataSourceName, vectorSize, token); + return await vectorStore.EnsureVectorStoreExists(collectionName, dataSourceName, vectorSize, token); } private async Task UpsertPointsAsync( diff --git a/runtime/src/qdrant_edge_database.rs b/runtime/src/qdrant_edge_database.rs index 67d725ad..b92ae490 100644 --- a/runtime/src/qdrant_edge_database.rs +++ b/runtime/src/qdrant_edge_database.rs @@ -132,6 +132,18 @@ pub struct QdrantEdgeOperationResponse { pub issue: String, } +#[derive(Serialize)] +pub struct QdrantEdgeEnsureStoreResponse { + pub success: bool, + pub issue: String, + pub data: Option, +} + +#[derive(Serialize)] +pub struct QdrantEdgeEnsureStoreResult { + pub created: bool, +} + #[derive(Serialize)] pub struct QdrantEdgeSearchResponse { pub success: bool, @@ -289,13 +301,16 @@ impl QdrantEdgeDatabase { }) } - fn ensure_store_exists(&mut self, store_name: &str, data_source_name: &str, vector_size: usize) -> QdrantEdgeResult<()> { + fn ensure_store_exists(&mut self, store_name: &str, data_source_name: &str, vector_size: usize) -> QdrantEdgeResult { validate_vector_size(vector_size)?; validate_data_source_name(data_source_name)?; let store_path = self.store_path(store_name)?; + let store_existed = store_is_initialized(&store_path, store_name)?; self.get_or_create_store(store_name, vector_size)?; write_store_display_name(&store_path, data_source_name)?; - Ok(()) + Ok(QdrantEdgeEnsureStoreResult { + created: !store_existed, + }) } fn insert_embedding(&mut self, store_name: &str, points: Vec) -> QdrantEdgeResult<()> { @@ -439,10 +454,32 @@ pub async fn qdrant_edge_info(_token: APIToken) -> Json { }) } -pub async fn ensure_qdrant_edge_store(_token: APIToken, Json(request): Json) -> Json { - execute_qdrant_edge_operation(|database| { - database.ensure_store_exists(&request.store_name, &request.data_source_name, request.vector_size) - }) +pub async fn ensure_qdrant_edge_store(_token: APIToken, Json(request): Json) -> Json { + let mut database_guard = QDRANT_EDGE_DATABASE.lock().unwrap(); + let Some(database) = database_guard.as_mut() else { + return Json(QdrantEdgeEnsureStoreResponse { + success: false, + issue: "Qdrant Edge is not available.".to_string(), + data: None, + }); + }; + + match database.ensure_store_exists(&request.store_name, &request.data_source_name, request.vector_size) { + Ok(result) => Json(QdrantEdgeEnsureStoreResponse { + success: true, + issue: String::new(), + data: Some(result), + }), + Err(error) => { + let issue = error.to_string(); + error!(Source = "Qdrant Edge"; "Qdrant Edge operation failed: {issue}"); + Json(QdrantEdgeEnsureStoreResponse { + success: false, + issue, + data: None, + }) + }, + } } pub async fn insert_qdrant_edge_embedding(_token: APIToken, Json(request): Json) -> Json { @@ -929,6 +966,31 @@ mod tests { assert!(validate_data_source_name(&"a".repeat(41)).is_err()); } + #[test] + fn ensure_store_reports_creation_and_updates_the_display_name() { + let test_directory = std::env::temp_dir().join(format!( + "ai-studio-qdrant-ensure-{}-{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + let store_name = "rag_6cc665a82b1e4d42bc748015b7b391ec"; + let mut database = QdrantEdgeDatabase::new(test_directory.clone()); + + let created = database.ensure_store_exists(store_name, "Original name", 3).unwrap(); + assert!(created.created); + + let existing = database.ensure_store_exists(store_name, "Renamed source", 3).unwrap(); + assert!(!existing.created); + let display_name_path = database.store_path(store_name).unwrap().join(STORE_DISPLAY_NAME_MARKER); + assert_eq!(fs::read_to_string(display_name_path).unwrap(), "Renamed source"); + + drop(database); + fs::remove_dir_all(test_directory).unwrap(); + } + #[test] fn point_ids_must_be_valid_uuids() { assert!(to_point_id("6cc665a8-2b1e-4d42-bc74-8015b7b391ec").is_ok());