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

This commit is contained in:
Samuele Lorefice
2024-12-26 19:47:07 +01:00
parent 000b32c41d
commit e90e0200e1
5 changed files with 213 additions and 100 deletions

8
TelegramBot/Actor.cs Normal file
View File

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

10
TelegramBot/Agent.cs Normal file
View File

@@ -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;
}

View File

@@ -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;
}

View File

@@ -1,31 +1,21 @@
using OpenAI;
using OpenAI.Chat;
using System.ClientModel;
using OpenAI.Chat;
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 krolikPrompt = File.ReadAllText($"prompt/{GetEnv("KROLIK_PROMPT_FILE")}");
Dictionary<long, List<ChatMessage>> oaiChats = new();
string nemesisPrompt = File.ReadAllText($"prompt/{Env.Get("NEMESIS_PROMPT_FILE")}");
string krolikPrompt = File.ReadAllText($"prompt/{Env.Get("KROLIK_PROMPT_FILE")}");
//Ratelimit
TimeSpan rateLimit = new(0, 0, 0, 30);
TimeSpan rateLimit = new(0, 0, 0, 10);
Dictionary<long, DateTime> lastMessage = new();
HashSet<long> unlimitedChats = new();
@@ -36,105 +26,94 @@ bool IsRateLimited(long chatId) {
return false;
}
#region TelegramBot Startup
var options = new OpenAIClientOptions() {
Endpoint = new(baseUrl),
NetworkTimeout = new TimeSpan(0, 0, 15)
};
string nemesisToken = Env.Get("NEMESIS_BOT_TOKEN");
string krolikToken = Env.Get("KROLIK_BOT_TOKEN");
var openAiApiKey = new ApiKeyCredential(apiKey);
var openAiClient = new OpenAIClient(openAiApiKey, options);
var chatClient = openAiClient.GetChatClient(model);
using var nemcts = new CancellationTokenSource();
using var krocts = new CancellationTokenSource();
string token = GetEnv("TELEGRAM_BOT_TOKEN") ?? string.Empty;
Console.WriteLine("OpenAI Chat Client created");
var nemesisBot = new TelegramBotClient(nemesisToken, cancellationToken:nemcts.Token);
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}
""");
""");
//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;
}
//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);
await AnswerChat(chatid, msg.Text, tokenlenght, agent.Actor);
}
}
async Task AnswerChat(long chatId, string input) {
//Check if the chat is already in the dictionary
if (!oaiChats.ContainsKey(chatId))
AddChatToDictionary(chatId);
async Task AnswerChat(long chatId, string input, int tokenLenght, Actor actor) {
//Limit the message to 1024 characters to avoid out of context jump
string text = input;
if (input.Length > 1024) text = input.Substring(0, 1024);
//Add the current message to the chat
ChatMessageRotate(chatId, new UserChatMessage(text));
//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));
//Get the response from the OpenAI API
var result = oaiAgent.GetChatResponse(chatId, actor);
//Send the response to the user
await bot.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);
await krolikBot.SendMessage(chatId, result);
}

5
TelegramBot/Tools.cs Normal file
View File

@@ -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");
}