mirror of
https://github.com/MindWorkAI/AI-Studio.git
synced 2026-10-16 14:53:48 +00:00
added sql integration instead of json files
This commit is contained in:
PaulKoudelka committed
1 parent
2e06af13e0
commit
c59e33e1b1
22 files changed
+700
-156
No files matched your search
+345
@@ -0,0 +1,345 @@
|
||||
using System.Globalization;
|
||||
|
||||
using AIStudio.Settings;
|
||||
using AIStudio.Tools.PluginSystem;
|
||||
using AIStudio.Tools.Services;
|
||||
|
||||
using Microsoft.Data.Sqlite;
|
||||
|
||||
namespace AIStudio.Tools.Databases.EmbeddingState;
|
||||
|
||||
public sealed class SqliteEmbeddingStateClientImplementation(
|
||||
string name,
|
||||
string databasePath,
|
||||
string basePath,
|
||||
string version) : EmbeddingStateClient(name, basePath)
|
||||
{
|
||||
private const string DATABASE_NAME = "SQLite";
|
||||
private const string DATABASE_FILENAME = "rag-embedding-state.sqlite3";
|
||||
|
||||
private readonly string databasePath = databasePath;
|
||||
private readonly string connectionString = new SqliteConnectionStringBuilder
|
||||
{
|
||||
DataSource = databasePath,
|
||||
Mode = SqliteOpenMode.ReadWriteCreate,
|
||||
Cache = SqliteCacheMode.Shared,
|
||||
ForeignKeys = true,
|
||||
DefaultTimeout = 30,
|
||||
}.ToString();
|
||||
|
||||
private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(SqliteEmbeddingStateClientImplementation).Namespace, nameof(SqliteEmbeddingStateClientImplementation));
|
||||
|
||||
public override string CacheKey => $"{this.Name}:{this.databasePath}:{version}";
|
||||
|
||||
public static async Task<DatabaseClient> CreateAsync(
|
||||
ILogger logger,
|
||||
ILogger<DatabaseClient> databaseClientLogger,
|
||||
CancellationToken cancellationToken)
|
||||
{
|
||||
if (string.IsNullOrWhiteSpace(SettingsManager.DataDirectory))
|
||||
return CreateNoEmbeddingStateClient(DATABASE_NAME, "The application data directory is not available yet.", DatabaseClientStatus.STARTING, databaseClientLogger);
|
||||
|
||||
try
|
||||
{
|
||||
SQLitePCL.Batteries_V2.Init();
|
||||
|
||||
var basePath = Path.Combine(SettingsManager.DataDirectory, "databases", "sqlite");
|
||||
Directory.CreateDirectory(basePath);
|
||||
|
||||
var databasePath = Path.Combine(basePath, DATABASE_FILENAME);
|
||||
var client = new SqliteEmbeddingStateClientImplementation(DATABASE_NAME, databasePath, basePath, string.Empty);
|
||||
await client.InitializeAsync(cancellationToken);
|
||||
var version = await client.GetSqliteVersionAsync(cancellationToken);
|
||||
|
||||
client = new SqliteEmbeddingStateClientImplementation(DATABASE_NAME, databasePath, basePath, version);
|
||||
client.SetLogger(databaseClientLogger);
|
||||
return client;
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
logger.LogWarning(exception, "{DatabaseName} is not available. Indexed file fingerprints are disabled.", DATABASE_NAME);
|
||||
return CreateNoEmbeddingStateClient(DATABASE_NAME, exception.Message, DatabaseClientStatus.UNAVAILABLE, databaseClientLogger);
|
||||
}
|
||||
}
|
||||
|
||||
public override async IAsyncEnumerable<(string Label, string Value)> GetDisplayInfo()
|
||||
{
|
||||
yield return (TB("Reported version"), version);
|
||||
yield return (TB("Database path"), this.databasePath);
|
||||
yield return (TB("Storage size"), this.GetStorageSize());
|
||||
yield return (TB("Indexed data sources"), (await this.CountAsync("data_sources", CancellationToken.None)).ToString(CultureInfo.InvariantCulture));
|
||||
yield return (TB("Indexed files"), (await this.CountAsync("embedded_files", CancellationToken.None)).ToString(CultureInfo.InvariantCulture));
|
||||
}
|
||||
|
||||
public override async Task<DataSourceEmbeddingManifest> GetManifestAsync(string dataSourceId, CancellationToken token)
|
||||
{
|
||||
await using var connection = await this.OpenConnectionAsync(token);
|
||||
var manifest = new DataSourceEmbeddingManifest();
|
||||
|
||||
await using (var command = CreateCommand(connection, """
|
||||
SELECT embedding_provider_id, embedding_signature, vector_size
|
||||
FROM data_sources
|
||||
WHERE data_source_id = $dataSourceId
|
||||
"""))
|
||||
{
|
||||
command.Parameters.AddWithValue("$dataSourceId", dataSourceId);
|
||||
await using var reader = await command.ExecuteReaderAsync(token);
|
||||
if (await reader.ReadAsync(token))
|
||||
{
|
||||
manifest.EmbeddingProviderId = reader.GetString(0);
|
||||
manifest.EmbeddingSignature = reader.GetString(1);
|
||||
manifest.VectorSize = reader.GetInt32(2);
|
||||
}
|
||||
}
|
||||
|
||||
await using (var command = CreateCommand(connection, """
|
||||
SELECT file_path, fingerprint, file_size, last_write_utc, embedded_at_utc, chunk_count
|
||||
FROM embedded_files
|
||||
WHERE data_source_id = $dataSourceId
|
||||
"""))
|
||||
{
|
||||
command.Parameters.AddWithValue("$dataSourceId", dataSourceId);
|
||||
await using var reader = await command.ExecuteReaderAsync(token);
|
||||
while (await reader.ReadAsync(token))
|
||||
{
|
||||
manifest.Files[reader.GetString(0)] = new EmbeddedFileRecord(
|
||||
reader.GetString(1),
|
||||
reader.GetInt64(2),
|
||||
ParseUtc(reader.GetString(3)),
|
||||
ParseUtc(reader.GetString(4)),
|
||||
reader.GetInt32(5));
|
||||
}
|
||||
}
|
||||
|
||||
return manifest;
|
||||
}
|
||||
|
||||
public override async Task UpsertDataSourceAsync(
|
||||
string dataSourceId,
|
||||
string dataSourceName,
|
||||
string dataSourceType,
|
||||
string embeddingProviderId,
|
||||
string embeddingSignature,
|
||||
int vectorSize,
|
||||
CancellationToken token)
|
||||
{
|
||||
await using var connection = await this.OpenConnectionAsync(token);
|
||||
await ExecuteNonQueryAsync(connection, """
|
||||
INSERT INTO data_sources (
|
||||
data_source_id,
|
||||
data_source_name,
|
||||
data_source_type,
|
||||
embedding_provider_id,
|
||||
embedding_signature,
|
||||
vector_size,
|
||||
updated_at_utc)
|
||||
VALUES (
|
||||
$dataSourceId,
|
||||
$dataSourceName,
|
||||
$dataSourceType,
|
||||
$embeddingProviderId,
|
||||
$embeddingSignature,
|
||||
$vectorSize,
|
||||
$updatedAtUtc)
|
||||
ON CONFLICT(data_source_id) DO UPDATE SET
|
||||
data_source_name = excluded.data_source_name,
|
||||
data_source_type = excluded.data_source_type,
|
||||
embedding_provider_id = excluded.embedding_provider_id,
|
||||
embedding_signature = excluded.embedding_signature,
|
||||
vector_size = excluded.vector_size,
|
||||
updated_at_utc = excluded.updated_at_utc
|
||||
""", token,
|
||||
("$dataSourceId", dataSourceId),
|
||||
("$dataSourceName", dataSourceName),
|
||||
("$dataSourceType", dataSourceType),
|
||||
("$embeddingProviderId", embeddingProviderId),
|
||||
("$embeddingSignature", embeddingSignature),
|
||||
("$vectorSize", vectorSize),
|
||||
("$updatedAtUtc", ToUtcText(DateTime.UtcNow)));
|
||||
}
|
||||
|
||||
public override async Task UpdateVectorSizeAsync(string dataSourceId, int vectorSize, CancellationToken token)
|
||||
{
|
||||
await using var connection = await this.OpenConnectionAsync(token);
|
||||
await ExecuteNonQueryAsync(connection, """
|
||||
UPDATE data_sources
|
||||
SET vector_size = $vectorSize,
|
||||
updated_at_utc = $updatedAtUtc
|
||||
WHERE data_source_id = $dataSourceId
|
||||
""", token,
|
||||
("$dataSourceId", dataSourceId),
|
||||
("$vectorSize", vectorSize),
|
||||
("$updatedAtUtc", ToUtcText(DateTime.UtcNow)));
|
||||
}
|
||||
|
||||
public override async Task UpsertFileAsync(string dataSourceId, EmbeddingStateFile file, CancellationToken token)
|
||||
{
|
||||
await using var connection = await this.OpenConnectionAsync(token);
|
||||
await ExecuteNonQueryAsync(connection, """
|
||||
INSERT INTO embedded_files (
|
||||
data_source_id,
|
||||
file_path,
|
||||
file_name,
|
||||
relative_path,
|
||||
fingerprint,
|
||||
file_size,
|
||||
last_write_utc,
|
||||
embedded_at_utc,
|
||||
chunk_count)
|
||||
VALUES (
|
||||
$dataSourceId,
|
||||
$filePath,
|
||||
$fileName,
|
||||
$relativePath,
|
||||
$fingerprint,
|
||||
$fileSize,
|
||||
$lastWriteUtc,
|
||||
$embeddedAtUtc,
|
||||
$chunkCount)
|
||||
ON CONFLICT(data_source_id, file_path) DO UPDATE SET
|
||||
file_name = excluded.file_name,
|
||||
relative_path = excluded.relative_path,
|
||||
fingerprint = excluded.fingerprint,
|
||||
file_size = excluded.file_size,
|
||||
last_write_utc = excluded.last_write_utc,
|
||||
embedded_at_utc = excluded.embedded_at_utc,
|
||||
chunk_count = excluded.chunk_count
|
||||
""", token,
|
||||
("$dataSourceId", dataSourceId),
|
||||
("$filePath", file.FilePath),
|
||||
("$fileName", file.FileName),
|
||||
("$relativePath", file.RelativePath),
|
||||
("$fingerprint", file.Fingerprint),
|
||||
("$fileSize", file.FileSize),
|
||||
("$lastWriteUtc", ToUtcText(file.LastWriteUtc)),
|
||||
("$embeddedAtUtc", ToUtcText(file.EmbeddedAtUtc)),
|
||||
("$chunkCount", file.ChunkCount));
|
||||
}
|
||||
|
||||
public override async Task DeleteFileAsync(string dataSourceId, string filePath, CancellationToken token)
|
||||
{
|
||||
await using var connection = await this.OpenConnectionAsync(token);
|
||||
await ExecuteNonQueryAsync(connection, """
|
||||
DELETE FROM embedded_files
|
||||
WHERE data_source_id = $dataSourceId
|
||||
AND file_path = $filePath
|
||||
""", token,
|
||||
("$dataSourceId", dataSourceId),
|
||||
("$filePath", filePath));
|
||||
}
|
||||
|
||||
public override async Task DeleteDataSourceAsync(string dataSourceId, CancellationToken token)
|
||||
{
|
||||
await using var connection = await this.OpenConnectionAsync(token);
|
||||
await ExecuteNonQueryAsync(connection, """
|
||||
DELETE FROM data_sources
|
||||
WHERE data_source_id = $dataSourceId
|
||||
""", token,
|
||||
("$dataSourceId", dataSourceId));
|
||||
}
|
||||
|
||||
public override void Dispose()
|
||||
{
|
||||
}
|
||||
|
||||
private async Task InitializeAsync(CancellationToken token)
|
||||
{
|
||||
await using var connection = await this.OpenConnectionAsync(token);
|
||||
await ExecuteNonQueryAsync(connection, """
|
||||
PRAGMA journal_mode=WAL;
|
||||
PRAGMA foreign_keys=ON;
|
||||
|
||||
CREATE TABLE IF NOT EXISTS data_sources (
|
||||
data_source_id TEXT PRIMARY KEY,
|
||||
data_source_name TEXT NOT NULL,
|
||||
data_source_type TEXT NOT NULL,
|
||||
embedding_provider_id TEXT NOT NULL,
|
||||
embedding_signature TEXT NOT NULL,
|
||||
vector_size INTEGER NOT NULL DEFAULT 0,
|
||||
updated_at_utc TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS embedded_files (
|
||||
data_source_id TEXT NOT NULL,
|
||||
file_path TEXT COLLATE NOCASE NOT NULL,
|
||||
file_name TEXT NOT NULL,
|
||||
relative_path TEXT NOT NULL,
|
||||
fingerprint TEXT NOT NULL,
|
||||
file_size INTEGER NOT NULL,
|
||||
last_write_utc TEXT NOT NULL,
|
||||
embedded_at_utc TEXT NOT NULL,
|
||||
chunk_count INTEGER NOT NULL,
|
||||
PRIMARY KEY (data_source_id, file_path),
|
||||
FOREIGN KEY (data_source_id)
|
||||
REFERENCES data_sources(data_source_id)
|
||||
ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_embedded_files_data_source
|
||||
ON embedded_files(data_source_id);
|
||||
""", token);
|
||||
}
|
||||
|
||||
private async Task<string> GetSqliteVersionAsync(CancellationToken token)
|
||||
{
|
||||
await using var connection = await this.OpenConnectionAsync(token);
|
||||
await using var command = CreateCommand(connection, "SELECT sqlite_version()");
|
||||
var versionObject = await command.ExecuteScalarAsync(token);
|
||||
return Convert.ToString(versionObject, CultureInfo.InvariantCulture) ?? string.Empty;
|
||||
}
|
||||
|
||||
private async Task<long> CountAsync(string tableName, CancellationToken token)
|
||||
{
|
||||
await using var connection = await this.OpenConnectionAsync(token);
|
||||
await using var command = CreateCommand(connection, $"SELECT COUNT(*) FROM {tableName}");
|
||||
var countObject = await command.ExecuteScalarAsync(token);
|
||||
return Convert.ToInt64(countObject, CultureInfo.InvariantCulture);
|
||||
}
|
||||
|
||||
private async Task<SqliteConnection> OpenConnectionAsync(CancellationToken token)
|
||||
{
|
||||
var connection = new SqliteConnection(this.connectionString);
|
||||
await connection.OpenAsync(token);
|
||||
return connection;
|
||||
}
|
||||
|
||||
private static SqliteCommand CreateCommand(SqliteConnection connection, string commandText)
|
||||
{
|
||||
var command = connection.CreateCommand();
|
||||
command.CommandText = commandText;
|
||||
return command;
|
||||
}
|
||||
|
||||
private static async Task ExecuteNonQueryAsync(
|
||||
SqliteConnection connection,
|
||||
string commandText,
|
||||
CancellationToken token,
|
||||
params (string Name, object? Value)[] parameters)
|
||||
{
|
||||
await using var command = CreateCommand(connection, commandText);
|
||||
foreach (var (name, value) in parameters)
|
||||
command.Parameters.AddWithValue(name, value ?? DBNull.Value);
|
||||
|
||||
await command.ExecuteNonQueryAsync(token);
|
||||
}
|
||||
|
||||
private static string ToUtcText(DateTime dateTime)
|
||||
{
|
||||
var utc = dateTime.Kind is DateTimeKind.Utc ? dateTime : dateTime.ToUniversalTime();
|
||||
return utc.ToString("O", CultureInfo.InvariantCulture);
|
||||
}
|
||||
|
||||
private static DateTime ParseUtc(string value)
|
||||
{
|
||||
return DateTime.TryParse(value, CultureInfo.InvariantCulture, DateTimeStyles.RoundtripKind, out var dateTime)
|
||||
? dateTime.ToUniversalTime()
|
||||
: DateTime.UnixEpoch;
|
||||
}
|
||||
|
||||
private static NoEmbeddingStateClient CreateNoEmbeddingStateClient(string name, string? unavailableReason, DatabaseClientStatus status, ILogger<DatabaseClient> databaseClientLogger)
|
||||
{
|
||||
var client = new NoEmbeddingStateClient(name, unavailableReason, status);
|
||||
client.SetLogger(databaseClientLogger);
|
||||
return client;
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user