Skip to content

Commit

Permalink
Add telemetry collection for the azure agent (#251)
Browse files Browse the repository at this point in the history
  • Loading branch information
@NoriZC
NoriZC authored Oct 29, 2024
1 parent 16e4c40 commit da3e6d5
Show file tree
Hide file tree
Showing 6 changed files with 382 additions and 66 deletions.
33 changes: 30 additions & 3 deletions shell/agents/Microsoft.Azure.Agent/AzureAgent.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ public sealed class AzureAgent : ILLMAgent
public string SettingFile { private set; get; }

internal ArgumentPlaceholder ArgPlaceholder { set; get; }
internal CopilotResponse CopilotResponse => _copilotResponse;

private const string SettingFileName = "az.config.json";
private const string LoggingFileName = "log..txt";
Expand Down Expand Up @@ -82,6 +83,7 @@ public void Dispose()
_httpClient.Dispose();

Log.CloseAndFlush();
Telemetry.CloseAndFlush();
}

public void Initialize(AgentConfig config)
Expand Down Expand Up @@ -110,11 +112,34 @@ public void Initialize(AgentConfig config)
.CreateLogger();
Log.Information("Azure agent initialized.");
}

if (_setting.Telemetry)
{
Telemetry.Initialize();
}
}

public IEnumerable<CommandBase> GetCommands() => [new ReplaceCommand(this)];
public bool CanAcceptFeedback(UserAction action) => false;
public void OnUserAction(UserActionPayload actionPayload) {}
public bool CanAcceptFeedback(UserAction action) => Telemetry.Enabled;
public void OnUserAction(UserActionPayload actionPayload) {
// Send telemetry about the user action.
bool isUserFeedback = false;
string details = null;
UserAction action = actionPayload.Action;

if (action is UserAction.Dislike)
{
var dislike = (DislikePayload) actionPayload;
isUserFeedback = true;
details = string.Format("{0} | {1}", dislike.ShortFeedback, dislike.LongFeedback);
}
else if (action is UserAction.Like)
{
isUserFeedback = true;
}

Telemetry.Trace(AzTrace.UserAction(action.ToString(), _copilotResponse, details, isUserFeedback));
}

public async Task RefreshChatAsync(IShell shell, bool force)
{
Expand Down Expand Up @@ -254,6 +279,8 @@ public async Task<bool> ChatAsync(string input, IShell shell)
host.WriteLine("\nYou've reached the maximum length of a conversation. To continue, please run '/refresh' to start a new conversation.\n");
}
}

Telemetry.Trace(AzTrace.Chat(_copilotResponse));
}
catch (Exception ex) when (ex is TokenRequestException or ConnectionDroppedException)
{
Expand Down Expand Up @@ -362,8 +389,8 @@ private ResponseData ParseCLIHandlerResponse(IShell shell)
else
{
// The placeholder section is not in the format as we've instructed ...
// TODO: send telemetry about this case.
Log.Error("Placeholder section not in expected format:\n{0}", text);
Telemetry.Trace(AzTrace.Exception(_copilotResponse, "Placeholder section not in expected format."));
}

ReplaceKnownPlaceholders(data);
Expand Down
117 changes: 69 additions & 48 deletions shell/agents/Microsoft.Azure.Agent/ChatSession.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -87,16 +87,8 @@ internal async Task<string> RefreshAsync(IStatusContext context, bool force, Can
}
}

try
{
_token = await GenerateTokenAsync(context, cancellationToken);
return await StartConversationAsync(context, cancellationToken);
}
catch (Exception)
{
Reset();
throw;
}
_token = await GenerateTokenAsync(context, cancellationToken);
return await OpenConversationAsync(context, cancellationToken);
}

private void Reset()
Expand All @@ -113,57 +105,83 @@ private void Reset()

