Prediction/Commands Service API v2 updates

tools/Az.Tools.Predictor/Az.Tools.Predictor.Test/AzPredictorServiceTests.cs

@@ -107,8 +107,6 @@ public void VerifyParameterValues()
         [InlineData("Get-AzKeyVault -VaultName")]
         [InlineData("GET-AZSTORAGEACCOUNTKEY -NAME ")]
         [InlineData("new-azresourcegroup -name hello")]
-        [InlineData("Get-AzContext -Name")]
-        [InlineData("Get-AzContext -ErrorAction")]
         public void VerifyUsingCommandBasedPredictor(string userInput)
         {
             var predictionContext = PredictionContext.Create(userInput);
@@ -150,8 +148,8 @@ public void VerifyUsingCommandBasedPredictor(string userInput)
         /// Verifies that when no prediction is in the command based list, we'll use the fallback list.
         /// </summary>
         [Theory]
-        [InlineData("Get-AzResource -Name hello -Pre")]
-        [InlineData("Get-AzADServicePrincipal -ApplicationObject")]
+        [InlineData("New-AzApiManagementContext -ResourceGroupName hello -Serv")]
+        [InlineData("Get-AzAlert -TimeRange '1h' -Incl")]
         public void VerifyUsingFallbackPredictor(string userInput)
         {
             var predictionContext = PredictionContext.Create(userInput);

tools/Az.Tools.Predictor/Az.Tools.Predictor.Test/AzPredictorTests.cs

@@ -137,8 +137,6 @@ public void VerifySupportedCommandMasked()
         [Theory]
         [InlineData("new-azresourcegroup -name hello")]
         [InlineData("Get-AzContext -Name")]
-        [InlineData("Get-AzContext -ErrorAction")]
-        [InlineData("Get-AzADServicePrincipal -ApplicationObject")]
         public void VerifySuggestion(string userInput)
         {
             var predictionContext = PredictionContext.Create(userInput);
@@ -164,7 +162,7 @@ public void VerifySuggestionOnIncompleteCommand()
             null);
 
             var userInput = "New-AzResourceGroup -Name 'ResourceGroup01' -Location 'Central US' -WhatIf -";
-            var expected = "New-AzResourceGroup -Name 'ResourceGroup01' -Location 'Central US' -WhatIf -Verbose ***";
+            var expected = "New-AzResourceGroup -Name 'ResourceGroup01' -Location 'Central US' -WhatIf -Tag value1";
 
             var predictionContext = PredictionContext.Create(userInput);
             var actual = localAzPredictor.GetSuggestion(predictionContext, CancellationToken.None);

tools/Az.Tools.Predictor/Az.Tools.Predictor.Test/CommandLinePredictorTests.cs

@@ -169,7 +169,6 @@ public void GetPredictionWithCommandName(string userInput)
         [InlineData("Get-AzKeyVault -VaultName")]
         [InlineData("GET-AZSTORAGEACCOUNTKEY -NAME ")]
         [InlineData("new-azresourcegroup -name hello")]
-        [InlineData("Get-AzContext -Name")]
         public void GetPredictionWithCommandNameParameters(string userInput)
         {
             var predictionContext = PredictionContext.Create(userInput);
@@ -237,7 +236,7 @@ public void VerifyPredictionForCommand()
                     1,
                     CancellationToken.None);
 
-            Assert.Equal("Connect-AzAccount -Credential <PSCredential> -ServicePrincipal -Tenant <>", result.PredictiveSuggestions.First().SuggestionText);
+            Assert.Equal("Connect-AzAccount -Identity", result.PredictiveSuggestions.First().SuggestionText);
         }
 
         /// <summary>
@@ -260,7 +259,7 @@ public void VerifyPredictionForCommandAndParameters()
                     1,
                     CancellationToken.None);
 
-            Assert.Equal("Get-AzStorageAccountKey -Name 'ContosoStorage' -ResourceGroupName 'ContosoGroup02'", result.PredictiveSuggestions.First().SuggestionText);
+            Assert.Equal("Get-AzStorageAccountKey -Name 'myStorageAccount' -ResourceGroupName 'ContosoGroup02'", result.PredictiveSuggestions.First().SuggestionText);
         }
     }
 }

tools/Az.Tools.Predictor/Az.Tools.Predictor.Test/Mocks/MockAzPredictorService.cs

@@ -12,6 +12,7 @@
 // limitations under the License.
 // ----------------------------------------------------------------------------------
 
+using System;
 using System.Collections.Generic;
 
 namespace Microsoft.Azure.PowerShell.Tools.AzPredictor.Test.Mocks
@@ -32,7 +33,7 @@ sealed class MockAzPredictorService : AzPredictorService
         /// <param name="history">The history that the suggestion is for</param>
         /// <param name="suggestions">The suggestions collection</param>
         /// <param name="commands">The commands collection</param>
