Enhanced Model Management and SystemPrompt Integration in OllamaAgent

shell/agents/AIShell.Ollama.Agent/Command.cs

@@ -0,0 +1,280 @@
+using System.CommandLine;
+using System.CommandLine.Completions;
+using System.Threading.Tasks;
+using AIShell.Abstraction;
+
+namespace AIShell.Ollama.Agent;
+
+internal sealed class ConfigCommand : CommandBase
+{
+    private readonly OllamaAgent _agnet;
+     public ConfigCommand(OllamaAgent agent)
+        : base("config", "Command for config management within the 'ollama' agent.")
+    {
+        _agnet = agent;
+
+        var use = new Command("use", "Specify a config to use.");
+        var useConfig = new Argument<string>(
+            name: "Config",
+            getDefaultValue: () => null,
+            description: "Name of a configuration.").AddCompletions(ConfigNameCompleter);
+        use.AddArgument(useConfig);
+        use.SetHandler(UseConfigAction, useConfig);
+
+        var list = new Command("list", "List a specific config, or all available configs.");
+        var listConfig = new Argument<string>(
+            name: "Config",
+            getDefaultValue: () => null,
+            description: "Name of a configuration.").AddCompletions(ConfigNameCompleter);
+        list.AddArgument(listConfig);
+        list.SetHandler(ListConfigAction, listConfig);
+
+        AddCommand(list);
+        AddCommand(use);
+    }
+
+    private void ListConfigAction(string name)
+    {
+        IHost host = Shell.Host;
+
+        // Reload the setting file if needed.
+        _agnet.ReloadSettings();
+
+        Settings settings = _agnet.Settings;
+
+        if (settings is null)
+        {
+            host.WriteErrorLine("Invalid configuration.");
+            return;
+        }
+
+        if (string.IsNullOrEmpty(name))
+        {
+            settings.ListAllConfigs(host);
+            return;
+        }
+
+        try
+        {
+            settings.ShowOneConfig(host, name);
+        }
+        catch (InvalidOperationException ex)
+        {
+            string availableConfigNames = ConfigNamesAsString();
+            host.WriteErrorLine($"{ex.Message} Available cofiguration(s): {availableConfigNames}.");
+        }
+    }
+
+    private async Task UseConfigAction(string name)
+    {
+        // Reload the setting file if needed.
+        _agnet.ReloadSettings();
+
+        var setting = _agnet.Settings;
+        var host = Shell.Host;
+
+        if (setting is null || setting.Configs.Count is 0)
+        {
+            host.WriteErrorLine("No configs configured.");
+            return;
+        }
+
+        try
+        {
+            ModelConfig chosenConfig = (string.IsNullOrEmpty(name)
+                ? host.PromptForSelectionAsync(
+                    title: "[orange1]Please select a [Blue]Configuration[/] to use[/]:",
+                    choices: setting.Configs,
+                    converter: ConfigName,
+                    CancellationToken.None).GetAwaiter().GetResult()
+                : setting.Configs.FirstOrDefault(c => c.Name == name)) ?? throw new InvalidOperationException($"The configuration '{name}' doesn't exist.");
+            await setting.UseConfg(host, chosenConfig);
+            host.MarkupLine($"Using the config [green]{chosenConfig.Name}[/]:");
+        }
+        catch (InvalidOperationException ex)
+        {
+            string availableConfigNames = ConfigNamesAsString();
+            host.WriteErrorLine($"{ex.Message} Available configurations: {availableConfigNames}.");
+        }
+    }
+
+    private static string ConfigName(ModelConfig config) => config.Name.Any(Char.IsWhiteSpace) ? $"\"{config.Name}\"" : config.Name;
+    private IEnumerable<string> ConfigNameCompleter(CompletionContext context) => _agnet.Settings?.Configs?.Select(ConfigName) ?? [];
+    private string ConfigNamesAsString() => string.Join(", ", ConfigNameCompleter(null));
+}
+
+internal sealed class SystemPromptCommand : CommandBase
+{
+    private readonly OllamaAgent _agnet;
+
+    public SystemPromptCommand(OllamaAgent agent)
+        : base("system-prompt", "Command for system prompt management within the 'ollama' agent.")
+    {
+        _agnet = agent;
+
+        var show = new Command("show", "Show the current system prompt.");
+        show.SetHandler(ShowSystemPromptAction);
+
+        var set = new Command("set", "Sets the system prompt.");
+        var systemPromptModel = new Argument<string>(
+            name: "System-Prompt",
+            getDefaultValue: () => null,
+            description: "The system prompt");
+        set.AddArgument(systemPromptModel);
+        set.SetHandler(SetSystemPromptAction, systemPromptModel);
+
+        AddCommand(show);
+        AddCommand(set);
+    }
+
+    private void ShowSystemPromptAction()
+    {
+        IHost host = Shell.Host;
+
+        // Reload the setting file if needed.
+        _agnet.ReloadSettings();
+
+        Settings settings = _agnet.Settings;
+
+        if (settings is null)
+        {
+            host.WriteErrorLine("Invalid configuration.");
+            return;
+        }
+
+        try
+        {
+            settings.ShowSystemPrompt(host);
+        }
+        catch (InvalidOperationException ex)
+        {
+            host.WriteErrorLine($"{ex.Message}");
+        }
+    }
+
+    private void SetSystemPromptAction(string prompt)
+    {
+        IHost host = Shell.Host;
+
+        // Reload the setting file if needed.
+        _agnet.ReloadSettings();
+        _agnet.ResetContext();
+
+        Settings settings = _agnet.Settings;
+
+        if (settings is null)
+        {
+            host.WriteErrorLine("Invalid configuration.");
+            return;
+        }
+
+        try
+        {
+            settings.SetSystemPrompt(host, prompt);
+        }
+        catch (InvalidOperationException ex)
+        {
+            host.WriteErrorLine($"{ex.Message}.");
+        }
+    }
+}
+
+internal sealed class ModelCommand : CommandBase
+{
+    private readonly OllamaAgent _agnet;
+
+    public ModelCommand(OllamaAgent agent)
+        : base("model", "Command for model management within the 'ollama' agent.")
+    {
+        _agnet = agent;
+
+        var use = new Command("use", "Specify a model to use, or choose one from the available models.");
+        var useModel = new Argument<string>(
+            name: "Model",
+            getDefaultValue: () => null,
+            description: "Name of a model.").AddCompletions(ModelNameCompleter);
+        use.AddArgument(useModel);
+        use.SetHandler(UseModelAction, useModel);
+
+        var list = new Command("list", "List a specific model, or all available models.");
+        var listModel = new Argument<string>(
+            name: "Model",
+            getDefaultValue: () => null,
+            description: "Name of a model.").AddCompletions(ModelNameCompleter);
+        list.AddArgument(listModel);
+        list.SetHandler(ListModelAction, listModel);
+
+        AddCommand(list);
+        AddCommand(use);
+    }
+
+    private void ListModelAction(string name)
+    {
+        IHost host = Shell.Host;
+
+        // Reload the setting file if needed.
+        _agnet.ReloadSettings();
+
+        Settings settings = _agnet.Settings;
+
+        if (settings is null)
+        {
+            host.WriteErrorLine("Invalid configuration.");
+            return;
+        }
+
+        if (string.IsNullOrEmpty(name))
+        {
+            settings.ListAllModels(host).GetAwaiter().GetResult();
+            return;
+        }
+
+        try
+        {
+            settings.ShowOneModel(host, name).GetAwaiter().GetResult();
+        }
+        catch (InvalidOperationException ex)
+        {
+            string availableModelNames = ModelNamesAsString();
+            host.WriteErrorLine($"{ex.Message} Available Models(s): {availableModelNames}.");
+        }
+    }
+
+    private void UseModelAction(string name)
+    {
+        // Reload the setting file if needed.
+        _agnet.ReloadSettings();
+
+        var setting = _agnet.Settings;
+        var host = Shell.Host;
+
+        if (setting is null || setting.GetAllModels().GetAwaiter().GetResult().Count is 0)
+        {
+            host.WriteErrorLine("No models configured.");
+            return;
+        }
+
+        try
+        {
+            OllamaModel chosenModel = string.IsNullOrEmpty(name)
+                ? host.PromptForSelectionAsync(
+                    title: "[orange1]Please select a [Blue]Model[/] to use[/]:",
+                    choices: setting.GetAllModels().GetAwaiter().GetResult(),
+                    converter: ModelName,
+                    CancellationToken.None).GetAwaiter().GetResult()
+                : setting.GetModelByName(name).GetAwaiter().GetResult();
+
+            setting.UseModel(chosenModel);
+            host.MarkupLine($"Using the model [green]{chosenModel.Name}[/]:");
+        }
+        catch (InvalidOperationException ex)
+        {
+            string availableModelNames = ModelNamesAsString();
+            host.WriteErrorLine($"{ex.Message} Available Modless: {availableModelNames}.");
+        }
+    }
+
+    private static string ModelName(OllamaModel model) => model.Name;
+    private IEnumerable<string> ModelNameCompleter(CompletionContext context) => _agnet.Settings?.GetAllModels().GetAwaiter().GetResult().Select(ModelName) ?? [];
+    private string ModelNamesAsString() => string.Join(", ", ModelNameCompleter(null));
+}

