|
| 1 | +using System.Diagnostics; |
| 2 | +using System.Text; |
| 3 | +using System.Text.Json; |
| 4 | +using Azure.Identity; |
| 5 | +using AIShell.Abstraction; |
| 6 | + |
| 7 | +namespace AIShell.Azure.CLI; |
| 8 | + |
| 9 | +public sealed class AzCLIAgent : ILLMAgent |
| 10 | +{ |
| 11 | + public string Name => "az-cli"; |
| 12 | + public string Description => "This AI assistant can help generate Azure CLI scripts or commands for managing Azure resources and end-to-end scenarios that involve multiple different Azure resources."; |
| 13 | + public string Company => "Microsoft"; |
| 14 | + public List<string> SampleQueries => [ |
| 15 | + "Create a VM with a public IP address", |
| 16 | + "How to create a web app?", |
| 17 | + "Backup an Azure SQL database to a storage container" |
| 18 | + ]; |
| 19 | + public Dictionary<string, string> LegalLinks { private set; get; } = null; |
| 20 | + public string SettingFile { private set; get; } = null; |
| 21 | + internal ArgumentPlaceholder ArgPlaceholder { set; get; } |
| 22 | + internal UserValueStore ValueStore { get; } = new(); |
| 23 | + |
| 24 | + private const string SettingFileName = "az-cli.agent.json"; |
| 25 | + private readonly Stopwatch _watch = new(); |
| 26 | + |
| 27 | + private AzCLIChatService _chatService; |
| 28 | + private StringBuilder _text; |
| 29 | + private MetricHelper _metricHelper; |
| 30 | + private LinkedList<HistoryMessage> _historyForTelemetry; |
| 31 | + |
| 32 | + public void Dispose() |
| 33 | + { |
| 34 | + _chatService?.Dispose(); |
| 35 | + } |
| 36 | + |
| 37 | + public void Initialize(AgentConfig config) |
| 38 | + { |
| 39 | + _text = new StringBuilder(); |
| 40 | + _chatService = new AzCLIChatService(); |
| 41 | + _historyForTelemetry = []; |
| 42 | + _metricHelper = new MetricHelper(AzCLIChatService.Endpoint); |
| 43 | + |
| 44 | + LegalLinks = new(StringComparer.OrdinalIgnoreCase) |
| 45 | + { |
| 46 | + ["Terms"] = "https://aka.ms/TermsofUseCopilot", |
| 47 | + ["Privacy"] = "https://aka.ms/privacy", |
| 48 | + ["FAQ"] = "https://aka.ms/CopilotforAzureClientToolsFAQ", |
| 49 | + ["Transparency"] = "https://aka.ms/CopilotAzCLIPSTransparency", |
| 50 | + }; |
| 51 | + |
| 52 | + SettingFile = Path.Combine(config.ConfigurationRoot, SettingFileName); |
| 53 | + } |
| 54 | + |
| 55 | + public IEnumerable<CommandBase> GetCommands() => [new ReplaceCommand(this)]; |
| 56 | + |
| 57 | + public bool CanAcceptFeedback(UserAction action) => !MetricHelper.TelemetryOptOut; |
| 58 | + |
| 59 | + public void OnUserAction(UserActionPayload actionPayload) |
| 60 | + { |
| 61 | + // Send telemetry about the user action. |
| 62 | + // DisLike Action |
| 63 | + string DetailedMessage = null; |
| 64 | + LinkedList<HistoryMessage> history = null; |
| 65 | + if (actionPayload.Action == UserAction.Dislike) |
| 66 | + { |
| 67 | + DislikePayload dislikePayload = (DislikePayload)actionPayload; |
| 68 | + DetailedMessage = string.Format("{0} | {1}", dislikePayload.ShortFeedback, dislikePayload.LongFeedback); |
| 69 | + if (dislikePayload.ShareConversation) |
| 70 | + { |
| 71 | + history = _historyForTelemetry; |
| 72 | + } |
| 73 | + else |
| 74 | + { |
| 75 | + _historyForTelemetry.Clear(); |
| 76 | + } |
| 77 | + } |
| 78 | + // Like Action |
| 79 | + else if (actionPayload.Action == UserAction.Like) |
| 80 | + { |
| 81 | + LikePayload likePayload = (LikePayload)actionPayload; |
| 82 | + if (likePayload.ShareConversation) |
| 83 | + { |
| 84 | + history = _historyForTelemetry; |
| 85 | + } |
| 86 | + else |
| 87 | + { |
| 88 | + _historyForTelemetry.Clear(); |
| 89 | + } |
| 90 | + } |
| 91 | + |
| 92 | + _metricHelper.LogTelemetry( |
| 93 | + new AzTrace() |
| 94 | + { |
| 95 | + Command = actionPayload.Action.ToString(), |
| 96 | + CorrelationID = _chatService.CorrelationID, |
| 97 | + EventType = "Feedback", |
| 98 | + Handler = "Azure CLI", |
| 99 | + DetailedMessage = DetailedMessage, |
| 100 | + HistoryMessage = history |
| 101 | + }); |
| 102 | + } |
| 103 | + |
| 104 | + public Task RefreshChatAsync(IShell shell, bool force) |
| 105 | + { |
| 106 | + // Reset the history so the subsequent chat can start fresh. |
| 107 | + _chatService.ChatHistory.Clear(); |
| 108 | + ArgPlaceholder = null; |
| 109 | + ValueStore.Clear(); |
| 110 | + |
| 111 | + return Task.CompletedTask; |
| 112 | + } |
| 113 | + |
| 114 | + public async Task<bool> ChatAsync(string input, IShell shell) |
| 115 | + { |
| 116 | + // Measure time spent |
| 117 | + _watch.Restart(); |
| 118 | + var startTime = DateTime.Now; |
| 119 | + |
| 120 | + IHost host = shell.Host; |
| 121 | + CancellationToken token = shell.CancellationToken; |
| 122 | + |
| 123 | + try |
| 124 | + { |
| 125 | + AzCliResponse azResponse = await host.RunWithSpinnerAsync( |
| 126 | + status: "Thinking ...", |
| 127 | + func: async context => await _chatService.GetChatResponseAsync(context, input, token) |
| 128 | + ).ConfigureAwait(false); |
| 129 | + |
| 130 | + if (azResponse is not null) |
| 131 | + { |
| 132 | + if (azResponse.Error is not null) |
| 133 | + { |
| 134 | + host.WriteLine($"\n{azResponse.Error}\n"); |
| 135 | + return true; |
| 136 | + } |
| 137 | + |
| 138 | + ResponseData data = azResponse.Data; |
| 139 | + AddMessageToHistory( |
| 140 | + JsonSerializer.Serialize(data, Utils.JsonOptions), |
| 141 | + fromUser: false); |
| 142 | + |
| 143 | + string answer = GenerateAnswer(input, data); |
| 144 | + host.RenderFullResponse(answer); |
| 145 | + |
| 146 | + // Measure time spent |
| 147 | + _watch.Stop(); |
| 148 | + |
| 149 | + if (!MetricHelper.TelemetryOptOut) |
| 150 | + { |
| 151 | + // TODO: extract into RecordQuestionTelemetry() : RecordTelemetry() |
| 152 | + var EndTime = DateTime.Now; |
| 153 | + var Duration = TimeSpan.FromTicks(_watch.ElapsedTicks); |
| 154 | + |
| 155 | + // Append last Q&A history in HistoryMessage |
| 156 | + _historyForTelemetry.AddLast(new HistoryMessage("user", input, _chatService.CorrelationID)); |
| 157 | + _historyForTelemetry.AddLast(new HistoryMessage("assistant", answer, _chatService.CorrelationID)); |
| 158 | + |
| 159 | + _metricHelper.LogTelemetry( |
| 160 | + new AzTrace() |
| 161 | + { |
| 162 | + CorrelationID = _chatService.CorrelationID, |
| 163 | + Duration = Duration, |
| 164 | + EndTime = EndTime, |
| 165 | + EventType = "Question", |
| 166 | + Handler = "Azure CLI", |
| 167 | + StartTime = startTime |
| 168 | + }); |
| 169 | + } |
| 170 | + } |
| 171 | + } |
| 172 | + catch (RefreshTokenException ex) |
| 173 | + { |
| 174 | + Exception inner = ex.InnerException; |
| 175 | + if (inner is CredentialUnavailableException) |
| 176 | + { |
| 177 | + host.WriteErrorLine($"Access token not available. Query cannot be served."); |
| 178 | + host.WriteErrorLine($"The '{Name}' agent depends on the Azure CLI credential to acquire access token. Please run 'az login' from a command-line shell to setup account."); |
| 179 | + } |
| 180 | + else |
| 181 | + { |
| 182 | + host.WriteErrorLine($"Failed to get the access token. {inner.Message}"); |
| 183 | + } |
| 184 | + |
| 185 | + return false; |
| 186 | + } |
| 187 | + finally |
| 188 | + { |
| 189 | + // Stop the watch in case of early return or exception. |
| 190 | + _watch.Stop(); |
| 191 | + } |
| 192 | + |
| 193 | + return true; |
| 194 | + } |
| 195 | + |
| 196 | + internal string GenerateAnswer(string input, ResponseData data) |
| 197 | + { |
| 198 | + _text.Clear(); |
| 199 | + _text.Append(data.Description).Append("\n\n"); |
| 200 | + |
| 201 | + // We keep 'ArgPlaceholder' unchanged when it's re-generating in '/replace' with only partial placeholders replaced. |
| 202 | + if (!ReferenceEquals(ArgPlaceholder?.ResponseData, data) || data.PlaceholderSet is null) |
| 203 | + { |
| 204 | + ArgPlaceholder?.DataRetriever?.Dispose(); |
| 205 | + ArgPlaceholder = null; |
| 206 | + } |
| 207 | + |
| 208 | + if (data.CommandSet.Count > 0) |
| 209 | + { |
| 210 | + // AzCLI handler incorrectly include pseudo values in the placeholder set, so we need to filter them out. |
| 211 | + UserValueStore.FilterOutPseudoValues(data); |
| 212 | + if (data.PlaceholderSet?.Count > 0) |
| 213 | + { |
| 214 | + // Create the data retriever for the placeholders ASAP, so it gets |
| 215 | + // more time to run in background. |
| 216 | + ArgPlaceholder ??= new ArgumentPlaceholder(input, data); |
| 217 | + } |
| 218 | + |
| 219 | + for (int i = 0; i < data.CommandSet.Count; i++) |
| 220 | + { |
| 221 | + CommandItem action = data.CommandSet[i]; |
| 222 | + // Replace the pseudo values with the real values. |
| 223 | + string script = ValueStore.ReplacePseudoValues(action.Script); |
| 224 | + |
| 225 | + _text.Append($"{i+1}. {action.Desc}") |
| 226 | + .Append("\n\n") |
| 227 | + .Append("```sh\n") |
| 228 | + .Append($"# {action.Desc}\n") |
| 229 | + .Append(script).Append('\n') |
| 230 | + .Append("```\n\n"); |
| 231 | + } |
| 232 | + |
| 233 | + if (ArgPlaceholder is not null) |
| 234 | + { |
| 235 | + _text.Append("Please provide values for the following placeholder variables:\n\n"); |
| 236 | + |
| 237 | + for (int i = 0; i < data.PlaceholderSet.Count; i++) |
| 238 | + { |
| 239 | + PlaceholderItem item = data.PlaceholderSet[i]; |
| 240 | + _text.Append($"- `{item.Name}`: {item.Desc}\n"); |
| 241 | + } |
| 242 | + |
| 243 | + _text.Append("\nRun `/replace` to get assistance in placeholder replacement.\n"); |
| 244 | + } |
| 245 | + } |
| 246 | + |
| 247 | + return _text.ToString(); |
| 248 | + } |
| 249 | + |
| 250 | + internal void AddMessageToHistory(string message, bool fromUser) |
| 251 | + { |
| 252 | + if (!string.IsNullOrEmpty(message)) |
| 253 | + { |
| 254 | + var history = _chatService.ChatHistory; |
| 255 | + while (history.Count > Utils.HistoryCount - 1) |
| 256 | + { |
| 257 | + history.RemoveAt(0); |
| 258 | + } |
| 259 | + |
| 260 | + history.Add(new ChatMessage() |
| 261 | + { |
| 262 | + Role = fromUser ? "user" : "assistant", |
| 263 | + Content = message |
| 264 | + }); |
| 265 | + } |
| 266 | + } |
| 267 | +} |
0 commit comments