-        public MockAzPredictorService(string history, IList<string> suggestions, IList<string> commands)
+        public MockAzPredictorService(string history, IList<PredictiveCommand> suggestions, IList<PredictiveCommand> commands)
         {
             if (history != null)
             {

tools/Az.Tools.Predictor/Az.Tools.Predictor.Test/ModelEntry.cs

@@ -0,0 +1,63 @@
+// ----------------------------------------------------------------------------------
+//
+// Copyright Microsoft Corporation
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ----------------------------------------------------------------------------------
+
+using Newtonsoft.Json;
+using Newtonsoft.Json.Serialization;
+
+namespace Microsoft.Azure.PowerShell.Tools.AzPredictor.Test
+{
+    /// <summary>
+    /// Represents a command entry in the model files.
+    /// </summary>
+    [JsonObject(NamingStrategyType = typeof(CamelCaseNamingStrategy))]
+    public class ModelEntry
+    {
+        /// <summary>
+        /// The command in the model.
+        /// </summary>
+        [JsonProperty("suggestion", Required = Required.Always)]
+        public string Command { get; set; }
+
+        /// <summary>
+        /// The description of the command in the model.
+        /// </summary>
+        [JsonProperty(Required = Required.Always)]
+        public string Description { get; set; }
+
+        /// <summary>
+        /// The prediction count in the model.
+        /// </summary>
+        [JsonProperty("suggestion count", Required = Required.Always)]
+        public int PredictionCount { get; set; }
+
+        /// <summary>
+        /// The history count in the model.
+        /// </summary>
+        [JsonProperty("history count", Required = Required.Always)]
+        public int HistoryCount { get; set; }
+
+        /// <summary>
+        /// Transforms the model entry into the client PredictiveCommand object.
+        /// </summary>
+        /// <returns>The PredictiveCommand object used on the client.</returns>
+        public PredictiveCommand TransformEntry()
+        {
+            return new PredictiveCommand()
+            {
+                Command = this.Command,
+                Description = this.Description
+            };
+        }
+    }
+}

tools/Az.Tools.Predictor/Az.Tools.Predictor.Test/ModelFixture.cs

@@ -12,12 +12,12 @@
 // limitations under the License.
 // ----------------------------------------------------------------------------------
 
-using Microsoft.Azure.PowerShell.Tools.AzPredictor.Utilities;
+using Newtonsoft.Json;
 using System;
 using System.Collections.Generic;
 using System.IO;
 using System.IO.Compression;
-using System.Text.Json;
+using System.Linq;
 using Xunit;
 
 namespace Microsoft.Azure.PowerShell.Tools.AzPredictor.Test
@@ -32,16 +32,18 @@ public sealed class ModelFixture : IDisposable
         private const string PredictionsModelZip = "PredictionsModel.zip";
         private const string PredictionsModelJson = "PredictionsModel.json";
         private const string DataDirectoryName = "Data";
+        private static readonly Version CommandsVersionToUse = new Version("5.1.0");
+        private static readonly Version PredictionsVersionToUse = new Version("5.1.0");
 
         /// <summary>
         /// Gets a list of string for the commands.
         /// </summary>
-        public IList<string> CommandCollection { get; private set; }
+        public IList<PredictiveCommand> CommandCollection { get; private set; }
 
         /// <summary>
         /// Gets a dictionary for the predictions.
         /// </summary>
-        public IDictionary<string, IList<string>> PredictionCollection { get; private set; }
+        public IDictionary<string, IList<PredictiveCommand>> PredictionCollection { get; private set; }
 
         /// <summary>
         /// Constructs a new instance of <see cref="ModelFixture" />
@@ -52,11 +54,24 @@ public ModelFixture()
             var fileInfo = new FileInfo(currentLocation);
             var directory = fileInfo.DirectoryName;
             var dataDirectory = Path.Join(directory, ModelFixture.DataDirectoryName);
-            var commandsModel = ModelFixture.ReadZipEntry(Path.Join(dataDirectory, ModelFixture.CommandsModelZip), ModelFixture.CommandsModelJson);
-            var predictionsModel = ModelFixture.ReadZipEntry(Path.Join(dataDirectory, ModelFixture.PredictionsModelZip), ModelFixture.PredictionsModelJson);
+            var commandsModelVersions= JsonConvert.DeserializeObject<IDictionary<Version, IList<ModelEntry>>>(ModelFixture.ReadZipEntry(Path.Join(dataDirectory, ModelFixture.CommandsModelZip), ModelFixture.CommandsModelJson));
+            var predictionsModelVersions = JsonConvert.DeserializeObject<IDictionary<Version, Dictionary<string, IList<ModelEntry>>>>(ModelFixture.ReadZipEntry(Path.Join(dataDirectory, ModelFixture.PredictionsModelZip), ModelFixture.PredictionsModelJson));
 
-            this.CommandCollection = JsonSerializer.Deserialize<IList<string>>(commandsModel, JsonUtilities.DefaultSerializerOptions);
-            this.PredictionCollection = JsonSerializer.Deserialize<IDictionary<string, IList<string>>>(predictionsModel, JsonUtilities.DefaultSerializerOptions);
+            var commandsModel = commandsModelVersions[CommandsVersionToUse];
+            var predictionsModel = predictionsModelVersions[PredictionsVersionToUse];
+
+            this.CommandCollection = commandsModel.Select(x => x.TransformEntry()).ToList();
+            var predictiveCollection = new Dictionary<string, IList<PredictiveCommand>>();
+            foreach (var command in predictionsModel)
+            {
+                var predictiveCommandEntries = new List<PredictiveCommand>();
+                foreach (var modelEntry in command.Value)
+                {
+                    predictiveCommandEntries.Add(modelEntry.TransformEntry());
+                }
+                predictiveCollection.Add(command.Key, predictiveCommandEntries);
+            }
+            this.PredictionCollection = predictiveCollection;
         }
 
         /// <inheritdoc/>

tools/Az.Tools.Predictor/Az.Tools.Predictor/AzContext.cs

@@ -32,6 +32,9 @@ internal sealed class AzContext : IAzContext
     {
         private static readonly Version DefaultVersion = new Version("0.0.0.0");
 
+        /// <inheritdoc/>
+        public Version AzVersion { get; private set; } = DefaultVersion;
+
         /// <inheritdoc/>
         public string UserId { get; private set; } = string.Empty;
 
@@ -100,6 +103,7 @@ public Version ModuleVersion
         /// <inheritdoc/>
         public void UpdateContext()
         {
+            AzVersion = GetAzVersion();
             UserId = GenerateSha256HashString(GetUserAccountId());
         }
 
@@ -120,6 +124,34 @@ private string GetUserAccountId()
             return string.Empty;
         }
 
+        /// <summary>
+        /// Gets the latest version from the loaded Az modules.
+        /// </summary>
+        private Version GetAzVersion()
+        {
+            Version latestAz = DefaultVersion;
+
+            try
+            {
+                var outputs = AzContext.ExecuteScript<PSObject>("Get-Module -Name Az -ListAvailable");
+                foreach (PSObject obj in outputs)
+                {
+                    string psVersion = obj.Properties["Version"].Value.ToString();
+                    int pos = psVersion.IndexOf('-');
+                    Version currentAz = (pos == -1) ? new Version(psVersion) : new Version(psVersion.Substring(0, pos));
+                    if (currentAz > latestAz)
+                    {
+                        latestAz = currentAz;
+                    }
+                }
+            }
+            catch (Exception)
+            {
+            }
+
+            return latestAz;
+        }
+
         /// <summary>
         /// Executes the PowerShell cmdlet in the current powershell session.
         /// </summary>

tools/Az.Tools.Predictor/Az.Tools.Predictor/AzPredictorService.cs

@@ -44,11 +44,11 @@ public sealed class RequestContext
                 public Version VersionNumber{ get; set; } = new Version(0, 0);
             }
 
