Skip to content

Commit

Permalink
Update how we send request and collect telemetry (#13327)
Browse files Browse the repository at this point in the history
* Fix the query

* Reduce the number of requests to the service.

- We request the prediction for the command history. When the command
  history isn't changed, we don't need to request the prediction again.

* Not to collect the parameter value in the telemetry.
  • Loading branch information
kceiw authored Nov 3, 2020
1 parent 4b1ff75 commit ec15095
Show file tree
Hide file tree
Showing 6 changed files with 151 additions and 68 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -66,8 +66,8 @@ public void VerifyUsingSuggestion(string userInput)
var actual = this._service.GetSuggestion(predictionContext.InputAst, 1, CancellationToken.None);
Assert.NotEmpty(actual);
Assert.NotNull(actual.First().Item1);
Assert.Equal(expected.First(), actual.First().Item1);
Assert.Equal(PredictionSource.CurrentCommand, actual.First().Item2);
Assert.Equal(expected.First().Key, actual.First().Item1);
Assert.Equal(PredictionSource.CurrentCommand, actual.First().Item3);
}

/// <summary>
Expand All @@ -83,8 +83,8 @@ public void VerifyUsingCommand(string userInput)
var actual = this._service.GetSuggestion(predictionContext.InputAst, 1, CancellationToken.None);
Assert.NotEmpty(actual);
Assert.NotNull(actual.First().Item1);
Assert.Equal(expected.First(), actual.First().Item1);
Assert.Equal(PredictionSource.StaticCommands, actual.First().Item2);
Assert.Equal(expected.First().Key, actual.First().Item1);
Assert.Equal(PredictionSource.StaticCommands, actual.First().Item3);
}

/// <summary>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,7 @@ public void VerifyPredictionForCommand()
var predictionContext = PredictionContext.Create("Connect-AzAccount");
var result = this._predictor.Query(predictionContext.InputAst, 1, CancellationToken.None);

Assert.Equal("Connect-AzAccount -Credential <PSCredential> -ServicePrincipal -Tenant <>", result.First());
Assert.Equal("Connect-AzAccount -Credential <PSCredential> -ServicePrincipal -Tenant <>", result.First().Key);
}

/// <summary>
Expand All @@ -123,7 +123,7 @@ public void VerifyPredictionForCommandAndParameters()
var predictionContext = PredictionContext.Create("GET-AZSTORAGEACCOUNTKEY -NAME");
var result = this._predictor.Query(predictionContext.InputAst, 1, CancellationToken.None);

Assert.Equal("Get-AzStorageAccountKey -Name 'ContosoStorage' -ResourceGroupName 'ContosoGroup02'", result.First());
Assert.Equal("Get-AzStorageAccountKey -Name 'ContosoStorage' -ResourceGroupName 'ContosoGroup02'", result.First().Key);
}
}
}
72 changes: 68 additions & 4 deletions tools/Az.Tools.Predictor/Az.Tools.Predictor/AzPredictor.cs
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,9 @@ internal sealed class AzPredictor : ICommandPredictor

private Queue<string> _lastTwoMaskedCommands = new Queue<string>(AzPredictorConstants.CommandHistoryCountToProcess);

// This contains the user modified texts and the original suggestion.
private Dictionary<string, string> _userAcceptedAndSuggestion = new Dictionary<string, string>();

