using MySql.Data.MySqlClient; using System.Data.Common; using Server.Migrations; using Server.Helper; using Server.Exceptions; using AdaptiveExpressions; using Shared.Models; using System.Text.Json; using Microsoft.Extensions.Options; using Server.Models; using Shared; using System.Diagnostics; namespace Server; public class SearchdomainManager : IDisposable { private Dictionary searchdomains = []; private readonly ILogger _logger; private readonly EmbeddingSearchOptions _options; public readonly AIProvider aIProvider; private readonly DatabaseHelper _databaseHelper; private readonly string connectionString; private MySqlConnection connection; public SQLHelper helper; public EnumerableLruCache> embeddingCache; public long EmbeddingCacheMaxCount; private bool disposed = false; public SearchdomainManager(ILogger logger, IOptions options, AIProvider aIProvider, DatabaseHelper databaseHelper) { _logger = logger; _options = options.Value; this.aIProvider = aIProvider; _databaseHelper = databaseHelper; EmbeddingCacheMaxCount = _options.Cache.CacheTopN; if (options.Value.Cache.StoreEmbeddingCache) { var stopwatch = Stopwatch.StartNew(); embeddingCache = CacheHelper.GetEmbeddingStore(options.Value); stopwatch.Stop(); _logger.LogInformation("GetEmbeddingStore completed in {ElapsedMilliseconds} ms", stopwatch.ElapsedMilliseconds); } else { embeddingCache = new((int)EmbeddingCacheMaxCount); } connectionString = _options.ConnectionStrings.SQL; connection = new MySqlConnection(connectionString); connection.Open(); helper = new SQLHelper(connection, connectionString); } public Searchdomain GetSearchdomain(string searchdomain) { if (searchdomains.TryGetValue(searchdomain, out Searchdomain? value)) { return value; } try { return SetSearchdomain(searchdomain, new Searchdomain(searchdomain, connectionString, helper, aIProvider, embeddingCache, _logger)); } catch (MySqlException) { _logger.LogError("Unable to find the searchdomain {searchdomain}", searchdomain); throw new SearchdomainNotFoundException(searchdomain); } catch (Exception ex) { _logger.LogError("Unable to load the searchdomain {searchdomain} due to the following exception: {ex}", [searchdomain, ex.Message]); throw; } } public void InvalidateSearchdomainCache(string searchdomainName) { var searchdomain = GetSearchdomain(searchdomainName); searchdomain.UpdateEntityCache(); searchdomain.InvalidateSearchCache(); } public async Task> ListSearchdomainsAsync() { return await helper.ExecuteQueryAsync("SELECT name FROM searchdomain", [], x => x.GetString(0)); } public async Task CreateSearchdomain(string searchdomain, SearchdomainSettings settings) { return await CreateSearchdomain(searchdomain, JsonSerializer.Serialize(settings)); } public async Task CreateSearchdomain(string searchdomain, string settings = "{}") { if (searchdomains.TryGetValue(searchdomain, out Searchdomain? value)) { _logger.LogError("Searchdomain {searchdomain} could not be created, as it already exists", [searchdomain]); throw new SearchdomainAlreadyExistsException(searchdomain); } Dictionary parameters = new() { { "name", searchdomain }, { "settings", settings} }; int id = await helper.ExecuteSQLCommandGetInsertedID("INSERT INTO searchdomain (name, settings) VALUES (@name, @settings)", parameters); searchdomains.Add(searchdomain, new(searchdomain, connectionString, helper, aIProvider, embeddingCache, _logger)); return id; } public async Task DeleteSearchdomain(string searchdomain) { int counter = await _databaseHelper.RemoveAllEntities(helper, searchdomain); _logger.LogDebug($"Number of entities deleted as part of deleting the searchdomain \"{searchdomain}\": {counter}"); await helper.ExecuteSQLNonQuery("DELETE FROM searchdomain WHERE name = @name", new() {{"name", searchdomain}}); searchdomains.Remove(searchdomain); _logger.LogDebug($"Searchdomain has been successfully removed"); return counter; } private Searchdomain SetSearchdomain(string name, Searchdomain searchdomain) { searchdomains[name] = searchdomain; return searchdomain; } public bool IsSearchdomainLoaded(string name) { return searchdomains.ContainsKey(name); } // Cleanup procedure private async Task Cleanup() { try { if (_options.Cache.StoreEmbeddingCache) { var stopwatch = Stopwatch.StartNew(); await CacheHelper.UpdateEmbeddingStore(embeddingCache, _options); stopwatch.Stop(); _logger.LogInformation("UpdateEmbeddingStore completed in {ElapsedMilliseconds} ms", stopwatch.ElapsedMilliseconds); } _logger.LogInformation("SearchdomainManager cleanup completed"); } catch (Exception ex) { _logger.LogError(ex, "Error during SearchdomainManager cleanup"); } } public void Dispose() { Dispose(true).Wait(); GC.SuppressFinalize(this); } protected virtual async Task Dispose(bool disposing) { if (!disposed && disposing) { await Cleanup(); disposed = true; } } }