Add Variety to Az.Predictor

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

@@ -62,11 +62,12 @@ public AzPredictorServiceTests(ModelFixture fixture)
         public void VerifyUsingSuggestion(string userInput)
         {
             var predictionContext = PredictionContext.Create(userInput);
-            var expected = this._suggestionsPredictor.Query(predictionContext.InputAst, 1, CancellationToken.None);
-            var actual = this._service.GetSuggestion(predictionContext.InputAst, 1, CancellationToken.None);
+            var presentCommands = new System.Collections.Generic.Dictionary<string, int>();
+            var expected = this._suggestionsPredictor.Query(predictionContext.InputAst, presentCommands, 1, 1, CancellationToken.None);
+            var actual = this._service.GetSuggestion(predictionContext.InputAst, 1, 1, CancellationToken.None);
             Assert.NotEmpty(actual);
             Assert.NotNull(actual.First().Item1);
-            Assert.Equal(expected.First().Key, actual.First().Item1);
+            Assert.Equal(expected.Item1.First().Key, actual.First().Item1);
             Assert.Equal(PredictionSource.CurrentCommand, actual.First().Item3);
         }
 
@@ -79,11 +80,12 @@ public void VerifyUsingSuggestion(string userInput)
         public void VerifyUsingCommand(string userInput)
         {
             var predictionContext = PredictionContext.Create(userInput);
-            var expected = this._commandsPredictor.Query(predictionContext.InputAst, 1, CancellationToken.None);
-            var actual = this._service.GetSuggestion(predictionContext.InputAst, 1, CancellationToken.None);
+            var presentCommands = new System.Collections.Generic.Dictionary<string, int>();
+            var expected = this._commandsPredictor.Query(predictionContext.InputAst, presentCommands, 1, 1, CancellationToken.None);
+            var actual = this._service.GetSuggestion(predictionContext.InputAst, 1, 1, CancellationToken.None);
             Assert.NotEmpty(actual);
             Assert.NotNull(actual.First().Item1);
-            Assert.Equal(expected.First().Key, actual.First().Item1);
+            Assert.Equal(expected.Item1.First().Key, actual.First().Item1);
             Assert.Equal(PredictionSource.StaticCommands, actual.First().Item3);
         }
 
@@ -101,7 +103,7 @@ public void VerifyUsingCommand(string userInput)
         public void VerifyNoPrediction(string userInput)
         {
             var predictionContext = PredictionContext.Create(userInput);
-            var actual = this._service.GetSuggestion(predictionContext.InputAst, 1, CancellationToken.None);
+            var actual = this._service.GetSuggestion(predictionContext.InputAst, 1, 1, CancellationToken.None);
             Assert.Empty(actual);
         }
     }

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

