1  // Copyright (c) Microsoft Corporation
 2  // The Microsoft Corporation licenses this file to you under the MIT license.
 3  // See the LICENSE file in the project root for more information.
 4  
 5  using AdvancedPaste.Models;
 6  using Microsoft.SemanticKernel;
 7  
 8  namespace AdvancedPaste.Helpers;
 9  
10  /// <summary>
11  /// Helper class for extracting AI service usage information from chat messages.
12  /// </summary>
13  public static class AIServiceUsageHelper
14  {
15      /// <summary>
16      /// Extracts AI service usage information from OpenAI chat message metadata.
17      /// </summary>
18      /// <param name="chatMessage">The chat message containing usage metadata.</param>
19      /// <returns>AI service usage information or AIServiceUsage.None if extraction fails.</returns>
20      public static AIServiceUsage GetOpenAIServiceUsage(ChatMessageContent chatMessage)
21      {
22          // Try to get usage information from metadata
23          if (chatMessage.Metadata?.TryGetValue("Usage", out var usageObj) == true)
24          {
25              // Handle different possible usage types through reflection to be version-agnostic
26              var usageType = usageObj.GetType();
27  
28              try
29              {
30                  // Try common property names for prompt tokens
31                  var promptTokensProp = usageType.GetProperty("PromptTokens") ??
32                                       usageType.GetProperty("InputTokens") ??
33                                       usageType.GetProperty("InputTokenCount");
34  
35                  var completionTokensProp = usageType.GetProperty("CompletionTokens") ??
36                                           usageType.GetProperty("OutputTokens") ??
37                                           usageType.GetProperty("OutputTokenCount");
38  
39                  if (promptTokensProp != null && completionTokensProp != null)
40                  {
41                      var promptTokens = (int)(promptTokensProp.GetValue(usageObj) ?? 0);
42                      var completionTokens = (int)(completionTokensProp.GetValue(usageObj) ?? 0);
43                      return new AIServiceUsage(promptTokens, completionTokens);
44                  }
45              }
46              catch
47              {
48                  // If reflection fails, fall back to no usage
49              }
50          }
51  
52          return AIServiceUsage.None;
53      }
54  }