mirror of
https://github.com/MindWorkAI/AI-Studio.git
synced 2026-08-13 14:22:11 +00:00
forced overlap on chunks
This commit is contained in:
parent
010f35b08a
commit
2a3a82820c
@ -3871,9 +3871,6 @@ UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALDIRECTORYDIALOG::T1948697886"
|
||||
-- The overlap must be smaller than the effective token limit.
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALDIRECTORYDIALOG::T2101951526"] = "The overlap must be smaller than the effective token limit."
|
||||
|
||||
-- Number of tokens repeated at the start of the next chunk. Use 0 to disable overlap.
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALDIRECTORYDIALOG::T2263831524"] = "Number of tokens repeated at the start of the next chunk. Use 0 to disable overlap."
|
||||
|
||||
-- Add
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALDIRECTORYDIALOG::T2646845972"] = "Add"
|
||||
|
||||
@ -3889,6 +3886,9 @@ UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALDIRECTORYDIALOG::T2838542994"
|
||||
-- Token limit
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALDIRECTORYDIALOG::T2961294165"] = "Token limit"
|
||||
|
||||
-- Number of tokens repeated at the start of the next chunk. Use 0 to use the default overlap.
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALDIRECTORYDIALOG::T3181127959"] = "Number of tokens repeated at the start of the next chunk. Use 0 to use the default overlap."
|
||||
|
||||
-- Please enter 0 or a positive overlap length.
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALDIRECTORYDIALOG::T3242265813"] = "Please enter 0 or a positive overlap length."
|
||||
|
||||
@ -4033,9 +4033,6 @@ UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALFILEDIALOG::T1948697886"] = "
|
||||
-- The overlap must be smaller than the effective token limit.
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALFILEDIALOG::T2101951526"] = "The overlap must be smaller than the effective token limit."
|
||||
|
||||
-- Number of tokens repeated at the start of the next chunk. Use 0 to disable overlap.
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALFILEDIALOG::T2263831524"] = "Number of tokens repeated at the start of the next chunk. Use 0 to disable overlap."
|
||||
|
||||
-- Add
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALFILEDIALOG::T2646845972"] = "Add"
|
||||
|
||||
@ -4054,6 +4051,9 @@ UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALFILEDIALOG::T2859265837"] = "
|
||||
-- Token limit
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALFILEDIALOG::T2961294165"] = "Token limit"
|
||||
|
||||
-- Number of tokens repeated at the start of the next chunk. Use 0 to use the default overlap.
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALFILEDIALOG::T3181127959"] = "Number of tokens repeated at the start of the next chunk. Use 0 to use the default overlap."
|
||||
|
||||
-- Please enter 0 or a positive overlap length.
|
||||
UI_TEXT_CONTENT["AISTUDIO::DIALOGS::DATASOURCELOCALFILEDIALOG::T3242265813"] = "Please enter 0 or a positive overlap length."
|
||||
|
||||
|
||||
@ -117,7 +117,7 @@
|
||||
Min="0"
|
||||
Immediate="@true"
|
||||
Validation="@this.ValidateChunkOverlapTokenLength"
|
||||
HelperText="@T("Number of tokens repeated at the start of the next chunk. Use 0 to disable overlap.")"
|
||||
HelperText="@T("Number of tokens repeated at the start of the next chunk. Use 0 to use the default overlap.")"
|
||||
Variant="Variant.Outlined"
|
||||
Adornment="Adornment.Start"
|
||||
AdornmentIcon="@Icons.Material.Filled.CompareArrows"
|
||||
|
||||
@ -117,7 +117,7 @@
|
||||
Min="0"
|
||||
Immediate="@true"
|
||||
Validation="@this.ValidateChunkOverlapTokenLength"
|
||||
HelperText="@T("Number of tokens repeated at the start of the next chunk. Use 0 to disable overlap.")"
|
||||
HelperText="@T("Number of tokens repeated at the start of the next chunk. Use 0 to use the default overlap.")"
|
||||
Variant="Variant.Outlined"
|
||||
Adornment="Adornment.Start"
|
||||
AdornmentIcon="@Icons.Material.Filled.CompareArrows"
|
||||
|
||||
@ -14,6 +14,7 @@ namespace AIStudio.Tools.Services;
|
||||
public sealed partial class DataSourceEmbeddingService
|
||||
{
|
||||
private const string OFFICE_LOCK_FILE_PREFIX = "~$";
|
||||
private const int DEFAULT_CHUNK_OVERLAP_TOKEN_LENGTH = 300;
|
||||
|
||||
private static readonly string[] RAG_DELIMITED_TABLE_FILE_EXTENSIONS = ["csv", "tsv"];
|
||||
private static readonly string[] RAG_SPREADSHEET_FILE_EXTENSIONS = ["ods", "xlsm", "xlsb"];
|
||||
@ -86,22 +87,23 @@ public sealed partial class DataSourceEmbeddingService
|
||||
int ruleIndex,
|
||||
ChunkingOptions options,
|
||||
EmbeddingProvider embeddingProvider,
|
||||
[System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token)
|
||||
[System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token,
|
||||
string requiredOverlapPrefix = "")
|
||||
{
|
||||
text = text.Trim();
|
||||
if (string.IsNullOrWhiteSpace(text))
|
||||
yield break;
|
||||
|
||||
var tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, text, token);
|
||||
var tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, AddOverlapPrefix(text, requiredOverlapPrefix), token);
|
||||
if (tokenCount <= options.MaxChunkTokenLength)
|
||||
{
|
||||
yield return text;
|
||||
yield return AddOverlapPrefix(text, requiredOverlapPrefix);
|
||||
yield break;
|
||||
}
|
||||
|
||||
if (ruleIndex >= strategy.Rules.Count)
|
||||
{
|
||||
await foreach (var hardChunk in this.SplitTextByHardCutAsync(text, options, embeddingProvider, token))
|
||||
await foreach (var hardChunk in this.SplitTextByHardCutAsync(text, options, embeddingProvider, token, requiredOverlapPrefix))
|
||||
yield return hardChunk;
|
||||
|
||||
yield break;
|
||||
@ -110,7 +112,7 @@ public sealed partial class DataSourceEmbeddingService
|
||||
var rule = strategy.Rules[ruleIndex];
|
||||
if (rule.Split is null)
|
||||
{
|
||||
await foreach (var hardChunk in this.SplitTextByHardCutAsync(text, options, embeddingProvider, token))
|
||||
await foreach (var hardChunk in this.SplitTextByHardCutAsync(text, options, embeddingProvider, token, requiredOverlapPrefix))
|
||||
yield return hardChunk;
|
||||
|
||||
yield break;
|
||||
@ -119,30 +121,34 @@ public sealed partial class DataSourceEmbeddingService
|
||||
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))
|
||||
await foreach (var chunk in this.SplitTextByRulesAsync(text, sourceSegments, strategy, ruleIndex + 1, options, embeddingProvider, token, requiredOverlapPrefix))
|
||||
yield return chunk;
|
||||
|
||||
yield break;
|
||||
}
|
||||
|
||||
logger.LogDebug(
|
||||
"Splitting content for embedding provider '{EmbeddingProviderName}' with strategy '{ChunkingStrategy}' and rule '{ChunkingRule}'. TokenCount={TokenCount}, MaxChunkTokenLength={MaxChunkTokenLength}.",
|
||||
"Splitting content for embedding provider '{EmbeddingProviderName}' with strategy '{ChunkingStrategy}' and rule '{ChunkingRule}'. TokenCount={TokenCount}, MaxChunkTokenLength={MaxChunkTokenLength}, OverlapTokenLength={OverlapTokenLength}.",
|
||||
embeddingProvider.Name,
|
||||
strategy.Name,
|
||||
rule.Name,
|
||||
tokenCount,
|
||||
options.MaxChunkTokenLength);
|
||||
options.MaxChunkTokenLength,
|
||||
options.OverlapTokenLength);
|
||||
|
||||
var index = 0;
|
||||
var overlapPrefix = requiredOverlapPrefix;
|
||||
|
||||
while (index < units.Count)
|
||||
{
|
||||
token.ThrowIfCancellationRequested();
|
||||
|
||||
var unitCount = await this.FindLargestUnitCountWithinMaxChunkLengthAsync(units, index, embeddingProvider, options.MaxChunkTokenLength, token);
|
||||
var unitCount = await this.FindLargestUnitCountWithinMaxChunkLengthAsync(units, index, embeddingProvider, options.MaxChunkTokenLength, token, overlapPrefix);
|
||||
if (unitCount > 0)
|
||||
{
|
||||
var chunk = string.Concat(units.Skip(index).Take(unitCount)).Trim();
|
||||
var rawChunk = string.Concat(units.Skip(index).Take(unitCount)).Trim();
|
||||
var chunk = AddOverlapPrefix(rawChunk, overlapPrefix);
|
||||
overlapPrefix = string.Empty;
|
||||
if (!string.IsNullOrWhiteSpace(chunk))
|
||||
yield return chunk;
|
||||
|
||||
@ -150,18 +156,45 @@ public sealed partial class DataSourceEmbeddingService
|
||||
if (nextIndex >= units.Count)
|
||||
yield break;
|
||||
|
||||
index = await this.CalculateNextStartIndexAsync(units, index, nextIndex, options, embeddingProvider, token);
|
||||
var nextStartIndex = await this.CalculateNextStartIndexAsync(units, index, nextIndex, options, embeddingProvider, token);
|
||||
if (nextStartIndex < nextIndex)
|
||||
{
|
||||
logger.LogDebug(
|
||||
"Applied delimiter overlap while chunking. Strategy='{ChunkingStrategy}', Rule='{ChunkingRule}', PreviousStartUnitIndex={PreviousStartUnitIndex}, PreviousEndUnitIndex={PreviousEndUnitIndex}, NextStartUnitIndex={NextStartUnitIndex}, OverlapUnits={OverlapUnits}, OverlapTokenLength={OverlapTokenLength}.",
|
||||
strategy.Name,
|
||||
rule.Name,
|
||||
index,
|
||||
nextIndex,
|
||||
nextStartIndex,
|
||||
nextIndex - nextStartIndex,
|
||||
options.OverlapTokenLength);
|
||||
|
||||
index = nextStartIndex;
|
||||
}
|
||||
else
|
||||
{
|
||||
overlapPrefix = await this.CreateOverlapPrefixAsync(chunk, strategy, rule, options, embeddingProvider, token);
|
||||
index = nextIndex;
|
||||
}
|
||||
|
||||
continue;
|
||||
}
|
||||
|
||||
await foreach (var splitUnit in this.SplitTextByRulesAsync(units[index], [units[index]], strategy, ruleIndex + 1, options, embeddingProvider, token))
|
||||
string? lastSplitUnit = null;
|
||||
await foreach (var splitUnit in this.SplitTextByRulesAsync(units[index], [units[index]], strategy, ruleIndex + 1, options, embeddingProvider, token, overlapPrefix))
|
||||
{
|
||||
lastSplitUnit = splitUnit;
|
||||
yield return splitUnit;
|
||||
}
|
||||
|
||||
overlapPrefix = lastSplitUnit is null
|
||||
? string.Empty
|
||||
: await this.CreateOverlapPrefixAsync(lastSplitUnit, strategy, rule, options, embeddingProvider, token);
|
||||
index++;
|
||||
}
|
||||
}
|
||||
|
||||
private async Task<int> FindLargestUnitCountWithinMaxChunkLengthAsync(IReadOnlyList<string> units, int startIndex, EmbeddingProvider embeddingProvider, int maxChunkTokenLength, CancellationToken token)
|
||||
private async Task<int> FindLargestUnitCountWithinMaxChunkLengthAsync(IReadOnlyList<string> units, int startIndex, EmbeddingProvider embeddingProvider, int maxChunkTokenLength, CancellationToken token, string overlapPrefix = "")
|
||||
{
|
||||
var low = 1;
|
||||
var high = units.Count - startIndex;
|
||||
@ -172,7 +205,7 @@ public sealed partial class DataSourceEmbeddingService
|
||||
token.ThrowIfCancellationRequested();
|
||||
|
||||
var mid = low + (high - low) / 2;
|
||||
var candidate = string.Concat(units.Skip(startIndex).Take(mid)).Trim();
|
||||
var candidate = AddOverlapPrefix(string.Concat(units.Skip(startIndex).Take(mid)).Trim(), overlapPrefix);
|
||||
var tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, candidate, token);
|
||||
if (tokenCount <= maxChunkTokenLength)
|
||||
{
|
||||
@ -186,6 +219,55 @@ public sealed partial class DataSourceEmbeddingService
|
||||
return best;
|
||||
}
|
||||
|
||||
private async Task<string> CreateOverlapPrefixAsync(string chunk, ChunkingStrategy strategy, ChunkingRule rule, ChunkingOptions options, EmbeddingProvider embeddingProvider, CancellationToken token)
|
||||
{
|
||||
return await this.CreateOverlapPrefixAsync(chunk, strategy.Name, rule.Name, options, embeddingProvider, token);
|
||||
}
|
||||
|
||||
private async Task<string> CreateOverlapPrefixAsync(string chunk, string strategyName, string ruleName, ChunkingOptions options, EmbeddingProvider embeddingProvider, CancellationToken token)
|
||||
{
|
||||
if (options.OverlapTokenLength <= 0)
|
||||
return string.Empty;
|
||||
|
||||
chunk = chunk.Trim();
|
||||
if (string.IsNullOrWhiteSpace(chunk))
|
||||
return string.Empty;
|
||||
|
||||
var chunkTokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, chunk, token);
|
||||
if (chunkTokenCount <= options.OverlapTokenLength)
|
||||
{
|
||||
logger.LogDebug(
|
||||
"Applied whole-chunk overlap while chunking because the previous chunk is smaller than the requested overlap. Strategy='{ChunkingStrategy}', Rule='{ChunkingRule}', RequestedOverlapTokenLength={RequestedOverlapTokenLength}, ActualOverlapTokenCount={ActualOverlapTokenCount}.",
|
||||
strategyName,
|
||||
ruleName,
|
||||
options.OverlapTokenLength,
|
||||
chunkTokenCount);
|
||||
|
||||
return chunk;
|
||||
}
|
||||
|
||||
var overlapStartIndex = await this.CalculateHardCutOverlapStartIndexAsync(chunk, 0, chunk.Length, options, embeddingProvider, token);
|
||||
if (overlapStartIndex >= chunk.Length)
|
||||
overlapStartIndex = FindLastNonWhitespaceStartIndex(chunk);
|
||||
|
||||
if (overlapStartIndex >= chunk.Length)
|
||||
return string.Empty;
|
||||
|
||||
var overlapPrefix = chunk[overlapStartIndex..].Trim();
|
||||
if (string.IsNullOrWhiteSpace(overlapPrefix))
|
||||
return string.Empty;
|
||||
|
||||
var tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, overlapPrefix, token);
|
||||
logger.LogDebug(
|
||||
"Applied hard-cut overlap while chunking because delimiter overlap was not available. Strategy='{ChunkingStrategy}', Rule='{ChunkingRule}', RequestedOverlapTokenLength={RequestedOverlapTokenLength}, ActualOverlapTokenCount={ActualOverlapTokenCount}.",
|
||||
strategyName,
|
||||
ruleName,
|
||||
options.OverlapTokenLength,
|
||||
tokenCount);
|
||||
|
||||
return overlapPrefix;
|
||||
}
|
||||
|
||||
private async Task<int> CalculateNextStartIndexAsync(IReadOnlyList<string> units, int chunkStartIndex, int chunkEndIndex, ChunkingOptions options, EmbeddingProvider embeddingProvider, CancellationToken token)
|
||||
{
|
||||
if (options.OverlapTokenLength <= 0)
|
||||
@ -217,13 +299,26 @@ public sealed partial class DataSourceEmbeddingService
|
||||
return bestStartIndex <= chunkStartIndex ? chunkEndIndex : bestStartIndex;
|
||||
}
|
||||
|
||||
private async IAsyncEnumerable<string> SplitTextByHardCutAsync(string text, ChunkingOptions options, EmbeddingProvider embeddingProvider, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token)
|
||||
private async IAsyncEnumerable<string> SplitTextByHardCutAsync(
|
||||
string text,
|
||||
ChunkingOptions options,
|
||||
EmbeddingProvider embeddingProvider,
|
||||
[System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token,
|
||||
string requiredOverlapPrefix = "")
|
||||
{
|
||||
text = text.Trim();
|
||||
var startIndex = 0;
|
||||
var overlapPrefix = requiredOverlapPrefix;
|
||||
while (startIndex < text.Length)
|
||||
{
|
||||
token.ThrowIfCancellationRequested();
|
||||
|
||||
while (startIndex < text.Length && char.IsWhiteSpace(text[startIndex]))
|
||||
startIndex++;
|
||||
|
||||
if (startIndex >= text.Length)
|
||||
yield break;
|
||||
|
||||
var low = startIndex + 1;
|
||||
var high = text.Length;
|
||||
var bestEndIndex = startIndex;
|
||||
@ -231,7 +326,7 @@ public sealed partial class DataSourceEmbeddingService
|
||||
while (low <= high)
|
||||
{
|
||||
var mid = low + (high - low) / 2;
|
||||
var candidate = text[startIndex..mid].Trim();
|
||||
var candidate = AddOverlapPrefix(text[startIndex..mid].Trim(), overlapPrefix);
|
||||
var tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, candidate, token);
|
||||
if (tokenCount <= options.MaxChunkTokenLength)
|
||||
{
|
||||
@ -244,19 +339,36 @@ public sealed partial class DataSourceEmbeddingService
|
||||
|
||||
if (bestEndIndex == startIndex)
|
||||
{
|
||||
var smallestCandidate = text[startIndex..Math.Min(startIndex + 1, text.Length)].Trim();
|
||||
if (!string.IsNullOrWhiteSpace(overlapPrefix))
|
||||
{
|
||||
var smallestOverlapPrefix = GetSmallestOverlapPrefix(overlapPrefix);
|
||||
if (!string.IsNullOrWhiteSpace(smallestOverlapPrefix) && !string.Equals(smallestOverlapPrefix, overlapPrefix, StringComparison.Ordinal))
|
||||
{
|
||||
logger.LogDebug(
|
||||
"Reduced hard-cut overlap because the configured overlap leaves no room for new content. RequestedOverlapTokenLength={RequestedOverlapTokenLength}, MaxChunkTokenLength={MaxChunkTokenLength}.",
|
||||
options.OverlapTokenLength,
|
||||
options.MaxChunkTokenLength);
|
||||
|
||||
overlapPrefix = smallestOverlapPrefix;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
var smallestCandidate = AddOverlapPrefix(text[startIndex..Math.Min(startIndex + 1, text.Length)].Trim(), overlapPrefix);
|
||||
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();
|
||||
var chunk = AddOverlapPrefix(text[startIndex..bestEndIndex].Trim(), overlapPrefix);
|
||||
overlapPrefix = string.Empty;
|
||||
if (!string.IsNullOrWhiteSpace(chunk))
|
||||
yield return chunk;
|
||||
|
||||
if (bestEndIndex >= text.Length)
|
||||
yield break;
|
||||
|
||||
startIndex = await this.CalculateHardCutOverlapStartIndexAsync(text, startIndex, bestEndIndex, options, embeddingProvider, token);
|
||||
overlapPrefix = await this.CreateOverlapPrefixAsync(chunk, "hard-cut", "Hard cut", options, embeddingProvider, token);
|
||||
startIndex = bestEndIndex;
|
||||
}
|
||||
}
|
||||
|
||||
@ -288,6 +400,31 @@ public sealed partial class DataSourceEmbeddingService
|
||||
return bestStartIndex <= chunkStartIndex ? chunkEndIndex : bestStartIndex;
|
||||
}
|
||||
|
||||
private static int FindLastNonWhitespaceStartIndex(string text)
|
||||
{
|
||||
for (var index = text.Length - 1; index >= 0; index--)
|
||||
{
|
||||
if (!char.IsWhiteSpace(text[index]))
|
||||
return index;
|
||||
}
|
||||
|
||||
return text.Length;
|
||||
}
|
||||
|
||||
private static string GetSmallestOverlapPrefix(string text)
|
||||
{
|
||||
var index = FindLastNonWhitespaceStartIndex(text);
|
||||
return index >= text.Length ? string.Empty : text[index..].Trim();
|
||||
}
|
||||
|
||||
private static string AddOverlapPrefix(string chunk, string overlapPrefix)
|
||||
{
|
||||
if (string.IsNullOrWhiteSpace(overlapPrefix))
|
||||
return chunk.Trim();
|
||||
|
||||
return $"{overlapPrefix.TrimEnd()}\n{chunk.TrimStart()}".Trim();
|
||||
}
|
||||
|
||||
private async Task<int> GetEmbeddingTokenCountAsync(EmbeddingProvider embeddingProvider, string text, CancellationToken token)
|
||||
{
|
||||
var response = await rustService.GetTokenCount(embeddingProvider.Name, embeddingProvider.TokenizerPath, text, token);
|
||||
@ -311,7 +448,10 @@ public sealed partial class DataSourceEmbeddingService
|
||||
var configuredOverlapTokenLength = dataSource is IInternalDataSource overlapDataSource
|
||||
? overlapDataSource.ChunkOverlapTokenLength
|
||||
: 0;
|
||||
var overlapTokenLength = Math.Clamp(configuredOverlapTokenLength, 0, Math.Max(0, maxChunkTokenLength - 1));
|
||||
var requestedOverlapTokenLength = configuredOverlapTokenLength > 0
|
||||
? configuredOverlapTokenLength
|
||||
: DEFAULT_CHUNK_OVERLAP_TOKEN_LENGTH;
|
||||
var overlapTokenLength = Math.Clamp(requestedOverlapTokenLength, 0, Math.Max(0, maxChunkTokenLength - 1));
|
||||
|
||||
return new(maxChunkTokenLength, overlapTokenLength);
|
||||
}
|
||||
|
||||
@ -58,10 +58,9 @@ public sealed class DataSourceLocalRetrievalService(
|
||||
if (maxMatches == 0)
|
||||
return [];
|
||||
|
||||
var candidateLimit = maxMatches * 2;
|
||||
var collectionName = DataSourceEmbeddingNames.GetCollectionName(dataSource.Name, dataSource.Id);
|
||||
var vectorTask = this.SearchVectorAsync(dataSource, query, candidateLimit, collectionName, token);
|
||||
var bm25Task = this.SearchBm25Async(dataSource, query, candidateLimit, token);
|
||||
var vectorTask = this.SearchVectorAsync(dataSource, query, maxMatches, collectionName, token);
|
||||
var bm25Task = this.SearchBm25Async(dataSource, query, maxMatches, token);
|
||||
|
||||
await Task.WhenAll(vectorTask, bm25Task);
|
||||
token.ThrowIfCancellationRequested();
|
||||
@ -118,7 +117,13 @@ public sealed class DataSourceLocalRetrievalService(
|
||||
return [];
|
||||
}
|
||||
|
||||
return await vectorStore.SearchEmbeddingAsync(collectionName, vector, maxMatches, token);
|
||||
var results = this.LimitSearchResults(
|
||||
dataSource,
|
||||
"vector",
|
||||
await vectorStore.SearchEmbeddingAsync(collectionName, vector, maxMatches, token),
|
||||
maxMatches);
|
||||
this.LogVectorResults(dataSource, results);
|
||||
return results;
|
||||
}
|
||||
catch (OperationCanceledException) when (token.IsCancellationRequested)
|
||||
{
|
||||
@ -146,7 +151,13 @@ public sealed class DataSourceLocalRetrievalService(
|
||||
return [];
|
||||
}
|
||||
|
||||
return await embeddingState.SearchChunksAsync(dataSource.Id, query, maxMatches, token);
|
||||
var results = this.LimitSearchResults(
|
||||
dataSource,
|
||||
"BM25",
|
||||
await embeddingState.SearchChunksAsync(dataSource.Id, query, maxMatches, token),
|
||||
maxMatches);
|
||||
this.LogBm25Results(dataSource, results);
|
||||
return results;
|
||||
}
|
||||
catch (OperationCanceledException) when (token.IsCancellationRequested)
|
||||
{
|
||||
@ -159,6 +170,22 @@ public sealed class DataSourceLocalRetrievalService(
|
||||
}
|
||||
}
|
||||
|
||||
private IReadOnlyList<T> LimitSearchResults<T>(IInternalDataSource dataSource, string searchName, IReadOnlyList<T> results, int maxMatches)
|
||||
{
|
||||
if (results.Count <= maxMatches)
|
||||
return results;
|
||||
|
||||
logger.LogWarning(
|
||||
"Local RAG {SearchName} search returned {ReturnedHits} chunks for data source '{DataSourceName}' ({DataSourceId}), which exceeds the configured maximum {MaxMatches}. Truncating to the datasource limit.",
|
||||
searchName,
|
||||
results.Count,
|
||||
dataSource.Name,
|
||||
dataSource.Id,
|
||||
maxMatches);
|
||||
|
||||
return results.Take(maxMatches).ToList();
|
||||
}
|
||||
|
||||
private static IReadOnlyList<LocalRetrievalHit> MergeResults(
|
||||
IReadOnlyList<VectorSearchResult> vectorResults,
|
||||
IReadOnlyList<EmbeddingStateSearchResult> bm25Results,
|
||||
@ -263,8 +290,13 @@ public sealed class DataSourceLocalRetrievalService(
|
||||
private static string BuildReferenceTitle(LocalRetrievalHit hit)
|
||||
{
|
||||
var sourceName = FirstNonEmpty(hit.FileName, hit.DataSourceName);
|
||||
var page = hit.PageNumber is > 0 ? $", page {hit.PageNumber}" : string.Empty;
|
||||
return $"{sourceName} (chunk {hit.ChunkIndex + 1}{page})";
|
||||
return BuildChunkTitle(sourceName, hit.ChunkIndex, hit.PageNumber);
|
||||
}
|
||||
|
||||
private static string BuildChunkTitle(string sourceName, int chunkIndex, int? pageNumber)
|
||||
{
|
||||
var page = pageNumber is > 0 ? $", page {pageNumber}" : string.Empty;
|
||||
return $"{sourceName} (chunk {chunkIndex + 1}{page})";
|
||||
}
|
||||
|
||||
private static string BuildReferenceLink(string path, LocalRetrievalHit hit)
|
||||
@ -304,4 +336,52 @@ public sealed class DataSourceLocalRetrievalService(
|
||||
|
||||
private static string FirstNonEmpty(params string[] values) =>
|
||||
values.FirstOrDefault(value => !string.IsNullOrWhiteSpace(value)) ?? string.Empty;
|
||||
|
||||
private void LogVectorResults(IInternalDataSource dataSource, IReadOnlyList<VectorSearchResult> results)
|
||||
{
|
||||
if (results.Count == 0)
|
||||
{
|
||||
logger.LogInformation("Local RAG vector search found no chunks for data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id);
|
||||
return;
|
||||
}
|
||||
|
||||
foreach (var result in results.Select((result, index) => (Result: result, Rank: index + 1)))
|
||||
{
|
||||
logger.LogInformation(
|
||||
"Local RAG vector search found chunk for data source '{DataSourceName}' ({DataSourceId}). Rank={Rank}, Score={Score}, ChunkId='{ChunkId}', ParentFileId='{ParentFileId}', File='{FileName}', Path='{Path}', Title='{Title}'.",
|
||||
dataSource.Name,
|
||||
dataSource.Id,
|
||||
result.Rank,
|
||||
result.Result.Score,
|
||||
result.Result.ChunkId,
|
||||
result.Result.ParentFileId,
|
||||
result.Result.FileName,
|
||||
FirstNonEmpty(result.Result.AbsolutePath, result.Result.FilePath),
|
||||
BuildChunkTitle(FirstNonEmpty(result.Result.FileName, dataSource.Name), result.Result.ChunkIndex, result.Result.PageNumber));
|
||||
}
|
||||
}
|
||||
|
||||
private void LogBm25Results(IInternalDataSource dataSource, IReadOnlyList<EmbeddingStateSearchResult> results)
|
||||
{
|
||||
if (results.Count == 0)
|
||||
{
|
||||
logger.LogInformation("Local RAG BM25 search found no chunks for data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id);
|
||||
return;
|
||||
}
|
||||
|
||||
foreach (var result in results.Select((result, index) => (Result: result, Rank: index + 1)))
|
||||
{
|
||||
logger.LogInformation(
|
||||
"Local RAG BM25 search found chunk for data source '{DataSourceName}' ({DataSourceId}). Rank={Rank}, Score={Score}, ChunkId='{ChunkId}', ParentFileId='{ParentFileId}', File='{FileName}', Path='{Path}', Title='{Title}'.",
|
||||
dataSource.Name,
|
||||
dataSource.Id,
|
||||
result.Rank,
|
||||
result.Result.Score,
|
||||
result.Result.ChunkId,
|
||||
result.Result.ParentFileId,
|
||||
result.Result.FileName,
|
||||
result.Result.AbsolutePath,
|
||||
BuildChunkTitle(FirstNonEmpty(result.Result.FileName, dataSource.Name), result.Result.ChunkIndex, result.Result.PageNumber));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Loading…
Reference in New Issue
Block a user