shell/agents/AIShell.Ollama.Agent/OllamaAgent.cs

@@ -110,7 +110,7 @@ public void Initialize(AgentConfig config)
     /// <summary>
     /// Get commands that an agent can register to the shell when being loaded.
     /// </summary>
-    public IEnumerable<CommandBase> GetCommands() => null;
+    public IEnumerable<CommandBase> GetCommands() => [new ConfigCommand(this), new ModelCommand(this), new SystemPromptCommand(this)];
 
     /// <summary>
     /// Gets the path to the setting file of the agent.
@@ -148,6 +148,11 @@ public Task RefreshChatAsync(IShell shell, bool force)
         return Task.CompletedTask;
     }
 
+    public void ResetContext()
+    {
+        _request.Context = null;
+    }
+
     /// <summary>
     /// Main chat function that takes the users input and passes it to the LLM and renders it.
     /// </summary>
@@ -171,11 +176,18 @@ public async Task<bool> ChatAsync(string input, IShell shell)
             return false;
         }
 
+        var activeModel = await _settings.GetActiveModel().ConfigureAwait(false);
+
         // Prepare request
         _request.Prompt = input;
-        _request.Model = _settings.Model;
+        _request.Model = activeModel.Name;
         _request.Stream = _settings.Stream;
 
+        if (!string.IsNullOrWhiteSpace(_settings.RunningConfig.SystemPrompt))
+        {
+            _request.System = _settings.RunningConfig.SystemPrompt;
+        }
+
         try
         {
             if (_request.Stream)
@@ -238,15 +250,23 @@ public async Task<bool> ChatAsync(string input, IShell shell)
         catch (HttpRequestException e)
         {
             host.WriteErrorLine($"{e.Message}");
-            host.WriteErrorLine($"Ollama model:    \"{_settings.Model}\"");
-            host.WriteErrorLine($"Ollama endpoint: \"{_settings.Endpoint}\"");
-            host.WriteErrorLine($"Ollama settings: \"{SettingFile}\"");
+            host.WriteErrorLine($"Ollama active model: \"{activeModel.Name}\"");
+            host.WriteErrorLine($"Ollama endpoint:     \"{_settings.Endpoint}\"");
+            host.WriteErrorLine($"Ollama settings:     \"{SettingFile}\"");
+        }
+        finally
+        {
+            if (_settings.RunningConfig.ResetContext)
+            {
+                // Reset the request context
+            	ResetContext();
+            }
         }
 
         return true;
     }
 