@@ -142,7 +142,7 @@ public void VerifySupportedCommandMasked()
         public void VerifySuggestion(string userInput)
         {
             var predictionContext = PredictionContext.Create(userInput);
-            var expected = this._service.GetSuggestion(predictionContext.InputAst, 1, CancellationToken.None);
+            var expected = this._service.GetSuggestion(predictionContext.InputAst, 1, 1, CancellationToken.None);
             var actual = this._azPredictor.GetSuggestion(predictionContext, CancellationToken.None);
 
             Assert.Equal(expected.Select(e => e.Item1), actual.Select(a => a.SuggestionText));

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

@@ -50,8 +50,9 @@ public PredictorTests(ModelFixture fixture)
         public void GetNullPredictionWithCommandName(string userInput)
         {
             var predictionContext = PredictionContext.Create(userInput);
-            var result = this._predictor.Query(predictionContext.InputAst, 1, CancellationToken.None);
-            Assert.Empty(result);
+            var presentCommands = new System.Collections.Generic.Dictionary<string, int>();
+            var result = this._predictor.Query(predictionContext.InputAst, presentCommands, 1, 1, CancellationToken.None);
+            Assert.Empty(result.Item1);
         }
 
         /// <summary>
@@ -66,8 +67,9 @@ public void GetNullPredictionWithCommandName(string userInput)
         public void GetPredictionWithCommandName(string userInput)
         {
             var predictionContext = PredictionContext.Create(userInput);
-            var result = this._predictor.Query(predictionContext.InputAst, 1, CancellationToken.None);
-            Assert.NotEmpty(result);
+            var presentCommands = new System.Collections.Generic.Dictionary<string, int>();
+            var result = this._predictor.Query(predictionContext.InputAst, presentCommands, 1, 1, CancellationToken.None);
+            Assert.NotEmpty(result.Item1);
         }
 
         /// <summary>
@@ -81,8 +83,9 @@ public void GetPredictionWithCommandName(string userInput)
         public void GetPredictionWithCommandNameParameters(string userInput)
         {
             var predictionContext = PredictionContext.Create(userInput);
-            var result = this._predictor.Query(predictionContext.InputAst, 1, CancellationToken.None);
-            Assert.NotEmpty(result);
+            var presentCommands = new System.Collections.Generic.Dictionary<string, int>();
+            var result = this._predictor.Query(predictionContext.InputAst, presentCommands, 1, 1, CancellationToken.None);
+            Assert.NotEmpty(result.Item1);
         }
 
         /// <summary>
@@ -97,8 +100,9 @@ public void GetPredictionWithCommandNameParameters(string userInput)
         public void GetNullPredictionWithCommandNameParameters(string userInput)
         {
             var predictionContext = PredictionContext.Create(userInput);
-            var result = this._predictor.Query(predictionContext.InputAst, 1, CancellationToken.None);
-            Assert.Empty(result);
+            var presentCommands = new System.Collections.Generic.Dictionary<string, int>();
+            var result = this._predictor.Query(predictionContext.InputAst, presentCommands, 1, 1, CancellationToken.None);
+            Assert.Empty(result.Item1);
         }
 
         /// <summary>
@@ -108,9 +112,10 @@ public void GetNullPredictionWithCommandNameParameters(string userInput)
         public void VerifyPredictionForCommand()
         {
             var predictionContext = PredictionContext.Create("Connect-AzAccount");
-            var result = this._predictor.Query(predictionContext.InputAst, 1, CancellationToken.None);
+            var presentCommands = new System.Collections.Generic.Dictionary<string, int>();
+            var result = this._predictor.Query(predictionContext.InputAst, presentCommands, 1, 1, CancellationToken.None);
 
-            Assert.Equal("Connect-AzAccount -Credential <PSCredential> -ServicePrincipal -Tenant <>", result.First().Key);
+            Assert.Equal("Connect-AzAccount -Credential <PSCredential> -ServicePrincipal -Tenant <>", result.Item1.First().Key);
         }
 
         /// <summary>
@@ -120,9 +125,10 @@ public void VerifyPredictionForCommand()
         public void VerifyPredictionForCommandAndParameters()
         {
             var predictionContext = PredictionContext.Create("GET-AZSTORAGEACCOUNTKEY -NAME");
-            var result = this._predictor.Query(predictionContext.InputAst, 1, CancellationToken.None);
+            var presentCommands = new System.Collections.Generic.Dictionary<string, int>();
+            var result = this._predictor.Query(predictionContext.InputAst, presentCommands, 1, 1, CancellationToken.None);
 
-            Assert.Equal("Get-AzStorageAccountKey -Name 'ContosoStorage' -ResourceGroupName 'ContosoGroup02'", result.First().Key);
+            Assert.Equal("Get-AzStorageAccountKey -Name 'ContosoStorage' -ResourceGroupName 'ContosoGroup02'", result.Item1.First().Key);
         }
     }
 }

tools/Az.Tools.Predictor/Az.Tools.Predictor/AzPredictor.cs

