Skip to content

Commit 2b67582

Browse files
cnupycnupy
committed
OllamaAgent can now handle and switch between multiple models seamlessly, listing available models dynamically via an API
* Modified the agent to include a model selection mechanism that allows it to iterate through predefined models. * Introduced a `model` command to manage the switching between models based on available models. It is possible to configure SystemPrompt with the selected model into OllamaAgent, providing users with a toolkit for task-oriented conversations. * Modified the agent to accommodate Ollama's SystemPrompt functionality. * Introduced a `system-prompt` command to manage the setting. OllamaAgent now supports the use of predefined model configurations, allowing users to easily switch between different models and system prompts based on specific requirements. * Introduced a `ModelConfig` Record to encapsulate data for each configuration. * Introduced a `config` command to manage the switching between configuration based on available predefined sets.
1 parent 765a0ca commit 2b67582

File tree

3 files changed

+509
-21
lines changed

3 files changed

+509
-21
lines changed

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

Lines changed: 280 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,280 @@
1+
using System.CommandLine;
2+
using System.CommandLine.Completions;
3+
using System.Threading.Tasks;
4+
using AIShell.Abstraction;
5+
6+
namespace AIShell.Ollama.Agent;
7+
8+
internal sealed class ConfigCommand : CommandBase
9+
{
10+
private readonly OllamaAgent _agnet;
11+
public ConfigCommand(OllamaAgent agent)
12+
: base("config", "Command for config management within the 'ollama' agent.")
13+
{
14+
_agnet = agent;
15+
16+
var use = new Command("use", "Specify a config to use.");
17+
var useConfig = new Argument<string>(
18+
name: "Config",
19+
getDefaultValue: () => null,
20+
description: "Name of a configuration.").AddCompletions(ConfigNameCompleter);
21+
use.AddArgument(useConfig);
22+
use.SetHandler(UseConfigAction, useConfig);
23+
24+
var list = new Command("list", "List a specific config, or all available configs.");
25+
var listConfig = new Argument<string>(
26+
name: "Config",
27+
getDefaultValue: () => null,
28+
description: "Name of a configuration.").AddCompletions(ConfigNameCompleter);
29+
list.AddArgument(listConfig);
30+
list.SetHandler(ListConfigAction, listConfig);
31+
32+
AddCommand(list);
33+
AddCommand(use);
34+
}
35+
36+
private void ListConfigAction(string name)
37+
{
38+
IHost host = Shell.Host;
39+
40+
// Reload the setting file if needed.
41+
_agnet.ReloadSettings();
42+
43+
Settings settings = _agnet.Settings;
44+
45+
if (settings is null)
46+
{
47+
host.WriteErrorLine("Invalid configuration.");
48+
return;
49+
}
50+
51+
if (string.IsNullOrEmpty(name))
52+
{
53+
settings.ListAllConfigs(host);
54+
return;
55+
}
56+
57+
try
58+
{
59+
settings.ShowOneConfig(host, name);
60+
}
61+
catch (InvalidOperationException ex)
62+
{
63+
string availableConfigNames = ConfigNamesAsString();
64+
host.WriteErrorLine($"{ex.Message} Available cofiguration(s): {availableConfigNames}.");
65+
}
66+
}
67+
68+
private async Task UseConfigAction(string name)
69+
{
70+
// Reload the setting file if needed.
71+
_agnet.ReloadSettings();
72+
73+
var setting = _agnet.Settings;
74+
var host = Shell.Host;
75+
76+
if (setting is null || setting.Configs.Count is 0)
77+
{
78+
host.WriteErrorLine("No configs configured.");
79+
return;
80+
}
81+
82+
try
83+
{
84+
ModelConfig chosenConfig = (string.IsNullOrEmpty(name)
85+
? host.PromptForSelectionAsync(
86+
title: "[orange1]Please select a [Blue]Configuration[/] to use[/]:",
87+
choices: setting.Configs,
88+
converter: ConfigName,
89+
CancellationToken.None).GetAwaiter().GetResult()
90+
: setting.Configs.FirstOrDefault(c => c.Name == name)) ?? throw new InvalidOperationException($"The configuration '{name}' doesn't exist.");
91+
await setting.UseConfg(host, chosenConfig);
92+
host.MarkupLine($"Using the config [green]{chosenConfig.Name}[/]:");
93+
}
94+
catch (InvalidOperationException ex)
95+
{
96+
string availableConfigNames = ConfigNamesAsString();
97+
host.WriteErrorLine($"{ex.Message} Available configurations: {availableConfigNames}.");
98+
}
99+
}
100+
101+
private static string ConfigName(ModelConfig config) => config.Name.Any(Char.IsWhiteSpace) ? $"\"{config.Name}\"" : config.Name;
102+
private IEnumerable<string> ConfigNameCompleter(CompletionContext context) => _agnet.Settings?.Configs?.Select(ConfigName) ?? [];
103+
private string ConfigNamesAsString() => string.Join(", ", ConfigNameCompleter(null));
104+
}
105+
106+
internal sealed class SystemPromptCommand : CommandBase
107+
{
108+
private readonly OllamaAgent _agnet;
109+
110+
public SystemPromptCommand(OllamaAgent agent)
111+
: base("system-prompt", "Command for system prompt management within the 'ollama' agent.")
112+
{
113+
_agnet = agent;
114+
115+
var show = new Command("show", "Show the current system prompt.");
116+
show.SetHandler(ShowSystemPromptAction);
117+
118+
var set = new Command("set", "Sets the system prompt.");
119+
var systemPromptModel = new Argument<string>(
120+
name: "System-Prompt",
121+
getDefaultValue: () => null,
122+
description: "The system prompt");
123+
set.AddArgument(systemPromptModel);
124+
set.SetHandler(SetSystemPromptAction, systemPromptModel);
125+
126+
AddCommand(show);
127+
AddCommand(set);
128+
}
129+
130+
private void ShowSystemPromptAction()
131+
{
132+
IHost host = Shell.Host;
133+
134+
// Reload the setting file if needed.
135+
_agnet.ReloadSettings();
136+
137+
Settings settings = _agnet.Settings;
138+
139+
if (settings is null)
140+
{
141+
host.WriteErrorLine("Invalid configuration.");
142+
return;
143+
}
144+
145+
try
146+
{
147+
settings.ShowSystemPrompt(host);
148+
}
149+
catch (InvalidOperationException ex)
150+
{
151+
host.WriteErrorLine($"{ex.Message}");
152+
}
153+
}
154+
155+
private void SetSystemPromptAction(string prompt)
156+
{
157+
IHost host = Shell.Host;
158+
159+
// Reload the setting file if needed.
160+
_agnet.ReloadSettings();
161+
_agnet.ResetContext();
162+
163+
Settings settings = _agnet.Settings;
164+
165+
if (settings is null)
166+
{
167+
host.WriteErrorLine("Invalid configuration.");
168+
return;
169+
}
170+
171+
try
172+
{
173+
settings.SetSystemPrompt(host, prompt);
174+
}
175+
catch (InvalidOperationException ex)
176+
{
177+
host.WriteErrorLine($"{ex.Message}.");
178+
}
179+
}
180+
}
181+
182+
internal sealed class ModelCommand : CommandBase
183+
{
184+
private readonly OllamaAgent _agnet;
185+
186+
public ModelCommand(OllamaAgent agent)
187+
: base("model", "Command for model management within the 'ollama' agent.")
188+
{
189+
_agnet = agent;
190+
191+
var use = new Command("use", "Specify a model to use, or choose one from the available models.");
192+
var useModel = new Argument<string>(
193+
name: "Model",
194+
getDefaultValue: () => null,
195+
description: "Name of a model.").AddCompletions(ModelNameCompleter);
196+
use.AddArgument(useModel);
197+
use.SetHandler(UseModelAction, useModel);
198+
199+
var list = new Command("list", "List a specific model, or all available models.");
200+
var listModel = new Argument<string>(
201+
name: "Model",
202+
getDefaultValue: () => null,
203+
description: "Name of a model.").AddCompletions(ModelNameCompleter);
204+
list.AddArgument(listModel);
205+
list.SetHandler(ListModelAction, listModel);
206+
207+
AddCommand(list);
208+
AddCommand(use);
209+
}
210+
211+
private void ListModelAction(string name)
212+
{
213+
IHost host = Shell.Host;
214+
215+
// Reload the setting file if needed.
216+
_agnet.ReloadSettings();
217+
218+
Settings settings = _agnet.Settings;
219+
220+
if (settings is null)
221+
{
222+
host.WriteErrorLine("Invalid configuration.");
223+
return;
224+
}
225+
226+
if (string.IsNullOrEmpty(name))
227+
{
228+
settings.ListAllModels(host).GetAwaiter().GetResult();
229+
return;
230+
}
231+
232+
try
233+
{
234+
settings.ShowOneModel(host, name).GetAwaiter().GetResult();
235+
}
236+
catch (InvalidOperationException ex)
237+
{
238+
string availableModelNames = ModelNamesAsString();
239+
host.WriteErrorLine($"{ex.Message} Available Models(s): {availableModelNames}.");
240+
}
241+
}
242+
243+
private void UseModelAction(string name)
244+
{
245+
// Reload the setting file if needed.
246+
_agnet.ReloadSettings();
247+
248+
var setting = _agnet.Settings;
249+
var host = Shell.Host;
250+
251+
if (setting is null || setting.GetAllModels().GetAwaiter().GetResult().Count is 0)
252+
{
253+
host.WriteErrorLine("No models configured.");
254+
return;
255+
}
256+
257+
try
258+
{
259+
OllamaModel chosenModel = string.IsNullOrEmpty(name)
260+
? host.PromptForSelectionAsync(
261+
title: "[orange1]Please select a [Blue]Model[/] to use[/]:",
262+
choices: setting.GetAllModels().GetAwaiter().GetResult(),
263+
converter: ModelName,
264+
CancellationToken.None).GetAwaiter().GetResult()
265+
: setting.GetModelByName(name).GetAwaiter().GetResult();
266+
267+
setting.UseModel(chosenModel);
268+
host.MarkupLine($"Using the model [green]{chosenModel.Name}[/]:");
269+
}
270+
catch (InvalidOperationException ex)
271+
{
272+
string availableModelNames = ModelNamesAsString();
273+
host.WriteErrorLine($"{ex.Message} Available Modless: {availableModelNames}.");
274+
}
275+
}
276+
277+
private static string ModelName(OllamaModel model) => model.Name;
278+
private IEnumerable<string> ModelNameCompleter(CompletionContext context) => _agnet.Settings?.GetAllModels().GetAwaiter().GetResult().Select(ModelName) ?? [];
279+
private string ModelNamesAsString() => string.Join(", ", ModelNameCompleter(null));
280+
}

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

