Added an agent to select the appropriate data sources for any prompt (#284)

app/MindWork AI Studio/Agents/AgentDataSourceSelection.cs

@@ -0,0 +1,416 @@
+using System.Text;
+using System.Text.Json;
+
+using AIStudio.Chat;
+using AIStudio.Provider;
+using AIStudio.Settings;
+using AIStudio.Settings.DataModel;
+using AIStudio.Tools.ERIClient;
+using AIStudio.Tools.Services;
+
+namespace AIStudio.Agents;
+
+public sealed class AgentDataSourceSelection (ILogger<AgentDataSourceSelection> logger, ILogger<AgentBase> baseLogger, SettingsManager settingsManager, DataSourceService dataSourceService, ThreadSafeRandom rng) : AgentBase(baseLogger, settingsManager, dataSourceService, rng)
+{
+    private static readonly ContentBlock EMPTY_BLOCK = new()
+    {
+        Content = null,
+        ContentType = ContentType.NONE,
+        Role = ChatRole.AGENT,
+        Time = DateTimeOffset.UtcNow,
+    };
+    
+    private readonly List<ContentBlock> answers = new();
+    
+    #region Overrides of AgentBase
+
+    /// <inheritdoc />
+    protected override Type Type => Type.SYSTEM;
+    
+    /// <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; }
+    
+    /// <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 time = this.AddUserRequest(thread, text.Text);
+        await this.AddAIResponseAsync(thread, 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 lastPrompt, 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 agentProvider = this.SettingsManager.ConfigurationData.Providers.FirstOrDefault(x => x.Id == provider.Id);
+
+        // If the user preselected an agent provider, we try to use this one:
+        if (this.SettingsManager.ConfigurationData.AgentDataSourceSelection.PreselectAgentOptions)
+        {
+            var configuredAgentProvider = this.SettingsManager.ConfigurationData.Providers.FirstOrDefault(x => x.Id == this.SettingsManager.ConfigurationData.AgentDataSourceSelection.PreselectedAgentProvider);
+
+            // If the configured agent provider is available, we use it:
+            if (configuredAgentProvider != default)
+                agentProvider = configuredAgentProvider;
+        }
+
+        // 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 = lastPrompt switch
+        {
+            ContentText text => text.Text,
+
+            // Image prompts may be empty, e.g., when the image is too large:
+            ContentImage image => await image.AsBase64(token),
+
+            // 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.");
+        
+        // Notice: We do not dispose the Rust service here. The Rust service is a singleton
+        // and will be disposed when the application shuts down:
+        var rustService = Program.SERVICE_PROVIDER.GetService<RustService>()!;
+        
+        var sb = new StringBuilder();
+        sb.AppendLine("The following data sources are available for selection:");
+        foreach (var ds in dataSources.AllowedDataSources)
+        {
+            switch (ds)
+            {
+                case DataSourceLocalDirectory localDirectory:
+                    sb.AppendLine($"- Id={ds.Id}, name='{localDirectory.Name}', type=local directory, path='{localDirectory.Path}'");
+                    break;
+
+                case DataSourceLocalFile localFile:
+                    sb.AppendLine($"- Id={ds.Id}, name='{localFile.Name}', type=local file, path='{localFile.FilePath}'");
+                    break;
+
+                case IERIDataSource eriDataSource:
+                    var eriServerDescription = string.Empty;
+
+                    try
+                    {
+                        //
+                        // Call the ERI server to get the server description:
+                        //
+                        using var eriClient = ERIClientFactory.Get(eriDataSource.Version, eriDataSource)!;
+                        var authResponse = await eriClient.AuthenticateAsync(eriDataSource, rustService, token);
+                        if (authResponse.Successful)
+                        {
+                            var serverDescriptionResponse = await eriClient.GetDataSourceInfoAsync(token);
+                            if (serverDescriptionResponse.Successful)
+                            {
+                                eriServerDescription = serverDescriptionResponse.Data.Description;
+
+                                // Remove all line breaks from the description:
+                                eriServerDescription = eriServerDescription.Replace("\n", " ").Replace("\r", " ");
+                            }
+                            else
+                                logger.LogWarning($"Was not able to retrieve the server description from the ERI data source '{eriDataSource.Name}'. Message: {serverDescriptionResponse.Message}");
+                        }
+                        else
+                            logger.LogWarning($"Was not able to authenticate with the ERI data source '{eriDataSource.Name}'. Message: {authResponse.Message}");
+                    }
+                    catch (Exception e)
+                    {
+                        logger.LogWarning($"The ERI data source '{eriDataSource.Name}' is not available. Thus, we cannot retrieve the server description. Error: {e.Message}");
+                    }
+
+                    //
+                    // Append the ERI data source to the list. Use the server description if available:
+                    //
+                    if (string.IsNullOrWhiteSpace(eriServerDescription))
+                        sb.AppendLine($"- Id={ds.Id}, name='{eriDataSource.Name}', type=external data source");
+                    else
+                        sb.AppendLine($"- Id={ds.Id}, name='{eriDataSource.Name}', type=external data source, description='{eriServerDescription}'");
+
+                    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 = this.ExtractJson(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 [];
+        }
+    }
+    
+    /// <summary>
+    /// Extracts the JSON list from the given text. The text may contain additional
+    /// information around the JSON list. The method tries to extract the JSON list
+    /// from the text.
+    /// </summary>
+    /// <remarks>
+    /// Algorithm: The method searches for the first line that contains only a '[' character.
+    /// Then, it searches for the first line that contains only a ']' character. The method
+    /// returns the text between these two lines (including the brackets). When the method
+    /// cannot find the JSON list, it returns an empty string.
+    /// <br/><br/>
+    /// This overload is using strings instead of spans. We can use this overload in any
+    /// async method. Thus, it is a wrapper around the span-based method. Yes, we are losing
+    /// the memory efficiency of the span-based method, but we still gain the performance
+    /// of the span-based method: the entire search algorithm is span-based.
+    /// </remarks>
+    /// <param name="text">The text that may contain the JSON list.</param>
+    /// <returns>The extracted JSON list.</returns>
+    private string ExtractJson(string text) => ExtractJson(text.AsSpan()).ToString();
+
+    /// <summary>
+    /// Extracts the JSON list from the given text. The text may contain additional
+    /// information around the JSON list. The method tries to extract the JSON list
+    /// from the text.
+    /// </summary>
+    /// <remarks>
+    /// Algorithm: The method searches for the first line that contains only a '[' character.
+    /// Then, it searches for the first line that contains only a ']' character. The method
+    /// returns the text between these two lines (including the brackets). When the method
+    /// cannot find the JSON list, it returns an empty string.
+    /// </remarks>
+    /// <param name="text">The text that may contain the JSON list.</param>
+    /// <returns>The extracted JSON list.</returns>
+    private static ReadOnlySpan<char> ExtractJson(ReadOnlySpan<char> text)
+    {
+        var startIndex = -1;
+        var endIndex = -1;
+        var foundStart = false;
+        var foundEnd = false;
+        var lineStart = 0;
+        
+        for (var i = 0; i <= text.Length; i++)
+        {
+            // Handle the end of the line or the end of the text:
+            if (i == text.Length || text[i] == '\n')
+            {
+                if (IsCharacterAloneInLine(text, lineStart, i, '[') && !foundStart)
+                {
+                    startIndex = lineStart;
+                    foundStart = true;
+                }
+                else if (IsCharacterAloneInLine(text, lineStart, i, ']') && foundStart && !foundEnd)
+                {
+                    endIndex = i;
+                    foundEnd = true;
+                    break;
+                }
+                
+                lineStart = i + 1;
+            }
+        }
+        
+        if (foundStart && foundEnd)
+        {
+            // Adjust endIndex for slicing, ensuring it's within bounds:
+            return text.Slice(startIndex, Math.Min(text.Length, endIndex + 1) - startIndex);
+        }
+        
+        return ReadOnlySpan<char>.Empty;
+    }
+    
+    private static bool IsCharacterAloneInLine(ReadOnlySpan<char> text, int lineStart, int lineEnd, char character)
+    {
+        for (var i = lineStart; i < lineEnd; i++)
+            if (!char.IsWhiteSpace(text[i]) && text[i] != character)
+                return false;
+
+        return true;
+    }
+}