Skip to content

Commit ec15095

Browse files
authored
Update how we send request and collect telemetry (#13327)
* 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.
1 parent 4b1ff75 commit ec15095

File tree

6 files changed

+151
-68
lines changed

6 files changed

+151
-68
lines changed

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -66,8 +66,8 @@ public void VerifyUsingSuggestion(string userInput)
6666
var actual = this._service.GetSuggestion(predictionContext.InputAst, 1, CancellationToken.None);
6767
Assert.NotEmpty(actual);
6868
Assert.NotNull(actual.First().Item1);
69-
Assert.Equal(expected.First(), actual.First().Item1);
70-
Assert.Equal(PredictionSource.CurrentCommand, actual.First().Item2);
69+
Assert.Equal(expected.First().Key, actual.First().Item1);
70+
Assert.Equal(PredictionSource.CurrentCommand, actual.First().Item3);
7171
}
7272

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

9090
/// <summary>

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -111,7 +111,7 @@ public void VerifyPredictionForCommand()
111111
var predictionContext = PredictionContext.Create("Connect-AzAccount");
112112
var result = this._predictor.Query(predictionContext.InputAst, 1, CancellationToken.None);
113113

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

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

126-
Assert.Equal("Get-AzStorageAccountKey -Name 'ContosoStorage' -ResourceGroupName 'ContosoGroup02'", result.First());
126+
Assert.Equal("Get-AzStorageAccountKey -Name 'ContosoStorage' -ResourceGroupName 'ContosoGroup02'", result.First().Key);
127127
}
128128
}
129129
}

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

Lines changed: 68 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,9 @@ internal sealed class AzPredictor : ICommandPredictor
6060

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

