Merge pull request #153 from hchen2020/master

Add TokenStatistics for cost control.

Author: Oceania2018

src/Infrastructure/BotSharp.Abstraction/Conversations/ITokenStatistics.cs

@@ -0,0 +1,10 @@
+namespace BotSharp.Abstraction.Conversations;
+
+public interface ITokenStatistics
+{
+    int Total { get; }
+    float AccumulatedCost { get; }
+    float Cost { get; }
+    void AddToken(int promptCount, int completionCount);
+    void PrintStatistics();
+}

src/Infrastructure/BotSharp.Abstraction/Routing/Models/RoutingArgs.cs

@@ -5,7 +5,7 @@ public class RoutingArgs
     [JsonPropertyName("reason")]
     public string Reason { get; set; } = string.Empty;
 
-    [JsonPropertyName("agent_name")]
+    [JsonPropertyName("agent")]
     public string AgentName { get; set; } = string.Empty;
 
     public override string ToString()

src/Infrastructure/BotSharp.Core/BotSharpServiceCollectionExtensions.cs

@@ -61,6 +61,7 @@ public static IServiceCollection AddBotSharp(this IServiceCollection services, I
         }
 
         services.AddScoped<IInstructService, InstructService>();
+        services.AddScoped<ITokenStatistics, TokenStatistics>();
 
         return services;
     }

src/Infrastructure/BotSharp.Core/Conversations/Services/ConversationService.SendMessage.cs

@@ -63,6 +63,9 @@ await routing.InstructLoop(agent) :
 
         await HandleAssistantMessage(response, onMessageReceived);
 
+        var statistics = _services.GetRequiredService<ITokenStatistics>();
+        statistics.PrintStatistics();
+
         return true;
     }
 

src/Infrastructure/BotSharp.Core/Conversations/Services/TokenStatistics.cs

@@ -0,0 +1,53 @@
+using System.Drawing;
+
+namespace BotSharp.Core.Conversations.Services;
+
+public class TokenStatistics : ITokenStatistics
+{
+    private int _promptTokenCount = 0;
+    private int _completionTokenCount = 0;
+    private readonly IServiceProvider _services;
+    private readonly ILogger _logger;
+    public int Total => _promptTokenCount + _completionTokenCount;
+
+    public float Cost => _promptTokenCount / 1000f * 0.0015f + _completionTokenCount / 1000f * 0.002f;
+
+    public float AccumulatedCost
+    {
+        get 
+        {
+            var stat = _services.GetRequiredService<IConversationStateService>();
+            var promptTokenCount = int.Parse(stat.GetState("prompt_total", "0"));
+            var completionTokenCount = int.Parse(stat.GetState("completion_total", "0"));
+            return promptTokenCount / 1000f * 0.0015f + completionTokenCount / 1000f * 0.002f;
+        }
+    }
+
+    public TokenStatistics(IServiceProvider services, ILogger<TokenStatistics> logger) 
+    { 
+        _services = services;
+        _logger = logger;
+    }
+
+    public void AddToken(int promptCount, int completionCount)
+    {
+        _promptTokenCount += promptCount;
+        _completionTokenCount += completionCount;
+
+        // Accumulated Token
+        var stat = _services.GetRequiredService<IConversationStateService>();
+        var count1 = int.Parse(stat.GetState("prompt_total", "0"));
+        stat.SetState("prompt_total", promptCount + count1);
+        var count2 = int.Parse(stat.GetState("completion_total", "0"));
+        stat.SetState("completion_total", completionCount + count2);
+    }
+
+    public void PrintStatistics()
+    {
+#if DEBUG
+        Console.WriteLine($"Token Usage: {_promptTokenCount} prompt + {_completionTokenCount} completion = {Total} total (${Cost}), accumulated cost: ${AccumulatedCost}", Color.DarkGray);
+#else
+        _logger.LogInformation($"Token Usage: {_promptTokenCount} prompt + {_completionTokenCount} completion = {Total} total (${Cost}), accumulated cost: ${AccumulatedCost}");
+#endif
+    }
+}

src/Infrastructure/BotSharp.Core/Routing/Handlers/ContinueExecuteTaskRoutingHandler.cs

@@ -14,7 +14,7 @@ public class ContinueExecuteTaskRoutingHandler : RoutingHandlerBase, IRoutingHan
 
     public List<NameDesc> Parameters => new List<NameDesc>
     {
-        new NameDesc("agent_name", "the name of the agent"),
+        new NameDesc("agent", "the name of the agent"),
         new NameDesc("args", "required parameters extracted from question"),
         new NameDesc("reason", "why continue to execute current task")
     };