Lines changed: 35 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -110,7 +110,7 @@ public void Initialize(AgentConfig config)
110110
/// <summary>
111111
/// Get commands that an agent can register to the shell when being loaded.
112112
/// </summary>
113-
public IEnumerable<CommandBase> GetCommands() => null;
113+
public IEnumerable<CommandBase> GetCommands() => [new ConfigCommand(this), new ModelCommand(this), new SystemPromptCommand(this)];
114114

115115
/// <summary>
116116
/// Gets the path to the setting file of the agent.
@@ -148,6 +148,11 @@ public Task RefreshChatAsync(IShell shell, bool force)
148148
return Task.CompletedTask;
149149
}
150150

151+
public void ResetContext()
152+
{
153+
_request.Context = null;
154+
}
155+
151156
/// <summary>
152157
/// Main chat function that takes the users input and passes it to the LLM and renders it.
153158
/// </summary>
@@ -171,11 +176,18 @@ public async Task<bool> ChatAsync(string input, IShell shell)
171176
return false;
172177
}
173178

179+
var activeModel = await _settings.GetActiveModel().ConfigureAwait(false);
180+
174181
// Prepare request
175182
_request.Prompt = input;
176-
_request.Model = _settings.Model;
183+
_request.Model = activeModel.Name;
177184
_request.Stream = _settings.Stream;
178185

