Skip to content

Commit 9b66a08

Browse files
daxian-dbwdaxian-dbw
authored
Enable parameter injection for Azure PowerShell response (#339)
See the PR description for detailed decision choices: #339
1 parent 119e6a8 commit 9b66a08

File tree

9 files changed

+517
-35
lines changed

9 files changed

+517
-35
lines changed

‎build.psm1

+37
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,43 @@ function Start-Build
9191
$app_csproj = GetProjectFile $app_dir
9292
dotnet publish $app_csproj -c $Configuration -o $app_out_dir -r $RID --sc
9393

94+
## Move the 'Modules' folder to the appbase folder.
95+
if ($LASTEXITCODE -eq 0) {
96+
## Remove the artifacts that are not for the current platform, to reduce size.
97+
$otherPlatDir = Join-Path $app_out_dir 'runtimes' ($IsWindows ? 'unix' : 'win')
98+
if (Test-Path $otherPlatDir -PathType Container) {
99+
Remove-Item $otherPlatDir -Recurse -Force -ErrorAction Stop
100+
}
101+
102+
## Move the 'Modules' folder and if possible, remove the 'runtimes' folder all together afterward.
103+
$platDir = Join-Path $app_out_dir 'runtimes' ($IsWindows ? 'win' : 'unix') 'lib'
104+
if (Test-Path $platDir -PathType Container) {
105+
$moduleDir = Get-ChildItem $platDir -Directory -Include 'Modules' -Recurse
106+
if ($moduleDir) {
107+
## Remove the existing 'Modules' folder if it already exists.
108+
$target = Join-Path $app_out_dir 'Modules'
109+
if (Test-Path $target -PathType Container) {
110+
Remove-Item $target -Recurse -Force -ErrorAction Stop
111+
}
112+
113+
## Move 'Modules' folder.
114+
Move-Item $moduleDir.FullName $app_out_dir -Force -ErrorAction Stop
115+
116+
## Remove the 'runtimes' folder if possible.
117+
$parent = $moduleDir.Parent
118+
while ($parent.FullName -ne $app_out_dir) {
119+
$files = Get-ChildItem $parent.FullName -File -Recurse
120+
if (-not $files) {
121+
Remove-Item $parent.FullName -Recurse -Force -ErrorAction Stop
122+
$parent = $parent.Parent
123+
} else {
124+
break
125+
}
126+
}
127+
}
128+
}
129+
}
130+
94131
if ($LASTEXITCODE -eq 0 -and $AgentToInclude -contains 'openai-gpt') {
95132
Write-Host "`n[Build the OpenAI agent ...]`n" -ForegroundColor Green
96133
$openai_csproj = GetProjectFile $openai_agent_dir

‎shell/AIShell.App/AIShell.App.csproj

+1
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
<ItemGroup>
1414
<ProjectReference Include="..\AIShell.Kernel\AIShell.Kernel.csproj" />
1515
<PackageReference Include="System.CommandLine" Version="2.0.0-beta4.22272.1" />
16+
<PackageReference Include="Microsoft.PowerShell.SDK" Version="7.4.7" />
1617
</ItemGroup>
1718

1819
</Project>

‎shell/AIShell.Kernel/Utility/LoadContext.cs

+106-8
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,17 @@
1-
using System.Reflection;
1+
using System.Collections.Concurrent;
2+
using System.Reflection;
3+
using System.Runtime.InteropServices;
24
using System.Runtime.Loader;
35

46
namespace AIShell.Kernel;
57

68
internal class AgentAssemblyLoadContext : AssemblyLoadContext
79
{
810
private readonly string _dependencyDir;
11+
private readonly string _nativeLibExt;
12+
private readonly List<string> _runtimeLibDir;
13+
private readonly List<string> _runtimeNativeDir;
14+
private readonly ConcurrentDictionary<string, Assembly> _cache;
915

1016
internal AgentAssemblyLoadContext(string name, string dependencyDir)
1117
: base($"{name.Replace(' ', '.')}-ALC", isCollectible: false)
@@ -17,20 +23,112 @@ internal AgentAssemblyLoadContext(string name, string dependencyDir)
1723

1824
// Save the full path to the dependencies directory when creating the context.
1925
_dependencyDir = dependencyDir;
26+
_runtimeLibDir = [];
27+
_runtimeNativeDir = [];
28+
_cache = [];
29+
30+
if (OperatingSystem.IsWindows())
31+
{
32+
_nativeLibExt = ".dll";
33+
AddToList(_runtimeLibDir, Path.Combine(dependencyDir, "runtimes", "win", "lib"));
34+
AddToList(_runtimeNativeDir, Path.Combine(dependencyDir, "runtimes", "win", "native"));
35+
}
36+
else if (OperatingSystem.IsLinux())
37+
{
38+
_nativeLibExt = ".so";
39+
AddToList(_runtimeLibDir, Path.Combine(dependencyDir, "runtimes", "unix", "lib"));
40+
AddToList(_runtimeLibDir, Path.Combine(dependencyDir, "runtimes", "linux", "lib"));
41+
42+
AddToList(_runtimeNativeDir, Path.Combine(dependencyDir, "runtimes", "unix", "native"));
43+
AddToList(_runtimeNativeDir, Path.Combine(dependencyDir, "runtimes", "linux", "native"));
44+
}
45+
else if (OperatingSystem.IsMacOS())
46+
{
47+
_nativeLibExt = ".dylib";
48+
AddToList(_runtimeLibDir, Path.Combine(dependencyDir, "runtimes", "unix", "lib"));
49+
AddToList(_runtimeLibDir, Path.Combine(dependencyDir, "runtimes", "osx", "lib"));
50+
51+
AddToList(_runtimeNativeDir, Path.Combine(dependencyDir, "runtimes", "unix", "native"));
52+
AddToList(_runtimeNativeDir, Path.Combine(dependencyDir, "runtimes", "osx", "native"));
53+
}
54+
55+
AddToList(_runtimeLibDir, Path.Combine(dependencyDir, "runtimes", RuntimeInformation.RuntimeIdentifier, "lib"));
56+
AddToList(_runtimeNativeDir, Path.Combine(dependencyDir, "runtimes", RuntimeInformation.RuntimeIdentifier, "native"));
57+
58+
ResolvingUnmanagedDll += ResolveUnmanagedDll;
2059
}
2160

2261
protected override Assembly Load(AssemblyName assemblyName)
2362
{
24-
// Create a path to the assembly in the dependencies directory.
25-
string path = Path.Combine(_dependencyDir, $"{assemblyName.Name}.dll");
63+
if (_cache.TryGetValue(assemblyName.Name, out Assembly assembly))
64+
{
65+
return assembly;
66+
}
67+
68+
lock (this)
69+
{
70+
if (_cache.TryGetValue(assemblyName.Name, out assembly))
71+
{
72+
return assembly;
73+
}
74+
75+
// Create a path to the assembly in the dependencies directory.
76+
string assemblyFile = $"{assemblyName.Name}.dll";
77+
string path = Path.Combine(_dependencyDir, assemblyFile);
78+
79+
if (File.Exists(path))
80+
{
81+
// If the assembly exists in our dependency directory, then load it into this load context.
82+
assembly = LoadFromAssemblyPath(path);
83+
}
84+
else
85+
{
86+
foreach (string dir in _runtimeLibDir)
87+
{
88+
IEnumerable<string> result = Directory.EnumerateFiles(dir, assemblyFile, SearchOption.AllDirectories);
89+
path = result.FirstOrDefault();
90+
91+
if (path is not null)
92+
{
93+
assembly = LoadFromAssemblyPath(path);
94+
break;
95+
}
96+
}
97+
}
98+
99+
// Add the probing result to cache, regardless of whether we found it.
100+
// If we didn't find it, we will add 'null' to the cache so that we don't probe
101+
// again in case another loading request comes for the same assembly.
102+
_cache.TryAdd(assemblyName.Name, assembly);
103+
104+
// Return the assembly if we found it, or return 'null' otherwise to depend on the default load context to resolve the request.
105+
return assembly;
106+
}
107+
}
108+
109+
private nint ResolveUnmanagedDll(Assembly assembly, string libraryName)
110+
{
111+
string libraryFile = $"{libraryName}{_nativeLibExt}";
26112

27-
if (File.Exists(path))
113+
foreach (string dir in _runtimeNativeDir)
28114
{
29-
// If the assembly exists in our dependency directory, then load it into this load context.
30-
return LoadFromAssemblyPath(path);
115+
IEnumerable<string> result = Directory.EnumerateFiles(dir, libraryFile, SearchOption.AllDirectories);
116+
string path = result.FirstOrDefault();
117+
118+
if (path is not null)
119+
{
120+
return NativeLibrary.Load(path);
121+
}
31122
}
32123

33-
// Otherwise we will depend on the default load context to resolve the request.
34-
return null;
124+
return nint.Zero;
125+
}
126+
127+
private static void AddToList(List<string> depList, string dirPath)
128+
{
129+
if (Directory.Exists(dirPath))
130+
{
131+
depList.Add(dirPath);
132+
}
35133
}
36134
}

‎shell/agents/Microsoft.Azure.Agent/AzureAgent.cs

+5-4
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ public sealed class AzureAgent : ILLMAgent
2727
2. DO NOT include the command for creating a new resource group unless the query explicitly asks for it. Otherwise, assume a resource group already exists.
2828
3. DO NOT include an additional example with made-up values unless it provides additional context or value beyond the initial command.
2929
4. DO NOT use the line continuation operator (backslash `\` in Bash) in the generated commands.
30-
5. Always represent a placeholder in the form of `<placeholder-name>`.
30+
5. Always represent a placeholder in the form of `<placeholder-name>` and enclose it within double quotes.
3131
6. Always use the consistent placeholder names across all your responses. For example, `<resourceGroupName>` should be used for all the places where a resource group name value is needed.
3232
7. When the commands contain placeholders, the placeholders should be summarized in markdown bullet points at the end of the response in the same order as they appear in the commands, following this format:
3333
```
@@ -260,9 +260,9 @@ public async Task<bool> ChatAsync(string input, IShell shell)
260260
{
261261
// Process CLI handler response specially to support parameter injection.
262262
ResponseData data = null;
263-
if (_copilotResponse.TopicName == CopilotActivity.CLIHandlerTopic)
263+
if (_copilotResponse.TopicName is CopilotActivity.CLIHandlerTopic or CopilotActivity.PSHandlerTopic)
264264
{
265-
data = ParseCLIHandlerResponse(shell);
265+
data = ParseCodeResponse(shell);
266266
}
267267

268268
if (data?.PlaceholderSet is not null)
@@ -349,7 +349,7 @@ public async Task<bool> ChatAsync(string input, IShell shell)
349349
return true;
350350
}
351351

352-
private ResponseData ParseCLIHandlerResponse(IShell shell)
352+
private ResponseData ParseCodeResponse(IShell shell)
353353
{
354354
string text = _copilotResponse.Text;
355355
List<CodeBlock> codeBlocks = shell.ExtractCodeBlocks(text, out List<SourceInfo> sourceInfos);
@@ -402,6 +402,7 @@ private ResponseData ParseCLIHandlerResponse(IShell shell)
402402
ResponseData data = new() {
403403
Text = text,
404404
CommandSet = commands,
405+
TopicName = _copilotResponse.TopicName,
405406
PlaceholderSet = placeholders,
406407
Locale = _copilotResponse.Locale,
407408
};

‎shell/agents/Microsoft.Azure.Agent/Command.cs

+14-10
Original file line numberDiff line numberDiff line change
@@ -33,9 +33,19 @@ private static string SyntaxHighlightAzCommand(string command, string parameter,
3333
const string vtReset = "\x1b[0m";
3434

3535
StringBuilder cStr = new(capacity: command.Length + parameter.Length + placeholder.Length + 50);
36-
cStr.Append(vtItalic)
37-
.Append(vtCommand).Append("az").Append(vtFgDefault).Append(command.AsSpan(2)).Append(' ')
38-
.Append(vtParameter).Append(parameter).Append(vtFgDefault).Append(' ')
36+
cStr.Append(vtItalic);
37+
38+
int index = command.IndexOf(' ');
39+
if (index is -1)
40+
{
41+
cStr.Append(vtCommand).Append(command).Append(vtFgDefault).Append(' ');
42+
}
43+
else
44+
{
45+
cStr.Append(vtCommand).Append(command.AsSpan(0, index)).Append(vtFgDefault).Append(command.AsSpan(index)).Append(' ');
46+
}
47+
48+
cStr.Append(vtParameter).Append(parameter).Append(vtFgDefault).Append(' ')
3949
.Append(vtVariable).Append(placeholder).Append(vtFgDefault)
4050
.Append(vtReset);
4151

@@ -125,15 +135,9 @@ private void ReplaceAction()
125135

126136
// Prompt for argument without printing captions again.
127137
string value = host.PromptForArgument(argInfo, printCaption: false);
138+
value = value?.Trim();
128139
if (!string.IsNullOrEmpty(value))
129140
{
130-
// Add quotes for the value if needed.
131-
value = value.Trim();
132-
if (value.StartsWith('-') || value.Contains(' ') || value.Contains('|'))
133-
{
134-
value = $"\"{value}\"";
135-
}
136-
137141
_values.Add(item.Name, value);
138142
_agent.SaveUserValue(item.Name, value);
139143

0 commit comments

Comments
 (0)