private async Task<string> GenerateTokenAsync(IStatusContext context, CancellationToken cancellationToken)
{
context.Status("Get Azure CLI login token ...");
// Get an access token from the AzCLI login, using the specific audience guid.
AccessToken accessToken = await new AzureCliCredential()
.GetTokenAsync(
new TokenRequestContext(["7000789f-b583-4714-ab18-aef39213018a/.default"]),
cancellationToken);

context.Status("Request for DirectLine token ...");
StringContent content = new("{\"conversationType\": \"Chat\"}", Encoding.UTF8, Utils.JsonContentType);
HttpRequestMessage request = new(HttpMethod.Post, DL_TOKEN_URL) { Content = content };
request.Headers.Authorization = new AuthenticationHeaderValue("Bearer", accessToken.Token);
try
{
context.Status("Get Azure CLI login token ...");
// Get an access token from the AzCLI login, using the specific audience guid.
AccessToken accessToken = await new AzureCliCredential()
.GetTokenAsync(
new TokenRequestContext(["7000789f-b583-4714-ab18-aef39213018a/.default"]),
cancellationToken);

context.Status("Request for DirectLine token ...");
StringContent content = new("{\"conversationType\": \"Chat\"}", Encoding.UTF8, Utils.JsonContentType);
HttpRequestMessage request = new(HttpMethod.Post, DL_TOKEN_URL) { Content = content };
request.Headers.Authorization = new AuthenticationHeaderValue("Bearer", accessToken.Token);

HttpResponseMessage response = await _httpClient.SendAsync(request, cancellationToken);
response.EnsureSuccessStatusCode();

HttpResponseMessage response = await _httpClient.SendAsync(request, cancellationToken);
response.EnsureSuccessStatusCode();
using Stream stream = await response.Content.ReadAsStreamAsync(cancellationToken);
var dlToken = JsonSerializer.Deserialize<DirectLineToken>(stream, Utils.JsonOptions);
return dlToken.DirectLine.Token;
}
catch (Exception e)
{
if (e is not OperationCanceledException)
{
Telemetry.Trace(AzTrace.Exception("Failed to generate the initial DL token."), e);
}

using Stream stream = await response.Content.ReadAsStreamAsync(cancellationToken);
var dlToken = JsonSerializer.Deserialize<DirectLineToken>(stream, Utils.JsonOptions);
return dlToken.DirectLine.Token;
Reset();
throw;
}
}

private async Task<string> StartConversationAsync(IStatusContext context, CancellationToken cancellationToken)
private async Task<string> OpenConversationAsync(IStatusContext context, CancellationToken cancellationToken)
{
context.Status("Start a new chat session ...");
HttpRequestMessage request = new(HttpMethod.Post, CONVERSATION_URL);
request.Headers.Authorization = new AuthenticationHeaderValue("Bearer", _token);
try
{
context.Status("Start a new chat session ...");
HttpRequestMessage request = new(HttpMethod.Post, CONVERSATION_URL);
request.Headers.Authorization = new AuthenticationHeaderValue("Bearer", _token);

HttpResponseMessage response = await _httpClient.SendAsync(request, cancellationToken);
response.EnsureSuccessStatusCode();
HttpResponseMessage response = await _httpClient.SendAsync(request, cancellationToken);
response.EnsureSuccessStatusCode();

using Stream content = await response.Content.ReadAsStreamAsync(cancellationToken);
SessionPayload spl = JsonSerializer.Deserialize<SessionPayload>(content, Utils.JsonOptions);
using Stream content = await response.Content.ReadAsStreamAsync(cancellationToken);
SessionPayload spl = JsonSerializer.Deserialize<SessionPayload>(content, Utils.JsonOptions);

_token = spl.Token;
_conversationId = spl.ConversationId;
_conversationUrl = $"{CONVERSATION_URL}/{_conversationId}/activities";
_streamUrl = spl.StreamUrl;
_expireOn = DateTime.UtcNow.AddSeconds(spl.ExpiresIn);
_copilotReceiver = await AzureCopilotReceiver.CreateAsync(_streamUrl);
_token = spl.Token;
_conversationId = spl.ConversationId;
_conversationUrl = $"{CONVERSATION_URL}/{_conversationId}/activities";
_streamUrl = spl.StreamUrl;
_expireOn = DateTime.UtcNow.AddSeconds(spl.ExpiresIn);
_copilotReceiver = await AzureCopilotReceiver.CreateAsync(_streamUrl);

Log.Debug("[ChatSession] Conversation started. Id: {0}", _conversationId);
Log.Debug("[ChatSession] Conversation started. Id: {0}", _conversationId);

while (true)
while (true)
{
CopilotActivity activity = _copilotReceiver.Take(cancellationToken);
if (activity.IsMessage && activity.IsFromCopilot && _copilotReceiver.Watermark is 0)
{
activity.ExtractMetadata(out _, out ConversationState conversationState);
int chatNumber = conversationState.DailyConversationNumber;
int requestNumber = conversationState.TurnNumber;
return $"{activity.Text}\nThis is chat #{chatNumber}, request #{requestNumber}.\n";
}
}
}
catch (Exception e)
{
CopilotActivity activity = _copilotReceiver.Take(cancellationToken);
if (activity.IsMessage && activity.IsFromCopilot && _copilotReceiver.Watermark is 0)
if (e is not OperationCanceledException)
{
activity.ExtractMetadata(out _, out ConversationState conversationState);
int chatNumber = conversationState.DailyConversationNumber;
int requestNumber = conversationState.TurnNumber;
return $"{activity.Text}\nThis is chat #{chatNumber}, request #{requestNumber}.\n";
Telemetry.Trace(AzTrace.Exception("Failed to open conversation with the initial DL token."), e);
}

Reset();
throw;
}
}

