Add Client side telemetry for Azure PowerShell Agent

shell/ShellCopilot.Azure.Agent/AzPS/AzPSAgent.cs

@@ -1,5 +1,6 @@
 using Azure.Identity;
 using ShellCopilot.Abstraction;
+using System.Diagnostics;
 
 namespace ShellCopilot.Azure.PowerShell;
 
@@ -21,6 +22,9 @@ public sealed class AzPSAgent : ILLMAgent
     private string _configRoot;
     private RenderingStyle _renderingStyle;
     private AzPSChatService _chatService;
+    private MetricHelper _metricHelper;
+    private List<HistoryMessage> _historyForTelemetry;
+    private Stopwatch _watch = new Stopwatch();
 
     public void Dispose()
     {
@@ -45,15 +49,66 @@ public void Initialize(AgentConfig config)
             ["Privacy statement"] = "https://aka.ms/privacy",
         };
 
+        _historyForTelemetry = [];
+        _metricHelper = new MetricHelper(AzPSChatService.Endpoint);
         _chatService = new AzPSChatService(config.IsInteractive, tenantId);
     }
 
     public IEnumerable<CommandBase> GetCommands() => null;
-    public bool CanAcceptFeedback(UserAction action) => false;
-    public void OnUserAction(UserActionPayload actionPayload) {}
+
+    public bool CanAcceptFeedback(UserAction action) => true;
+
+    public void OnUserAction(UserActionPayload actionPayload) 
+    {
+        // DisLike Action
+        string DetailedMessage = null;
+        List<HistoryMessage> history = null;
+        if (actionPayload.Action == UserAction.Dislike)
+        {
+            DislikePayload dislikePayload = (DislikePayload)actionPayload;
+            DetailedMessage = string.Format("{0} | {1}", dislikePayload.ShortFeedback, dislikePayload.LongFeedback);
+            if (dislikePayload.ShareConversation)
+            {
+                history = _historyForTelemetry;
+            }
+            else
+            {
+                _historyForTelemetry.Clear();
+            }
+        }
+        // Like Action
+        else if (actionPayload.Action == UserAction.Like)
+        {
+            LikePayload likePayload = (LikePayload)actionPayload;
+            if (likePayload.ShareConversation)
+            {
+                history = _historyForTelemetry;
+            }
+            else
+            {
+                _historyForTelemetry.Clear();
+            }
+        }
+
+        // TODO: Extract into RecrodActionTelemetry : RecordTelemetry()
+        _metricHelper.LogTelemetry(
+            new AzTrace()
+            {
+                Command = actionPayload.Action.ToString(),
+                CorrelationID = _chatService.CorrelationID,
+                EventType = "Feedback",
+                Handler = "Azure PowerShell",
+                DetailedMessage = DetailedMessage,
+                HistoryMessage = history
+            });
+    }
 
     public async Task<bool> Chat(string input, IShell shell)
     {
+        // Measure time spent
+        _watch.Restart();
+        var startTime = DateTime.Now;
+
         IHost host = shell.Host;
         CancellationToken token = shell.CancellationToken;
 
@@ -69,6 +124,7 @@ public async Task<bool> Chat(string input, IShell shell)
             if (chunkReader is null)
             {
                 // Operation was cancelled by user.
+                _watch.Stop();
                 return true;
             }
 
@@ -92,10 +148,34 @@ public async Task<bool> Chat(string input, IShell shell)
                 // Operation was cancelled by user.
             }
 