-            public string History { get; set; }
+            public IEnumerable<string> History { get; set; }
             public string ClientType { get; set; } = AzPredictorService.ClientType;
             public RequestContext Context { get; set; } = new RequestContext();
 
-            public PredictionRequestBody(string command) => History = command;
+            public PredictionRequestBody(IEnumerable<string> commands) => History = commands;
         };
 
         private sealed class CommandRequestContext
@@ -98,7 +98,7 @@ public AzPredictorService(string serviceUri, ITelemetryClient telemetryClient, I
             Validation.CheckArgument(telemetryClient, $"{nameof(telemetryClient)} cannot be null.");
             Validation.CheckArgument(azContext, $"{nameof(azContext)} cannot be null.");
 
-            _commandsEndpoint = $"{serviceUri}{AzPredictorConstants.CommandsEndpoint}?clientType={AzPredictorService.ClientType}&context={JsonSerializer.Serialize(new CommandRequestContext(), JsonUtilities.DefaultSerializerOptions)}";
+            _commandsEndpoint = $"{serviceUri}{AzPredictorConstants.CommandsEndpoint}?clientType={AzPredictorService.ClientType}&context.versionNumber={azContext.AzVersion}";
             _predictionsEndpoint = serviceUri + AzPredictorConstants.PredictionsEndpoint;
             _telemetryClient = telemetryClient;
             _azContext = azContext;
@@ -264,9 +264,10 @@ public virtual void RequestPredictions(IEnumerable<string> commands)
                             {
                                 SessionId = _telemetryClient.SessionId,
                                 CorrelationId = _telemetryClient.CorrelationId,
+                                VersionNumber = this._azContext.AzVersion
                             };
 
-                            var requestBody = new PredictionRequestBody(localCommands)
+                            var requestBody = new PredictionRequestBody(commands)
                             {
                                 Context = requestContext,
                             };
@@ -277,7 +278,7 @@ public virtual void RequestPredictions(IEnumerable<string> commands)
 
                             httpResponseMessage.EnsureSuccessStatusCode();
                             var reply = await httpResponseMessage.Content.ReadAsStreamAsync(cancellationToken);
-                            var suggestionsList = await JsonSerializer.DeserializeAsync<IList<string>>(reply, JsonUtilities.DefaultSerializerOptions);
+                            var suggestionsList = await JsonSerializer.DeserializeAsync<IList<PredictiveCommand>>(reply, JsonUtilities.DefaultSerializerOptions);
 
                             SetCommandBasedPreditor(localCommands, suggestionsList);
                         }
