Add local parameter value map to Az.Predictor

tools/Az.Tools.Predictor/Az.Tools.Predictor/Az.Tools.Predictor.csproj

@@ -43,5 +43,6 @@ For more information on Az Predictor, please visit the following: https://aka.ms
   <ItemGroup>
     <None Include="Az.Tools.Predictor.psd1" CopyToOutputDirectory="PreserveNewest" />
     <None Include="AzPredictorSettings.json" CopyToOutputDirectory="PreserveNewest" />
+    <None Include="command_param_to_resource_map.json" CopyToOutputDirectory="PreserveNewest" />  
   </ItemGroup>
 </Project>

tools/Az.Tools.Predictor/Az.Tools.Predictor/CommandLinePredictor.cs

@@ -122,10 +122,11 @@ public CommandLineSuggestion GetSuggestion(string inputCommandName,
                     resultBuilder.Clear();
                     resultBuilder.Append(_commandLinePredictions[i].Name);
                     usedParams.Clear();
+                    string commandNoun = ParameterValuePredictor.GetAzCommandNoun(_commandLinePredictions[i].Name).ToLower();
 
-                    if (DoesPredictionParameterSetMatchInput(resultBuilder, inputParameterSet, _commandLinePredictions[i].ParameterSet, usedParams))
+                    if (DoesPredictionParameterSetMatchInput(resultBuilder, inputParameterSet, commandNoun, _commandLinePredictions[i].ParameterSet, usedParams))
                     {
-                        PredictRestOfParameters(resultBuilder, _commandLinePredictions[i].ParameterSet.Parameters, usedParams);
+                        PredictRestOfParameters(resultBuilder, commandNoun, _commandLinePredictions[i].ParameterSet.Parameters, usedParams);
 
                         if (resultBuilder.Length <= rawUserInput.Length)
                         {
@@ -177,15 +178,16 @@ public CommandLineSuggestion GetSuggestion(string inputCommandName,
         /// Appends unused parameters to the builder.
         /// </summary>
         /// <param name="builder">StringBuilder that aggregates the prediction text output.</param>
+        /// <param name="commandNoun">Command Noun.</param>
         /// <param name="parameters">Chosen prediction parameters.</param>
         /// <param name="usedParams">Set of used parameters for set.</param>
-        private void PredictRestOfParameters(StringBuilder builder, IReadOnlyList<Parameter> parameters, HashSet<int> usedParams)
+        private void PredictRestOfParameters(StringBuilder builder, string commandNoun, IReadOnlyList<Parameter> parameters, HashSet<int> usedParams)
         {
             for (var j = 0; j < parameters.Count; j++)
             {
                 if (!usedParams.Contains(j))
                 {
-                    BuildParameterValue(builder, parameters[j]);
+                    BuildParameterValue(builder, commandNoun, parameters[j]);
                 }
             }
         }
@@ -195,9 +197,10 @@ private void PredictRestOfParameters(StringBuilder builder, IReadOnlyList<Parame
         /// </summary>
         /// <param name="builder">StringBuilder that aggregates the prediction text output.</param>
         /// <param name="inputParameters">Parsed ParameterSet from the user input AST.</param>
+        /// <param name="commandNoun">Command Noun.</param>
         /// <param name="predictionParameters">Candidate prediction parameter set.</param>
         /// <param name="usedParams">Set of used parameters for set.</param>
-        private bool DoesPredictionParameterSetMatchInput(StringBuilder builder, ParameterSet inputParameters, ParameterSet predictionParameters, HashSet<int> usedParams)
+        private bool DoesPredictionParameterSetMatchInput(StringBuilder builder, ParameterSet inputParameters, string commandNoun,ParameterSet predictionParameters, HashSet<int> usedParams)
         {
             foreach (var inputParameter in inputParameters.Parameters)
             {
@@ -215,7 +218,7 @@ private bool DoesPredictionParameterSetMatchInput(StringBuilder builder, Paramet
                     }
                     else
                     {
-                        BuildParameterValue(builder, predictionParameters.Parameters[matchIndex]);
+                        BuildParameterValue(builder, commandNoun, predictionParameters.Parameters[matchIndex]);
                     }
                 }
             }
@@ -234,11 +237,12 @@ private bool DoesPredictionParameterSetMatchInput(StringBuilder builder, Paramet
         /// "TestVM" is predicted for Get-AzVM.
         /// </summary>
         /// <param name="builder">The string builder to create the whole predicted command line.</param>
+        /// <param name="commandNoun">Command Noun.</param>
         /// <param name="parameter">The parameter name and value from prediction.</param>
-        private void BuildParameterValue(StringBuilder builder, Parameter parameter)
+        private void BuildParameterValue(StringBuilder builder, string commandNoun, Parameter parameter)
         {
             var parameterName = parameter.Name;
-            var parameterValue = this._parameterValuePredictor?.GetParameterValueFromAzCommand(parameterName);
+            string parameterValue = this._parameterValuePredictor?.GetParameterValueFromAzCommand(commandNoun, parameterName);
 
             if (string.IsNullOrWhiteSpace(parameterValue))
             {

tools/Az.Tools.Predictor/Az.Tools.Predictor/ParameterValuePredictor.cs

@@ -1,23 +1,35 @@
-using System;
+using Microsoft.Azure.PowerShell.Tools.AzPredictor.Utilities;
+using System;
+using System.IO;
+using System.Text;
+using System.Text.Json;
+using System.Text.Json.Serialization;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Linq;
 using System.Management.Automation.Language;
 
+
+
 namespace Microsoft.Azure.PowerShell.Tools.AzPredictor
 {
     /// <summary>
     /// A predictor to learn and provide values for Azure PowerShell commands' parameters values.
     /// </summary>
     sealed class ParameterValuePredictor
     {
-        /// <summary>
-        /// The collections of the parameter names that is used directly as the key in local parameter collection.
-        /// </summary>
-        private static readonly IReadOnlyCollection<string> _specialLocalParameterNames = new HashSet<string>(StringComparer.OrdinalIgnoreCase) { "location", "credential", "addressprefix" };
-
         private readonly ConcurrentDictionary<string, string> _localParameterValues = new ConcurrentDictionary<string, string>();
 
+        private readonly Dictionary<string, Dictionary<string, string>> _command_param_to_resource_map;
+
+        public ParameterValuePredictor()
+        {
+            var fileInfo = new FileInfo(typeof(Settings).Assembly.Location);
+            var directory = fileInfo.DirectoryName;
+            var mappingFilePath = Path.Join(directory, "command_param_to_resource_map.json");
+            _command_param_to_resource_map = JsonSerializer.Deserialize<Dictionary<string, Dictionary<string, string>>>(File.ReadAllText(mappingFilePath), JsonUtilities.DefaultSerializerOptions);
+        }
+
         /// <summary>
         /// Process the command from history
         /// </summary>
@@ -38,30 +50,28 @@ public void ProcessHistoryCommand(CommandAst command)
         /// > Get-AzVM -VMName &lt;TestVM&gt;
         /// "TestVM" is predicted for Get-AzVM.
         /// </summary>
+        /// <param name="commandNoun">The command noun</param>
         /// <param name="parameterName">The parameter name</param>
         /// <returns>The parameter value from the history command. Null if that is not available.</returns>
-        public string GetParameterValueFromAzCommand(string parameterName)
+        public string GetParameterValueFromAzCommand(string commandNoun, string parameterName)
         {
-            if (_localParameterValues.TryGetValue(parameterName.ToUpper(), out var value))
+            if (_command_param_to_resource_map.ContainsKey(commandNoun))
             {
-                return value;
+                parameterName = parameterName.ToLower();
+                if (_command_param_to_resource_map[commandNoun].ContainsKey(parameterName))
+                {
+                    var key = _command_param_to_resource_map[commandNoun][parameterName];
+                    if (_localParameterValues.TryGetValue(key, out var value))
+                    {
+                        return value;
+                    }
+                }
             }
-
             return null;
         }
 
-        /// <summary>
-        /// Gets the key to the local parameter dictionary from the command noun and the parameter name.
-        /// </summary>
-        /// <param name="commandNoun">The noun in the PowerShell command, e.g. the noun for command New-AzVM is VM.</param>
-        /// <param name="parameterName">The command's parameter name, e.g. "New-AzVM -Name" the parameter name is Name</param>
-        /// <returns></returns>
-        private static string GetLocalParameterKey(string commandNoun, string parameterName)
-        {
-            return _specialLocalParameterNames.Contains(parameterName) ? parameterName.ToUpper() : string.Concat(commandNoun, parameterName).ToUpper();
-        }
-
-        private static string GetAzCommandNoun(string commandName)
+
+        public static string GetAzCommandNoun(string commandName)
         {
             var monikerIndex = commandName?.IndexOf(AzPredictorConstants.AzCommandMoniker, StringComparison.OrdinalIgnoreCase);
 
@@ -91,7 +101,7 @@ private void ExtractLocalParameters(System.Collections.ObjectModel.ReadOnlyColle
             // We need to extract the noun to construct the parameter name.
 
             var commandName = command.FirstOrDefault()?.ToString();
-            var commandNoun = ParameterValuePredictor.GetAzCommandNoun(commandName);
+            var commandNoun = ParameterValuePredictor.GetAzCommandNoun(commandName).ToLower();
             if (commandNoun == null)
             {
                 return;
@@ -101,11 +111,17 @@ private void ExtractLocalParameters(System.Collections.ObjectModel.ReadOnlyColle
             {
                 if (command[i - 1] is CommandParameterAst parameterAst && command[i] is StringConstantExpressionAst)
                 {
-                    var parameterName = parameterAst.ParameterName;
-                    var key = ParameterValuePredictor.GetLocalParameterKey(commandNoun, parameterName);
-                    var parameterValue = command[i].ToString();
-                    this._localParameterValues.AddOrUpdate(key, parameterValue, (k, v) => parameterValue);
-                }
+                    var parameterName = command[i - 1].ToString().ToLower().Trim('-');
+                    if (_command_param_to_resource_map.ContainsKey(commandNoun))
+                    {
+                        if (_command_param_to_resource_map[commandNoun].ContainsKey(parameterName))
+                        {
+                            var key = _command_param_to_resource_map[commandNoun][parameterName];
+                            var parameterValue = command[i].ToString();
+                            this._localParameterValues.AddOrUpdate(key, parameterValue, (k, v) => parameterValue);
+                        }
+                    }
+                }   
             }
         }
     }