Removes ratelimit, refactors everything in more sections, adds tokenization calculation

TelegramBot/Actor.cs

@@ -0,0 +1,8 @@
+namespace TelegramBot;
+
+public enum Actor {
+    User,
+    Krolik,
+    Nemesis,
+    System
+}

TelegramBot/Agent.cs

@@ -0,0 +1,10 @@
+using Telegram.Bot;
+
+namespace TelegramBot;
+
+public class Agent(Actor actor, long telegramId, string name, TelegramBotClient bot) {
+    public Actor Actor { get; } = actor;
+    public long TelegramId { get; } = telegramId;
+    public string Name { get; } = name;
+    public TelegramBotClient Bot { get; } = bot;
+}

TelegramBot/OpenAIAgent.cs

@@ -1,6 +1,117 @@
-namespace TelegramBot;
+using System.ClientModel;
+using System.Net.Http.Json;
+using System.Text.Json.Serialization;
+using OpenAI;
+using OpenAI.Chat;
 
-public class OpenAIAgent
+namespace TelegramBot;
+
+public class OpenAiAgent
 {
+    private ApiKeyCredential apikey;
+    private OpenAIClient oaiClient;
+    private ChatClient chatClient;
+    private HttpClient httpClient;
+    private Dictionary<(long, Actor), List<(ChatMessage, int)>> oaiChats = new();
     
+    public OpenAiAgent(string baseUrl, string apiKey, string model) {
+        OpenAIClientOptions options = new OpenAIClientOptions() {
+            Endpoint = new(baseUrl),
+            NetworkTimeout = new(0, 0, 15)
+        };
+
+        apikey = new(apiKey);
+        oaiClient = new(apikey, options);
+        chatClient = oaiClient.GetChatClient(model);
+
+        httpClient = new() {
+            BaseAddress = new(baseUrl),
+            Timeout = new(0, 0, 15)
+        };
+        
+        Console.WriteLine(
+            $"""
+                Base URL: {baseUrl}
+                Model: {model}
+                API Key: {apiKey}
+             """);
+    }
+
+    public int GetTokenLenght(string message) {
+        TokenizeRequest request = new(message);
+        var req = httpClient.PostAsync("tokenize", JsonContent.Create(request));
+        req.Wait();
+        var response = req.Result.Content.ReadFromJsonAsync<TokenizeResponseBase>();
+        response.Wait();
+        return response.Result.Tokens.Length;
+    }
+
+    public void ChatHistoryAppend(Actor actor, long chatId, string message) {
+        //get the chat from the dictionary
+        var chat = GetChatHistory((chatId, actor));
+        //get the token lenght of the message
+        var tokenLenght = GetTokenLenght(message);
+        //create a new chat message
+        switch (actor) {
+            case Actor.User:
+                chat.Add((new UserChatMessage(message), tokenLenght));
+                break;
+            case Actor.Krolik:
+                chat.Add((new AssistantChatMessage(message), tokenLenght));
+                break;
+            case Actor.Nemesis:
+                chat.Add((new AssistantChatMessage(message), tokenLenght));
+                break;
+            case Actor.System:
+                chat.Add((new SystemChatMessage(message), tokenLenght));
+                break;
+        }
+    }
+
+    public List<(ChatMessage, int)> GetChatHistory((long, Actor) key) => GetChatHistory(key.Item1, key.Item2);
+    
+    public List<(ChatMessage, int)> GetChatHistory(long chatId, Actor actor) {
+        oaiChats.TryGetValue((chatId, actor), out var chat);
+        if(chat == null) {
+            AddChatToDictionary(chatId, actor);
+            chat = oaiChats[(chatId, actor)];
+        }
+        return chat!;
+    }
+
+    public string GetChatResponse(long chatId, Actor actor) {
+        return String.Empty;
+    }
+
+    public void AddChatToDictionary(long id) {
+        AddChatToDictionary(id, Actor.Krolik);
+        AddChatToDictionary(id, Actor.Nemesis);
+    }
+
+    public void AddChatToDictionary(long id, Actor actor) {
+        //Check if the chat already exists
+        if (oaiChats.ContainsKey((id, actor))) return;
+        //Create a new chat object
+        var chat = new List<(ChatMessage, int)>();
+        //chat.Add(new SystemChatMessage(nemesisPrompt));
+        //add the entry to the dictionary
+        oaiChats.Add((id, actor), chat);
+    }
+    
+    public void ResetChat(long chatId) {
+        //Remove the chat from the dictionary
+        oaiChats.Where(x => x.Key.Item1 == chatId).ToList().ForEach(x => oaiChats.Remove(x.Key));
+        //Add the chat back to the dictionary
+        AddChatToDictionary(chatId);
+    }
+}
+
+public struct TokenizeRequest(string content) {
+    [JsonPropertyName("content")] public string Content { get; set; } = content;
+    [JsonPropertyName("add_special")] public bool AddSpecial { get; set; } = false;
+    [JsonPropertyName("with_pieces")] public bool WithPieces { get; set; } = false;
+}
+
+public struct TokenizeResponseBase(int[] tokens) {
+    [JsonPropertyName("tokens")] public int[] Tokens { get; set; } = tokens;
 }