186+
if (!string.IsNullOrWhiteSpace(_settings.RunningConfig.SystemPrompt))
187+
{
188+
_request.System = _settings.RunningConfig.SystemPrompt;
189+
}
190+
179191
try
180192
{
181193
if (_request.Stream)
@@ -238,15 +250,23 @@ public async Task<bool> ChatAsync(string input, IShell shell)
238250
catch (HttpRequestException e)
239251
{
240252
host.WriteErrorLine($"{e.Message}");
241-
host.WriteErrorLine($"Ollama model: \"{_settings.Model}\"");
253+
host.WriteErrorLine($"Ollama active model: \"{activeModel.Name}\"");
242254
host.WriteErrorLine($"Ollama endpoint: \"{_settings.Endpoint}\"");
243255
host.WriteErrorLine($"Ollama settings: \"{SettingFile}\"");
244256
}
257+
finally
258+
{
259+
if (_settings.RunningConfig.ResetContext)
260+
{
261+
// Reset the request context
262+
ResetContext();
263+
}
264+
}
245265

246266
return true;
247267
}
248268

249-
private void ReloadSettings()
269+
internal void ReloadSettings()
250270
{
251271
if (_reloadSettings)
252272
{
@@ -308,13 +328,20 @@ private void NewExampleSettingFile()
308328
// 1. Install Ollama: `winget install Ollama.Ollama`
309329
// 2. Start Ollama API server: `ollama serve`
310330
// 3. Install Ollama model: `ollama pull phi3`
311-
312-
// Declare Ollama model
313-
"Model": "phi3",
331+
"Configs": [
332+
{
333+
"Name": "PowerShell Expert",
334+
"Description": "A ollama agent with expertise in PowerShell scripting and command line utilities.",
335+
"ModelName": "phi3",
336+
"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."
337+
}
338+
],
314339
// Declare Ollama endpoint
315340
"Endpoint": "http://localhost:11434",
316341
// Enable Ollama streaming
317-
"Stream": false
342+
"Stream": false,
343+
// Specify the default model to use
344+
"DefaultConfig": "PowerShell Expert",
318345
}
319346
""";
320347
File.WriteAllText(SettingFile, SampleContent, Encoding.UTF8);

0 commit comments

Comments
 (0)