@@ -336,7 +337,7 @@ protected virtual void RequestAllPredictiveCommands()
 
                             httpResponseMessage.EnsureSuccessStatusCode();
                             var reply = await httpResponseMessage.Content.ReadAsStringAsync();
-                            var commandsReply = JsonSerializer.Deserialize<IList<string>>(reply, JsonUtilities.DefaultSerializerOptions);
+                            var commandsReply = JsonSerializer.Deserialize<IList<PredictiveCommand>>(reply, JsonUtilities.DefaultSerializerOptions);
                             SetFallbackPredictor(commandsReply);
                         }
                         catch (Exception e)
@@ -359,20 +360,20 @@ protected virtual void RequestAllPredictiveCommands()
         /// Sets the fallback predictor.
         /// </summary>
         /// <param name="commands">The command collection to set the predictor</param>
-        protected void SetFallbackPredictor(IList<string> commands)
+        protected void SetFallbackPredictor(IList<PredictiveCommand> commands)
         {
             Validation.CheckArgument(commands, $"{nameof(commands)} cannot be null.");
 
             _fallbackPredictor = new CommandLinePredictor(commands, _parameterValuePredictor);
-            _allPredictiveCommands = commands.Select(x => AzPredictorService.GetCommandName(x)).ToHashSet<string>(StringComparer.OrdinalIgnoreCase); // this could be slow
+            _allPredictiveCommands = commands.Select(x => AzPredictorService.GetCommandName(x.Command)).ToHashSet<string>(StringComparer.OrdinalIgnoreCase); // this could be slow
         }
 
         /// <summary>
         /// Sets the predictor based on the command history.
         /// </summary>
         /// <param name="commands">The commands that the suggestions are for</param>
         /// <param name="suggestions">The suggestion collection to set the predictor</param>
-        protected void SetCommandBasedPreditor(string commands, IList<string> suggestions)
+        protected void SetCommandBasedPreditor(string commands, IList<PredictiveCommand> suggestions)
         {
             Validation.CheckArgument(!string.IsNullOrWhiteSpace(commands), $"{nameof(commands)} cannot be null or whitespace.");
             Validation.CheckArgument(suggestions, $"{nameof(suggestions)} cannot be null.");

tools/Az.Tools.Predictor/Az.Tools.Predictor/AzPredictorSettings.json

@@ -1,5 +1,5 @@
 {
   "maxAllowedCommandDuplicate": 1,
-  "serviceUri": "https://app.aladdin.microsoft.com/api/v1",
+  "serviceUri": "https://app.aladdin.microsoft.com/api/v2.0",
   "suggestionCount": 7
 }

tools/Az.Tools.Predictor/Az.Tools.Predictor/CommandLine.cs

@@ -27,6 +27,11 @@ sealed class CommandLine
         /// </summary>
         public string Name { get; }
 
+        /// <summary>
+        /// Gets the description text for the command.
+        /// </summary>
+        public string Description { get; }
+
         /// <summary>
         /// Gets the <see cref="ParameterSet "/>.
         /// </summary>
@@ -36,12 +41,14 @@ sealed class CommandLine
         /// Create a new instance of <see cref="CommandLine"/> with the command name and parameter set.
         /// </summary>
         /// <param name="name">The command name.</param>
+        /// <param name="description">The command's description</param>
         /// <param name="parameterSet">The parameter set.</param>
-        public CommandLine(string name, ParameterSet parameterSet)
+        public CommandLine(string name, string description, ParameterSet parameterSet)
         {
             Validation.CheckArgument(!string.IsNullOrWhiteSpace(name), $"{nameof(name)} must not be null or whitespace.");
 
             Name = name;
+            Description = description;
             ParameterSet = parameterSet;
         }
     }