using System.Security.Cryptography; using System.Text; using AIStudio.Settings; using AIStudio.Settings.DataModel; using AIStudio.Tools.PluginSystem; using AIStudio.Tools.Rust; namespace AIStudio.Tools.Services; public sealed partial class DataSourceEmbeddingService { private const string OFFICE_LOCK_FILE_PREFIX = "~$"; private static readonly string[] RAG_DELIMITED_TABLE_FILE_EXTENSIONS = ["csv", "tsv"]; private static readonly string[] RAG_SPREADSHEET_FILE_EXTENSIONS = ["ods", "xlsm", "xlsb"]; private static readonly string[] RAG_SPREADSHEET_ADD_IN_FILE_EXTENSIONS = ["xla", "xlam"]; private static readonly string[] SKIPPED_RAG_FILE_EXTENSIONS = ["lnk"]; private enum RagFileIndexingDecision { INDEXABLE, EXCLUDED, UNSUPPORTED, } private async IAsyncEnumerable StreamEmbeddingChunksAsync(string filePath, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token) { if (this.IsImageFilePath(filePath)) { yield return this.BuildImageIndexText(filePath); yield break; } var currentChunk = new StringBuilder(); await foreach (var segment in rustService.StreamArbitraryFileData(filePath, token: token)) { var normalized = NormalizeChunkSegment(segment); if (string.IsNullOrWhiteSpace(normalized)) continue; if (currentChunk.Length > 0 && currentChunk.Length + normalized.Length + Environment.NewLine.Length > MAX_CHUNK_LENGTH) { if (currentChunk.Length >= MIN_CHUNK_LENGTH) { var chunk = currentChunk.ToString().Trim(); if (!string.IsNullOrWhiteSpace(chunk)) yield return chunk; var overlap = chunk.Length > CHUNK_OVERLAP_LENGTH ? chunk[^CHUNK_OVERLAP_LENGTH..] : chunk; currentChunk.Clear(); currentChunk.Append(overlap); currentChunk.AppendLine(); } else { currentChunk.AppendLine(); } } currentChunk.Append(normalized); currentChunk.AppendLine(); } var finalChunk = currentChunk.ToString().Trim(); if (!string.IsNullOrWhiteSpace(finalChunk)) yield return finalChunk; } private FileEnumerationResult GetInputFiles(IDataSource dataSource) { var result = new FileEnumerationResult(); switch (dataSource) { case DataSourceLocalFile localFile when File.Exists(localFile.FilePath): var file = new FileInfo(localFile.FilePath); switch (this.GetRagFileIndexingDecision(file)) { case RagFileIndexingDecision.INDEXABLE: result.Files.Add(file); break; case RagFileIndexingDecision.EXCLUDED: logger.LogDebug("Skipping excluded file '{FilePath}' while indexing.", file.FullName); break; default: result.FailedFiles = 1; result.LastError = $"The selected file '{localFile.FilePath}' is not supported for background embeddings."; break; } return result; case DataSourceLocalDirectory localDirectory when Directory.Exists(localDirectory.Path): this.EnumerateAccessibleFiles(localDirectory.Path, result); return result; } switch (dataSource) { case DataSourceLocalFile localFile: result.FailedFiles = 1; result.LastError = $"The selected file '{localFile.FilePath}' does not exist."; break; case DataSourceLocalDirectory localDirectory: result.FailedFiles = 1; result.LastError = $"The selected directory '{localDirectory.Path}' does not exist."; break; } return result; } private void EnumerateAccessibleFiles(string rootPath, FileEnumerationResult result) { var pendingDirectories = new Stack(); pendingDirectories.Push(rootPath); while (pendingDirectories.Count > 0) { var currentPath = pendingDirectories.Pop(); IEnumerable subDirectories; IEnumerable files; try { subDirectories = Directory.EnumerateDirectories(currentPath); files = Directory.EnumerateFiles(currentPath); } catch (Exception exception) { logger.LogWarning(exception, "Cannot access directory '{DirectoryPath}' while indexing.", currentPath); result.FailedFiles++; result.LastError = $"The directory '{currentPath}' could not be accessed."; continue; } foreach (var filePath in files) { FileInfo fileInfo; try { fileInfo = new FileInfo(filePath); if (!fileInfo.Exists) continue; } catch (Exception exception) { logger.LogWarning(exception, "Cannot inspect file '{FilePath}' while indexing.", filePath); result.FailedFiles++; result.LastError = $"The file '{filePath}' could not be inspected."; continue; } switch (this.GetRagFileIndexingDecision(fileInfo)) { case RagFileIndexingDecision.INDEXABLE: result.Files.Add(fileInfo); break; case RagFileIndexingDecision.EXCLUDED: logger.LogDebug("Skipping excluded file '{FilePath}' while indexing.", fileInfo.FullName); break; } } foreach (var subDirectory in subDirectories) { if (this.IsSkippedRagDirectory(subDirectory)) continue; pendingDirectories.Push(subDirectory); } } } private string TryGetRelativePath(IDataSource dataSource, FileInfo file) => dataSource switch { DataSourceLocalDirectory localDirectory => Path.GetRelativePath(localDirectory.Path, file.FullName), _ => file.Name }; private static string NormalizeChunkSegment(string input) { return input .Replace("\r\n", "\n", StringComparison.Ordinal) .Replace('\r', '\n') .Trim(); } private bool IsImageFilePath(string filePath) { return FileTypes.IsAllowedPath(filePath, FileTypes.IMAGE); } private bool IsSupportedRagFilePath(string filePath) { var extension = Path.GetExtension(filePath).TrimStart('.'); return FileTypes.IsAllowedPath(filePath, FileTypes.DOCUMENT, FileTypes.IMAGE) || RAG_DELIMITED_TABLE_FILE_EXTENSIONS.Contains(extension, StringComparer.OrdinalIgnoreCase) || RAG_SPREADSHEET_FILE_EXTENSIONS.Contains(extension, StringComparer.OrdinalIgnoreCase) || RAG_SPREADSHEET_ADD_IN_FILE_EXTENSIONS.Contains(extension, StringComparer.OrdinalIgnoreCase); } private RagFileIndexingDecision GetRagFileIndexingDecision(FileInfo file) { if (this.IsSkippedRagFile(file)) return RagFileIndexingDecision.EXCLUDED; return this.IsSupportedRagFilePath(file.FullName) ? RagFileIndexingDecision.INDEXABLE : RagFileIndexingDecision.UNSUPPORTED; } private bool IsSkippedRagFile(FileInfo file) { if (IsSkippedRagFileName(file.Name)) return true; try { return file.Attributes.HasFlag(FileAttributes.ReparsePoint) || file.Attributes.HasFlag(FileAttributes.Offline) || file.Attributes.HasFlag(FileAttributes.Temporary) || file.Attributes.HasFlag(FileAttributes.System); } catch (Exception exception) { logger.LogWarning(exception, "Cannot inspect file '{FilePath}' while indexing.", file.FullName); return true; } } private static bool IsSkippedRagFileName(string fileName) { var extension = Path.GetExtension(fileName).TrimStart('.'); return SKIPPED_RAG_FILE_EXTENSIONS.Contains(extension, StringComparer.OrdinalIgnoreCase) || fileName.StartsWith(OFFICE_LOCK_FILE_PREFIX, StringComparison.Ordinal); } private bool IsSkippedRagDirectory(string path) { try { var directory = new DirectoryInfo(path); return directory.Attributes.HasFlag(FileAttributes.ReparsePoint) || directory.Attributes.HasFlag(FileAttributes.Offline) || directory.Attributes.HasFlag(FileAttributes.System); } catch (Exception exception) { logger.LogWarning(exception, "Cannot inspect directory '{DirectoryPath}' while indexing.", path); return true; } } private string BuildImageIndexText(string filePath) { var fileName = Path.GetFileName(filePath); var fileNameWithoutExtension = Path.GetFileNameWithoutExtension(filePath); var extension = Path.GetExtension(filePath).TrimStart('.'); var normalizedName = fileNameWithoutExtension .Replace('_', ' ') .Replace('-', ' ') .Trim(); return $$""" Image asset File name: {{fileName}} Type: {{extension}} Search terms: {{normalizedName}} Path: {{filePath}} Note: The current RAG embedding pipeline stores image files by metadata only. Visual content is not OCRed or captioned yet. """; } private string BuildEmbeddingSignature(EmbeddingProvider embeddingProvider) { return string.Join('|', embeddingProvider.Id, embeddingProvider.UsedLLMProvider, embeddingProvider.Model.Id, embeddingProvider.Host, embeddingProvider.Hostname, embeddingProvider.TokenizerPath); } private async Task BuildFingerprintAsync(FileInfo file, CancellationToken token) { await using var stream = new FileStream( file.FullName, FileMode.Open, FileAccess.Read, FileShare.ReadWrite | FileShare.Delete, 1024 * 128, FileOptions.Asynchronous | FileOptions.SequentialScan); var contentHash = await SHA256.HashDataAsync(stream, token); var fingerprintSource = $"{file.FullName}|{Convert.ToHexString(contentHash)}"; var bytes = SHA256.HashData(Encoding.UTF8.GetBytes(fingerprintSource)); return Convert.ToHexString(bytes); } private string GetCollectionName(string dataSourceName, string dataSourceId) { var safeId = dataSourceId .ToLowerInvariant() .Replace("-", string.Empty, StringComparison.Ordinal); var safeName = new string(dataSourceName .ToLowerInvariant() .Where(c => c is >= 'a' and <= 'z' or >= '0' and <= '9') .Take(32) .ToArray()); safeName = string.IsNullOrWhiteSpace(safeName) ? "datasource" : safeName; return $"rag_{safeName}_{safeId}"; } private string CreatePointId(string dataSourceId, string fingerprint, int chunkIndex) { var source = $"{dataSourceId}:{fingerprint}:{chunkIndex}"; var hash = SHA256.HashData(Encoding.UTF8.GetBytes(source)); var guidBytes = hash[..16].ToArray(); guidBytes[6] = (byte)((guidBytes[6] & 0x0F) | 0x40); guidBytes[8] = (byte)((guidBytes[8] & 0x3F) | 0x80); return new Guid(guidBytes).ToString(); } }