TelegramBot/Program.cs

@@ -1,140 +1,119 @@
-using OpenAI;
+using OpenAI.Chat;
-using OpenAI.Chat;
-using System.ClientModel;
 using Telegram.Bot;
 using Telegram.Bot.Types;
-using Telegram.Bot.Types.Enums;
+using Telegram.Bot.Types.Enums;using TelegramBot;
 using File = System.IO.File;
 
-string GetEnv(string name) => Environment.GetEnvironmentVariable(name) ?? throw new Exception($"Environment variable {name} is not set");
+string baseUrl  = Env.Get("OPENAI_BASE_URL");
+string apiKey   = Env.Get("OPENAI_API_KEY");
+string model    = Env.Get("OPENAI_MODEL");
+var oaiAgent    = new OpenAiAgent(baseUrl, apiKey, model);
 
-string baseUrl  = GetEnv("OPENAI_BASE_URL");
-string model    = GetEnv("OPENAI_MODEL");
-string apiKey   = GetEnv("OPENAI_API_KEY");
 Console.WriteLine("Starting the bot...");
-Console.WriteLine(
-    $"""
-        Base URL: {baseUrl}
-        Model: {model}
-        API Key: {apiKey}
-     """);
 
-string nemesisPrompt = File.ReadAllText($"prompt/{GetEnv("NEMESIS_PROMPT_FILE")}");
+string nemesisPrompt = File.ReadAllText($"prompt/{Env.Get("NEMESIS_PROMPT_FILE")}");
-string krolikPrompt = File.ReadAllText($"prompt/{GetEnv("KROLIK_PROMPT_FILE")}");
+string krolikPrompt = File.ReadAllText($"prompt/{Env.Get("KROLIK_PROMPT_FILE")}");
-
-Dictionary<long, List<ChatMessage>> oaiChats = new();
 
 //Ratelimit
-TimeSpan rateLimit = new(0, 0, 0, 30);
+TimeSpan rateLimit = new(0, 0, 0, 10);
 Dictionary<long, DateTime> lastMessage = new();
 HashSet<long> unlimitedChats = new();
 
 bool IsRateLimited(long chatId) {
     if (lastMessage.ContainsKey(chatId) && DateTime.Now - lastMessage[chatId] < rateLimit) return true;
-    
+
     lastMessage[chatId] = DateTime.Now;
     return false;
 }
 
+#region TelegramBot Startup
 
-var options = new OpenAIClientOptions() {
+string nemesisToken = Env.Get("NEMESIS_BOT_TOKEN");
-    Endpoint = new(baseUrl),
+string krolikToken = Env.Get("KROLIK_BOT_TOKEN");
-    NetworkTimeout = new TimeSpan(0, 0, 15)
-};
 
-var openAiApiKey    = new ApiKeyCredential(apiKey);
+using var nemcts = new CancellationTokenSource();
-var openAiClient    = new OpenAIClient(openAiApiKey, options);
+using var krocts = new CancellationTokenSource();
-var chatClient      = openAiClient.GetChatClient(model);
 
-string token = GetEnv("TELEGRAM_BOT_TOKEN") ?? string.Empty;
+var nemesisBot = new TelegramBotClient(nemesisToken, cancellationToken:nemcts.Token);
-Console.WriteLine("OpenAI Chat Client created");
+var krolikBot = new TelegramBotClient(krolikToken, cancellationToken:krocts.Token);
+
+var nemProfile = nemesisBot.GetMe();
+Agent Nemesis = new(Actor.Nemesis, nemProfile.Result.Id, nemProfile.Result.FirstName, nemesisBot);
+
+var kroProfile = krolikBot.GetMe();
+Agent Krolik = new(Actor.Krolik, kroProfile.Result.Id, kroProfile.Result.FirstName, krolikBot);
+
+nemesisBot.OnMessage += OnNemMessage;
+krolikBot.OnMessage += OnKroMessage;
+
+await nemesisBot.DropPendingUpdates();
+Console.WriteLine("Nemesis Bot running");
+await krolikBot.DropPendingUpdates();
+Console.WriteLine("Krolik Bot running");
 
-using var cts = new CancellationTokenSource();
-var bot = new TelegramBotClient(token, cancellationToken:cts.Token);
-await bot.DropPendingUpdates();
-var me = bot.GetMe();
-bot.OnMessage += OnMessage;
-Console.WriteLine("Bot running");
 Thread.Sleep(Timeout.Infinite);
-cts.Cancel(); // stop the bot
+nemcts.Cancel(); // stop nembot
+krocts.Cancel(); // stop krobot
 