src/Infrastructure/BotSharp.Core/Routing/Handlers/RetrieveDataFromAgentRoutingHandler.cs

@@ -14,7 +14,7 @@ public class RetrieveDataFromAgentRoutingHandler : RoutingHandlerBase, IRoutingH
 
     public List<NameDesc> Parameters => new List<NameDesc>
     {
-        new NameDesc("agent_name", "the name of the agent"),
+        new NameDesc("agent", "the name of the agent"),
         new NameDesc("question", "the question you will ask the agent to get the necessary data"),
         new NameDesc("reason", "why retrieve data"),
         new NameDesc("args", "required parameters extracted from question and hand over to the next agent")

src/Infrastructure/BotSharp.Core/Routing/Handlers/RouteToAgentRoutingHandler.cs

@@ -15,7 +15,7 @@ public class RouteToAgentRoutingHandler : RoutingHandlerBase, IRoutingHandler
 
     public List<NameDesc> Parameters => new List<NameDesc>
     {
-        new NameDesc("agent_name", "the name of the agent from AGENTS"),
+        new NameDesc("agent", "the name of the agent from AGENTS"),
         new NameDesc("reason", "why route to this agent"),
         new NameDesc("args", "parameters extracted from context")
     };

src/Infrastructure/BotSharp.Core/Routing/RoutingService.cs

@@ -4,7 +4,6 @@
 using BotSharp.Abstraction.Routing;
 using BotSharp.Abstraction.Routing.Models;
 using BotSharp.Abstraction.Routing.Settings;
-using BotSharp.Abstraction.Templating;
 namespace BotSharp.Core.Routing;
 
 public class RoutingService : IRoutingService
@@ -116,7 +115,8 @@ public Agent LoadRouter()
         var prompt = @"You're a Router with reasoning. Follow these steps to handle user's request:
 1. Read the CONVERSATION context.
 2. Select a appropriate function from FUNCTIONS.
-3. Determine which agent from AGENTS is suitable for the current task.";
+3. Determine which agent from AGENTS is suitable for the current task.
+4. Re-think about selected function is from FUNCTIONS to handle the request.";
 
         // Append function
         prompt += "\r\n";
@@ -160,7 +160,7 @@ 2. Select a appropriate function from FUNCTIONS.
             // Append parameters
             if (agent.RequiredFields.Any())
             {
-                prompt += $"\r\nRequired: {string.Join(',', agent.RequiredFields)}.";
+                prompt += $"\r\nRequired: {string.Join(", ", agent.RequiredFields)}.";
             }
             return agent;
         }).ToList();

src/Plugins/BotSharp.Plugin.AzureOpenAI/Providers/ChatCompletionProvider.cs

@@ -12,7 +12,6 @@
 using Microsoft.Extensions.Logging;
 using System;
 using System.Collections.Generic;
-using System.Drawing;
 using System.Linq;
 using System.Text.Json;
 using System.Threading.Tasks;
@@ -24,17 +23,20 @@ public class ChatCompletionProvider : IChatCompletion
     private readonly AzureOpenAiSettings _settings;
     private readonly IServiceProvider _services;
     private readonly ILogger _logger;
+    private readonly ITokenStatistics _tokenStatistics;
     private string _model;
 
     public virtual string Provider => "azure-openai";
 
     public ChatCompletionProvider(AzureOpenAiSettings settings, 
         ILogger<ChatCompletionProvider> logger,
-        IServiceProvider services)
+        IServiceProvider services,
+        ITokenStatistics tokenStatistics)
     {
         _settings = settings;
         _logger = logger;
         _services = services;
+        _tokenStatistics = tokenStatistics;
     }
 
     protected virtual (OpenAIClient, string) GetClient()
@@ -105,7 +107,7 @@ public async Task<bool> GetChatCompletionsAsync(Agent agent,
         var choice = response.Value.Choices[0];
         var message = choice.Message;
 
-        _logger.LogInformation($"Token Usage: {response.Value.Usage.PromptTokens} prompt + {response.Value.Usage.CompletionTokens} completion = {response.Value.Usage.TotalTokens} total");
+        _tokenStatistics.AddToken(response.Value.Usage.PromptTokens, response.Value.Usage.CompletionTokens);
 
         if (choice.FinishReason == CompletionsFinishReason.FunctionCall)
         {