63+
// This contains the user modified texts and the original suggestion.
64+
private Dictionary<string, string> _userAcceptedAndSuggestion = new Dictionary<string, string>();
65+
6366
/// <summary>
6467
/// Constructs a new instance of <see cref="AzPredictor"/>
6568
/// </summary>
@@ -76,6 +79,11 @@ public AzPredictor(IAzPredictorService service, ITelemetryClient telemetryClient
7679
/// <inhericdoc />
7780
public void StartEarlyProcessing(IReadOnlyList<string> history)
7881
{
82+
lock (_userAcceptedAndSuggestion)
83+
{
84+
_userAcceptedAndSuggestion.Clear();
85+
}
86+
7987
if (history.Count > 0)
8088
{
8189
if (_lastTwoMaskedCommands.Any())
@@ -141,7 +149,20 @@ ValueTuple<CommandAst, string> GetAstAndMaskedCommandLine(string commandLine)
141149
/// <inhericdoc />
142150
public void OnSuggestionAccepted(string acceptedSuggestion)
143151
{
144-
_telemetryClient.OnSuggestionAccepted(acceptedSuggestion);
152+
IDictionary<string, string> localSuggestedTexts = null;
153+
lock (_userAcceptedAndSuggestion)
154+
{
155+
localSuggestedTexts = _userAcceptedAndSuggestion;
156+
}
157+
158+
if (localSuggestedTexts.TryGetValue(acceptedSuggestion, out var suggestedText))
159+
{
160+
_telemetryClient.OnSuggestionAccepted(suggestedText);
161+
}
162+
else
163+
{
164+
_telemetryClient.OnSuggestionAccepted("NoRecord");
165+
}
145166
}
146167

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

157-
IEnumerable<ValueTuple<string, PredictionSource>> suggestions = Enumerable.Empty<ValueTuple<string, PredictionSource>>();
178+
IEnumerable<ValueTuple<string, IList<Tuple<string, string>>, PredictionSource>> suggestions = Enumerable.Empty<ValueTuple<string, IList<Tuple<string, string>>, PredictionSource>>();
158179

159180
try
160181
{
@@ -178,8 +199,51 @@ public List<PredictiveSuggestion> GetSuggestion(PredictionContext context, Cance
178199
}
179200
finally
180201
{
181-
var maskedCommandLine = MaskCommandLine(context.InputAst.FindAll((ast) => ast is CommandAst, true).LastOrDefault() as CommandAst);
182-
_telemetryClient.OnGetSuggestion(maskedCommandLine, suggestions, cancellationToken.IsCancellationRequested);
202+
var maskedCommandLine = AzPredictor.MaskCommandLine(context.InputAst.FindAll((ast) => ast is CommandAst, true).LastOrDefault() as CommandAst);
203+
var sb = new StringBuilder();
204+
// This is the list of records of the original suggestion and the prediction source.
205+
var telemetryData = new List<ValueTuple<string, PredictionSource>>();
206+
var userAcceptedAndSuggestion = new Dictionary<string, string>();
207+
208+
foreach (var s in suggestions)
209+
{
210+
sb.Clear();
211+
sb.Append(s.Item1.Split(' ')[0])
212+
.Append(AzPredictorConstants.CommandParameterSeperator);
213+
214+
foreach (var p in s.Item2)
215+
{
216+
sb.Append(p.Item1);
217+
if (p.Item2 != null)
218+
{
219+
sb.Append(AzPredictorConstants.CommandParameterSeperator)
220+
.Append(p.Item2);
221+
}
222+
223+
sb.Append(AzPredictorConstants.CommandParameterSeperator);
224+
}
225+
226+
if (sb[sb.Length - 1] == AzPredictorConstants.CommandParameterSeperator)
227+
{
228+
sb.Remove(sb.Length - 1, 1);
229+
}
230+
231+
var suggestedText = sb.ToString();
232+
telemetryData.Add(ValueTuple.Create(suggestedText, s.Item3));
233+
userAcceptedAndSuggestion[s.Item1] = suggestedText;
234+
}
235+
236+
lock (_userAcceptedAndSuggestion)
237+
{
238+
foreach (var u in userAcceptedAndSuggestion)
239+
{
240+
_userAcceptedAndSuggestion[u.Key] = u.Value;
241+
}
242+
}
243+
244+
_telemetryClient.OnGetSuggestion(maskedCommandLine,
245+
telemetryData,
246+
cancellationToken.IsCancellationRequested);
183247
}
184248

185249
return new List<PredictiveSuggestion>();

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

Lines changed: 67 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,8 @@ namespace Microsoft.Azure.PowerShell.Tools.AzPredictor
3030
/// </summary>
3131
internal class AzPredictorService : IAzPredictorService, IDisposable
3232
{
33+
private const string ClientType = "AzurePowerShell";
34+
3335
[JsonObject(NamingStrategyType = typeof(CamelCaseNamingStrategy))]
3436
private sealed class PredictionRequestBody
3537
{
@@ -38,16 +40,22 @@ public sealed class RequestContext
3840
public string CorrelationId { get; set; } = Guid.Empty.ToString();
3941
public string SessionId { get; set; } = Guid.Empty.ToString();
4042
public string SubscriptionId { get; set; } = Guid.Empty.ToString();
41-
public Version VersionNumber{ get; set; } = new Version(1, 0);
43+
public Version VersionNumber{ get; set; } = new Version(0, 0);
4244
}
4345

4446
public string History { get; set; }
45-
public string ClientType { get; set; } = "AzurePowerShell";
47+
public string ClientType { get; set; } = AzPredictorService.ClientType;
4648
public RequestContext Context { get; set; } = new RequestContext();
4749

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

53+
[JsonObject(NamingStrategyType = typeof(CamelCaseNamingStrategy))]
54+
private sealed class CommandRequestContext
55+
{
56+
public Version VersionNumber{ get; set; } = new Version(0, 0);
57+
}
58+
5159
private static readonly HttpClient _client = new HttpClient();
5260
private readonly string _commandsEndpoint;
5361
private readonly string _predictionsEndpoint;
@@ -56,7 +64,7 @@ public sealed class RequestContext
5664
private volatile string _commandForPrediction;
5765
private HashSet<string> _commandSet;
5866
private CancellationTokenSource _predictionRequestCancellationSource;
59-
private ParameterValuePredictor _parameterValuePredictor = new ParameterValuePredictor();
67+
private readonly ParameterValuePredictor _parameterValuePredictor = new ParameterValuePredictor();
6068

6169
private readonly ITelemetryClient _telemetryClient;
6270

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

@@ -109,16 +117,12 @@ protected virtual void Dispose(bool disposing)
109117
/// <remarks>
110118
/// Queries the Predictor with the user input if predictions are available, otherwise uses commands
111119
/// </remarks>
112-
public IEnumerable<ValueTuple<string, PredictionSource>> GetSuggestion(Ast input, int suggestionCount, CancellationToken cancellationToken)
120+
public IEnumerable<ValueTuple<string, IList<Tuple<string, string>>, PredictionSource>> GetSuggestion(Ast input, int suggestionCount, CancellationToken cancellationToken)
113121
{
114122
var commandSuggestions = this._commandSuggestions;
115123
var command = this._commandForPrediction;
116124

117-
// We've already used _commandSuggestions. There is no need to wait the request to complete at this point.
118-
// Cancel it.
119-
this._predictionRequestCancellationSource?.Cancel();
120-
121-
IList<ValueTuple<string, PredictionSource>> results = new List<ValueTuple<string, PredictionSource>>();
125+
IList<ValueTuple<string, IList<Tuple<string, string>>, PredictionSource>> results = new List<ValueTuple<string, IList<Tuple<string, string>>, PredictionSource>>();
122126

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

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

138-
foreach (var r in resultsFromSuggestion)
142+
if (resultsFromSuggestion != null)
139143
{
140-
results.Add(ValueTuple.Create(r, predictionSource));
144+
foreach (var r in resultsFromSuggestion)
145+
{
146+
results.Add(ValueTuple.Create(r.Key, r.Value, predictionSource));
147+
}
141148
}
142149
}
143150

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

149-
resultsFromCommands?.ExceptWith(resultsFromSuggestion);
150-
151156
if (resultsFromCommands != null)
152157
{
153158
foreach (var r in resultsFromCommands)
154159
{
155-
results.Add(ValueTuple.Create(r, PredictionSource.StaticCommands));
160+
if (resultsFromCommands?.ContainsKey(r.Key) == true)
161+
{
162+
continue;
163+
}
164+
165+
results.Add(ValueTuple.Create(r.Key, r.Value, PredictionSource.StaticCommands));
156166
}
157167
}
158168
}
@@ -163,45 +173,54 @@ public IEnumerable<ValueTuple<string, PredictionSource>> GetSuggestion(Ast input
163173
/// <inheritdoc/>
164174
public virtual void RequestPredictions(IEnumerable<string> commands)
165175
{
166-
// Even if it's called multiple times, we only need to keep the one for the latest command.
167-
168-
this._predictionRequestCancellationSource?.Cancel();
169-
this._predictionRequestCancellationSource = new CancellationTokenSource();
170-
171-
var cancellationToken = this._predictionRequestCancellationSource.Token;
172-
173176
var localCommands= string.Join(AzPredictorConstants.CommandConcatenator, commands);
174177
this._telemetryClient.OnRequestPrediction(localCommands);
175-
this.SetPredictionCommand(localCommands);
176178

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

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

194-
var reply = await httpResponseMessage.Content.ReadAsStringAsync(cancellationToken);
195-
var suggestionsList = JsonConvert.DeserializeObject<List<string>>(reply);
190+
this._predictionRequestCancellationSource?.Cancel();
191+
this._predictionRequestCancellationSource = new CancellationTokenSource();
196192

197-
this.SetSuggestionPredictor(localCommands, suggestionsList);
198-
}
199-
catch (Exception e) when (!(e is OperationCanceledException))
200-
{
201-
this._telemetryClient.OnRequestPredictionError(localCommands, e);
202-
}
203-
},
204-
cancellationToken);
193+
var cancellationToken = this._predictionRequestCancellationSource.Token;
194+
195+
// We don't need to block on the task. We send the HTTP request and update prediction list at the background.
196+
Task.Run(async () => {
197+
try
198+
{
199+
var requestContext = new PredictionRequestBody.RequestContext()
200+
{
201+
SessionId = this._telemetryClient.SessionId,
202+
CorrelationId = this._telemetryClient.CorrelationId,
203+
};
204+
var requestBody = new PredictionRequestBody(localCommands)
205+
{
206+
Context = requestContext,
207+
};
208+
209+
var requestBodyString = JsonConvert.SerializeObject(requestBody);
210+
var httpResponseMessage = await _client.PostAsync(this._predictionsEndpoint, new StringContent(requestBodyString, Encoding.UTF8, "application/json"), cancellationToken);
211+
212+
var reply = await httpResponseMessage.Content.ReadAsStringAsync(cancellationToken);
213+
var suggestionsList = JsonConvert.DeserializeObject<List<string>>(reply);
214+
215+
this.SetSuggestionPredictor(localCommands, suggestionsList);
216+
}
217+
catch (Exception e) when (!(e is OperationCanceledException))
218+
{
219+
this._telemetryClient.OnRequestPredictionError(localCommands, e);
220+
}
221+
},
222+
cancellationToken);
223+
}
205224
}
206225

207226
/// <inheritdoc/>

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,8 +30,8 @@ public interface IAzPredictorService
3030
/// <param name="input">User input from PSReadLine</param>
3131
/// <param name="suggestionCount">The number of suggestion to return.</param>
3232
/// <param name="cancellationToken">The cancellation token</param>
33-
/// <returns>The list of suggestions for <paramref name="input"/>. The maximum number of suggestion is <paramref name="suggestionCount"/></returns>
34-
public IEnumerable<ValueTuple<string, PredictionSource>> GetSuggestion(Ast input, int suggestionCount, CancellationToken cancellationToken);
33+
/// <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>
34+
public IEnumerable<ValueTuple<string, IList<Tuple<string, string>>, PredictionSource>> GetSuggestion(Ast input, int suggestionCount, CancellationToken cancellationToken);
3535

3636
/// <summary>
3737
/// Requests predictions, given a command string.

0 commit comments

Comments
 (0)