AI-Studio/app/MindWork AI Studio/Tools/Databases/DatabaseClientProvider.cs

using AIStudio.Tools.Databases.IndexStore;
using AIStudio.Tools.Databases.VectorStore;
using AIStudio.Tools.Services;

namespace AIStudio.Tools.Databases;

public sealed class DatabaseClientProvider(RustService rustService, ILoggerFactory loggerFactory) : IDisposable
{
    private readonly Dictionary<DatabaseRole, DatabaseClient> clients = new();
    private readonly Dictionary<DatabaseRole, SemaphoreSlim> locks = new();
    private readonly Lock locksLock = new();
    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)
    {
        var databaseLock = this.GetLock(databaseRole);
        await databaseLock.WaitAsync(cancellationToken);
        try
        {
            if (this.clients.TryGetValue(databaseRole, out var cachedClient) && cachedClient.IsAvailable)
                return cachedClient;

            var client = await this.CreateClientAsync(databaseRole, cancellationToken);
            return this.CacheIfAvailable(databaseRole, client);
        }
        finally
        {
            databaseLock.Release();
        }
    }

    public async Task<DatabaseClient> RefreshClientAsync(DatabaseRole databaseRole, CancellationToken cancellationToken = default)
    {
        var databaseLock = this.GetLock(databaseRole);
        await databaseLock.WaitAsync(cancellationToken);
        try
        {
            var client = await this.CreateClientAsync(databaseRole, cancellationToken);
            return this.CacheIfAvailable(databaseRole, client);
        }
        finally
        {
            databaseLock.Release();
        }
    }

    public async Task<VectorStoreClient> GetVectorStoreAsync(CancellationToken cancellationToken = default)
    {
        var client = await this.GetClientAsync(DatabaseRole.VECTOR_STORE, cancellationToken);
        if (client is VectorStoreClient vectorStore)
            return vectorStore;

        return new NoVectorStoreClient(
            client.Name,
            "The configured database client does not support vector store operations.",
            client.Status);
    }

    public async Task<IndexStoreClient> GetIndexStoreAsync(CancellationToken cancellationToken = default)
    {
        var client = await this.GetClientAsync(DatabaseRole.INDEX_STORE, cancellationToken);
        if (client is IndexStoreClient indexStore)
            return indexStore;

        return new NoIndexStoreClient(
            client.Name,
            "The configured database client does not support local RAG index operations.",
            client.Status);
    }

    /// <summary>
    /// Builds the client which stands in for a database role that cannot serve right now.
    /// </summary>
    /// <remarks>
    /// Callers outside this namespace get their stand-in from here instead of naming the concrete
    /// type themselves, so a new role does not have to be spelled out in every one of them.
    /// </remarks>
    /// <param name="databaseRole">The role the stand-in has to fill.</param>
    /// <param name="name">The name to show for the database.</param>
    /// <param name="reason">Why the database is not available.</param>
    /// <param name="status">Whether the database is starting or unavailable.</param>
    /// <returns>A client which answers every operation without a database behind it.</returns>
    public static DatabaseClient CreateUnavailableClient(DatabaseRole databaseRole, string name, string? reason, DatabaseClientStatus status) => databaseRole switch
    {
        DatabaseRole.VECTOR_STORE => new NoVectorStoreClient(name, reason, status),
        DatabaseRole.INDEX_STORE => new NoIndexStoreClient(name, reason, status),
        _ => new NoDatabaseClient(name, reason, status)
    };

    private DatabaseClient CacheIfAvailable(DatabaseRole databaseRole, DatabaseClient client)
    {
        if (!client.IsAvailable)
            return client;

        if (this.clients.TryGetValue(databaseRole, out var cachedClient))
        {
            if (IsSameClient(cachedClient, client))
            {
                client.Dispose();
                return cachedClient;
            }

            cachedClient.Dispose();
        }

        this.clients[databaseRole] = client;
        return client;
    }

    private SemaphoreSlim GetLock(DatabaseRole databaseRole)
    {
        lock (this.locksLock)
        {
            if (this.locks.TryGetValue(databaseRole, out var databaseLock))
                return databaseLock;

            databaseLock = new SemaphoreSlim(1, 1);
            this.locks[databaseRole] = databaseLock;
            return databaseLock;
        }
    }

    private async Task<DatabaseClient> CreateClientAsync(DatabaseRole databaseRole, CancellationToken cancellationToken) => databaseRole switch
    {
        DatabaseRole.VECTOR_STORE => await QdrantEdgeClientImplementation.CreateAsync(rustService, this.GetIndexStoreAsync, this.logger, this.databaseClientLogger, cancellationToken),
        DatabaseRole.INDEX_STORE => await SqliteIndexStoreClientImplementation.CreateAsync(this.logger, this.databaseClientLogger, cancellationToken),
        _ => new NoDatabaseClient(databaseRole.ToString(), "The requested database role is not supported.")
    };

    private static bool IsSameClient(DatabaseClient left, DatabaseClient right) =>
        left.IsAvailable
        && right.IsAvailable
        && left.CacheKey == right.CacheKey;

    public void Dispose()
    {
        foreach (var client in this.clients.Values)
            client.Dispose();

        foreach (var databaseLock in this.locks.Values)
            databaseLock.Dispose();
    }
}