mirror of
https://github.com/MindWorkAI/AI-Studio.git
synced 2026-10-13 20:33:49 +00:00
added metadata for bm25 and also confidence levels
This commit is contained in:
PaulKoudelka committed
1 parent
b3ba945b98
commit
c78a8f05ef
27 files changed
+511
-110
No files matched your search
@@ -65,7 +65,7 @@ public sealed partial class DatabaseClientProvider(RustService rustService, ILog
|
||||
|
||||
return new NoEmbeddingStateClient(
|
||||
client.Name,
|
||||
"The configured database client does not support embedding state operations.",
|
||||
"The configured database client does not support local RAG index operations.",
|
||||
client.Status);
|
||||
}
|
||||
|
||||
|
||||
@@ -24,15 +24,30 @@ public abstract class EmbeddingStateClient(string name, string path) : DatabaseC
|
||||
|
||||
public abstract Task DeleteFileAsync(string dataSourceId, string filePath, CancellationToken token);
|
||||
|
||||
public abstract Task UpsertChunksAsync(string dataSourceId, IReadOnlyList<EmbeddingStateChunk> chunks, CancellationToken token);
|
||||
|
||||
public abstract Task DeleteDataSourceAsync(string dataSourceId, CancellationToken token);
|
||||
}
|
||||
|
||||
public sealed record EmbeddingStateFile(
|
||||
string FilePath,
|
||||
string ParentFileId,
|
||||
string AbsolutePath,
|
||||
string FileName,
|
||||
string RelativePath,
|
||||
string FileType,
|
||||
string Fingerprint,
|
||||
long FileSize,
|
||||
DateTime CreationUtc,
|
||||
DateTime LastWriteUtc,
|
||||
DateTime EmbeddedAtUtc,
|
||||
int ChunkCount);
|
||||
int ChunkCount,
|
||||
string ComplianceLevel,
|
||||
int ComplianceLevelRank);
|
||||
|
||||
public sealed record EmbeddingStateChunk(
|
||||
string ChunkId,
|
||||
string ParentFileId,
|
||||
int? PageNumber,
|
||||
int ChunkIndex,
|
||||
string ChunkText,
|
||||
DateTime EmbeddedAtUtc);
|
||||
@@ -44,6 +44,8 @@ public sealed class NoEmbeddingStateClient(string name, string? unavailableReaso
|
||||
|
||||
public override Task DeleteFileAsync(string dataSourceId, string filePath, CancellationToken token) => Task.CompletedTask;
|
||||
|
||||
public override Task UpsertChunksAsync(string dataSourceId, IReadOnlyList<EmbeddingStateChunk> chunks, CancellationToken token) => Task.CompletedTask;
|
||||
|
||||
public override Task DeleteDataSourceAsync(string dataSourceId, CancellationToken token) => Task.CompletedTask;
|
||||
|
||||
public override void Dispose()
|
||||
|
||||
+184
-33
@@ -14,8 +14,8 @@ public sealed class SqliteEmbeddingStateClientImplementation(
|
||||
string basePath,
|
||||
string version) : EmbeddingStateClient(name, basePath)
|
||||
{
|
||||
private const string DATABASE_NAME = "SQLite";
|
||||
private const string DATABASE_FILENAME = "rag-embedding-state.sqlite3";
|
||||
private const string DATABASE_NAME = "Local RAG Index";
|
||||
private const string DATABASE_FILENAME = "rag-index.sqlite3";
|
||||
|
||||
private readonly string databasePath = databasePath;
|
||||
private readonly string connectionString = new SqliteConnectionStringBuilder
|
||||
@@ -57,7 +57,7 @@ public sealed class SqliteEmbeddingStateClientImplementation(
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
logger.LogWarning(exception, "{DatabaseName} is not available. Indexed file fingerprints are disabled.", DATABASE_NAME);
|
||||
logger.LogWarning(exception, "{DatabaseName} is not available. Indexed file fingerprints and search chunks are disabled.", DATABASE_NAME);
|
||||
return CreateNoEmbeddingStateClient(DATABASE_NAME, exception.Message, DatabaseClientStatus.UNAVAILABLE, databaseClientLogger);
|
||||
}
|
||||
}
|
||||
@@ -69,6 +69,7 @@ public sealed class SqliteEmbeddingStateClientImplementation(
|
||||
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));
|
||||
yield return (TB("Search chunks"), (await this.CountAsync("embedding_chunks", CancellationToken.None)).ToString(CultureInfo.InvariantCulture));
|
||||
}
|
||||
|
||||
public override async Task<DataSourceEmbeddingManifest> GetManifestAsync(string dataSourceId, CancellationToken token)
|
||||
@@ -94,9 +95,10 @@ public sealed class SqliteEmbeddingStateClientImplementation(
|
||||
}
|
||||
|
||||
await using (var command = CreateCommand(connection, """
|
||||
SELECT file_path, fingerprint, file_size, last_write_utc, embedded_at_utc, chunk_count
|
||||
SELECT absolute_path, fingerprint, file_size, last_write_utc, embedded_at_utc, chunk_count
|
||||
FROM embedded_files
|
||||
WHERE data_source_id = $dataSourceId
|
||||
AND chunk_count > 0
|
||||
"""))
|
||||
{
|
||||
command.Parameters.AddWithValue("$dataSourceId", dataSourceId);
|
||||
@@ -197,43 +199,64 @@ public sealed class SqliteEmbeddingStateClientImplementation(
|
||||
await using var connection = await this.OpenConnectionAsync(token);
|
||||
await ExecuteNonQueryAsync(connection, """
|
||||
INSERT INTO embedded_files (
|
||||
parent_file_id,
|
||||
data_source_id,
|
||||
file_path,
|
||||
absolute_path,
|
||||
file_name,
|
||||
relative_path,
|
||||
file_type,
|
||||
fingerprint,
|
||||
file_size,
|
||||
creation_utc,
|
||||
last_write_utc,
|
||||
embedded_at_utc,
|
||||
chunk_count)
|
||||
chunk_count,
|
||||
compliance_level,
|
||||
compliance_level_rank)
|
||||
VALUES (
|
||||
$parentFileId,
|
||||
$dataSourceId,
|
||||
$filePath,
|
||||
$absolutePath,
|
||||
$fileName,
|
||||
$relativePath,
|
||||
$fileType,
|
||||
$fingerprint,
|
||||
$fileSize,
|
||||
$creationUtc,
|
||||
$lastWriteUtc,
|
||||
$embeddedAtUtc,
|
||||
$chunkCount)
|
||||
ON CONFLICT(data_source_id, file_path) DO UPDATE SET
|
||||
$chunkCount,
|
||||
$complianceLevel,
|
||||
$complianceLevelRank)
|
||||
ON CONFLICT(parent_file_id) DO UPDATE SET
|
||||
data_source_id = excluded.data_source_id,
|
||||
absolute_path = excluded.absolute_path,
|
||||
file_name = excluded.file_name,
|
||||
relative_path = excluded.relative_path,
|
||||
file_type = excluded.file_type,
|
||||
fingerprint = excluded.fingerprint,
|
||||
file_size = excluded.file_size,
|
||||
creation_utc = excluded.creation_utc,
|
||||
last_write_utc = excluded.last_write_utc,
|
||||
embedded_at_utc = excluded.embedded_at_utc,
|
||||
chunk_count = excluded.chunk_count
|
||||
chunk_count = excluded.chunk_count,
|
||||
compliance_level = excluded.compliance_level,
|
||||
compliance_level_rank = excluded.compliance_level_rank
|
||||
""", token,
|
||||
("$parentFileId", file.ParentFileId),
|
||||
("$dataSourceId", dataSourceId),
|
||||
("$filePath", file.FilePath),
|
||||
("$absolutePath", file.AbsolutePath),
|
||||
("$fileName", file.FileName),
|
||||
("$relativePath", file.RelativePath),
|
||||
("$fileType", file.FileType),
|
||||
("$fingerprint", file.Fingerprint),
|
||||
("$fileSize", file.FileSize),
|
||||
("$creationUtc", ToUtcText(file.CreationUtc)),
|
||||
("$lastWriteUtc", ToUtcText(file.LastWriteUtc)),
|
||||
("$embeddedAtUtc", ToUtcText(file.EmbeddedAtUtc)),
|
||||
("$chunkCount", file.ChunkCount));
|
||||
("$chunkCount", file.ChunkCount),
|
||||
("$complianceLevel", file.ComplianceLevel),
|
||||
("$complianceLevelRank", file.ComplianceLevelRank));
|
||||
}
|
||||
|
||||
public override async Task DeleteFileAsync(string dataSourceId, string filePath, CancellationToken token)
|
||||
@@ -242,12 +265,57 @@ public sealed class SqliteEmbeddingStateClientImplementation(
|
||||
await ExecuteNonQueryAsync(connection, """
|
||||
DELETE FROM embedded_files
|
||||
WHERE data_source_id = $dataSourceId
|
||||
AND file_path = $filePath
|
||||
AND absolute_path = $filePath
|
||||
""", token,
|
||||
("$dataSourceId", dataSourceId),
|
||||
("$filePath", filePath));
|
||||
}
|
||||
|
||||
public override async Task UpsertChunksAsync(string dataSourceId, IReadOnlyList<EmbeddingStateChunk> chunks, CancellationToken token)
|
||||
{
|
||||
if (chunks.Count == 0)
|
||||
return;
|
||||
|
||||
await using var connection = await this.OpenConnectionAsync(token);
|
||||
await using var transaction = (SqliteTransaction)await connection.BeginTransactionAsync(token);
|
||||
|
||||
foreach (var chunk in chunks)
|
||||
{
|
||||
token.ThrowIfCancellationRequested();
|
||||
|
||||
await ExecuteNonQueryAsync(connection, transaction, """
|
||||
INSERT INTO embedding_chunks (
|
||||
chunk_id,
|
||||
parent_file_id,
|
||||
page_number,
|
||||
chunk_index,
|
||||
chunk_text,
|
||||
embedded_at_utc)
|
||||
VALUES (
|
||||
$chunkId,
|
||||
$parentFileId,
|
||||
$pageNumber,
|
||||
$chunkIndex,
|
||||
$chunkText,
|
||||
$embeddedAtUtc)
|
||||
ON CONFLICT(chunk_id) DO UPDATE SET
|
||||
parent_file_id = excluded.parent_file_id,
|
||||
page_number = excluded.page_number,
|
||||
chunk_index = excluded.chunk_index,
|
||||
chunk_text = excluded.chunk_text,
|
||||
embedded_at_utc = excluded.embedded_at_utc
|
||||
""", token,
|
||||
("$chunkId", chunk.ChunkId),
|
||||
("$parentFileId", chunk.ParentFileId),
|
||||
("$pageNumber", chunk.PageNumber),
|
||||
("$chunkIndex", chunk.ChunkIndex),
|
||||
("$chunkText", chunk.ChunkText),
|
||||
("$embeddedAtUtc", ToUtcText(chunk.EmbeddedAtUtc)));
|
||||
}
|
||||
|
||||
await transaction.CommitAsync(token);
|
||||
}
|
||||
|
||||
public override async Task DeleteDataSourceAsync(string dataSourceId, CancellationToken token)
|
||||
{
|
||||
await using var connection = await this.OpenConnectionAsync(token);
|
||||
@@ -281,26 +349,109 @@ public sealed class SqliteEmbeddingStateClientImplementation(
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS embedded_files (
|
||||
parent_file_id TEXT PRIMARY KEY,
|
||||
data_source_id TEXT NOT NULL,
|
||||
file_path TEXT COLLATE NOCASE NOT NULL,
|
||||
absolute_path TEXT COLLATE NOCASE NOT NULL,
|
||||
file_name TEXT NOT NULL,
|
||||
relative_path TEXT NOT NULL,
|
||||
file_type TEXT NOT NULL,
|
||||
fingerprint TEXT NOT NULL,
|
||||
file_size INTEGER NOT NULL,
|
||||
creation_utc TEXT 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),
|
||||
compliance_level TEXT NOT NULL,
|
||||
compliance_level_rank INTEGER NOT NULL,
|
||||
FOREIGN KEY (data_source_id)
|
||||
REFERENCES data_sources(data_source_id)
|
||||
ON DELETE CASCADE
|
||||
ON DELETE CASCADE,
|
||||
UNIQUE(data_source_id, absolute_path)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS embedding_chunks (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
chunk_id TEXT NOT NULL UNIQUE,
|
||||
parent_file_id TEXT NOT NULL,
|
||||
page_number INTEGER NULL,
|
||||
chunk_index INTEGER NOT NULL,
|
||||
chunk_text TEXT NOT NULL,
|
||||
embedded_at_utc TEXT NOT NULL,
|
||||
FOREIGN KEY (parent_file_id)
|
||||
REFERENCES embedded_files(parent_file_id)
|
||||
ON DELETE CASCADE,
|
||||
UNIQUE(parent_file_id, chunk_index)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_embedded_files_data_source
|
||||
ON embedded_files(data_source_id);
|
||||
""", token);
|
||||
|
||||
await EnsureColumnAsync(connection, "data_sources", "source_hash", "TEXT NOT NULL DEFAULT ''", token);
|
||||
CREATE INDEX IF NOT EXISTS idx_embedded_files_absolute_path
|
||||
ON embedded_files(absolute_path);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_embedded_files_file_type
|
||||
ON embedded_files(file_type);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_embedded_files_compliance
|
||||
ON embedded_files(compliance_level_rank);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_embedding_chunks_parent_file
|
||||
ON embedding_chunks(parent_file_id);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_embedding_chunks_page
|
||||
ON embedding_chunks(page_number);
|
||||
|
||||
CREATE VIRTUAL TABLE IF NOT EXISTS embedding_chunks_fts
|
||||
USING fts5(chunk_id UNINDEXED, file_name, chunk_text);
|
||||
|
||||
CREATE TRIGGER IF NOT EXISTS embedding_chunks_ai
|
||||
AFTER INSERT ON embedding_chunks
|
||||
BEGIN
|
||||
INSERT INTO embedding_chunks_fts(rowid, chunk_id, file_name, chunk_text)
|
||||
VALUES (
|
||||
new.id,
|
||||
new.chunk_id,
|
||||
(SELECT file_name FROM embedded_files WHERE parent_file_id = new.parent_file_id),
|
||||
new.chunk_text);
|
||||
END;
|
||||
|
||||
CREATE TRIGGER IF NOT EXISTS embedding_chunks_ad
|
||||
AFTER DELETE ON embedding_chunks
|
||||
BEGIN
|
||||
DELETE FROM embedding_chunks_fts
|
||||
WHERE rowid = old.id;
|
||||
END;
|
||||
|
||||
CREATE TRIGGER IF NOT EXISTS embedding_chunks_au
|
||||
AFTER UPDATE ON embedding_chunks
|
||||
BEGIN
|
||||
DELETE FROM embedding_chunks_fts
|
||||
WHERE rowid = old.id;
|
||||
|
||||
INSERT INTO embedding_chunks_fts(rowid, chunk_id, file_name, chunk_text)
|
||||
VALUES (
|
||||
new.id,
|
||||
new.chunk_id,
|
||||
(SELECT file_name FROM embedded_files WHERE parent_file_id = new.parent_file_id),
|
||||
new.chunk_text);
|
||||
END;
|
||||
|
||||
CREATE TRIGGER IF NOT EXISTS embedded_files_file_name_au
|
||||
AFTER UPDATE OF file_name ON embedded_files
|
||||
BEGIN
|
||||
DELETE FROM embedding_chunks_fts
|
||||
WHERE rowid IN (
|
||||
SELECT id
|
||||
FROM embedding_chunks
|
||||
WHERE parent_file_id = new.parent_file_id
|
||||
);
|
||||
|
||||
INSERT INTO embedding_chunks_fts(rowid, chunk_id, file_name, chunk_text)
|
||||
SELECT id, chunk_id, new.file_name, chunk_text
|
||||
FROM embedding_chunks
|
||||
WHERE parent_file_id = new.parent_file_id;
|
||||
END;
|
||||
""", token);
|
||||
}
|
||||
|
||||
private async Task<string> GetSqliteVersionAsync(CancellationToken token)
|
||||
@@ -319,21 +470,6 @@ public sealed class SqliteEmbeddingStateClientImplementation(
|
||||
return Convert.ToInt64(countObject, CultureInfo.InvariantCulture);
|
||||
}
|
||||
|
||||
private static async Task EnsureColumnAsync(SqliteConnection connection, string tableName, string columnName, string columnDefinition, CancellationToken token)
|
||||
{
|
||||
await using (var command = CreateCommand(connection, $"PRAGMA table_info({tableName})"))
|
||||
{
|
||||
await using var reader = await command.ExecuteReaderAsync(token);
|
||||
while (await reader.ReadAsync(token))
|
||||
{
|
||||
if (string.Equals(reader.GetString(1), columnName, StringComparison.OrdinalIgnoreCase))
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
await ExecuteNonQueryAsync(connection, $"ALTER TABLE {tableName} ADD COLUMN {columnName} {columnDefinition}", token);
|
||||
}
|
||||
|
||||
private async Task<SqliteConnection> OpenConnectionAsync(CancellationToken token)
|
||||
{
|
||||
var connection = new SqliteConnection(this.connectionString);
|
||||
@@ -361,6 +497,21 @@ public sealed class SqliteEmbeddingStateClientImplementation(
|
||||
await command.ExecuteNonQueryAsync(token);
|
||||
}
|
||||
|
||||
private static async Task ExecuteNonQueryAsync(
|
||||
SqliteConnection connection,
|
||||
SqliteTransaction transaction,
|
||||
string commandText,
|
||||
CancellationToken token,
|
||||
params (string Name, object? Value)[] parameters)
|
||||
{
|
||||
await using var command = CreateCommand(connection, commandText);
|
||||
command.Transaction = transaction;
|
||||
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();
|
||||
|
||||
@@ -6,11 +6,19 @@ public sealed record VectorStoragePoint(
|
||||
string DataSourceId,
|
||||
string DataSourceName,
|
||||
string DataSourceType,
|
||||
string ChunkId,
|
||||
string ParentFileId,
|
||||
string FilePath,
|
||||
string AbsolutePath,
|
||||
string FileName,
|
||||
string RelativePath,
|
||||
string FileType,
|
||||
int? PageNumber,
|
||||
int ChunkIndex,
|
||||
string Text,
|
||||
string Fingerprint,
|
||||
DateTime CreationUtc,
|
||||
DateTime LastWriteUtc,
|
||||
DateTime EmbeddedAtUtc);
|
||||
DateTime EmbeddedAtUtc,
|
||||
string ComplianceLevel,
|
||||
int ComplianceLevelRank);
|
||||
@@ -2,8 +2,10 @@ using System.Security.Cryptography;
|
||||
using System.Text;
|
||||
using System.Text.RegularExpressions;
|
||||
|
||||
using AIStudio.Provider;
|
||||
using AIStudio.Settings;
|
||||
using AIStudio.Settings.DataModel;
|
||||
using AIStudio.Tools.Databases.EmbeddingState;
|
||||
using AIStudio.Tools.PluginSystem;
|
||||
using AIStudio.Tools.Rust;
|
||||
|
||||
@@ -27,6 +29,8 @@ public sealed partial class DataSourceEmbeddingService
|
||||
|
||||
private sealed record ExtractedFileContent(string Text, IReadOnlyList<string> SourceSegments);
|
||||
|
||||
private sealed record EmbeddingChunkDraft(string ChunkId, string Text, int ChunkIndex, int? PageNumber);
|
||||
|
||||
private sealed record ChunkingOptions(int MaxChunkTokenLength, int OverlapTokenLength);
|
||||
|
||||
private sealed record ChunkingStrategy(string Name, IReadOnlyList<ChunkingRule> Rules);
|
||||
@@ -748,6 +752,7 @@ public sealed partial class DataSourceEmbeddingService
|
||||
embeddingProvider.Hostname,
|
||||
embeddingProvider.TokenizerPath,
|
||||
embeddingProvider.EffectiveTokenLimit,
|
||||
GetDataSourceComplianceLevel(dataSource).ToString(),
|
||||
dataSource is IInternalDataSource internalDataSource ? internalDataSource.MaxChunkTokenLength : 0,
|
||||
dataSource is IInternalDataSource overlapDataSource ? overlapDataSource.ChunkOverlapTokenLength : 0,
|
||||
chunkingOptions.MaxChunkTokenLength,
|
||||
@@ -847,6 +852,59 @@ public sealed partial class DataSourceEmbeddingService
|
||||
return Convert.ToHexString(bytes);
|
||||
}
|
||||
|
||||
private EmbeddingStateFile CreateEmbeddingStateFile(IDataSource dataSource, FileInfo file, string fingerprint, int chunkCount, DateTime embeddedAtUtc)
|
||||
{
|
||||
file.Refresh();
|
||||
var absolutePath = Path.GetFullPath(file.FullName);
|
||||
var complianceLevel = GetDataSourceComplianceLevel(dataSource);
|
||||
return new(
|
||||
this.CreateParentFileId(dataSource.Id, absolutePath),
|
||||
absolutePath,
|
||||
file.Name,
|
||||
this.TryGetRelativePath(dataSource, file),
|
||||
GetFileType(file),
|
||||
fingerprint,
|
||||
file.Exists ? file.Length : 0,
|
||||
file.Exists ? file.CreationTimeUtc : DateTime.UnixEpoch,
|
||||
file.Exists ? file.LastWriteTimeUtc : DateTime.UnixEpoch,
|
||||
embeddedAtUtc,
|
||||
chunkCount,
|
||||
complianceLevel.ToString(),
|
||||
(int)complianceLevel);
|
||||
}
|
||||
|
||||
private IReadOnlyList<EmbeddingStateChunk> CreateEmbeddingStateChunks(EmbeddingStateFile parentFile, IReadOnlyList<EmbeddingChunkDraft> batch, DateTime embeddedAtUtc)
|
||||
{
|
||||
return batch
|
||||
.Select(chunk => new EmbeddingStateChunk(
|
||||
chunk.ChunkId,
|
||||
parentFile.ParentFileId,
|
||||
chunk.PageNumber,
|
||||
chunk.ChunkIndex,
|
||||
chunk.Text,
|
||||
embeddedAtUtc))
|
||||
.ToList();
|
||||
}
|
||||
|
||||
private static ConfidenceLevel GetDataSourceComplianceLevel(IDataSource dataSource) =>
|
||||
dataSource.ComplianceLevel is ConfidenceLevel.NONE
|
||||
? ConfidenceLevel.UNKNOWN
|
||||
: dataSource.ComplianceLevel;
|
||||
|
||||
private static string GetFileType(FileInfo file)
|
||||
{
|
||||
var extension = file.Extension.TrimStart('.').ToLowerInvariant();
|
||||
return string.IsNullOrWhiteSpace(extension) ? "unknown" : extension;
|
||||
}
|
||||
|
||||
private static int? TryExtractPageNumber(string chunk)
|
||||
{
|
||||
var match = Regex.Match(chunk, @"^\s*#\s+Page\s+(\d+)\b", RegexOptions.CultureInvariant | RegexOptions.IgnoreCase);
|
||||
return match.Success && int.TryParse(match.Groups[1].Value, out var pageNumber) && pageNumber > 0
|
||||
? pageNumber
|
||||
: null;
|
||||
}
|
||||
|
||||
private string GetCollectionName(string dataSourceName, string dataSourceId)
|
||||
{
|
||||
var safeId = dataSourceId
|
||||
@@ -864,9 +922,14 @@ public sealed partial class DataSourceEmbeddingService
|
||||
return $"rag_{safeName}_{safeId}";
|
||||
}
|
||||
|
||||
private string CreatePointId(string dataSourceId, string fingerprint, int chunkIndex)
|
||||
private string CreatePointId(string dataSourceId, string fingerprint, int chunkIndex) =>
|
||||
CreateStableGuid($"{dataSourceId}:chunk:{fingerprint}:{chunkIndex}");
|
||||
|
||||
private string CreateParentFileId(string dataSourceId, string absolutePath) =>
|
||||
CreateStableGuid($"{dataSourceId}:parent-file:{absolutePath}");
|
||||
|
||||
private static string CreateStableGuid(string source)
|
||||
{
|
||||
var source = $"{dataSourceId}:{fingerprint}:{chunkIndex}";
|
||||
var hash = SHA256.HashData(Encoding.UTF8.GetBytes(source));
|
||||
var guidBytes = hash[..16].ToArray();
|
||||
|
||||
|
||||
@@ -17,11 +17,11 @@ public sealed partial class DataSourceEmbeddingService
|
||||
embeddingState ??= await databaseClientProvider.GetEmbeddingStateAsync(token);
|
||||
if (!embeddingState.IsAvailable)
|
||||
{
|
||||
logger.LogWarning("Could not delete SQLite embedding state for data source '{DataSourceId}' because the embedding state database '{DatabaseName}' is unavailable.", dataSourceId, embeddingState.Name);
|
||||
logger.LogWarning("Could not delete local RAG index state for data source '{DataSourceId}' because the database '{DatabaseName}' is unavailable.", dataSourceId, embeddingState.Name);
|
||||
return;
|
||||
}
|
||||
|
||||
await embeddingState.DeleteDataSourceAsync(dataSourceId, token);
|
||||
logger.LogInformation("Reset persisted embedding state for data source '{DataSourceId}'.", dataSourceId);
|
||||
logger.LogInformation("Reset persisted local RAG index state for data source '{DataSourceId}'.", dataSourceId);
|
||||
}
|
||||
}
|
||||
@@ -162,7 +162,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
|
||||
var embeddingState = await databaseClientProvider.GetEmbeddingStateAsync(token);
|
||||
if (!embeddingState.IsAvailable)
|
||||
{
|
||||
logger.LogWarning("Locking identity settings for data source '{DataSourceId}' because the embedding state database '{DatabaseName}' is unavailable.", dataSourceId, embeddingState.Name);
|
||||
logger.LogWarning("Locking identity settings for data source '{DataSourceId}' because the local RAG index database '{DatabaseName}' is unavailable.", dataSourceId, embeddingState.Name);
|
||||
return true;
|
||||
}
|
||||
|
||||
@@ -395,7 +395,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
|
||||
dataSource.Id,
|
||||
embeddingState.Name);
|
||||
token.ThrowIfCancellationRequested();
|
||||
this.UpsertStatus(this.GetFallbackStatus(dataSource, "The SQLite embedding state database is not available."));
|
||||
this.UpsertStatus(this.GetFallbackStatus(dataSource, "The local RAG index database is not available."));
|
||||
return;
|
||||
}
|
||||
|
||||
@@ -522,15 +522,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
|
||||
chunkCount);
|
||||
await embeddingState.UpsertFileAsync(
|
||||
dataSource.Id,
|
||||
new EmbeddingStateFile(
|
||||
file.FullName,
|
||||
file.Name,
|
||||
this.TryGetRelativePath(dataSource, file),
|
||||
fingerprint,
|
||||
file.Length,
|
||||
file.LastWriteTimeUtc,
|
||||
embeddedAtUtc,
|
||||
chunkCount),
|
||||
this.CreateEmbeddingStateFile(dataSource, file, fingerprint, chunkCount, embeddedAtUtc),
|
||||
token);
|
||||
manifest.Files[file.FullName] = record;
|
||||
completedFiles++;
|
||||
@@ -603,22 +595,26 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
|
||||
file.FullName,
|
||||
collectionName);
|
||||
await this.DeleteFilePointsAsync(vectorStore, collectionName, file.FullName, token);
|
||||
await embeddingState.DeleteFileAsync(dataSource.Id, file.FullName, token);
|
||||
|
||||
var parentFile = this.CreateEmbeddingStateFile(dataSource, file, fingerprint, 0, DateTime.UtcNow);
|
||||
await embeddingState.UpsertFileAsync(dataSource.Id, parentFile, token);
|
||||
|
||||
var embeddingBatchSize = embeddingProvider.EffectiveEmbeddingBatchSize;
|
||||
var batch = new List<(string Text, int ChunkIndex)>(embeddingBatchSize);
|
||||
var batch = new List<EmbeddingChunkDraft>(embeddingBatchSize);
|
||||
var totalChunkCount = 0;
|
||||
|
||||
await foreach (var chunk in this.StreamEmbeddingChunksAsync(file.FullName, dataSource, embeddingProvider, token))
|
||||
{
|
||||
batch.Add((chunk, totalChunkCount));
|
||||
batch.Add(new(this.CreatePointId(dataSource.Id, fingerprint, totalChunkCount), chunk, totalChunkCount, TryExtractPageNumber(chunk)));
|
||||
totalChunkCount++;
|
||||
|
||||
if (batch.Count >= embeddingBatchSize)
|
||||
await this.FlushBatchAsync(embeddingState, vectorStore, dataSource, file, fingerprint, embeddingProvider, provider, manifest, collectionName, batch, token);
|
||||
await this.FlushBatchAsync(embeddingState, vectorStore, dataSource, file, fingerprint, parentFile, embeddingProvider, provider, manifest, collectionName, batch, token);
|
||||
}
|
||||
|
||||
if (batch.Count > 0)
|
||||
await this.FlushBatchAsync(embeddingState, vectorStore, dataSource, file, fingerprint, embeddingProvider, provider, manifest, collectionName, batch, token);
|
||||
await this.FlushBatchAsync(embeddingState, vectorStore, dataSource, file, fingerprint, parentFile, embeddingProvider, provider, manifest, collectionName, batch, token);
|
||||
|
||||
if (totalChunkCount == 0)
|
||||
throw new InvalidOperationException($"The file '{file.Name}' did not yield any text chunks.");
|
||||
@@ -639,11 +635,12 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
|
||||
IDataSource dataSource,
|
||||
FileInfo file,
|
||||
string fingerprint,
|
||||
EmbeddingStateFile parentFile,
|
||||
EmbeddingProvider embeddingProvider,
|
||||
IProvider provider,
|
||||
DataSourceEmbeddingManifest manifest,
|
||||
string collectionName,
|
||||
List<(string Text, int ChunkIndex)> batch,
|
||||
List<EmbeddingChunkDraft> batch,
|
||||
CancellationToken token)
|
||||
{
|
||||
logger.LogDebug(
|
||||
@@ -694,15 +691,22 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
|
||||
}
|
||||
|
||||
token.ThrowIfCancellationRequested();
|
||||
var embeddedAtUtc = DateTime.UtcNow;
|
||||
await this.UpsertPointsAsync(
|
||||
vectorStore,
|
||||
collectionName,
|
||||
dataSource,
|
||||
file,
|
||||
fingerprint,
|
||||
parentFile,
|
||||
batch,
|
||||
vectors,
|
||||
this.TryGetRelativePath(dataSource, file),
|
||||
embeddedAtUtc,
|
||||
token);
|
||||
token.ThrowIfCancellationRequested();
|
||||
await embeddingState.UpsertChunksAsync(
|
||||
dataSource.Id,
|
||||
this.CreateEmbeddingStateChunks(parentFile, batch, embeddedAtUtc),
|
||||
token);
|
||||
|
||||
logger.LogDebug(
|
||||
@@ -725,26 +729,34 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
|
||||
IDataSource dataSource,
|
||||
FileInfo file,
|
||||
string fingerprint,
|
||||
IReadOnlyList<(string Text, int ChunkIndex)> batch,
|
||||
EmbeddingStateFile parentFile,
|
||||
IReadOnlyList<EmbeddingChunkDraft> batch,
|
||||
IReadOnlyList<IReadOnlyList<float>> vectors,
|
||||
string relativePath,
|
||||
DateTime embeddedAtUtc,
|
||||
CancellationToken token)
|
||||
{
|
||||
var embeddedAtUtc = DateTime.UtcNow;
|
||||
var points = batch.Select((item, index) => new VectorStoragePoint(
|
||||
this.CreatePointId(dataSource.Id, fingerprint, item.ChunkIndex),
|
||||
item.ChunkId,
|
||||
vectors[index],
|
||||
dataSource.Id,
|
||||
dataSource.Name,
|
||||
dataSource.Type.ToString(),
|
||||
item.ChunkId,
|
||||
parentFile.ParentFileId,
|
||||
file.FullName,
|
||||
file.Name,
|
||||
relativePath,
|
||||
parentFile.AbsolutePath,
|
||||
parentFile.FileName,
|
||||
parentFile.RelativePath,
|
||||
parentFile.FileType,
|
||||
item.PageNumber,
|
||||
item.ChunkIndex,
|
||||
item.Text,
|
||||
fingerprint,
|
||||
file.LastWriteTimeUtc,
|
||||
embeddedAtUtc)).ToList();
|
||||
parentFile.CreationUtc,
|
||||
parentFile.LastWriteUtc,
|
||||
embeddedAtUtc,
|
||||
parentFile.ComplianceLevel,
|
||||
parentFile.ComplianceLevelRank)).ToList();
|
||||
|
||||
await vectorStore.InsertEmbedding(collectionName, points, token);
|
||||
}
|
||||
@@ -805,7 +817,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
|
||||
.ToList();
|
||||
|
||||
logger.LogInformation(
|
||||
"Starting initial persisted hash check for {DataSourceCount} supported internal data source(s). Incomplete or failed embedding state will be retried during this pass. File watchers will be activated after this check completes.",
|
||||
"Starting initial persisted hash check for {DataSourceCount} supported internal data source(s). Incomplete or failed local RAG index state will be retried during this pass. File watchers will be activated after this check completes.",
|
||||
supportedDataSources.Count);
|
||||
|
||||
foreach (var dataSource in supportedDataSources)
|
||||
@@ -875,7 +887,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
|
||||
var manifest = await embeddingState.GetManifestAsync(dataSource.Id, token);
|
||||
|
||||
logger.LogInformation(
|
||||
"Loaded persisted embedding manifest for data source '{DataSourceName}' ({DataSourceId}). StoredFiles={StoredFiles}, StoredSourceHashPrefix={StoredSourceHashPrefix}, StoredSignaturePrefix={StoredSignaturePrefix}, CurrentSignaturePrefix={CurrentSignaturePrefix}.",
|
||||
"Loaded persisted local RAG index manifest for data source '{DataSourceName}' ({DataSourceId}). StoredFiles={StoredFiles}, StoredSourceHashPrefix={StoredSourceHashPrefix}, StoredSignaturePrefix={StoredSignaturePrefix}, CurrentSignaturePrefix={CurrentSignaturePrefix}.",
|
||||
dataSource.Name,
|
||||
dataSource.Id,
|
||||
manifest.Files.Count,
|
||||
|
||||
Reference in new issue
Block a user