/// <summary>
/// Constructs a new instance of <see cref="AzPredictor"/>
/// </summary>
Expand All @@ -76,6 +79,11 @@ public AzPredictor(IAzPredictorService service, ITelemetryClient telemetryClient
/// <inhericdoc />
public void StartEarlyProcessing(IReadOnlyList<string> history)
{
lock (_userAcceptedAndSuggestion)
{
_userAcceptedAndSuggestion.Clear();
}

if (history.Count > 0)
{
if (_lastTwoMaskedCommands.Any())
Expand Down Expand Up @@ -141,7 +149,20 @@ ValueTuple<CommandAst, string> GetAstAndMaskedCommandLine(string commandLine)
/// <inhericdoc />
public void OnSuggestionAccepted(string acceptedSuggestion)
{
_telemetryClient.OnSuggestionAccepted(acceptedSuggestion);
IDictionary<string, string> localSuggestedTexts = null;
lock (_userAcceptedAndSuggestion)
{
localSuggestedTexts = _userAcceptedAndSuggestion;
}

if (localSuggestedTexts.TryGetValue(acceptedSuggestion, out var suggestedText))
{
_telemetryClient.OnSuggestionAccepted(suggestedText);
}
else
{
_telemetryClient.OnSuggestionAccepted("NoRecord");
}
}

/// <inhericdoc />
Expand All @@ -154,7 +175,7 @@ public List<PredictiveSuggestion> GetSuggestion(PredictionContext context, Cance
// is prefixed with `userInput`, it should be ordered before result that is not prefixed
// with `userInput`.

IEnumerable<ValueTuple<string, PredictionSource>> suggestions = Enumerable.Empty<ValueTuple<string, PredictionSource>>();
IEnumerable<ValueTuple<string, IList<Tuple<string, string>>, PredictionSource>> suggestions = Enumerable.Empty<ValueTuple<string, IList<Tuple<string, string>>, PredictionSource>>();

try
{
Expand All @@ -178,8 +199,51 @@ public List<PredictiveSuggestion> GetSuggestion(PredictionContext context, Cance
}
finally
{
var maskedCommandLine = MaskCommandLine(context.InputAst.FindAll((ast) => ast is CommandAst, true).LastOrDefault() as CommandAst);
_telemetryClient.OnGetSuggestion(maskedCommandLine, suggestions, cancellationToken.IsCancellationRequested);
var maskedCommandLine = AzPredictor.MaskCommandLine(context.InputAst.FindAll((ast) => ast is CommandAst, true).LastOrDefault() as CommandAst);
var sb = new StringBuilder();
// This is the list of records of the original suggestion and the prediction source.
var telemetryData = new List<ValueTuple<string, PredictionSource>>();
var userAcceptedAndSuggestion = new Dictionary<string, string>();

foreach (var s in suggestions)
{
sb.Clear();
sb.Append(s.Item1.Split(' ')[0])
.Append(AzPredictorConstants.CommandParameterSeperator);

foreach (var p in s.Item2)
{
sb.Append(p.Item1);
if (p.Item2 != null)
{
sb.Append(AzPredictorConstants.CommandParameterSeperator)
.Append(p.Item2);
}

sb.Append(AzPredictorConstants.CommandParameterSeperator);
}

if (sb[sb.Length - 1] == AzPredictorConstants.CommandParameterSeperator)
{
sb.Remove(sb.Length - 1, 1);
}

var suggestedText = sb.ToString();
telemetryData.Add(ValueTuple.Create(suggestedText, s.Item3));
userAcceptedAndSuggestion[s.Item1] = suggestedText;
}

lock (_userAcceptedAndSuggestion)
{
foreach (var u in userAcceptedAndSuggestion)
{
_userAcceptedAndSuggestion[u.Key] = u.Value;
}
}

_telemetryClient.OnGetSuggestion(maskedCommandLine,
telemetryData,
cancellationToken.IsCancellationRequested);
}

return new List<PredictiveSuggestion>();
Expand Down
115 changes: 67 additions & 48 deletions tools/Az.Tools.Predictor/Az.Tools.Predictor/AzPredictorService.cs
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,8 @@ namespace Microsoft.Azure.PowerShell.Tools.AzPredictor
/// </summary>
internal class AzPredictorService : IAzPredictorService, IDisposable
{
private const string ClientType = "AzurePowerShell";

[JsonObject(NamingStrategyType = typeof(CamelCaseNamingStrategy))]
private sealed class PredictionRequestBody
{
Expand All @@ -38,16 +40,22 @@ public sealed class RequestContext
public string CorrelationId { get; set; } = Guid.Empty.ToString();
public string SessionId { get; set; } = Guid.Empty.ToString();
public string SubscriptionId { get; set; } = Guid.Empty.ToString();
public Version VersionNumber{ get; set; } = new Version(1, 0);
public Version VersionNumber{ get; set; } = new Version(0, 0);
}

public string History { get; set; }
public string ClientType { get; set; } = "AzurePowerShell";
public string ClientType { get; set; } = AzPredictorService.ClientType;
public RequestContext Context { get; set; } = new RequestContext();

public PredictionRequestBody(string command) => this.History = command;
};

[JsonObject(NamingStrategyType = typeof(CamelCaseNamingStrategy))]
private sealed class CommandRequestContext
{
public Version VersionNumber{ get; set; } = new Version(0, 0);
}

private static readonly HttpClient _client = new HttpClient();
private readonly string _commandsEndpoint;
private readonly string _predictionsEndpoint;
Expand All @@ -56,7 +64,7 @@ public sealed class RequestContext
private volatile string _commandForPrediction;
private HashSet<string> _commandSet;
private CancellationTokenSource _predictionRequestCancellationSource;
private ParameterValuePredictor _parameterValuePredictor = new ParameterValuePredictor();
private readonly ParameterValuePredictor _parameterValuePredictor = new ParameterValuePredictor();

private readonly ITelemetryClient _telemetryClient;

Expand All @@ -68,7 +76,7 @@ public sealed class RequestContext
/// <param name="telemetryClient">The telemetry client.</param>
public AzPredictorService(string serviceUri, ITelemetryClient telemetryClient)
{
this._commandsEndpoint = serviceUri + AzPredictorConstants.CommandsEndpoint;
this._commandsEndpoint = $"{serviceUri}{AzPredictorConstants.CommandsEndpoint}?clientType={AzPredictorService.ClientType}&context={JsonConvert.SerializeObject(new CommandRequestContext())}";
this._predictionsEndpoint = serviceUri + AzPredictorConstants.PredictionsEndpoint;
this._telemetryClient = telemetryClient;

Expand Down Expand Up @@ -109,16 +117,12 @@ 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, PredictionSource>> GetSuggestion(Ast input, int suggestionCount, CancellationToken cancellationToken)
public IEnumerable<ValueTuple<string, IList<Tuple<string, string>>, PredictionSource>> GetSuggestion(Ast input, int suggestionCount, CancellationToken cancellationToken)
{
var commandSuggestions = this._commandSuggestions;
var command = this._commandForPrediction;

// We've already used _commandSuggestions. There is no need to wait the request to complete at this point.
// Cancel it.
this._predictionRequestCancellationSource?.Cancel();

IList<ValueTuple<string, PredictionSource>> results = new List<ValueTuple<string, PredictionSource>>();
IList<ValueTuple<string, IList<Tuple<string, string>>, PredictionSource>> results = new List<ValueTuple<string, IList<Tuple<string, string>>, PredictionSource>>();

var resultsFromSuggestion = commandSuggestions?.Item2?.Query(input, suggestionCount, cancellationToken);

Expand All @@ -135,9 +139,12 @@ public IEnumerable<ValueTuple<string, PredictionSource>> GetSuggestion(Ast input
predictionSource = PredictionSource.PreviousCommand;
}

foreach (var r in resultsFromSuggestion)
if (resultsFromSuggestion != null)
{
results.Add(ValueTuple.Create(r, predictionSource));
foreach (var r in resultsFromSuggestion)
{
results.Add(ValueTuple.Create(r.Key, r.Value, predictionSource));
}
}
}

Expand All @@ -146,13 +153,16 @@ public IEnumerable<ValueTuple<string, PredictionSource>> GetSuggestion(Ast input
var commands = this._commands;
var resultsFromCommands = commands?.Query(input, suggestionCount - resultsFromSuggestion.Count(), cancellationToken);

resultsFromCommands?.ExceptWith(resultsFromSuggestion);

if (resultsFromCommands != null)
{
foreach (var r in resultsFromCommands)
{
results.Add(ValueTuple.Create(r, PredictionSource.StaticCommands));
if (resultsFromCommands?.ContainsKey(r.Key) == true)
{
continue;
}

results.Add(ValueTuple.Create(r.Key, r.Value, PredictionSource.StaticCommands));
}
}
}
Expand All @@ -163,45 +173,54 @@ public IEnumerable<ValueTuple<string, PredictionSource>> GetSuggestion(Ast input
/// <inheritdoc/>
public virtual void RequestPredictions(IEnumerable<string> commands)
{
// Even if it's called multiple times, we only need to keep the one for the latest command.

this._predictionRequestCancellationSource?.Cancel();
this._predictionRequestCancellationSource = new CancellationTokenSource();

var cancellationToken = this._predictionRequestCancellationSource.Token;

var localCommands= string.Join(AzPredictorConstants.CommandConcatenator, commands);
this._telemetryClient.OnRequestPrediction(localCommands);
this.SetPredictionCommand(localCommands);

// We don't need to block on the task. We send the HTTP request and update prediction list at the background.
Task.Run(async () => {
try
{
var requestContext = new PredictionRequestBody.RequestContext()
{
SessionId = this._telemetryClient.SessionId,
CorrelationId = this._telemetryClient.CorrelationId,
};
var requestBody = new PredictionRequestBody(localCommands)
{
Context = requestContext,
};
if (string.Equals(localCommands, this._commandForPrediction, StringComparison.Ordinal))
{
// It's the same history we've already requested the prediction for last time, skip it.
return;
}
else
{
this.SetPredictionCommand(localCommands);

var requestBodyString = JsonConvert.SerializeObject(requestBody);
var httpResponseMessage = await _client.PostAsync(this._predictionsEndpoint, new StringContent(requestBodyString, Encoding.UTF8, "application/json"), cancellationToken);
// When it's called multiple times, we only need to keep the one for the latest command.

var reply = await httpResponseMessage.Content.ReadAsStringAsync(cancellationToken);
var suggestionsList = JsonConvert.DeserializeObject<List<string>>(reply);
this._predictionRequestCancellationSource?.Cancel();
this._predictionRequestCancellationSource = new CancellationTokenSource();

this.SetSuggestionPredictor(localCommands, suggestionsList);
}
catch (Exception e) when (!(e is OperationCanceledException))
{
this._telemetryClient.OnRequestPredictionError(localCommands, e);
}
},
cancellationToken);
var cancellationToken = this._predictionRequestCancellationSource.Token;

// We don't need to block on the task. We send the HTTP request and update prediction list at the background.
Task.Run(async () => {
try
{
var requestContext = new PredictionRequestBody.RequestContext()
{
SessionId = this._telemetryClient.SessionId,
CorrelationId = this._telemetryClient.CorrelationId,
};
var requestBody = new PredictionRequestBody(localCommands)
{
Context = requestContext,
};

var requestBodyString = JsonConvert.SerializeObject(requestBody);
var httpResponseMessage = await _client.PostAsync(this._predictionsEndpoint, new StringContent(requestBodyString, Encoding.UTF8, "application/json"), cancellationToken);

var reply = await httpResponseMessage.Content.ReadAsStringAsync(cancellationToken);
var suggestionsList = JsonConvert.DeserializeObject<List<string>>(reply);

this.SetSuggestionPredictor(localCommands, suggestionsList);
}
catch (Exception e) when (!(e is OperationCanceledException))
{
this._telemetryClient.OnRequestPredictionError(localCommands, e);
}
},
cancellationToken);
}
}

/// <inheritdoc/>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -30,8 +30,8 @@ public interface IAzPredictorService
/// <param name="input">User input from PSReadLine</param>
/// <param name="suggestionCount">The number of suggestion to return.</param>
/// <param name="cancellationToken">The cancellation token</param>
/// <returns>The list of suggestions for <paramref name="input"/>. The maximum number of suggestion is <paramref name="suggestionCount"/></returns>
public IEnumerable<ValueTuple<string, PredictionSource>> GetSuggestion(Ast input, int suggestionCount, CancellationToken cancellationToken);
/// <returns>The list of suggestions and the parameter set that construct the suggestion for <paramref name="input"/>. The maximum number of suggestion is <paramref name="suggestionCount"/></returns>
public IEnumerable<ValueTuple<string, IList<Tuple<string, string>>, PredictionSource>> GetSuggestion(Ast input, int suggestionCount, CancellationToken cancellationToken);

/// <summary>
/// Requests predictions, given a command string.
Expand Down
Loading

0 comments on commit ec15095

Please sign in to comment.