@@ -184,7 +184,7 @@ public List<PredictiveSuggestion> GetSuggestion(PredictionContext context, Cance
             {
                 maskedUserInput = AzPredictor.MaskCommandLine(context.InputAst.FindAll((ast) => ast is CommandAst, true).LastOrDefault() as CommandAst);
 
-                suggestions = _service.GetSuggestion(context.InputAst, _settings.SuggestionCount.Value, localCancellationToken);
+                suggestions = _service.GetSuggestion(context.InputAst, _settings.SuggestionCount.Value, _settings.MaxAllowedCommandDuplicate.Value, localCancellationToken);
 
                 localCancellationToken.ThrowIfCancellationRequested();
 

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

@@ -125,14 +125,17 @@ protected virtual void Dispose(bool disposing)
         /// <remarks>
         /// Queries the Predictor with the user input if predictions are available, otherwise uses commands
         /// </remarks>
-        public IEnumerable<ValueTuple<string, string, PredictionSource>> GetSuggestion(Ast input, int suggestionCount, CancellationToken cancellationToken)
+        public IEnumerable<ValueTuple<string, string, PredictionSource>> GetSuggestion(Ast input, int suggestionCount, int maxAllowedCommandDuplicate, CancellationToken cancellationToken)
         {
             var commandSuggestions = this._commandSuggestions;
             var command = this._commandForPrediction;
 
             IList<ValueTuple<string, string, PredictionSource>> results = new List<ValueTuple<string, string, PredictionSource>>();
+            var presentCommands = new System.Collections.Generic.Dictionary<string, int>(); 
+            var resultsFromSuggestionTuple = commandSuggestions?.Item2?.Query(input, presentCommands, suggestionCount, maxAllowedCommandDuplicate, cancellationToken);
+            var resultsFromSuggestion = resultsFromSuggestionTuple.Item1;
+            presentCommands = resultsFromSuggestionTuple.Item2.ToDictionary(kvp => kvp.Key, kvp => kvp.Value);
 
-            var resultsFromSuggestion = commandSuggestions?.Item2?.Query(input, suggestionCount, cancellationToken);
 
             if (resultsFromSuggestion != null)
             {
@@ -159,7 +162,9 @@ public IEnumerable<ValueTuple<string, string, PredictionSource>> GetSuggestion(A
             if ((resultsFromSuggestion == null) || (resultsFromSuggestion.Count() < suggestionCount))
             {
                 var commands = this._commands;
-                var resultsFromCommands = commands?.Query(input, suggestionCount - resultsFromSuggestion.Count(), cancellationToken);
+                var resultsFromCommandsTuple = commands?.Query(input, presentCommands,suggestionCount - resultsFromSuggestion.Count(), maxAllowedCommandDuplicate, cancellationToken);
+                var resultsFromCommands = resultsFromCommandsTuple.Item1;
+                presentCommands = resultsFromCommandsTuple.Item2.ToDictionary(kvp => kvp.Key, kvp => kvp.Value);
 
                 if (resultsFromCommands != null)
                 {

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

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

tools/Az.Tools.Predictor/Az.Tools.Predictor/IAzPredictorService.cs

@@ -29,9 +29,10 @@ public interface IAzPredictorService
         /// </summary>
         /// <param name="input">User input from PSReadLine</param>
         /// <param name="suggestionCount">The number of suggestion to return.</param>
+        /// <param name="maxAllowedCommandDuplicate">The maximum amount of the same commnds in the list of predictions.</param>
         /// <param name="cancellationToken">The cancellation token</param>
         /// <returns>The list of suggestions for <paramref name="input"/> and the source that create the suggestion. The maximum number of suggestion is <paramref name="suggestionCount"/></returns>
-        public IEnumerable<ValueTuple<string, string, PredictionSource>> GetSuggestion(Ast input, int suggestionCount, CancellationToken cancellationToken);
+        public IEnumerable<ValueTuple<string, string, PredictionSource>> GetSuggestion(Ast input, int suggestionCount, int maxAllowedCommandDuplicate, CancellationToken cancellationToken);
 
         /// <summary>
         /// Requests predictions, given a command string.

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

@@ -57,11 +57,13 @@ public Predictor(IList<string> modelPredictions, ParameterValuePredictor paramet
         /// Given a user input PowerShell AST, returns prediction text.
         /// </summary>
         /// <param name="input">PowerShell AST input of the user, generated by PSReadLine</param>
+        /// <param name="presentCommands">Commands already present.</param>
         /// <param name="suggestionCount">The number of suggestion to return.</param>
+        /// <param name="maxAllowedCommandDuplicate">The maximum amount of the same commnds in the list of predictions.</param>
         /// <param name="cancellationToken">The cancellation token</param>
         /// <returns>The collection of suggestions. The key is the predicted text adjusted based on <paramref name="input"/>. The
         /// value is the original text to create the adjusted text. </returns>
-        public IDictionary<string, string> Query(Ast input, int suggestionCount, CancellationToken cancellationToken)
+        public Tuple<IDictionary<string, string>, IDictionary<string, int>> Query(Ast input, IDictionary<string, int> presentCommands, int suggestionCount, int maxAllowedCommandDuplicate, CancellationToken cancellationToken)
         {
             if (suggestionCount <= 0)
             {
@@ -77,6 +79,7 @@ public IDictionary<string, string> Query(Ast input, int suggestionCount, Cancell
             }
 
             var results = new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase);
+            var resultsTemp = new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase);
 
             try
             {
@@ -102,9 +105,7 @@ public IDictionary<string, string> Query(Ast input, int suggestionCount, Cancell
                 var resultBuilder = new StringBuilder();
                 var usedParams = new HashSet<int>();
                 var sourceBuilder = new StringBuilder();
-                var presentCommands = new System.Collections.Generic.Dictionary<string, int>(); // Added
-                int maxAllowedCommandDupl = 1;
-
+
                 for (var i = 0; i < _predictions.Count && results.Count < suggestionCount; ++i)
                 {
                     if (commandNameQuery(_predictions[i].Command))
@@ -148,11 +149,15 @@ public IDictionary<string, string> Query(Ast input, int suggestionCount, Cancell
                                     results.Add(prediction.ToString(), sourceBuilder.ToString());
                                     presentCommands.Add(_predictions[i].Command, 1);
                                 }
-                                else if (presentCommands[_predictions[i].Command] < maxAllowedCommandDupl)
+                                else if (presentCommands[_predictions[i].Command] < maxAllowedCommandDuplicate)
                                 {
                                     results.Add(prediction.ToString(), sourceBuilder.ToString());
                                     presentCommands[_predictions[i].Command] += 1;
                                 }
+                                else
+                                {
+                                    resultsTemp.Add(prediction.ToString(), sourceBuilder.ToString());
+                                }
 
                                 if (results.Count == suggestionCount)
                                 {
@@ -166,8 +171,11 @@ public IDictionary<string, string> Query(Ast input, int suggestionCount, Cancell
             catch
             { 
             }
-
-            return results;
+            if ((results.Count < suggestionCount) && (resultsTemp.Count >0))
+            {
+                resultsTemp.ToList().GetRange(0, suggestionCount - results.Count).ForEach(x => results.Add(x.Key,x.Value));
+            }
+            return new Tuple<IDictionary<string, string>, IDictionary<string, int>>(results, presentCommands);
         }
 
         /// <summary>

tools/Az.Tools.Predictor/Az.Tools.Predictor/Settings.cs

@@ -40,6 +40,7 @@ sealed class Settings
         /// The number of suggestions to return to PSReadLine
         /// </summary>
         public int? SuggestionCount { get; set; }
+        public int? MaxAllowedCommandDuplicate { get; set; }
 
         private static bool? _isContinueOnTimeout;
         /// <summary>
@@ -121,6 +122,11 @@ private void OverrideSettingsFromProfile()
                     {
                         this.SuggestionCount = profileSettings.SuggestionCount;
                     }
+
+                    if (profileSettings.MaxAllowedCommandDuplicate.HasValue && (profileSettings.MaxAllowedCommandDuplicate.Value > 0))
+                    {
+                        this.MaxAllowedCommandDuplicate = profileSettings.MaxAllowedCommandDuplicate;
+                    }
                 }
                 catch
                 {