mirror of
https://github.com/MindWorkAI/AI-Studio.git
synced 2026-11-01 20:07:21 +00:00
296 lines
12 KiB
C#
296 lines
12 KiB
C#
using System.Text;
|
|
using System.Text.Json;
|
|
|
|
using AIStudio.Chat;
|
|
using AIStudio.Provider;
|
|
using AIStudio.Settings;
|
|
using AIStudio.Settings.DataModel;
|
|
using AIStudio.Tools.Services;
|
|
|
|
namespace AIStudio.Agents;
|
|
|
|
public sealed class AgentDataSourceSelection (ILogger<AgentDataSourceSelection> logger, ILogger<AgentBase> baseLogger, SettingsManager settingsManager, DataSourceService dataSourceService, DataSourceDescriptionService descriptionService, ThreadSafeRandom rng) : AgentBase(baseLogger, settingsManager, dataSourceService, rng)
|
|
{
|
|
private readonly List<ContentBlock> answers = new();
|
|
|
|
#region Overrides of AgentBase
|
|
|
|
/// <inheritdoc />
|
|
protected override Type Type => Type.SYSTEM;
|
|
|
|
/// <inheritdoc />
|
|
protected override Tools.Components Component => Tools.Components.AGENT_DATA_SOURCE_SELECTION;
|
|
|
|
/// <inheritdoc />
|
|
public override string Id => "Data Source Selection";
|
|
|
|
/// <inheritdoc />
|
|
protected override string JobDescription =>
|
|
"""
|
|
You receive a system and a user prompt, as well as a list of possible data sources as input.
|
|
Your task is to select the appropriate data sources for the given task. You may choose none,
|
|
one, or multiple sources, depending on what best fits the system and user prompt. You need
|
|
to estimate and assess which source, based on its description, might be helpful in
|
|
processing the prompts.
|
|
|
|
Your response is a JSON list in the following format:
|
|
|
|
```
|
|
[
|
|
{"id": "The data source ID", "reason": "Why did you choose this source?", "confidence": 0.87},
|
|
{"id": "The data source ID", "reason": "Why did you choose this source?", "confidence": 0.54}
|
|
]
|
|
```
|
|
|
|
You express your confidence as a floating-point number between 0.0 (maximum uncertainty) and
|
|
1.0 (you are absolutely certain that this source is needed).
|
|
|
|
The JSON schema is:
|
|
|
|
```
|
|
{
|
|
"$schema": "http://json-schema.org/draft-04/schema#",
|
|
"type": "array",
|
|
"items": [
|
|
{
|
|
"type": "object",
|
|
"properties": {
|
|
"id": {
|
|
"type": "string"
|
|
},
|
|
"reason": {
|
|
"type": "string"
|
|
},
|
|
"confidence": {
|
|
"type": "number"
|
|
}
|
|
},
|
|
"required": [
|
|
"id",
|
|
"reason",
|
|
"confidence"
|
|
]
|
|
}
|
|
]
|
|
}
|
|
```
|
|
|
|
When no data source is needed, you return an empty JSON list `[]`. You do not ask any
|
|
follow-up questions. You do not address the user. Your response consists solely of
|
|
the JSON list.
|
|
""";
|
|
|
|
/// <inheritdoc />
|
|
protected override string SystemPrompt(string availableDataSources) => $"""
|
|
{this.JobDescription}
|
|
|
|
{availableDataSources}
|
|
""";
|
|
|
|
/// <inheritdoc />
|
|
public override Settings.Provider ProviderSettings { get; set; } = Settings.Provider.NONE;
|
|
|
|
/// <summary>
|
|
/// The data source selection agent does not work with context. Use
|
|
/// the process input method instead.
|
|
/// </summary>
|
|
/// <returns>The chat thread without any changes.</returns>
|
|
public override Task<ChatThread> ProcessContext(ChatThread chatThread, IDictionary<string, string> additionalData) => Task.FromResult(chatThread);
|
|
|
|
/// <inheritdoc />
|
|
public override async Task<ContentBlock> ProcessInput(ContentBlock input, IDictionary<string, string> additionalData)
|
|
{
|
|
if (input.Content is not ContentText text)
|
|
return EMPTY_BLOCK;
|
|
|
|
if(text.InitialRemoteWait || text.IsStreaming)
|
|
return EMPTY_BLOCK;
|
|
|
|
if(string.IsNullOrWhiteSpace(text.Text))
|
|
return EMPTY_BLOCK;
|
|
|
|
if(!additionalData.TryGetValue("availableDataSources", out var availableDataSources) || string.IsNullOrWhiteSpace(availableDataSources))
|
|
return EMPTY_BLOCK;
|
|
|
|
var thread = this.CreateChatThread(this.SystemPrompt(availableDataSources));
|
|
var userRequest = AddUserRequest(thread, text.Text);
|
|
await this.AddAIResponseAsync(thread, userRequest.UserPrompt, userRequest.Time);
|
|
|
|
var answer = thread.Blocks[^1];
|
|
|
|
this.answers.Add(answer);
|
|
return answer;
|
|
}
|
|
|
|
// <inheritdoc />
|
|
public override Task<bool> MadeDecision(ContentBlock input) => Task.FromResult(true);
|
|
|
|
// <inheritdoc />
|
|
public override IReadOnlyCollection<ContentBlock> GetContext() => [];
|
|
|
|
// <inheritdoc />
|
|
public override IReadOnlyCollection<ContentBlock> GetAnswers() => this.answers;
|
|
|
|
#endregion
|
|
|
|
public async Task<List<SelectedDataSource>> PerformSelectionAsync(IProvider provider, IContent lastUserPrompt, ChatThread chatThread, AllowedSelectedDataSources dataSources, CancellationToken token = default)
|
|
{
|
|
logger.LogInformation("The AI should select the appropriate data sources.");
|
|
|
|
//
|
|
// 1. Which LLM provider should the agent use?
|
|
//
|
|
|
|
// We start with the provider currently selected by the user:
|
|
var requiredDataSecurity = dataSources.AllowedDataSources.GetRequiredSecurityPolicy();
|
|
var requiredConfidenceLevel = dataSources.AllowedDataSources.GetRequiredConfidenceLevel();
|
|
var agentProvider = this.SettingsManager.GetPreselectedProvider(Tools.Components.AGENT_DATA_SOURCE_SELECTION, provider.ConfiguredProviderId, true);
|
|
if (agentProvider == Settings.Provider.NONE)
|
|
{
|
|
logger.LogWarning("No provider is selected for the agent. The agent cannot select data sources.");
|
|
return [];
|
|
}
|
|
|
|
if (!agentProvider.AllowsDataSourceAccess(this.SettingsManager, requiredDataSecurity, requiredConfidenceLevel))
|
|
{
|
|
logger.LogWarning($"The agent for data source selection uses provider '{agentProvider.InstanceName}' with confidence '{agentProvider.GetConfidenceLevel(this.SettingsManager).GetName()}', but the available data sources require data security '{requiredDataSecurity}' and provider confidence '{requiredConfidenceLevel.GetName()}'. The agent cannot select data sources.");
|
|
return [];
|
|
}
|
|
|
|
// Assign the provider settings to the agent:
|
|
logger.LogInformation($"The agent for the data source selection uses the provider '{agentProvider.InstanceName}' ({agentProvider.UsedLLMProvider.ToName()}, confidence={agentProvider.UsedLLMProvider.GetConfidence(this.SettingsManager).Level.GetName()}).");
|
|
this.ProviderSettings = agentProvider;
|
|
|
|
//
|
|
// 2. Prepare the current system and user prompts as input for the agent:
|
|
//
|
|
var lastPromptContent = lastUserPrompt switch
|
|
{
|
|
ContentText text => text.Text,
|
|
|
|
// Image prompts may be empty, e.g., when the image is too large:
|
|
ContentImage image => await image.TryAsBase64(token) is (success: true, { } base64Image)
|
|
? base64Image
|
|
: string.Empty,
|
|
|
|
// Other content types are not supported yet:
|
|
_ => string.Empty,
|
|
};
|
|
|
|
if (string.IsNullOrWhiteSpace(lastPromptContent))
|
|
{
|
|
logger.LogWarning("The last prompt is empty. The AI cannot select data sources.");
|
|
return [];
|
|
}
|
|
|
|
//
|
|
// 3. Prepare the allowed data sources as input for the agent:
|
|
//
|
|
var additionalData = new Dictionary<string, string>();
|
|
logger.LogInformation("Preparing the list of allowed data sources for the agent to choose from.");
|
|
|
|
var sb = new StringBuilder();
|
|
sb.AppendLine("The following data sources are available for selection:");
|
|
foreach (var ds in dataSources.AllowedDataSources)
|
|
{
|
|
var description = await descriptionService.GetDescriptionAsync(ds, token);
|
|
var descriptionPart = string.IsNullOrWhiteSpace(description) ? string.Empty : $", description='{description}'";
|
|
switch (ds)
|
|
{
|
|
case DataSourceLocalDirectory localDirectory:
|
|
sb.AppendLine($"- Id={ds.Id}, name='{localDirectory.Name}', type=local directory, path='{localDirectory.Path}'{descriptionPart}");
|
|
break;
|
|
|
|
case DataSourceLocalFile localFile:
|
|
sb.AppendLine($"- Id={ds.Id}, name='{localFile.Name}', type=local file, path='{localFile.FilePath}'{descriptionPart}");
|
|
break;
|
|
|
|
case IERIDataSource eriDataSource:
|
|
sb.AppendLine($"- Id={ds.Id}, name='{eriDataSource.Name}', type=external data source{descriptionPart}");
|
|
break;
|
|
}
|
|
}
|
|
|
|
logger.LogInformation("Prepared the list of allowed data sources for the agent.");
|
|
additionalData.Add("availableDataSources", sb.ToString());
|
|
|
|
//
|
|
// 4. Let the agent select the data sources:
|
|
//
|
|
var prompt = $"""
|
|
The system prompt is:
|
|
|
|
```
|
|
{chatThread.SystemPrompt}
|
|
```
|
|
|
|
The user prompt is:
|
|
|
|
```
|
|
{lastPromptContent}
|
|
```
|
|
""";
|
|
|
|
// Call the agent:
|
|
var aiResponse = await this.ProcessInput(new ContentBlock
|
|
{
|
|
Time = DateTimeOffset.UtcNow,
|
|
ContentType = ContentType.TEXT,
|
|
Role = ChatRole.USER,
|
|
Content = new ContentText
|
|
{
|
|
Text = prompt,
|
|
},
|
|
}, additionalData);
|
|
|
|
if(aiResponse.Content is null)
|
|
{
|
|
logger.LogWarning("The agent did not return a response.");
|
|
return [];
|
|
}
|
|
|
|
switch (aiResponse)
|
|
{
|
|
|
|
//
|
|
// 5. Parse the agent response:
|
|
//
|
|
case { ContentType: ContentType.TEXT, Content: ContentText textContent }:
|
|
{
|
|
//
|
|
// What we expect is a JSON list of SelectedDataSource objects:
|
|
//
|
|
var selectedDataSourcesJson = textContent.Text;
|
|
|
|
//
|
|
// We know how bad LLM may be in generating JSON without surrounding text.
|
|
// Thus, we expect the worst and try to extract the JSON list from the text:
|
|
//
|
|
var json = ModelAnswerJson.ExtractObjectList(selectedDataSourcesJson);
|
|
|
|
try
|
|
{
|
|
var aiSelectedDataSources = JsonSerializer.Deserialize<List<SelectedDataSource>>(json, JSON_SERIALIZER_OPTIONS);
|
|
return aiSelectedDataSources ?? [];
|
|
}
|
|
catch
|
|
{
|
|
logger.LogWarning("The agent answered with an invalid or unexpected JSON format.");
|
|
return [];
|
|
}
|
|
}
|
|
|
|
case { ContentType: ContentType.TEXT }:
|
|
logger.LogWarning("The agent answered with an unexpected inner content type.");
|
|
return [];
|
|
|
|
case { ContentType: ContentType.NONE }:
|
|
logger.LogWarning("The agent did not return a response.");
|
|
return [];
|
|
|
|
default:
|
|
logger.LogWarning($"The agent answered with an unexpected content type '{aiResponse.ContentType}'.");
|
|
return [];
|
|
}
|
|
}
|
|
} |