using System.Diagnostics;
using System.Text;
using System.Text.Json.Serialization;
using AIStudio.Provider;
using AIStudio.Settings;
using AIStudio.Tools.Databases.ProvenanceStore;
using AIStudio.Tools.PluginSystem;
using AIStudio.Tools.RAG;
using AIStudio.Tools.RAG.RAGProcesses;
using AIStudio.Tools.Rust;
using AIStudio.Tools.Security;
using AIStudio.Tools.Services;
using AIStudio.Tools.ToolCallingSystem;
using AIStudio.Tools.Web;
namespace AIStudio.Chat;
///
/// Text content in the chat.
///
public sealed class ContentText : IContent
{
private static readonly ILogger LOGGER = Program.LOGGER_FACTORY.CreateLogger();
private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(ContentText).Namespace, nameof(ContentText));
///
/// The minimum time between two streaming events, when the user
/// enables the energy saving mode.
///
private static readonly TimeSpan MIN_TIME = TimeSpan.FromSeconds(3);
#region Implementation of IContent
///
[JsonIgnore]
public bool InitialRemoteWait { get; set; }
///
[JsonIgnore]
public bool IsStreaming { get; set; }
///
[JsonIgnore]
public Func StreamingDone { get; set; } = IContent.NO_STREAMING_HANDLER;
///
[JsonIgnore]
public Func StreamingEvent { get; set; } = IContent.NO_STREAMING_HANDLER;
///
public List Sources { get; set; } = [];
///
public List FileAttachments { get; set; } = [];
public List ToolInvocations { get; set; } = [];
///
/// Optional structured data for a specialized visual presentation of this text.
///
///
/// The text remains the canonical content sent to providers, copied, and exported. A renderer
/// which understands the presentation may use its data instead, while every other consumer
/// safely continues to use the text.
///
public ContentPresentation? Presentation { get; set; }
///
/// The passages the data sources handed to the AI for this answer, as the AI received them.
///
///
/// Stored with the answer, so the user can still read what it rests on after the chat was
/// reloaded. Each passage belongs to one of the sources, see SourcePassageExtensions.GetPassagesOf.
/// Empty for every answer written before passages were kept, and for every answer no data
/// source contributed to.
///
public List SourcePassages { get; set; } = [];
///
/// Adds sources to this answer, together with the passages the AI received from them.
///
///
/// Both at once, because the title of a source names where its passages are, see
/// SourcePassageExtensions.NamePassagesInTitles. A passage added on its own would leave a title
/// behind which names only the pages known before it. Both lists are replaced as a whole rather
/// than changed, as in StoreReceivedContentAsync: the user interface may read them meanwhile.
///
/// The sources to add.
/// The passages to add.
public void AddSources(IEnumerable sources, IEnumerable passages)
{
var mergedSources = new List(this.Sources);
mergedSources.MergeSources(sources);
var mergedPassages = new List(this.SourcePassages);
mergedPassages.MergePassages(passages);
mergedSources.NamePassagesInTitles(mergedPassages);
// The passages first, so that every source the user interface finds has its passages:
this.SourcePassages = mergedPassages;
this.Sources = mergedSources;
}
///
/// Whether the user was told that a text stays in the chat file, since the provenance store could not take it.
///
///
/// Once per run of the app: a store which cannot be written stays so for a while, and a message
/// with every answer would tell the user nothing new.
///
private static int UNSTORED_CONTENT_REPORTED;
///
/// Moves what the AI received for this answer into the provenance store, so the chat keeps only references to it.
///
///
/// Called once the answer is complete, also when it was stopped or failed: until then, a passage
/// found twice is told apart by its text, see SourcePassageExtensions.MergePassages, and only the
/// passages which made it into the answer get stored. Nothing here may stop the answer from being
/// stored, so every failure ends in the passages keeping their texts.
///
/// The chat this answer belongs to.
public async Task StoreReceivedContentAsync(Guid chatId)
{
try
{
// Most answers, those of the agents among them, received nothing to store, and need no database:
if (!this.SourcePassages.Any(passage => passage.HoldsTextsToStore()))
return;
if (await ChatProvenance.GetStoreAsync() is { } store && await this.StoreReceivedContentAsync(store, chatId))
return;
}
catch (Exception exception)
{
LOGGER.LogError(exception, "Failed to store what the AI received for an answer of the chat '{ChatId}'. It stays in the chat file.", chatId);
}
if (Interlocked.Exchange(ref UNSTORED_CONTENT_REPORTED, 1) is 0)
await MessageBus.INSTANCE.SendWarning(new(Icons.Material.Filled.Storage, TB("AI Studio could not store what the AI received separately from your chat. It remains in the chat itself, so you can still read it.")));
}
///
/// Moves what the AI received for this answer into the given provenance store.
///
///
/// The list of passages is replaced as a whole rather than changed: the user interface may read
/// it meanwhile. A passage which the store could not take keeps its texts.
///
/// The provenance store.
/// The chat this answer belongs to.
/// True when everything was stored, false when something stays in the chat itself.
internal async Task StoreReceivedContentAsync(ProvenanceStoreClient store, Guid chatId)
{
var allStored = true;
var passages = new List(this.SourcePassages.Count);
foreach (var passage in this.SourcePassages)
{
if (!passage.HoldsTextsToStore())
{
passages.Add(passage);
continue;
}
Guid? provenanceId = null;
try
{
// Not cancellable: the answer is complete, also when the user stopped it.
provenanceId = await store.StoreAsync(chatId, passage.GetProvenanceKind(), passage.Text, passage.SurroundingContent, CancellationToken.None);
}
catch (Exception exception)
{
LOGGER.LogWarning(exception, "Failed to store a passage of the source '{SourceUrl}'. It stays in the chat file.", passage.SourceUrl);
}
if (provenanceId is null)
{
allStored = false;
passages.Add(passage);
continue;
}
passages.Add(passage with { ProvenanceId = provenanceId, Text = string.Empty, SurroundingContent = [] });
}
this.SourcePassages = passages;
return allStored;
}
///
/// What the provider said everything sent along with this answer cost, where it said anything.
///
///
/// Kept on the answer rather than beside the chat, so that it is stored, loaded, and exported
/// with the message it belongs to -- and so that it goes away when the message does. An edited
/// or regenerated answer removes its block, which removes these numbers with it. Null for every
/// answer written before this was recorded, and at every provider which reports nothing.
///
/// Having a report is not the same as the report still being true. Whether it still describes
/// the conversation is decided by ChatThread.ReportedHistoryFor, which looks at the thread around
/// the answer, not just at the answer.
///
public ReportedTokenUsage? ReportedUsage { get; set; }
[JsonIgnore]
public ToolRuntimeStatus ToolRuntimeStatus { get; set; } = new();
///
/// What the tool conversation of the running request adds to it, as far as it has got.
///
///
/// A model which calls tools asks several times before it answers, and every one of those
/// requests carries everything the tools returned so far -- up to three hundred thousand
/// characters of it. None of that is in this block's text, and none of it is in the traces
/// either: those say what happened, not what it costs. So it is kept here, where whoever
/// counts the conversation walks past anyway.
/// Replaced as a whole, never appended to: it is written by the thread which runs the tools
/// and read by the one which renders, and an exchange leaves the reader with a list which was
/// true at some moment rather than with one being rewritten under it.
/// Gone when the answer is there, and never persisted. The accumulated tool conversation lives
/// in the provider adapter, which is created for one request and dropped with it -- so the next
/// request does not carry it, and a number which still counted it would promise a cost nobody
/// is going to pay.
///
[JsonIgnore]
public IReadOnlyList PendingToolConversation { get; set; } = [];
///
/// Clears what the previous run of the tools left behind.
///
///
/// Both parts at once, because both belong to one request: the traces the user reads and the
/// payload the counting needs. They were cleared separately for exactly as long as there was
/// only one of them.
///
public void BeginToolRun()
{
this.ToolInvocations.Clear();
this.PendingToolConversation = [];
}
///
/// Says that no request is running anymore.
///
///
/// The traces stay -- they are what the user reads afterwards to see how the answer came
/// about. What goes is the payload, which belonged to a request that is over.
///
public void EndToolRun() => this.PendingToolConversation = [];
///
/// Keeps what the provider says the request behind this answer cost.
///
///
/// The one place where a stream's usage becomes part of the answer, for every path which writes
/// one. It arrives on one line of the stream, usually the last one, and only where the provider
/// reports it at all -- so a usage which states nothing leaves what was reported before alone.
///
/// What the current chunk of the stream says, which is mostly nothing.
/// The model the request went to.
/// How many blocks the conversation had when the request went out, this answer included.
public void RecordReportedUsage(TokenUsage usage, string modelId, int blockCount)
{
if (!usage.IsKnown)
return;
this.ReportedUsage = new()
{
PromptTokens = usage.PromptTokens,
ModelId = modelId,
BlockCount = blockCount,
};
}
///
/// Adds what one chunk of the stream brings to this answer: its text, its sources together with
/// their passages, and what the provider says the request cost.
///
///
/// The chat and the job service both write answers from a stream, and they used to do so step
/// by step each. A step one of them missed, such as the passages, would have been missing from
/// every answer written that way.
///
/// The chunk of the stream.
/// The model the request went to.
/// How many blocks the conversation had when the request went out, this answer included.
public void ApplyStreamChunk(ContentStreamChunk chunk, string modelId, int blockCount)
{
// The first chunk ends the waiting animation:
this.InitialRemoteWait = false;
this.IsStreaming = true;
this.Text += chunk;
if (chunk.Sources.Count > 0 || chunk.Passages.Count > 0)
this.AddSources(chunk.Sources, chunk.Passages);
this.RecordReportedUsage(chunk.Usage, modelId, blockCount);
}
///
public async Task CreateFromProviderAsync(IProvider provider, Model chatModel, IContent? lastUserPrompt, ChatThread? chatThread, CancellationToken token = default)
{
if(chatThread is null)
{
await this.CompleteWithoutStreaming();
return new();
}
if(!chatThread.IsLLMProviderAllowed(provider))
{
LOGGER.LogError("The provider is not allowed for this chat thread due to a prohibition of the organization, data security, or confidence-level requirements. Skipping the AI process.");
await this.CompleteWithoutStreaming();
return chatThread;
}
if(!await this.CheckSelectedModelAvailability(provider, chatModel, token))
{
await this.CompleteWithoutStreaming();
return chatThread;
}
// Call the RAG process. Right now, we only have one RAG process:
if (lastUserPrompt is not null)
{
try
{
var rag = new AISrcSelWithRetCtxVal();
chatThread = await rag.ProcessAsync(provider, chatModel, lastUserPrompt, chatThread, token);
}
catch (OperationCanceledException) when (token.IsCancellationRequested)
{
//
// The user canceled the request. That is not an error, and it must not reach the
// user as one. We do not rethrow here: the streaming task below observes the same
// token and ends the request itself, which keeps its finally block intact. That
// block is what tells the UI that the streaming is over.
//
LOGGER.LogInformation("The RAG process was canceled before the answer was requested.");
}
catch (Exception e)
{
LOGGER.LogError(e, "Skipping the RAG process due to an error.");
//
// The answer is about to be created without the data the user expected it to use.
// Without this message, that answer is indistinguishable from one that did use it:
//
await MessageBus.INSTANCE.SendWarning(new(Icons.Material.Filled.Source, TB("Your data sources could not be used. This answer was created without them.")));
}
}
// Store the last time we got a response. We use this later
// to determine whether we should notify the UI about the
// new content or not. Depends on the energy saving mode
// the user chose.
var last = DateTimeOffset.Now;
// Get the settings manager:
var settings = Program.SERVICE_PROVIDER.GetService()!;
// What the request carries, for telling later whether a report still describes this chat:
var blocksSent = chatThread.Blocks.Count;
//
// How long the answer takes, in three parts: the wait for its first chunk, which is the
// model reading the request and thinking; the streaming afterward; and how much of that
// the UI took to show it. Only these parts tell a slow model from a slow UI:
//
var requestStartedAt = Stopwatch.GetTimestamp();
TimeSpan? timeToFirstChunk = null;
var chunkCount = 0;
var uiDuration = TimeSpan.Zero;
async Task NotifyStreamingAsync()
{
var notifyingStartedAt = Stopwatch.GetTimestamp();
await this.StreamingEvent();
uiDuration += Stopwatch.GetElapsedTime(notifyingStartedAt);
}
// Start another thread by using a task to uncouple
// the UI thread from the AI processing:
try
{
await Task.Run(async () =>
{
try
{
// We show the waiting animation until we get the first response:
this.InitialRemoteWait = true;
// Iterate over the responses from the AI:
await foreach (var contentStreamChunk in provider.StreamChatCompletion(chatModel, chatThread, settings, token))
{
// When the user cancels the request, we stop the loop:
if (token.IsCancellationRequested)
break;
chunkCount++;
timeToFirstChunk ??= Stopwatch.GetElapsedTime(requestStartedAt);
// Add the response, its sources, and what it cost:
this.ApplyStreamChunk(contentStreamChunk, chatModel.Id, blocksSent);
// Notify the UI that the content has changed,
// depending on the energy saving mode:
var now = DateTimeOffset.Now;
switch (settings.ConfigurationData.App.IsSavingEnergy)
{
// Energy saving mode is off. We notify the UI
// as fast as possible -- no matter the odds:
case false:
await NotifyStreamingAsync();
break;
// Energy saving mode is on. We notify the UI
// only when the time between two events is
// greater than the minimum time:
case true when now - last > MIN_TIME:
last = now;
await NotifyStreamingAsync();
break;
}
}
}
finally
{
// Stop the waiting animation (in case the loop
// was stopped, or no content was received):
this.InitialRemoteWait = false;
this.IsStreaming = false;
}
}, token);
}
finally
{
var requestDuration = Stopwatch.GetElapsedTime(requestStartedAt);
if (timeToFirstChunk is { } firstChunkAfter)
LOGGER.LogInformation("The answer of model '{ModelId}' from '{ProviderInstanceName}' took {RequestSeconds:F1} s: {FirstChunkSeconds:F1} s until its first chunk, then {ChunkCount} chunks with {CharacterCount} characters, which the UI took {UISeconds:F1} s to show.", chatModel.Id, provider.InstanceName, requestDuration.TotalSeconds, firstChunkAfter.TotalSeconds, chunkCount, this.Text.Length, uiDuration.TotalSeconds);
else
LOGGER.LogInformation("The model '{ModelId}' from '{ProviderInstanceName}' sent no answer within {RequestSeconds:F1} s.", chatModel.Id, provider.InstanceName, requestDuration.TotalSeconds);
this.Text = this.Text.RemoveThinkTags().Trim();
this.EndToolRun();
// Before whoever is told next stores the chat, so that it refers to what the AI received instead of holding it:
await this.StoreReceivedContentAsync(chatThread.ChatId);
// Inform the UI that the streaming is done:
await this.StreamingDone();
}
return chatThread;
}
private async Task CompleteWithoutStreaming()
{
this.InitialRemoteWait = false;
this.IsStreaming = false;
await this.StreamingDone();
}
private static bool ModelsMatch(Model modelA, Model modelB)
{
var idA = modelA.Id.Trim();
var idB = modelB.Id.Trim();
return string.Equals(idA, idB, StringComparison.OrdinalIgnoreCase);
}
private async Task CheckSelectedModelAvailability(IProvider provider, Model chatModel, CancellationToken token = default)
{
if(chatModel.IsSystemModel)
return true;
if (string.IsNullOrWhiteSpace(chatModel.Id))
{
LOGGER.LogWarning("Skipping AI request because model ID is null or white space.");
return false;
}
if (!provider.HasModelLoadingCapability)
return true;
IReadOnlyList loadedModels;
try
{
var modelLoadResult = await provider.GetTextModels(token: token);
if (!modelLoadResult.Success)
{
var userMessage = modelLoadResult.FailureReason.ToUserMessage(provider.InstanceName);
if (!string.IsNullOrWhiteSpace(userMessage))
await MessageBus.INSTANCE.SendError(new(Icons.Material.Filled.CloudOff, userMessage));
LOGGER.LogWarning("Skipping selected model availability check for '{ProviderInstanceName}' (provider={ProviderType}) because loading the model list failed with reason {FailureReason}.", provider.InstanceName, provider.Provider, modelLoadResult.FailureReason);
return false;
}
loadedModels = modelLoadResult.Models;
}
catch (OperationCanceledException)
{
return false;
}
catch (Exception e)
{
LOGGER.LogWarning(e, "Skipping selected model availability check for '{ProviderInstanceName}' (provider={ProviderType}) because the model list could not be loaded.", provider.InstanceName, provider.Provider);
return true;
}
var availableModels = loadedModels.Where(model => !string.IsNullOrWhiteSpace(model.Id)).ToList();
if (availableModels.Count == 0)
{
var emptyModelsMessage = string.Format(
TB("We could load models from '{0}', but the provider did not return any usable text models."),
provider.InstanceName);
await MessageBus.INSTANCE.SendError(new(Icons.Material.Filled.CloudOff, emptyModelsMessage));
LOGGER.LogWarning("Skipping AI request because there are no models available from '{ProviderInstanceName}' (provider={ProviderType}).", provider.InstanceName, provider.Provider);
return false;
}
if(availableModels.Any(model => ModelsMatch(model, chatModel)))
return true;
var message = string.Format(
TB("The selected model '{0}' is no longer available from '{1}' (provider={2}). Please adapt your provider settings."),
chatModel.Id,
provider.InstanceName,
provider.Provider);
await MessageBus.INSTANCE.SendError(new(Icons.Material.Filled.CloudOff, message));
LOGGER.LogWarning("Skipping AI request because model '{ModelId}' is not available from '{ProviderInstanceName}' (provider={ProviderType}).", chatModel.Id, provider.InstanceName, provider.Provider);
return false;
}
///
///
/// The reported usage stays behind on purpose. A clone continues somewhere else -- as the
/// example conversation of a chat template, or as an assistant's conversation carried over into
/// a chat -- with another system prompt around it, and what the provider stated was for the
/// request this answer came out of.
///
public IContent DeepClone() => new ContentText
{
Text = this.Text,
InitialRemoteWait = this.InitialRemoteWait,
IsStreaming = this.IsStreaming,
Sources = [..this.Sources],
SourcePassages = [..this.SourcePassages],
FileAttachments = [..this.FileAttachments],
Presentation = this.Presentation?.DeepClone(),
ToolInvocations = [..this.ToolInvocations.Select(x => new ToolInvocationTrace
{
Order = x.Order,
ToolId = x.ToolId,
ToolName = x.ToolName,
ToolIcon = x.ToolIcon,
ToolCallId = x.ToolCallId,
Status = x.Status,
WasExecuted = x.WasExecuted,
StatusMessage = x.StatusMessage,
Arguments = new Dictionary(x.Arguments, StringComparer.Ordinal),
Result = x.Result,
JsonResult = x.JsonResult?.DeepClone(),
})],
};
#endregion
///
/// The web addresses in the documents attached to this message, as request keys, see
/// WebAddresses.CreateRequestKey.
///
///
/// The texts of the documents are known only when a message is sent, from the provenance store
/// or from the disk, so their addresses are collected right then, see PrepareTextContentForAI,
/// and count as given to the model, see ChatThread.IsWebAddressGivenToTheModel. Not stored:
/// every request collects them anew from what it sends.
///
[JsonIgnore]
public IReadOnlySet RuntimeAttachmentWebAddresses { get; internal set; } = new HashSet(StringComparer.Ordinal);
///
/// Whether this message holds attached documents which were never stored, see StoreAttachedDocumentsAsync.
///
/// Whether documents whose files changed since they were stored count as well, see DataChat.ReadChangedAttachmentsAgain.
internal bool HasDocumentsToStore(bool readChangedAgain) => this.FileAttachments.Any(attachment => IsToStore(attachment, readChangedAgain));
///
/// Whether an attachment is a document to read and store: one which was never stored, or, when
/// the user chose so, one whose file changed since.
///
private static bool IsToStore(FileAttachment attachment, bool readChangedAgain) => attachment.Type is FileAttachmentType.DOCUMENT && (attachment.ProvenanceId is null || readChangedAgain && attachment.ChangedSinceStored());
///
/// Stores the attached documents of this message which were never sent, reading each one from the disk once.
///
///
/// A document which cannot be read or stored right now keeps no entry: the message reads it when
/// it is built, as every message did before, and tells the user then what went wrong. The list of
/// attachments is replaced as a whole rather than changed, since the user interface may read it
/// meanwhile.
///
/// The provenance store.
/// The chat of this message.
/// The runtime, which reads the documents.
/// The request the documents are read for, which tells the user once about everything filtered out of them.
/// Whether documents whose files changed since they were stored are read again, see DataChat.ReadChangedAttachmentsAgain.
internal async Task StoreAttachedDocumentsAsync(ProvenanceStoreClient store, Guid chatId, RustService rustService, PromptInjectionReportingScope? reportingScope, bool readChangedAgain)
{
var attachments = new List(this.FileAttachments.Count);
var changed = false;
foreach (var attachment in this.FileAttachments)
{
var storedAttachment = await StoreAttachedDocumentAsync(attachment, store, chatId, rustService, reportingScope, readChangedAgain);
changed |= !ReferenceEquals(storedAttachment, attachment);
attachments.Add(storedAttachment);
}
if (changed)
this.FileAttachments = attachments;
}
///
/// Stores one attached document which was never sent, or whose file changed since, when the user chose so.
///
///
/// A document read again gets an entry of its own: the one before belongs to what the AI
/// received then, and the cleanup at the next start removes it once nothing refers to it.
///
/// The attachment with the entry of its text, or the attachment as it was when there is nothing to store, or nothing could be.
private static async Task StoreAttachedDocumentAsync(FileAttachment attachment, ProvenanceStoreClient store, Guid chatId, RustService rustService, PromptInjectionReportingScope? reportingScope, bool readChangedAgain)
{
if (!IsToStore(attachment, readChangedAgain))
return attachment;
var document = attachment.Normalize();
if (document.Type is not FileAttachmentType.DOCUMENT || !document.Exists)
return attachment;
try
{
// Quietly: building the message tells the user about a missing Pandoc, when it reads the document after all.
if (FileTypes.RequiresPandoc(document.FilePath))
{
var pandocState = await Pandoc.CheckAvailabilityAsync(rustService, showMessages: false, showSuccessMessage: false);
if (pandocState is not { IsAvailable: true, CheckWasSuccessful: true })
return attachment;
}
// Taken before reading, so that a change while reading counts as a change:
var version = FileVersion.Of(document.FilePath);
var extraction = await rustService.ReadArbitraryFileData(document.FilePath, int.MaxValue, reportingScope: reportingScope);
if (!extraction.HasUsableContent)
return attachment;
// Told here, since no later request reads the document again:
await ReportIncompleteReadingAsync(document, extraction);
// Not cancellable: what was read is what the request sends.
var provenanceId = await store.StoreAsync(chatId, ProvenanceKind.ATTACHMENT, extraction.Content, null, CancellationToken.None);
return provenanceId is null ? attachment : attachment with { ProvenanceId = provenanceId, Version = version };
}
catch (Exception exception)
{
LOGGER.LogWarning(exception, "Failed to store the file attachment '{FilePath}'. It is read from the disk with every request.", document.FilePath);
return attachment;
}
}
///
/// Tells the user when a document could be read, but not as it should have been.
///
/// The document.
/// What reading it gave.
private static async Task ReportIncompleteReadingAsync(FileAttachment document, FileExtractionResult extraction)
{
//
// The file is usable, but we lost parts of it. The user has to know which
// parts are missing, because the answer will be based on the rest.
//
if (extraction.Outcome is FileExtractionOutcome.PARTIAL)
{
LOGGER.LogWarning("Parts of the file attachment '{FilePath}' could not be read: pages={FailedPages}.", document.FilePath, string.Join(", ", extraction.FailedPages));
await MessageBus.INSTANCE.SendWarning(new(Icons.Material.Filled.Description, extraction.ToPartialUserMessage(document.FileName)));
}
// The file was read correctly, but its extension lies about what it contains:
if (extraction.HasExtensionMismatch)
{
LOGGER.LogWarning("The file attachment '{FilePath}' is actually a '{DetectedFormat}'.", document.FilePath, extraction.DetectedFormat);
await MessageBus.INSTANCE.SendWarning(new(Icons.Material.Filled.RuleFolder, extraction.ToExtensionMismatchUserMessage(document.FileName)));
}
}
///
/// Reads the stored texts of the attached documents.
///
/// The attachments.
/// Gets the provenance store, or null when there is none. It is asked only when a document has an entry.
/// The texts by the ids of their entries. An entry the store lacks, or cannot read, is missing.
private static async Task> ReadStoredDocumentTextsAsync(IEnumerable attachments, Func> getProvenanceStore)
{
var provenanceIds = attachments.Select(attachment => attachment.ProvenanceId).OfType().Distinct().ToList();
if (provenanceIds.Count is 0)
return new Dictionary();
try
{
if (await getProvenanceStore() is not { } store)
return new Dictionary();
var entries = await store.GetEntriesAsync(provenanceIds, CancellationToken.None);
return entries.ToDictionary(entry => entry.Key, entry => entry.Value.Content);
}
catch (Exception exception)
{
LOGGER.LogWarning(exception, "Failed to read the stored texts of the file attachments. Their files are read instead.");
return new Dictionary();
}
}
///
/// Filters the stored texts of the attached documents for prompt injections once more.
///
///
/// The texts were filtered when they were read from their files. Filtering them again with every
/// request applies the rules of today, as when every request read the files anew. All documents of
/// the message go in one request to the runtime, and the user hears about them with the request.
///
/// The stored texts by the ids of their entries.
/// The attachments, which name the files of the texts.
/// Gets the prompt injection filter.
/// The request the texts are sent with.
/// The filtered texts by the ids of their entries.
private static async Task> FilterStoredTextsAsync(IReadOnlyDictionary storedTexts, IReadOnlyList attachments, Func getPromptInjectionGuard, PromptInjectionReportingScope? reportingScope)
{
if (storedTexts.Count is 0)
return storedTexts;
var filePaths = attachments
.Where(attachment => attachment.ProvenanceId is not null)
.GroupBy(attachment => attachment.ProvenanceId!.Value)
.ToDictionary(group => group.Key, group => group.First().FilePath);
var ids = storedTexts.Keys.ToList();
var filteredTexts = await getPromptInjectionGuard().SanitizeAsync(ids.Select(id => new PromptInjectionText(storedTexts[id], PromptInjectionSource.ChatAttachment(filePaths[id]))).ToList(), reportingScope);
return ids.Zip(filteredTexts).ToDictionary(pair => pair.First, pair => pair.Second);
}
///
/// Builds the text the AI gets for this message, with the content of its attached documents.
///
/// The request this message is part of, which tells the user once about everything filtered out of its attachments, or null when the request reads no documents.
/// The text for the AI.
public Task PrepareTextContentForAI(PromptInjectionReportingScope? reportingScope) => this.PrepareTextContentForAI(reportingScope, Program.RUST_SERVICE, ChatProvenance.GetStoreAsync, () => Program.SERVICE_PROVIDER.GetRequiredService());
///
/// Builds the text the AI gets for this message, with the content of its attached documents.
///
///
/// A document stored at its first send comes from the provenance store, also when its file
/// changed or is gone by now, see FileAttachment.ProvenanceId: the AI gets the same text with
/// every request. Only a document without a stored text is read from the disk, as before. A stored
/// text passes the prompt injection filter once more, see FilterStoredTextsAsync.
///
/// The request this message is part of, which tells the user once about everything filtered out of its attachments, or null when the request reads no documents.
/// The runtime, which reads the documents without a stored text.
/// Gets the provenance store, or null when there is none. It is asked only when a document has an entry.
/// Gets the prompt injection filter, which is asked only when a document has a stored text.
/// The text for the AI.
internal async Task PrepareTextContentForAI(PromptInjectionReportingScope? reportingScope, RustService rustService, Func> getProvenanceStore, Func getPromptInjectionGuard)
{
var sb = new StringBuilder();
var attachmentWebAddresses = new HashSet(StringComparer.Ordinal);
sb.AppendLine(this.Text);
if(this.FileAttachments.Count > 0)
{
var normalizedAttachments = this.FileAttachments
.Select(attachment => attachment.Normalize())
.ToList();
var storedTexts = await FilterStoredTextsAsync(await ReadStoredDocumentTextsAsync(normalizedAttachments, getProvenanceStore), normalizedAttachments, getPromptInjectionGuard, reportingScope);
// Get the list of existing documents, those with a stored text included, whether their file is still there or not:
var existingDocuments = normalizedAttachments.Where(x => x.Type is FileAttachmentType.DOCUMENT && (StoredTextOf(x) is not null || x.Exists)).ToList();
//
// Report missing files. We tell the user about them instead of only logging: on a
// network drive, a file which is temporarily unreachable looks exactly like a deleted
// one, and silently dropping it would let the AI answer without that document.
//
var missingDocuments = normalizedAttachments.Except(existingDocuments).Where(x => x.Type is FileAttachmentType.DOCUMENT).ToList();
foreach (var missingDocument in missingDocuments)
{
LOGGER.LogWarning("File attachment no longer exists and will be skipped: '{MissingDocument}'.", missingDocument.FilePath);
await MessageBus.INSTANCE.SendError(new(Icons.Material.Filled.FindInPage, string.Format(TB("The file '{0}' is currently not available and was not sent."), missingDocument.FileName)));
}
// Only proceed if there are existing, allowed documents:
if (existingDocuments.Count > 0)
{
//
// Pandoc is only needed for the few formats we convert with it. PDFs, text files,
// spreadsheets, and presentations are read by the runtime itself, so a missing
// Pandoc installation must not stop them.
//
var pandocIsUsable = true;
if (existingDocuments.Any(document => StoredTextOf(document) is null && FileTypes.RequiresPandoc(document.FilePath)))
{
var pandocState = await Pandoc.CheckAvailabilityAsync(rustService, showMessages: true, showSuccessMessage: false);
pandocIsUsable = pandocState is { IsAvailable: true, CheckWasSuccessful: true };
if (!pandocState.IsAvailable)
LOGGER.LogWarning("File attachments which need Pandoc could not be processed because Pandoc is not available.");
else if (!pandocState.CheckWasSuccessful)
LOGGER.LogWarning("File attachments which need Pandoc could not be processed because the Pandoc version check failed.");
}
//
// The document blocks are collected separately, so we only announce attached
// files when at least one of them could actually be read. Announcing files we
// then hand over as empty blocks makes the AI answer about an empty document.
//
var documentBlocks = new StringBuilder();
foreach(var document in existingDocuments)
{
var content = StoredTextOf(document);
if (content is null)
{
if (document.IsForbidden)
{
LOGGER.LogWarning("File attachment '{FilePath}' has a forbidden file type and will be skipped.", document.FilePath);
continue;
}
if (!pandocIsUsable && FileTypes.RequiresPandoc(document.FilePath))
{
LOGGER.LogWarning("The file attachment '{FilePath}' needs Pandoc and will be skipped.", document.FilePath);
await MessageBus.INSTANCE.SendError(new(Icons.Material.Filled.Description, FileExtractionErrorCode.PANDOC_UNAVAILABLE.ToUserMessage(document.FileName)));
continue;
}
var extraction = await rustService.ReadArbitraryFileData(document.FilePath, int.MaxValue, reportingScope: reportingScope);
if (!extraction.HasUsableContent)
{
LOGGER.LogError("Reading the file attachment '{FilePath}' failed and it will not be sent: code={ErrorCode}, message='{ErrorMessage}'.", document.FilePath, extraction.ErrorCode, extraction.ErrorMessage);
await MessageBus.INSTANCE.SendError(new(Icons.Material.Filled.Description, extraction.ToUserMessage(document.FileName)));
continue;
}
await ReportIncompleteReadingAsync(document, extraction);
content = extraction.Content;
}
documentBlocks.AppendLine();
documentBlocks.AppendLine("---------------------------------------");
documentBlocks.AppendLine($"File path: {document.FilePath}");
documentBlocks.AppendLine("File content:");
documentBlocks.AppendLine("````");
documentBlocks.AppendLine(content);
documentBlocks.AppendLine("````");
foreach (var address in WebAddresses.Find(content))
if (WebAddresses.TryCreateRequestKey(address, out var requestKey))
attachmentWebAddresses.Add(requestKey);
}
if (documentBlocks.Length > 0)
{
sb.AppendLine();
sb.AppendLine("The following files are attached to this message:");
sb.Append(documentBlocks);
}
var numImages = normalizedAttachments.Count(x => x is { IsImage: true, Exists: true });
if (numImages > 0)
{
sb.AppendLine();
sb.AppendLine($"Additionally, there are {numImages} image file(s) attached to this message. ");
sb.AppendLine("Please consider them as part of the message content and use them to answer accordingly.");
}
}
string? StoredTextOf(FileAttachment document) => document.ProvenanceId is { } provenanceId ? storedTexts.GetValueOrDefault(provenanceId) : null;
}
this.RuntimeAttachmentWebAddresses = attachmentWebAddresses;
return sb.ToString();
}
///
/// The text content.
///
public string Text { get; set; } = string.Empty;
}