-    private void ReloadSettings()
+    internal void ReloadSettings()
     {
         if (_reloadSettings)
         {
@@ -308,13 +328,20 @@ private void NewExampleSettingFile()
             // 1. Install Ollama: `winget install Ollama.Ollama`
             // 2. Start Ollama API server: `ollama serve`
             // 3. Install Ollama model: `ollama pull phi3`
-
-            // Declare Ollama model
-            "Model": "phi3",
+            "Configs": [
+                {
+                    "Name": "PowerShell Expert",
+                    "Description": "A ollama agent with expertise in PowerShell scripting and command line utilities.",
+                    "ModelName": "phi3",
+                    "SystemPrompt": "1. You are a helpful and friendly assistant with expertise in PowerShell scripting and command line.\n2. Assume user is using the operating system `Windows 11` unless otherwise specified.\n3. Use the `code block` syntax in markdown to encapsulate any part in responses that is code, YAML, JSON or XML, but not table.\n4. When encapsulating command line code, use '```powershell' if it's PowerShell command; use '```sh' if it's non-PowerShell CLI command.\n5. When generating CLI commands, never ever break a command into multiple lines. Instead, always list all parameters and arguments of the command on the same line.\n6. Please keep the response concise but to the point. Do not overexplain."
+                }
+            ],
             // Declare Ollama endpoint
             "Endpoint": "http://localhost:11434",
             // Enable Ollama streaming
-            "Stream": false
+            "Stream": false,
+            // Specify the default model to use
+            "DefaultConfig": "PowerShell Expert",
         }
         """;
         File.WriteAllText(SettingFile, SampleContent, Encoding.UTF8);