Expand Down Expand Up @@ -216,6 +234,7 @@ private async Task RenewTokenAsync(CancellationToken cancellationToken)
catch (Exception e) when (e is not OperationCanceledException)
{
Reset();
Telemetry.Trace(AzTrace.Exception("Failed to refresh the DL token."), e);
throw new TokenRequestException($"Failed to refresh the 'DirectLine' token: {e.Message}.", e);
}
}
Expand Down Expand Up @@ -262,6 +281,8 @@ private HttpRequestMessage PrepareForChat(string input)
var request = new HttpRequestMessage(HttpMethod.Post, _conversationUrl) { Content = content };

request.Headers.Authorization = new AuthenticationHeaderValue("Bearer", _token);
// This header is for server side telemetry to identify where the request comes from.
request.Headers.Add("ClientType", "AIShell");
return request;
}

Expand Down
13 changes: 13 additions & 0 deletions shell/agents/Microsoft.Azure.Agent/Command.cs
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
using System.CommandLine;
using System.Text;

using AIShell.Abstraction;

namespace Microsoft.Azure.Agent;
Expand Down Expand Up @@ -158,6 +159,18 @@ private void ReplaceAction()
host.RenderDivider("Regenerate", DividerAlignment.Left);
host.MarkupLine($"\nQuery: [teal]{ap.Query}[/]");

if (Telemetry.Enabled)
{
Dictionary<string, bool> details = new(items.Count);
foreach (var item in items)
{
string name = item.Name;
details.Add(name, _values.ContainsKey(name));
}

Telemetry.Trace(AzTrace.UserAction("Replace", _agent.CopilotResponse, details));
}

try
{
string answer = host.RunWithSpinnerAsync(RegenerateAsync).GetAwaiter().GetResult();
Expand Down
31 changes: 16 additions & 15 deletions shell/agents/Microsoft.Azure.Agent/DataRetriever.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -562,29 +562,30 @@ private AzCLICommand QueryForMetadata(string azCommand)
{
using var cts = new CancellationTokenSource(1200);
var response = _httpClient.Send(request, HttpCompletionOption.ResponseHeadersRead, cts.Token);
response.EnsureSuccessStatusCode();

if (response.IsSuccessStatusCode)
{
using Stream stream = response.Content.ReadAsStream(cts.Token);
using JsonDocument document = JsonDocument.Parse(stream);
using Stream stream = response.Content.ReadAsStream(cts.Token);
using JsonDocument document = JsonDocument.Parse(stream);

JsonElement root = document.RootElement;
if (root.TryGetProperty("data", out JsonElement data) &&
data.TryGetProperty("metadata", out JsonElement metadata))
{
command = metadata.Deserialize<AzCLICommand>(Utils.JsonOptions);
}
}
else
JsonElement root = document.RootElement;
if (root.TryGetProperty("data", out JsonElement data) &&
data.TryGetProperty("metadata", out JsonElement metadata))
{
// TODO: telemetry.
Log.Error("[QueryForMetadata] Received status code '{0}' for command '{1}'", response.StatusCode, azCommand);
command = metadata.Deserialize<AzCLICommand>(Utils.JsonOptions);
}
}
catch (Exception e)
{
// TODO: telemetry.
Log.Error(e, "[QueryForMetadata] Exception while processing command: {0}", azCommand);
if (Telemetry.Enabled)
{
Dictionary<string, string> details = new()
{
["Command"] = azCommand,
["Message"] = "AzCLI metadata query and process raised an exception."
};
Telemetry.Trace(AzTrace.Exception(details), e);
}
}

return command;
Expand Down
Loading

0 comments on commit da3e6d5

Please sign in to comment.