-async Task OnMessage(Message msg, UpdateType type)
+#endregion
-{
+
+async Task OnNemMessage(Message msg, UpdateType type) {
     //Discard any message that is not a text message
     if (msg.Type != MessageType.Text) return;
+    
+    await OnMessage(msg, Nemesis);
+}
+
+async Task OnKroMessage(Message msg, UpdateType type) {
+    //Discard any message that is not a text message
+    if (msg.Type != MessageType.Text) return;
+    
+    await OnMessage(msg, Krolik);
+}
+
+//TODO: currently we only take in account private messages and messages directed to the bot/mentioning them.
+// We should also take in account the last x messages in groups to add more context
+async Task OnMessage(Message msg, Agent agent) {
     //Check if the message contains the bot's username or a reply to a message sent by the bot or a private chat
-    if (msg.Text!.Contains(me.Result.FirstName!, StringComparison.OrdinalIgnoreCase) ||
+        
-        msg.ReplyToMessage != null && msg.ReplyToMessage.From!.Id == me.Result.Id ||
-        msg.Chat.Type == ChatType.Private) {
         var chatid = msg.Chat.Id;
+        var tokenlenght = oaiAgent.GetTokenLenght(msg.Text);
         Console.WriteLine(
             $"""
-             Received message from {chatid} Type: {type}
+             {agent.Name} has received message from {chatid} TokenLenght: {tokenlenght}
-             Message: {msg.Text} 
+             Message: {msg.Text}
-            """);
+             """);
+        //Add the message to the chat history
+        oaiAgent.ChatHistoryAppend(Actor.User, chatid, "User: "+msg.Text);
+
+        //Check if the message is a reset command
+        if (msg.Text == "/reset" || msg.Text == "/reset@"+agent.Name) {
+            oaiAgent.ResetChat(chatid);
+            await agent.Bot.SendMessage(chatid, "Chat context has been reset");
+            return;
+        }
+        
+        // Otherwise process it normally
+    if (msg.Text!.Contains(agent.Name, StringComparison.OrdinalIgnoreCase) ||
+        msg.ReplyToMessage?.From?.Id == agent.TelegramId || msg.Chat.Type == ChatType.Private) {
+                
         //Check if the chat (group) is rate limited
         if (IsRateLimited(chatid)) {
             Console.WriteLine("No response due to ratelimit.");
             return;
         }
-
+        await AnswerChat(chatid, msg.Text, tokenlenght, agent.Actor);
-        //Check if the message is a reset command
-        if (msg.Text.StartsWith("/reset")) {
-            ResetChat(chatid);
-            await bot.SendMessage(chatid, "Chat context has been reset");
-            return;
-        }
-        // Otherwise process it normally
-        await AnswerChat(chatid, msg.Text);
     }
 }
 
-async Task AnswerChat(long chatId, string input) {
+
-    //Check if the chat is already in the dictionary
+async Task AnswerChat(long chatId, string input, int tokenLenght, Actor actor) {
-    if (!oaiChats.ContainsKey(chatId))
-        AddChatToDictionary(chatId);
-        
     //Limit the message to 1024 characters to avoid out of context jump
     string text = input;
-    if (input.Length > 1024) text = input.Substring(0, 1024); 
+    if (input.Length > 1024) text = input.Substring(0, 1024);
-        
+
-    //Add the current message to the chat
+    //Get the response from the OpenAI API
-    ChatMessageRotate(chatId, new UserChatMessage(text));
+    var result = oaiAgent.GetChatResponse(chatId, actor);
-    
-    //fetch existing messages history, append hint of the speaker to the message:
-    var messages = oaiChats[chatId];
-    messages.Add(new AssistantChatMessage("Nemesis: "));
-    
-    //Fetch the response from the model
-    var result = chatClient.CompleteChat(messages).Value.Content[0].Text;
-        
-    //Add the response to the chat
-    Console.WriteLine("Replying with: " + result);
-    oaiChats[chatId].Add(new AssistantChatMessage(result));
-        
     //Send the response to the user
-    await bot.SendMessage(chatId, result);
+    await krolikBot.SendMessage(chatId, result);
-}
-
-void AddChatToDictionary(long id) {
-    //Create a new chat object
-    var chat = new List<ChatMessage>();
-    chat.Add(new SystemChatMessage(nemesisPrompt));
-    //add the entry to the dictionary
-    oaiChats.Add(id, chat);
-}
-
-void ChatMessageRotate(long chatId, ChatMessage message){
-    //Remove the first message from the chat if the chat has more than 5 couples of messages (0 is our prompt)
-    if (oaiChats[chatId].Count > 10) oaiChats[chatId].RemoveRange(1, 2);
-    
-    //Add the new message to the chat
-    oaiChats[chatId].Add(message);
-}
-
-void ResetChat(long chatId) {
-    //Remove the chat from the dictionary
-    oaiChats.Remove(chatId);
-    //Add the chat back to the dictionary
-    AddChatToDictionary(chatId);
 }

TelegramBot/Tools.cs

@@ -0,0 +1,5 @@
+namespace TelegramBot;
+
+public class Env {
+    public static string Get(string var) => Environment.GetEnvironmentVariable(var) ?? throw new ($"Env var {var} is not set");
+}