Skip to content

Commit ad02261

Browse files
cnupycnupy
authored
Enhanced model management and system prompt integration in OllamaAgent (#351)
- Add the `model` command for model management - Add the `system-prompt` command to configure `SystemPrompt` with the selected model. - Add the `preset` command for preset management
1 parent 06836b2 commit ad02261

File tree

4 files changed

+615
-40
lines changed

4 files changed

+615
-40
lines changed

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

Lines changed: 309 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,309 @@
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 PresetCommand : CommandBase
9+
{
10+
private readonly OllamaAgent _agnet;
11+
12+
public PresetCommand(OllamaAgent agent)
13+
: base("preset", "Command for preset management within the 'ollama' agent.")
14+
{
15+
_agnet = agent;
16+
17+
var use = new Command("use", "Specify a preset to use.");
18+
var usePreset = new Argument<string>(
19+
name: "Preset",
20+
getDefaultValue: () => null,
21+
description: "Name of a preset.").AddCompletions(PresetNameCompleter);
22+
use.AddArgument(usePreset);
23+
use.SetHandler(UsePresetAction, usePreset);
24+
25+
var list = new Command("list", "List a specific preset, or all configured presets.");
26+
var listPreset = new Argument<string>(
27+
name: "Preset",
28+
getDefaultValue: () => null,
29+
description: "Name of a preset.").AddCompletions(PresetNameCompleter);
30+
list.AddArgument(listPreset);
31+
list.SetHandler(ListPresetAction, listPreset);
32+
33+
AddCommand(list);
34+
AddCommand(use);
35+
}
36+
37+
private void ListPresetAction(string name)
38+
{
39+
IHost host = Shell.Host;
40+
41+
// Reload the setting file if needed.
42+
_agnet.ReloadSettings();
43+
44+
Settings settings = _agnet.Settings;
45+
46+
if (settings is null)
47+
{
48+
host.WriteErrorLine("Error loading the configuration.");
49+
return;
50+
}
51+
52+
try
53+
{
54+
if (string.IsNullOrEmpty(name))
55+
{
56+
settings.ListAllPresets(host);
57+
return;
58+
}
59+
60+
settings.ShowOnePreset(host, name);
61+
}
62+
catch (Exception ex)
63+
{
64+
string availablePresetNames = PresetNamesAsString();
65+
host.WriteErrorLine($"{ex.Message} Available preset(s): {availablePresetNames}.");
66+
}
67+
}
68+
69+
private async Task UsePresetAction(string name)
70+
{
71+
// Reload the setting file if needed.
72+
_agnet.ReloadSettings();
73+
74+
var setting = _agnet.Settings;
75+
var host = Shell.Host;
76+
77+
if (setting is null)
78+
{
79+
host.WriteErrorLine("Error loading the configuration.");
80+
return;
81+
}
82+
83+
if (setting.Presets.Count is 0)
84+
{
85+
host.WriteErrorLine("There are no presets configured.");
86+
return;
87+
}
88+
89+
try
90+
{
91+
ModelConfig chosenPreset = (string.IsNullOrEmpty(name)
92+
? await host.PromptForSelectionAsync(
93+
title: "[orange1]Please select a [Blue]Preset[/] to use[/]:",
94+
choices: setting.Presets,
95+
converter: PresetName,
96+
CancellationToken.None)
97+
: setting.Presets.FirstOrDefault(c => c.Name == name)) ?? throw new InvalidOperationException($"The preset '{name}' doesn't exist.");
98+
await setting.UsePreset(host, chosenPreset);
99+
host.MarkupLine($"Using the preset [green]{chosenPreset.Name}[/]:");
100+
}
101+
catch (Exception ex)
102+
{
103+
string availablePresetNames = PresetNamesAsString();
104+
host.WriteErrorLine($"{ex.Message} Available presets: {availablePresetNames}.");
105+
}
106+
}
107+
108+
private static string PresetName(ModelConfig preset) => preset.Name.Any(Char.IsWhiteSpace) ? $"\"{preset.Name}\"" : preset.Name;
109+
private IEnumerable<string> PresetNameCompleter(CompletionContext context) => _agnet.Settings?.Presets?.Select(PresetName) ?? [];
110+
private string PresetNamesAsString() => string.Join(", ", PresetNameCompleter(null));
111+
}
112+
113+
internal sealed class SystemPromptCommand : CommandBase
114+
{
115+
private readonly OllamaAgent _agnet;
116+
117+
public SystemPromptCommand(OllamaAgent agent)
118+
: base("system-prompt", "Command for system prompt management within the 'ollama' agent.")
119+
{
120+
_agnet = agent;
121+
122+
var show = new Command("show", "Show the current system prompt.");
123+
show.SetHandler(ShowSystemPromptAction);
124+
125+
var set = new Command("set", "Sets the system prompt.");
126+
var systemPromptModel = new Argument<string>(
127+
name: "System-Prompt",
128+
getDefaultValue: () => null,
129+
description: "The system prompt");
130+
set.AddArgument(systemPromptModel);
131+
set.SetHandler(SetSystemPromptAction, systemPromptModel);
132+
133+
AddCommand(show);
134+
AddCommand(set);
135+
}
136+
137+
private void ShowSystemPromptAction()
138+
{
139+
IHost host = Shell.Host;
140+
141+
// Reload the setting file if needed.
142+
_agnet.ReloadSettings();
143+
144+
Settings settings = _agnet.Settings;
145+
146+
if (settings is null)
147+
{
148+
host.WriteErrorLine("Error loading the configuration.");
149+
return;
150+
}
151+
152+
try
153+
{
154+
settings.ShowSystemPrompt(host);
155+
}
156+
catch (Exception ex)
157+
{
158+
host.WriteErrorLine(ex.Message);
159+
}
160+
}
161+
162+
private void SetSystemPromptAction(string prompt)
163+
{
164+
IHost host = Shell.Host;
165+
166+
// Reload the setting file if needed.
167+
_agnet.ReloadSettings();
168+
_agnet.ResetContext();
169+
170+
Settings settings = _agnet.Settings;
171+
172+
if (settings is null)
173+
{
174+
host.WriteErrorLine("Error loading the configuration.");
175+
return;
176+
}
177+
178+
try
179+
{
180+
settings.SetSystemPrompt(host, prompt);
181+
}
182+
catch (Exception ex)
183+
{
184+
host.WriteErrorLine(ex.Message);
185+
}
186+
}
187+
}
188+
189+
internal sealed class ModelCommand : CommandBase
190+
{
191+
private readonly OllamaAgent _agnet;
192+
193+
public ModelCommand(OllamaAgent agent)
194+
: base("model", "Command for model management within the 'ollama' agent.")
195+
{
196+
_agnet = agent;
197+
198+
var use = new Command("use", "Specify a model to use, or choose one from the available models.");
199+
var useModel = new Argument<string>(
200+
name: "Model",
201+
getDefaultValue: () => null,
202+
description: "Name of a model.").AddCompletions(ModelNameCompleter);
203+
use.AddArgument(useModel);
204+
use.SetHandler(UseModelAction, useModel);
205+
206+
var list = new Command("list", "List a specific model, or all available models.");
207+
var listModel = new Argument<string>(
208+
name: "Model",
209+
getDefaultValue: () => null,
210+
description: "Name of a model.").AddCompletions(ModelNameCompleter);
211+
list.AddArgument(listModel);
212+
list.SetHandler(ListModelAction, listModel);
213+
214+
AddCommand(list);
215+
AddCommand(use);
216+
}
217+
218+
private async Task ListModelAction(string name)
219+
{
220+
IHost host = Shell.Host;
221+
222+
// Reload the setting file if needed.
223+
_agnet.ReloadSettings();
224+
225+
Settings settings = _agnet.Settings;
226+
227+
if (settings is null)
228+
{
229+
host.WriteErrorLine("Error loading the configuration.");
230+
return;
231+
}
232+
try
233+
{
234+
if (string.IsNullOrEmpty(name))
235+
{
236+
await settings.ListAllModels(host);
237+
return;
238+
}
239+
240+
await settings.ShowOneModel(host, name);
241+
}
242+
catch (Exception ex)
243+
{
244+
host.WriteErrorLine(ex.Message);
245+
}
246+
}
247+
248+
private async Task UseModelAction(string name)
249+
{
250+
// Reload the setting file if needed.
251+
_agnet.ReloadSettings();
252+
253+
var settings = _agnet.Settings;
254+
var host = Shell.Host;
255+
256+
if (settings is null)
257+
{
258+
host.WriteErrorLine("Error loading the configuration.");
259+
return;
260+
}
261+
262+
try
263+
{
264+
bool success = await settings.PerformSelfcheck(host, checkEndpointOnly: true);
265+
if (!success)
266+
{
267+
return;
268+
}
269+
270+
var allModels = await settings.GetAllModels();
271+
if (allModels.Count is 0)
272+
{
273+
host.WriteErrorLine($"No models found from '{settings.Endpoint}'.");
274+
return;
275+
}
276+
277+
if (string.IsNullOrEmpty(name))
278+
{
279+
name = await host.PromptForSelectionAsync(
280+
title: "[orange1]Please select a [Blue]Model[/] to use[/]:",
281+
choices: allModels,
282+
CancellationToken.None);
283+
}
284+
285+
await settings.UseModel(host, name);
286+
host.MarkupLine($"Using the model [green]{name}[/]");
287+
}
288+
catch (Exception ex)
289+
{
290+
host.WriteErrorLine(ex.Message);
291+
}
292+
}
293+
294+
private IEnumerable<string> ModelNameCompleter(CompletionContext context)
295+
{
296+
try
297+
{
298+
// Model retrieval may throw.
299+
var results = _agnet.Settings?.GetAllModels().Result;
300+
if (results is not null)
301+
{
302+
return results;
303+
}
304+
}
305+
catch (Exception) { }
306+
307+
return [];
308+
}
309+
}

0 commit comments

Comments
 (0)