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)
             {
@@ -51,13 +52,13 @@ public MockAzPredictorService(string history, IList<string> suggestions, IList<s
         }
 
         /// <inheritdoc/>
-        public override void RequestPredictions(IEnumerable<string> history)
+        public override void RequestPredictions(IEnumerable<string> history, Version version)
         {
             this.IsPredictionRequested = true;
         }
 
         /// <inheritdoc/>
-        protected override void RequestAllPredictiveCommands()
+        protected override void RequestAllPredictiveCommands(Version version)
         {
             // Do nothing since we've set the command and suggestion predictors.
         }

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; } = new Version("0.0.0.0");
+
         /// <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,36 @@ private string GetUserAccountId()
             return string.Empty;
         }
 
+        /// <summary>
+        /// Gets the latest version from the loaded Az modules.
+        /// </summary>
+        private Version GetAzVersion()
+        {
+            Version defaultVersion = new Version("0.0.0");
+
+            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/AzPredictor.cs

@@ -118,7 +118,7 @@ public void StartEarlyProcessing(IReadOnlyList<string> history)
                 }
 
                 _telemetryClient.OnHistory(new HistoryTelemetryData(lastCommand.Item2));
-                _service.RequestPredictions(_lastTwoMaskedCommands);
+                _service.RequestPredictions(_lastTwoMaskedCommands, _azContext?.AzVersion);
             }
 
             ValueTuple<CommandAst, string> GetAstAndMaskedCommandLine(string commandLine)