diff --git a/README.md b/README.md index 94ada2d..08e5366 100644 --- a/README.md +++ b/README.md @@ -52,6 +52,7 @@ The bot can be configured by setting the following environment variables: - `REDIS_HOST`: The host for the Redis server. - `REDIS_PORT`: The port for the Redis server. - `ADMIN_ID`: The Discord ID of the admin user. This user will have the ability to reset the chat. +- `CHAT_CHANNEL_ID`: The Discord ID of the specific channel where the bot will chat. Leave this empty to allow the bot to chat in all channels. (default is empty) - `BOT_NAME`: The name of the bot. This is used for the bot to recognize when it is mentioned in a message. - `CHAT_MAX_LENGTH`: The maximum length of the chat history to store in Redis. This is used to limit the amount of memory used by the bot. - `CTX`: The context length for the Ollama API. This determines how much of the chat history the bot will consider when generating a response. diff --git a/bot.env.example b/bot.env.example index 3a10020..37cd17f 100644 --- a/bot.env.example +++ b/bot.env.example @@ -1,9 +1,10 @@ DISCORD_TOKEN= ADMIN_ID=209692688415457282 +CHAT_CHANNEL_ID= OLLAMA_HOST=ollama OLLAMA_PORT=11434 OLLAMA_MODEL=llama3 REDIS_HOST=redis BOT_NAME=assistant CHAT_MAX_LENGTH=500 -CTX=2048 +CTX=2048 \ No newline at end of file diff --git a/bot.py b/bot.py index 42875ce..8ded4bd 100644 --- a/bot.py +++ b/bot.py @@ -46,12 +46,13 @@ class DiscordResponse: class Bot: - def __init__(self, ollama, discord, redis, model, admin_id, bot_name, chat_max_length=500, ctx=4096): + def __init__(self, ollama, discord, redis, model, admin_id, chat_channel_id, bot_name, chat_max_length=500, ctx=4096): self.ollama = ollama self.discord = discord self.redis = redis self.model = model self.admin_id = admin_id + self.chat_channel_id = chat_channel_id self.bot_name = bot_name self.chat_max_length = chat_max_length self.ctx = ctx @@ -106,6 +107,11 @@ class Bot: string_channel_id = str(message.channel.id) + # If CHAT_CHANNEL_ID is set, only respond in that channel. + if self.chat_channel_id: + if string_channel_id != self.chat_channel_id: + return + if self.discord.user == message.author: # don't respond to ourselves return @@ -242,6 +248,7 @@ def main(): parser.add_argument('--redis-port', default=os.getenv('REDIS_PORT', 6379), type=int) parser.add_argument('--admin-id', default=os.getenv('ADMIN_ID', ''), type=str) + parser.add_argument('--chat-channel-id', default=os.getenv('CHAT_CHANNEL_ID', ''), type=str) parser.add_argument('--bot-name', default=os.getenv('BOT_NAME', 'assistant'), type=str) parser.add_argument('--chat-max-length', default=os.getenv('CHAT_MAX_LENGTH', 500), type=int) @@ -260,6 +267,7 @@ def main(): redis.Redis(host=args.redis_host, port=args.redis_port, db=0, decode_responses=True), model=args.ollama_model, admin_id = args.admin_id, + chat_channel_id = args.chat_channel_id, bot_name = args.bot_name, chat_max_length=args.chat_max_length, ctx=args.ctx,