mirror of
https://github.com/MindWorkAI/AI-Studio.git
synced 2026-08-11 19:32:10 +00:00
942 lines
35 KiB
C#
942 lines
35 KiB
C#
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;
|
|
|
|
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 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);
|
|
|
|
private sealed record ChunkingRule(string Name, Func<string, IReadOnlyList<string>, IReadOnlyList<string>>? Split);
|
|
|
|
private sealed record DataSourceMetadataSnapshot(string SourceHash, IReadOnlyDictionary<string, string> FileHashes);
|
|
|
|
private async IAsyncEnumerable<string> StreamEmbeddingChunksAsync(string filePath, IDataSource dataSource, EmbeddingProvider embeddingProvider, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token)
|
|
{
|
|
var options = this.GetChunkingOptions(dataSource, embeddingProvider);
|
|
var strategy = this.GetChunkingStrategy(filePath);
|
|
ExtractedFileContent content;
|
|
|
|
if (this.IsImageFilePath(filePath))
|
|
{
|
|
var imageIndexText = this.BuildImageIndexText(filePath);
|
|
content = new(imageIndexText, [imageIndexText]);
|
|
}
|
|
else
|
|
{
|
|
content = await this.ReadExtractedFileContentAsync(filePath, token);
|
|
}
|
|
|
|
await foreach (var chunk in this.SplitByChunkingStrategyAsync(content, strategy, options, embeddingProvider, token))
|
|
yield return chunk;
|
|
}
|
|
|
|
private async Task<ExtractedFileContent> ReadExtractedFileContentAsync(string filePath, CancellationToken token)
|
|
{
|
|
var segments = new List<string>();
|
|
|
|
await foreach (var segment in rustService.StreamArbitraryFileData(filePath, token: token))
|
|
{
|
|
var normalized = NormalizeChunkSegment(segment);
|
|
if (!string.IsNullOrWhiteSpace(normalized))
|
|
segments.Add(normalized);
|
|
}
|
|
|
|
return new(string.Join("\n", segments).Trim(), segments);
|
|
}
|
|
|
|
private async IAsyncEnumerable<string> SplitByChunkingStrategyAsync(ExtractedFileContent content, ChunkingStrategy strategy, ChunkingOptions options, EmbeddingProvider embeddingProvider, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token)
|
|
{
|
|
await foreach (var chunk in this.SplitTextByRulesAsync(content.Text, content.SourceSegments, strategy, 0, options, embeddingProvider, token))
|
|
yield return chunk;
|
|
}
|
|
|
|
private async IAsyncEnumerable<string> SplitTextByRulesAsync(
|
|
string text,
|
|
IReadOnlyList<string> sourceSegments,
|
|
ChunkingStrategy strategy,
|
|
int ruleIndex,
|
|
ChunkingOptions options,
|
|
EmbeddingProvider embeddingProvider,
|
|
[System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token)
|
|
{
|
|
text = text.Trim();
|
|
if (string.IsNullOrWhiteSpace(text))
|
|
yield break;
|
|
|
|
var tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, text, token);
|
|
if (tokenCount <= options.MaxChunkTokenLength)
|
|
{
|
|
yield return text;
|
|
yield break;
|
|
}
|
|
|
|
if (ruleIndex >= strategy.Rules.Count)
|
|
{
|
|
await foreach (var hardChunk in this.SplitTextByHardCutAsync(text, options, embeddingProvider, token))
|
|
yield return hardChunk;
|
|
|
|
yield break;
|
|
}
|
|
|
|
var rule = strategy.Rules[ruleIndex];
|
|
if (rule.Split is null)
|
|
{
|
|
await foreach (var hardChunk in this.SplitTextByHardCutAsync(text, options, embeddingProvider, token))
|
|
yield return hardChunk;
|
|
|
|
yield break;
|
|
}
|
|
|
|
var units = NormalizeSplitUnits(rule.Split(text, sourceSegments), text);
|
|
if (units.Count <= 1)
|
|
{
|
|
await foreach (var chunk in this.SplitTextByRulesAsync(text, sourceSegments, strategy, ruleIndex + 1, options, embeddingProvider, token))
|
|
yield return chunk;
|
|
|
|
yield break;
|
|
}
|
|
|
|
logger.LogDebug(
|
|
"Splitting content for embedding provider '{EmbeddingProviderName}' with strategy '{ChunkingStrategy}' and rule '{ChunkingRule}'. TokenCount={TokenCount}, MaxChunkTokenLength={MaxChunkTokenLength}.",
|
|
embeddingProvider.Name,
|
|
strategy.Name,
|
|
rule.Name,
|
|
tokenCount,
|
|
options.MaxChunkTokenLength);
|
|
|
|
var index = 0;
|
|
|
|
while (index < units.Count)
|
|
{
|
|
token.ThrowIfCancellationRequested();
|
|
|
|
var unitCount = await this.FindLargestUnitCountWithinMaxChunkLengthAsync(units, index, embeddingProvider, options.MaxChunkTokenLength, token);
|
|
if (unitCount > 0)
|
|
{
|
|
var chunk = string.Concat(units.Skip(index).Take(unitCount)).Trim();
|
|
if (!string.IsNullOrWhiteSpace(chunk))
|
|
yield return chunk;
|
|
|
|
var nextIndex = index + unitCount;
|
|
if (nextIndex >= units.Count)
|
|
yield break;
|
|
|
|
index = await this.CalculateNextStartIndexAsync(units, index, nextIndex, options, embeddingProvider, token);
|
|
continue;
|
|
}
|
|
|
|
await foreach (var splitUnit in this.SplitTextByRulesAsync(units[index], [units[index]], strategy, ruleIndex + 1, options, embeddingProvider, token))
|
|
yield return splitUnit;
|
|
|
|
index++;
|
|
}
|
|
}
|
|
|
|
private async Task<int> FindLargestUnitCountWithinMaxChunkLengthAsync(IReadOnlyList<string> units, int startIndex, EmbeddingProvider embeddingProvider, int maxChunkTokenLength, CancellationToken token)
|
|
{
|
|
var low = 1;
|
|
var high = units.Count - startIndex;
|
|
var best = 0;
|
|
|
|
while (low <= high)
|
|
{
|
|
token.ThrowIfCancellationRequested();
|
|
|
|
var mid = low + (high - low) / 2;
|
|
var candidate = string.Concat(units.Skip(startIndex).Take(mid)).Trim();
|
|
var tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, candidate, token);
|
|
if (tokenCount <= maxChunkTokenLength)
|
|
{
|
|
best = mid;
|
|
low = mid + 1;
|
|
}
|
|
else
|
|
high = mid - 1;
|
|
}
|
|
|
|
return best;
|
|
}
|
|
|
|
private async Task<int> CalculateNextStartIndexAsync(IReadOnlyList<string> units, int chunkStartIndex, int chunkEndIndex, ChunkingOptions options, EmbeddingProvider embeddingProvider, CancellationToken token)
|
|
{
|
|
if (options.OverlapTokenLength <= 0)
|
|
return chunkEndIndex;
|
|
|
|
var bestStartIndex = chunkEndIndex;
|
|
var bestDistance = int.MaxValue;
|
|
|
|
for (var candidateStartIndex = chunkEndIndex - 1; candidateStartIndex > chunkStartIndex; candidateStartIndex--)
|
|
{
|
|
token.ThrowIfCancellationRequested();
|
|
|
|
var candidate = string.Concat(units.Skip(candidateStartIndex).Take(chunkEndIndex - candidateStartIndex)).Trim();
|
|
if (string.IsNullOrWhiteSpace(candidate))
|
|
continue;
|
|
|
|
var tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, candidate, token);
|
|
var distance = Math.Abs(tokenCount - options.OverlapTokenLength);
|
|
if (distance < bestDistance)
|
|
{
|
|
bestStartIndex = candidateStartIndex;
|
|
bestDistance = distance;
|
|
}
|
|
|
|
if (tokenCount >= options.OverlapTokenLength && bestStartIndex < chunkEndIndex)
|
|
break;
|
|
}
|
|
|
|
return bestStartIndex <= chunkStartIndex ? chunkEndIndex : bestStartIndex;
|
|
}
|
|
|
|
private async IAsyncEnumerable<string> SplitTextByHardCutAsync(string text, ChunkingOptions options, EmbeddingProvider embeddingProvider, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token)
|
|
{
|
|
var startIndex = 0;
|
|
while (startIndex < text.Length)
|
|
{
|
|
token.ThrowIfCancellationRequested();
|
|
|
|
var low = startIndex + 1;
|
|
var high = text.Length;
|
|
var bestEndIndex = startIndex;
|
|
|
|
while (low <= high)
|
|
{
|
|
var mid = low + (high - low) / 2;
|
|
var candidate = text[startIndex..mid].Trim();
|
|
var tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, candidate, token);
|
|
if (tokenCount <= options.MaxChunkTokenLength)
|
|
{
|
|
bestEndIndex = mid;
|
|
low = mid + 1;
|
|
}
|
|
else
|
|
high = mid - 1;
|
|
}
|
|
|
|
if (bestEndIndex == startIndex)
|
|
{
|
|
var smallestCandidate = text[startIndex..Math.Min(startIndex + 1, text.Length)].Trim();
|
|
var smallestCandidateTokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, smallestCandidate, token);
|
|
throw new InvalidOperationException($"The max chunk length for embedding provider '{embeddingProvider.Name}' is too low. The smallest possible split still has {smallestCandidateTokenCount} tokens, but the configured limit is {options.MaxChunkTokenLength}.");
|
|
}
|
|
|
|
var chunk = text[startIndex..bestEndIndex].Trim();
|
|
if (!string.IsNullOrWhiteSpace(chunk))
|
|
yield return chunk;
|
|
|
|
if (bestEndIndex >= text.Length)
|
|
yield break;
|
|
|
|
startIndex = await this.CalculateHardCutOverlapStartIndexAsync(text, startIndex, bestEndIndex, options, embeddingProvider, token);
|
|
}
|
|
}
|
|
|
|
private async Task<int> CalculateHardCutOverlapStartIndexAsync(string text, int chunkStartIndex, int chunkEndIndex, ChunkingOptions options, EmbeddingProvider embeddingProvider, CancellationToken token)
|
|
{
|
|
if (options.OverlapTokenLength <= 0 || chunkEndIndex - chunkStartIndex <= 1)
|
|
return chunkEndIndex;
|
|
|
|
var low = chunkStartIndex + 1;
|
|
var high = chunkEndIndex - 1;
|
|
var bestStartIndex = chunkEndIndex;
|
|
|
|
while (low <= high)
|
|
{
|
|
token.ThrowIfCancellationRequested();
|
|
|
|
var mid = low + (high - low) / 2;
|
|
var candidate = text[mid..chunkEndIndex].Trim();
|
|
var tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, candidate, token);
|
|
if (tokenCount <= options.OverlapTokenLength)
|
|
{
|
|
bestStartIndex = mid;
|
|
high = mid - 1;
|
|
}
|
|
else
|
|
low = mid + 1;
|
|
}
|
|
|
|
return bestStartIndex <= chunkStartIndex ? chunkEndIndex : bestStartIndex;
|
|
}
|
|
|
|
private async Task<int> GetEmbeddingTokenCountAsync(EmbeddingProvider embeddingProvider, string text, CancellationToken token)
|
|
{
|
|
var response = await rustService.GetTokenCount(embeddingProvider.Name, embeddingProvider.TokenizerPath, text, token);
|
|
if (response is { Success: true, Status: TokenizerStatus.AVAILABLE })
|
|
return response.Value.TokenCount;
|
|
|
|
var message = response?.Message ?? "No response was returned by the tokenizer service.";
|
|
throw new InvalidOperationException($"Could not count tokens for embedding provider '{embeddingProvider.Name}'. {message}");
|
|
}
|
|
|
|
private ChunkingOptions GetChunkingOptions(IDataSource dataSource, EmbeddingProvider embeddingProvider)
|
|
{
|
|
var providerMaxChunkTokenLength = Math.Max(1, embeddingProvider.EffectiveTokenLimit);
|
|
var dataSourceMaxChunkTokenLength = dataSource is IInternalDataSource { MaxChunkTokenLength: > 0 } internalDataSource
|
|
? internalDataSource.MaxChunkTokenLength
|
|
: 0;
|
|
var maxChunkTokenLength = dataSourceMaxChunkTokenLength > 0
|
|
? Math.Min(dataSourceMaxChunkTokenLength, providerMaxChunkTokenLength)
|
|
: providerMaxChunkTokenLength;
|
|
|
|
var configuredOverlapTokenLength = dataSource is IInternalDataSource overlapDataSource
|
|
? overlapDataSource.ChunkOverlapTokenLength
|
|
: 0;
|
|
var overlapTokenLength = Math.Clamp(configuredOverlapTokenLength, 0, Math.Max(0, maxChunkTokenLength - 1));
|
|
|
|
return new(maxChunkTokenLength, overlapTokenLength);
|
|
}
|
|
|
|
private ChunkingStrategy GetChunkingStrategy(string filePath)
|
|
{
|
|
if (this.IsImageFilePath(filePath))
|
|
return new("image", [
|
|
new("Whitespace", SplitByWhitespace),
|
|
new("Hard cut", null),
|
|
]);
|
|
|
|
if (this.IsPresentationFilePath(filePath))
|
|
return new("presentation", [
|
|
new("Slide", SplitBySourceSegments),
|
|
new("Line break", SplitByLineBreaks),
|
|
new("Whitespace", SplitByWhitespace),
|
|
new("Hard cut", null),
|
|
]);
|
|
|
|
if (this.IsDelimitedTableFilePath(filePath) || this.IsSpreadsheetFilePath(filePath))
|
|
return new("table", [
|
|
new("Row or sheet", SplitBySourceSegments),
|
|
new("Line break", SplitByLineBreaks),
|
|
new("Whitespace", SplitByWhitespace),
|
|
new("Hard cut", null),
|
|
]);
|
|
|
|
if (this.IsSourceCodeFilePath(filePath))
|
|
return GetSourceCodeChunkingStrategy(filePath);
|
|
|
|
return new("document", [
|
|
new("Heading", SplitByDocumentHeadings),
|
|
new("Page or extracted section", SplitBySourceSegments),
|
|
new("Paragraph", SplitByParagraphs),
|
|
new("Line break", SplitByLineBreaks),
|
|
new("Whitespace", SplitByWhitespace),
|
|
new("Hard cut", null),
|
|
]);
|
|
}
|
|
|
|
private static ChunkingStrategy GetSourceCodeChunkingStrategy(string filePath)
|
|
{
|
|
var rules = GetSourceCodeDelimiterRules(filePath).ToList();
|
|
rules.Add(new("Line break", SplitByLineBreaks));
|
|
rules.Add(new("Whitespace", SplitByWhitespace));
|
|
rules.Add(new("Hard cut", null));
|
|
|
|
return new("source-code", rules);
|
|
}
|
|
|
|
private static IReadOnlyList<ChunkingRule> GetSourceCodeDelimiterRules(string filePath) => Path.GetExtension(filePath).TrimStart('.') switch
|
|
{
|
|
_ => [],
|
|
};
|
|
|
|
private static List<string> NormalizeSplitUnits(IReadOnlyList<string> units, string fallbackText)
|
|
{
|
|
var result = units
|
|
.Where(unit => !string.IsNullOrWhiteSpace(unit))
|
|
.ToList();
|
|
|
|
return result.Count == 0 ? [fallbackText] : result;
|
|
}
|
|
|
|
private static IReadOnlyList<string> SplitBySourceSegments(string text, IReadOnlyList<string> sourceSegments)
|
|
{
|
|
return sourceSegments.Count > 1
|
|
? sourceSegments.Select(segment => segment + "\n").ToList()
|
|
: [text];
|
|
}
|
|
|
|
private static IReadOnlyList<string> SplitByDocumentHeadings(string text, IReadOnlyList<string> sourceSegments)
|
|
{
|
|
var lines = ReadLines(text);
|
|
if (lines.Count < 2)
|
|
return [text];
|
|
|
|
var result = new List<string>();
|
|
var segmentStart = 0;
|
|
|
|
for (var i = 0; i < lines.Count; i++)
|
|
{
|
|
var (lineStart, _, lineText) = lines[i];
|
|
if (lineStart == 0)
|
|
continue;
|
|
|
|
var previousLine = i > 0 ? lines[i - 1].Text : string.Empty;
|
|
var nextLine = i + 1 < lines.Count ? lines[i + 1].Text : string.Empty;
|
|
if (!IsDocumentHeadingLine(lineText, previousLine, nextLine))
|
|
continue;
|
|
|
|
result.Add(text[segmentStart..lineStart]);
|
|
segmentStart = lineStart;
|
|
}
|
|
|
|
if (segmentStart == 0)
|
|
return [text];
|
|
|
|
result.Add(text[segmentStart..]);
|
|
return result;
|
|
}
|
|
|
|
private static IReadOnlyList<string> SplitByParagraphs(string text, IReadOnlyList<string> sourceSegments)
|
|
{
|
|
var matches = Regex.Matches(text, @"\n[ \t]*\n", RegexOptions.CultureInvariant);
|
|
if (matches.Count == 0)
|
|
return [text];
|
|
|
|
var result = new List<string>();
|
|
var start = 0;
|
|
foreach (Match match in matches)
|
|
{
|
|
var end = match.Index + match.Length;
|
|
result.Add(text[start..end]);
|
|
start = end;
|
|
}
|
|
|
|
if (start < text.Length)
|
|
result.Add(text[start..]);
|
|
|
|
return result;
|
|
}
|
|
|
|
private static IReadOnlyList<string> SplitByLineBreaks(string text, IReadOnlyList<string> sourceSegments)
|
|
{
|
|
var result = new List<string>();
|
|
var start = 0;
|
|
|
|
for (var i = 0; i < text.Length; i++)
|
|
{
|
|
if (text[i] != '\n')
|
|
continue;
|
|
|
|
result.Add(text[start..(i + 1)]);
|
|
start = i + 1;
|
|
}
|
|
|
|
if (start < text.Length)
|
|
result.Add(text[start..]);
|
|
|
|
return result.Count == 0 ? [text] : result;
|
|
}
|
|
|
|
private static IReadOnlyList<string> SplitByWhitespace(string text, IReadOnlyList<string> sourceSegments)
|
|
{
|
|
var matches = Regex.Matches(text, @"\S+\s*", RegexOptions.CultureInvariant);
|
|
if (matches.Count == 0)
|
|
return [text];
|
|
|
|
return matches.Cast<Match>().Select(match => match.Value).ToList();
|
|
}
|
|
|
|
private static List<(int Start, int End, string Text)> ReadLines(string text)
|
|
{
|
|
var result = new List<(int Start, int End, string Text)>();
|
|
var start = 0;
|
|
|
|
for (var i = 0; i < text.Length; i++)
|
|
{
|
|
if (text[i] != '\n')
|
|
continue;
|
|
|
|
result.Add((start, i + 1, text[start..(i + 1)]));
|
|
start = i + 1;
|
|
}
|
|
|
|
if (start < text.Length)
|
|
result.Add((start, text.Length, text[start..]));
|
|
|
|
return result;
|
|
}
|
|
|
|
private static bool IsDocumentHeadingLine(string line, string previousLine, string nextLine)
|
|
{
|
|
var trimmed = line.Trim();
|
|
if (string.IsNullOrWhiteSpace(trimmed))
|
|
return false;
|
|
|
|
if (Regex.IsMatch(trimmed, @"^#{1,6}\s+\S", RegexOptions.CultureInvariant))
|
|
return true;
|
|
|
|
if (!string.IsNullOrWhiteSpace(previousLine) || !string.IsNullOrWhiteSpace(nextLine))
|
|
return false;
|
|
|
|
if (trimmed.Length is < 3 or > 120)
|
|
return false;
|
|
|
|
if (trimmed.Contains("|", StringComparison.Ordinal) || trimmed.EndsWith(".", StringComparison.Ordinal))
|
|
return false;
|
|
|
|
return Regex.IsMatch(trimmed, @"^(\d+(\.\d+)*\.?\s+\S|(?i:chapter|section)\s+\S|[A-Z0-9][A-Z0-9 ,:;'/&()_-]{2,})$", RegexOptions.CultureInvariant);
|
|
}
|
|
|
|
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.AddFailure(localFile.FilePath, $"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.AddFailure(localFile.FilePath, $"The selected file '{localFile.FilePath}' does not exist.");
|
|
break;
|
|
|
|
case DataSourceLocalDirectory localDirectory:
|
|
result.AddFailure(localDirectory.Path, $"The selected directory '{localDirectory.Path}' does not exist.");
|
|
break;
|
|
}
|
|
|
|
return result;
|
|
}
|
|
|
|
private void EnumerateAccessibleFiles(string rootPath, FileEnumerationResult result)
|
|
{
|
|
var pendingDirectories = new Stack<string>();
|
|
pendingDirectories.Push(rootPath);
|
|
|
|
while (pendingDirectories.Count > 0)
|
|
{
|
|
var currentPath = pendingDirectories.Pop();
|
|
IEnumerable<string> subDirectories;
|
|
IEnumerable<string> files;
|
|
|
|
try
|
|
{
|
|
subDirectories = Directory.EnumerateDirectories(currentPath);
|
|
files = Directory.EnumerateFiles(currentPath);
|
|
}
|
|
catch (Exception exception)
|
|
{
|
|
logger.LogWarning(exception, "Cannot access directory '{DirectoryPath}' while indexing.", currentPath);
|
|
result.AddFailure(currentPath, $"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.AddFailure(filePath, $"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 IsPresentationFilePath(string filePath)
|
|
{
|
|
return FileTypes.IsAllowedPath(filePath, FileTypes.POWER_POINT);
|
|
}
|
|
|
|
private bool IsDelimitedTableFilePath(string filePath)
|
|
{
|
|
var extension = Path.GetExtension(filePath).TrimStart('.');
|
|
return RAG_DELIMITED_TABLE_FILE_EXTENSIONS.Contains(extension, StringComparer.OrdinalIgnoreCase);
|
|
}
|
|
|
|
private bool IsSpreadsheetFilePath(string filePath)
|
|
{
|
|
var extension = Path.GetExtension(filePath).TrimStart('.');
|
|
return FileTypes.IsAllowedPath(filePath, FileTypes.EXCEL)
|
|
|| RAG_SPREADSHEET_FILE_EXTENSIONS.Contains(extension, StringComparer.OrdinalIgnoreCase)
|
|
|| RAG_SPREADSHEET_ADD_IN_FILE_EXTENSIONS.Contains(extension, StringComparer.OrdinalIgnoreCase);
|
|
}
|
|
|
|
private bool IsSourceCodeFilePath(string filePath)
|
|
{
|
|
return !this.IsHtmlFilePath(filePath) && FileTypes.IsAllowedPath(filePath, FileTypes.SOURCE_CODE);
|
|
}
|
|
|
|
private bool IsHtmlFilePath(string filePath)
|
|
{
|
|
var extension = Path.GetExtension(filePath).TrimStart('.');
|
|
return extension.Equals("html", StringComparison.OrdinalIgnoreCase)
|
|
|| extension.Equals("htm", StringComparison.OrdinalIgnoreCase);
|
|
}
|
|
|
|
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(IDataSource dataSource, EmbeddingProvider embeddingProvider, ChunkingOptions chunkingOptions)
|
|
{
|
|
return string.Join('|',
|
|
embeddingProvider.Id,
|
|
embeddingProvider.UsedLLMProvider,
|
|
embeddingProvider.Model.Id,
|
|
embeddingProvider.Host,
|
|
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,
|
|
chunkingOptions.OverlapTokenLength);
|
|
}
|
|
|
|
private DataSourceMetadataSnapshot BuildDataSourceMetadataSnapshot(IDataSource dataSource, IReadOnlyList<FileInfo> indexedFiles)
|
|
{
|
|
var fileHashes = indexedFiles
|
|
.OrderBy(file => file.FullName, StringComparer.OrdinalIgnoreCase)
|
|
.ToDictionary(file => file.FullName, BuildFileMetadataHash, StringComparer.OrdinalIgnoreCase);
|
|
|
|
var sourceHash = dataSource switch
|
|
{
|
|
DataSourceLocalFile localFile => indexedFiles.Count > 0
|
|
? fileHashes[indexedFiles[0].FullName]
|
|
: BuildMetadataHash("file", localFile.FilePath, Path.GetFileName(localFile.FilePath) ?? string.Empty, "missing", "0"),
|
|
|
|
DataSourceLocalDirectory localDirectory => this.BuildDirectoryMetadataHash(localDirectory, indexedFiles, fileHashes),
|
|
|
|
_ => BuildMetadataHash(dataSource.Type.ToString(), dataSource.Id, dataSource.Name)
|
|
};
|
|
|
|
return new(sourceHash, fileHashes);
|
|
}
|
|
|
|
private string BuildDirectoryMetadataHash(DataSourceLocalDirectory dataSource, IReadOnlyList<FileInfo> indexedFiles, IReadOnlyDictionary<string, string> fileHashes)
|
|
{
|
|
var directory = new DirectoryInfo(dataSource.Path);
|
|
directory.Refresh();
|
|
|
|
var totalSize = 0L;
|
|
var latestFileWriteTicks = 0L;
|
|
foreach (var file in indexedFiles)
|
|
{
|
|
file.Refresh();
|
|
if (!file.Exists)
|
|
continue;
|
|
|
|
totalSize += file.Length;
|
|
latestFileWriteTicks = Math.Max(latestFileWriteTicks, file.LastWriteTimeUtc.Ticks);
|
|
}
|
|
|
|
var latestWriteTicks = Math.Max(directory.LastWriteTimeUtc.Ticks, latestFileWriteTicks);
|
|
var parts = new List<string>
|
|
{
|
|
"directory",
|
|
directory.FullName,
|
|
directory.Name,
|
|
latestWriteTicks.ToString(),
|
|
totalSize.ToString(),
|
|
indexedFiles.Count.ToString()
|
|
};
|
|
|
|
foreach (var file in indexedFiles.OrderBy(file => file.FullName, StringComparer.OrdinalIgnoreCase))
|
|
{
|
|
parts.Add(this.TryGetRelativePath(dataSource, file));
|
|
parts.Add(fileHashes[file.FullName]);
|
|
}
|
|
|
|
return BuildMetadataHash(parts);
|
|
}
|
|
|
|
private static string BuildFileMetadataHash(FileInfo file)
|
|
{
|
|
file.Refresh();
|
|
if (!file.Exists)
|
|
{
|
|
return BuildMetadataHash(
|
|
"file",
|
|
file.FullName,
|
|
file.Name,
|
|
"missing",
|
|
"0");
|
|
}
|
|
|
|
return BuildMetadataHash(
|
|
"file",
|
|
file.FullName,
|
|
file.Name,
|
|
file.LastWriteTimeUtc.Ticks.ToString(),
|
|
file.Length.ToString());
|
|
}
|
|
|
|
private static string BuildMetadataHash(params string[] parts)
|
|
{
|
|
return BuildMetadataHash((IEnumerable<string>)parts);
|
|
}
|
|
|
|
private static string BuildMetadataHash(IEnumerable<string> parts)
|
|
{
|
|
var fingerprintSource = new StringBuilder();
|
|
foreach (var part in parts)
|
|
fingerprintSource.Append(part.Length).Append(':').Append(part).Append('|');
|
|
|
|
var bytes = SHA256.HashData(Encoding.UTF8.GetBytes(fingerprintSource.ToString()));
|
|
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
|
|
.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) =>
|
|
CreateStableGuid($"{dataSourceId}:chunk:{fingerprint}:{chunkIndex}");
|
|
|
|
private string CreateParentFileId(string dataSourceId, string absolutePath) =>
|
|
CreateStableGuid($"{dataSourceId}:parent-file:{absolutePath}");
|
|
|
|
private static string CreateStableGuid(string source)
|
|
{
|
|
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();
|
|
}
|
|
}
|