mirror of
https://github.com/MindWorkAI/AI-Studio.git
synced 2026-08-24 20:52:11 +00:00
made check for vectordb more robust
This commit is contained in:
parent
92e986d9b9
commit
987f493848
@ -22,8 +22,8 @@ public sealed class NoVectorStoreClient(string name, string? unavailableReason,
|
|||||||
await Task.CompletedTask;
|
await Task.CompletedTask;
|
||||||
}
|
}
|
||||||
|
|
||||||
public override Task EnsureVectorStoreExists(string storeName, string dataSourceName, int vectorSize, CancellationToken token) =>
|
public override Task<VectorStoreEnsureResult> EnsureVectorStoreExists(string storeName, string dataSourceName, int vectorSize, CancellationToken token) =>
|
||||||
Task.FromException(this.CreateUnavailableException());
|
Task.FromException<VectorStoreEnsureResult>(this.CreateUnavailableException());
|
||||||
|
|
||||||
public override Task InsertEmbedding(string storeName, IReadOnlyList<VectorStoragePoint> points, CancellationToken token) =>
|
public override Task InsertEmbedding(string storeName, IReadOnlyList<VectorStoragePoint> points, CancellationToken token) =>
|
||||||
Task.FromException(this.CreateUnavailableException());
|
Task.FromException(this.CreateUnavailableException());
|
||||||
|
|||||||
@ -82,8 +82,9 @@ public sealed class QdrantEdgeClientImplementation(
|
|||||||
yield return (TB("Number of vector stores"), displayStoresCount.ToString());
|
yield return (TB("Number of vector stores"), displayStoresCount.ToString());
|
||||||
}
|
}
|
||||||
|
|
||||||
public override Task EnsureVectorStoreExists(string storeName, string dataSourceName, int vectorSize, CancellationToken token) =>
|
public override async Task<VectorStoreEnsureResult> EnsureVectorStoreExists(string storeName, string dataSourceName, int vectorSize, CancellationToken token) =>
|
||||||
rustService.ExecuteDatabaseOperation(DATABASE_NAME, ENSURE_PATH, new EnsureVectorStoreRequest(storeName, dataSourceName, vectorSize), token);
|
await rustService.ExecuteDatabaseQuery<EnsureVectorStoreRequest, VectorStoreEnsureResult>( 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<VectorStoragePoint> points, CancellationToken token) =>
|
public override Task InsertEmbedding(string storeName, IReadOnlyList<VectorStoragePoint> points, CancellationToken token) =>
|
||||||
rustService.ExecuteDatabaseOperation(DATABASE_NAME, INSERT_PATH, new InsertEmbeddingRequest(storeName, points), token);
|
rustService.ExecuteDatabaseOperation(DATABASE_NAME, INSERT_PATH, new InsertEmbeddingRequest(storeName, points), token);
|
||||||
|
|||||||
@ -2,7 +2,7 @@
|
|||||||
|
|
||||||
public abstract class VectorStoreClient(string name, string path): DatabaseClient(name, path)
|
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<VectorStoreEnsureResult> EnsureVectorStoreExists(string storeName, string dataSourceName, int vectorSize, CancellationToken token);
|
||||||
|
|
||||||
public abstract Task InsertEmbedding(string storeName, IReadOnlyList<VectorStoragePoint> points, CancellationToken token);
|
public abstract Task InsertEmbedding(string storeName, IReadOnlyList<VectorStoragePoint> 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 abstract Task DeleteVectorStore(string storeName, CancellationToken token);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
public sealed record VectorStoreEnsureResult(bool Created);
|
||||||
|
|||||||
@ -428,6 +428,22 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
|
|||||||
return;
|
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))
|
if (!this.TryResolveEmbeddingProvider(dataSource, out var embeddingProvider))
|
||||||
{
|
{
|
||||||
token.ThrowIfCancellationRequested();
|
token.ThrowIfCancellationRequested();
|
||||||
@ -460,11 +476,8 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
|
|||||||
dataSource.Name,
|
dataSource.Name,
|
||||||
dataSource.Id);
|
dataSource.Id);
|
||||||
|
|
||||||
var collectionName = this.GetCollectionName(dataSource.Id);
|
|
||||||
var manifest = await this.EnsureCompatibleManifestAsync(dataSource, embeddingProvider, collectionName, vectorStore, embeddingState, token);
|
var manifest = await this.EnsureCompatibleManifestAsync(dataSource, embeddingProvider, collectionName, vectorStore, embeddingState, token);
|
||||||
token.ThrowIfCancellationRequested();
|
token.ThrowIfCancellationRequested();
|
||||||
if (manifest.VectorSize > 0)
|
|
||||||
await this.EnsureCollectionExistsAsync(vectorStore, collectionName, dataSource.Name, manifest.VectorSize, token);
|
|
||||||
|
|
||||||
var inputFiles = this.GetInputFiles(dataSource);
|
var inputFiles = this.GetInputFiles(dataSource);
|
||||||
var indexedFiles = inputFiles.Files;
|
var indexedFiles = inputFiles.Files;
|
||||||
@ -770,7 +783,20 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
|
|||||||
if (manifest.VectorSize == 0)
|
if (manifest.VectorSize == 0)
|
||||||
{
|
{
|
||||||
token.ThrowIfCancellationRequested();
|
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);
|
await embeddingState.UpdateVectorSizeAsync(dataSource.Id, vectorSize, token);
|
||||||
manifest.VectorSize = vectorSize;
|
manifest.VectorSize = vectorSize;
|
||||||
logger.LogInformation(
|
logger.LogInformation(
|
||||||
@ -819,9 +845,9 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
|
|||||||
batch.Clear();
|
batch.Clear();
|
||||||
}
|
}
|
||||||
|
|
||||||
private async Task EnsureCollectionExistsAsync(VectorStoreClient vectorStore, string collectionName, string dataSourceName, int vectorSize, CancellationToken token)
|
private async Task<VectorStoreEnsureResult> 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(
|
private async Task UpsertPointsAsync(
|
||||||
|
|||||||
@ -132,6 +132,18 @@ pub struct QdrantEdgeOperationResponse {
|
|||||||
pub issue: String,
|
pub issue: String,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Serialize)]
|
||||||
|
pub struct QdrantEdgeEnsureStoreResponse {
|
||||||
|
pub success: bool,
|
||||||
|
pub issue: String,
|
||||||
|
pub data: Option<QdrantEdgeEnsureStoreResult>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Serialize)]
|
||||||
|
pub struct QdrantEdgeEnsureStoreResult {
|
||||||
|
pub created: bool,
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Serialize)]
|
#[derive(Serialize)]
|
||||||
pub struct QdrantEdgeSearchResponse {
|
pub struct QdrantEdgeSearchResponse {
|
||||||
pub success: bool,
|
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<QdrantEdgeEnsureStoreResult> {
|
||||||
validate_vector_size(vector_size)?;
|
validate_vector_size(vector_size)?;
|
||||||
validate_data_source_name(data_source_name)?;
|
validate_data_source_name(data_source_name)?;
|
||||||
let store_path = self.store_path(store_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)?;
|
self.get_or_create_store(store_name, vector_size)?;
|
||||||
write_store_display_name(&store_path, data_source_name)?;
|
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<QdrantEdgeStoragePoint>) -> QdrantEdgeResult<()> {
|
fn insert_embedding(&mut self, store_name: &str, points: Vec<QdrantEdgeStoragePoint>) -> QdrantEdgeResult<()> {
|
||||||
@ -439,10 +454,32 @@ pub async fn qdrant_edge_info(_token: APIToken) -> Json<QdrantEdgeServiceInfo> {
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn ensure_qdrant_edge_store(_token: APIToken, Json(request): Json<EnsureQdrantEdgeStoreRequest>) -> Json<QdrantEdgeOperationResponse> {
|
pub async fn ensure_qdrant_edge_store(_token: APIToken, Json(request): Json<EnsureQdrantEdgeStoreRequest>) -> Json<QdrantEdgeEnsureStoreResponse> {
|
||||||
execute_qdrant_edge_operation(|database| {
|
let mut database_guard = QDRANT_EDGE_DATABASE.lock().unwrap();
|
||||||
database.ensure_store_exists(&request.store_name, &request.data_source_name, request.vector_size)
|
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<InsertQdrantEdgeEmbeddingRequest>) -> Json<QdrantEdgeOperationResponse> {
|
pub async fn insert_qdrant_edge_embedding(_token: APIToken, Json(request): Json<InsertQdrantEdgeEmbeddingRequest>) -> Json<QdrantEdgeOperationResponse> {
|
||||||
@ -929,6 +966,31 @@ mod tests {
|
|||||||
assert!(validate_data_source_name(&"a".repeat(41)).is_err());
|
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]
|
#[test]
|
||||||
fn point_ids_must_be_valid_uuids() {
|
fn point_ids_must_be_valid_uuids() {
|
||||||
assert!(to_point_id("6cc665a8-2b1e-4d42-bc74-8015b7b391ec").is_ok());
|
assert!(to_point_id("6cc665a8-2b1e-4d42-bc74-8015b7b391ec").is_ok());
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user