-            _chatService.AddResponseToHistory(streamingRender.AccumulatedContent);
+            string accumulatedContent = streamingRender.AccumulatedContent;
+            _chatService.AddResponseToHistory(accumulatedContent);
+
+            // Measure time spent
+            _watch.Stop();
+            // TODO: extract into RecordQuestionTelemetry() : RecordTelemetry()
+            var EndTime = DateTime.Now;
+            var Duration = TimeSpan.FromTicks(_watch.ElapsedTicks);
+
+            // Append last Q&A history in HistoryMessage
+            _historyForTelemetry.Add(new HistoryMessage("user", input, _chatService.CorrelationID));
+            _historyForTelemetry.Add(new HistoryMessage("assistant", accumulatedContent, _chatService.CorrelationID));
+
+            _metricHelper.LogTelemetry(
+                new AzTrace() {
+                    CorrelationID = _chatService.CorrelationID,
+                    Duration = Duration,
+                    EndTime = EndTime,
+                    EventType = "Question",
+                    Handler = "Azure PowerShell",
+                    StartTime = startTime
+                });
         }
         catch (RefreshTokenException ex)
         {
+            // Stop the watch in case it was not when exception happened.
+            _watch.Stop();
+
             Exception inner = ex.InnerException;
             if (inner is CredentialUnavailableException)
             {

shell/ShellCopilot.Azure.Agent/AzPS/AzPSChatService.cs

@@ -9,7 +9,7 @@ namespace ShellCopilot.Azure.PowerShell;
 
 internal class AzPSChatService : IDisposable
 {
-    private const string Endpoint = "https://azclitools-copilot.azure-api.net/azps/api/azure-powershell/copilot/streaming";
+    internal const string Endpoint = "https://azclitools-copilot.azure-api.net/azps/api/azure-powershell/copilot/streaming";
 
     private readonly bool _interactive;
     private readonly string[] _scopes;
@@ -18,6 +18,9 @@ internal class AzPSChatService : IDisposable
     private readonly AzurePowerShellCredentialOptions _credOptions;
 
     private AccessToken? _accessToken;
+    private string _correlationID;
+
+    internal string CorrelationID => _correlationID;
 
     internal AzPSChatService(bool isInteractive, string tenant)
     {
@@ -30,6 +33,7 @@ internal AzPSChatService(bool isInteractive, string tenant)
             : new() { TenantId = tenant };
 
         _accessToken = null;
+        _correlationID = null;
     }
 
     public void Dispose()
@@ -49,6 +53,12 @@ internal void AddResponseToHistory(string response)
         }
     }
 
+    private string NewCorrelationID()
+    {
+        _correlationID = Guid.NewGuid().ToString();
+        return _correlationID;
+    }
+
     private void RefreshToken(CancellationToken cancellationToken)
     {
         try
@@ -83,6 +93,11 @@ private HttpRequestMessage PrepareForChat(string input, bool streaming)
         var request = new HttpRequestMessage(HttpMethod.Post, Endpoint) { Content = content };
 
         request.Headers.Authorization = new AuthenticationHeaderValue("Bearer", _accessToken.Value.Token);
+
+        // These headers are for telemetry. We refresh correlation ID for each query.
+        request.Headers.Add("CorrelationId", NewCorrelationID());
+        request.Headers.Add("ClientType", "Copilot for client tools");
+
         return request;
     }
 

shell/ShellCopilot.Azure.Agent/ShellCopilot.Azure.Agent.csproj

@@ -17,6 +17,7 @@
 
   <ItemGroup>
     <PackageReference Include="Azure.Identity" Version="1.10.4" />
+    <PackageReference Include="Microsoft.ApplicationInsights.WorkerService" Version="2.18.0" />
     <PackageReference Include="ShellCopilot.Abstraction" Version="0.1.0-beta.7">
       <ExcludeAssets>contentFiles</ExcludeAssets>
       <PrivateAssets>all</PrivateAssets>

shell/ShellCopilot.Azure.Agent/Telemetry/AzTrace.cs

@@ -0,0 +1,81 @@
+using System.Text.Json;
+
+namespace ShellCopilot.Azure;
+
+public class AzTrace
+{
+    private static string s_installationId;
+    private static string GetInstallationID()
+    {
+        string userProfile = Environment.GetFolderPath(Environment.SpecialFolder.UserProfile);
+        string userProfilePath = Path.Combine(userProfile, ".Azure", "azureProfile.json");
+        FileStream jsonStream;
+        JsonElement array;
+        string installationID;
+
+        if (File.Exists(userProfilePath))
+        {
+            jsonStream = new FileStream(userProfilePath, FileMode.Open, FileAccess.Read);
+            array = JsonSerializer.Deserialize<JsonElement>(jsonStream);
+            installationID = array.GetProperty("installationId").GetString();
+        }
+        else
+        {
+            userProfilePath = Path.Combine(userProfile, ".Azure", "AzureRmContextSettings.json");
+            try
+            {
+                jsonStream = new FileStream(userProfilePath, FileMode.Open, FileAccess.Read);
+                array = JsonSerializer.Deserialize<JsonElement>(jsonStream);
+                installationID = array.GetProperty("Settings").GetProperty("InstallationId").GetString();
+            }
+            catch
+            {
+                // If finally no installation id found, just return null.
+                return null;
+            }
+        }
+
+        return installationID;
+    }
+
+    // "Azure PowerShell / Azure CLI"
+    public string Handler;
+    // CorrelationId from client side.
+    public string CorrelationID; 
+    // private bool _enableAzureDataCollection = null;
+    public TimeSpan? Duration;
+    public DateTime? StartTime;
+    public DateTime? EndTime;
+    public string InstallationID = s_installationId;
+    public string EventType;
+    public string Command;
+    /// <summary>
+    /// Detailed information containing additional Information - may contain:
+    /// Reason of dislike
+    /// </summary>
+    public string DetailedMessage;
+    internal List<HistoryMessage> HistoryMessage;
+    /// <summary>
+    /// Agent Information - may contain:
+    /// Handler Version
+    /// Product Version
+    /// .net/python Version
+    /// </summary>
+    public Dictionary<string, string> ExtendedProperties;
+    static AzTrace() => s_installationId = GetInstallationID();
+}
+
+// TODO: inherit from ChatMessage in PSSchema
+internal class HistoryMessage
+{
+    internal HistoryMessage(string role, string content, string correlationID)
+    {
+        Role = role;
+        Content = content;
+        CorrelationID = correlationID;
+    }
+
+    public string Role { get; }
+    public string Content { get; }
+    public string CorrelationID { get; }
+}

shell/ShellCopilot.Azure.Agent/Telemetry/MetricHelper.cs

@@ -0,0 +1,109 @@
+using Microsoft.ApplicationInsights;
+using Microsoft.ApplicationInsights.Channel;
+using Microsoft.ApplicationInsights.DataContracts;
+using Microsoft.ApplicationInsights.Extensibility;
+using Microsoft.ApplicationInsights.Extensibility.PerfCounterCollector.QuickPulse;
+using Microsoft.ApplicationInsights.WorkerService;
+using Microsoft.Extensions.DependencyInjection;
+using System.Text.Json;
+
+namespace ShellCopilot.Azure;
+
+public class MetricHelper
+{
+    private string _endpoint;
+
+    private TelemetryClient _telemetryClient;
+
+    public MetricHelper(string endpoint) 
+    {
+        _endpoint = endpoint;
+        InitializeTelemetryClient();
+    }
+
+    private void InitializeTelemetryClient()
+    {
+        // Create the DI container.
+        IServiceCollection services = new ServiceCollection();
+
+        // Add custom TelemetryInitializer
+        services.AddSingleton(typeof(ITelemetryInitializer), new MyCustomTelemetryInitializer(_endpoint));
+
+        // Configure TelemetryConfiguration
+        services.Configure<TelemetryConfiguration>(config =>
+        {
+            // Optionally configure AAD
+            // var credential = new DefaultAzureCredential();
+            // config.SetAzureTokenCredential(credential);
+        });
+
+        // Being a regular console app, there is no appsettings.json or configuration providers enabled by default.
+        // Hence connection string must be specified here.
+        services.AddApplicationInsightsTelemetryWorkerService((ApplicationInsightsServiceOptions options) => options.ConnectionString = "InstrumentationKey=c7d054ff-9f40-43e8-bf8e-7d76c58cc1af");
+
+        // Add custom TelemetryProcessor
+        services.AddApplicationInsightsTelemetryProcessor<MyCustomTelemetryProcessor>();
+
+        // Build ServiceProvider.
+        IServiceProvider serviceProvider = services.BuildServiceProvider();
+
+        // Obtain TelemetryClient instance from DI, for additional manual tracking or to flush.
+        _telemetryClient = serviceProvider.GetRequiredService<TelemetryClient>();
+    }
+
+    public void LogTelemetry(AzTrace trace)
+    {
+        Dictionary<string, string> eventProperties = new()
+        {
+            { "CorrelationID", trace.CorrelationID },
+            { "InstallationID", trace.InstallationID },
+            { "Handler", trace.Handler },
+            { "EventType", trace.EventType },
+            { "Duration", trace.Duration?.ToString() },
+            { "Command", trace.Command },
+            { "DetailedMessage", trace.DetailedMessage },
+            { "HistoryMessage", JsonSerializer.Serialize(trace.HistoryMessage) },
+            { "StartTime", trace.StartTime?.ToString() },
+            { "EndTime", trace.EndTime?.ToString() },
+        };
+
+        _telemetryClient.TrackTrace("shellCopilot", eventProperties);
+
+        // Explicitly call Flush() followed by sleep is required in Console Apps.
+        // This is to ensure that even if application terminates, telemetry is sent to the back-end.
+        _telemetryClient.Flush();
+        // Task.Delay(500000).Wait();
+    }
+}
+
+internal class MyCustomTelemetryInitializer : ITelemetryInitializer
+{
+    private string _endpoint;
+    public void Initialize(ITelemetry telemetry)
+    {
+        // Replace with actual properties.
+        (telemetry as ISupportProperties).Properties["Endpoint"] = _endpoint;
+    }
+
+    public MyCustomTelemetryInitializer(string endpoint) 
+    {
+        _endpoint = endpoint;
+    }
+}
+
+internal class MyCustomTelemetryProcessor : ITelemetryProcessor
+{
+    ITelemetryProcessor next;
+
+    public MyCustomTelemetryProcessor(ITelemetryProcessor next)
+    {
+        this.next = next;
+
+    }
+    public void Process(ITelemetry item)
+    {
+        // Example processor - not filtering out anything.
+        // This should be replaced with actual logic.
+        this.next.Process(item);
+    }
+}

shell/shell.sln

@@ -21,6 +21,14 @@ Global
 		Release|Any CPU = Release|Any CPU
 	EndGlobalSection
 	GlobalSection(ProjectConfigurationPlatforms) = postSolution
+		{1C02CDBD-D89A-4945-9FD7-4551FEE287D4}.CodeCoverage|Any CPU.ActiveCfg = Debug|Any CPU
+		{1C02CDBD-D89A-4945-9FD7-4551FEE287D4}.CodeCoverage|Any CPU.Build.0 = Debug|Any CPU
+		{1C02CDBD-D89A-4945-9FD7-4551FEE287D4}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
+		{1C02CDBD-D89A-4945-9FD7-4551FEE287D4}.Debug|Any CPU.Build.0 = Debug|Any CPU
+		{1C02CDBD-D89A-4945-9FD7-4551FEE287D4}.Linux|Any CPU.ActiveCfg = Debug|Any CPU
+		{1C02CDBD-D89A-4945-9FD7-4551FEE287D4}.Linux|Any CPU.Build.0 = Debug|Any CPU
+		{1C02CDBD-D89A-4945-9FD7-4551FEE287D4}.Release|Any CPU.ActiveCfg = Release|Any CPU
+		{1C02CDBD-D89A-4945-9FD7-4551FEE287D4}.Release|Any CPU.Build.0 = Release|Any CPU
 		{B0DAD6B9-06E4-4811-9367-60174BEDC475}.CodeCoverage|Any CPU.ActiveCfg = Debug|Any CPU
 		{B0DAD6B9-06E4-4811-9367-60174BEDC475}.CodeCoverage|Any CPU.Build.0 = Debug|Any CPU
 		{B0DAD6B9-06E4-4811-9367-60174BEDC475}.Debug|Any CPU.ActiveCfg = Debug|Any CPU