···11-import random
22-import os
33-import io
44-import discord
55-import pytesseract
66-import PIL
77-import pypdf
88-from discord.ext import commands
99-from discord.message import Message
1010-from pathlib import Path
1111-from ..lib.markov import MarkovChain
1212-from ..lib.permissions import require_owner
1313-1414-1515-class Markov(commands.Cog):
1616- def __init__(self, bot: discord.bot.Bot):
1717- self.bot = bot
1818- self.reply_channels = [
1919- int(x) for x in os.getenv("REPLY_CHANNELS", "0").split(",")
2020- ]
2121- self.chain_file = Path(os.getenv("BRAIN_FILE", "brain.msgpackz"))
2222- if self.chain_file.is_file():
2323- print(f"Attempting load from {self.chain_file}...")
2424- try:
2525- self.markov = MarkovChain.load_from_file(self.chain_file)
2626- print("Load Complete")
2727- except Exception as E:
2828- print(f"Error while loading\n{E}")
2929- else:
3030- self.markov = MarkovChain({})
3131-3232- async def update_words(self):
3333- amount = len(self.markov.edges.keys())
3434- try:
3535- self.markov.save_to_file(self.chain_file)
3636- except Exception as E:
3737- print(f"Error while saving\n{E}")
3838-3939- await self.bot.change_presence(
4040- activity=discord.CustomActivity(name=f"I know {amount} words!")
4141- )
4242-4343- @require_owner
4444- @commands.slash_command()
4545- async def dump_chain(self, ctx: discord.ApplicationContext):
4646- o = self.markov.dumpb()
4747- fd = io.BytesIO(o)
4848- await ctx.respond(
4949- ephemeral=True, file=discord.File(fd, filename="brain.msgpackz")
5050- )
5151-5252- @require_owner
5353- @commands.slash_command()
5454- async def load_chain(
5555- self, ctx: discord.ApplicationContext, raw: discord.Option(discord.Attachment)
5656- ):
5757- new = MarkovChain.loadb(await raw.read())
5858- self.markov.merge(new)
5959- await ctx.respond("Imported", ephemeral=True)
6060- await self.update_words()
6161-6262- @require_owner
6363- @commands.slash_command()
6464- async def scan_history(self, ctx: discord.ApplicationContext):
6565- await ctx.defer(ephemeral=True)
6666- async for msg in ctx.history(limit=None):
6767- if msg.author.id != self.bot.application_id:
6868- self.markov.learn(msg.content)
6969- await ctx.respond("> Bingus Learned!", ephemeral=True)
7070- await self.update_words()
7171-7272- @commands.slash_command()
7373- async def markov(
7474- self, ctx: discord.ApplicationContext, prompt: discord.Option(str)
7575- ):
7676- print("Bingus is responding!")
7777- response = self.markov.respond(prompt)
7878- if response is not None and len(response) != 0:
7979- await ctx.respond(f"{prompt} {response[:1000]}")
8080- else:
8181- await ctx.respond("> Bingus couldn't think of what to say!")
8282- print("Bingus Responded!")
8383-8484- @commands.slash_command()
8585- async def weights(
8686- self, ctx: discord.ApplicationContext, token: discord.Option(str)
8787- ):
8888- edges = self.markov.get_edges(token)
8989- if edges is None:
9090- await ctx.respond("> Bingus doesn't know that word!")
9191- else:
9292- head = f"Weights for **{token}**"
9393- msg = "\n".join([f"{str(k)}: {v}" for k, v in edges.items()])
9494- if len(msg) > 1750:
9595- fd = io.BytesIO(msg.encode())
9696- await ctx.respond(head, file=discord.File(fd, filename="weights.txt"))
9797- else:
9898- await ctx.respond(f"{head}:\n{msg}")
9999-100100- @require_owner
101101- @commands.slash_command()
102102- async def pdf(
103103- self, ctx: discord.ApplicationContext, file: discord.Option(discord.Attachment)
104104- ):
105105- await ctx.defer(ephemeral=True)
106106- raw = await file.read()
107107- try:
108108- pdf = pypdf.PdfReader(io.BytesIO(raw))
109109- # i = 0
110110- for page in pdf.pages:
111111- # i += 1
112112- # printf("Bingus learned a page! {i}/{pdf.get_num_pages()}")
113113- text = page.extract_text()
114114- self.markov.learn(text)
115115- await self.update_words()
116116- await ctx.respond(
117117- "> Bingus learned something from the pdf!", ephemeral=True
118118- )
119119- except pypdf.errors.PdfReadError:
120120- await ctx.respond("> Bingus only understands pdf files!", ephemeral=True)
121121-122122- @require_owner
123123- @commands.slash_command()
124124- async def ocr(
125125- self, ctx: discord.ApplicationContext, file: discord.Option(discord.Attachment)
126126- ):
127127- raw = await file.read()
128128- try:
129129- image = PIL.Image.open(io.BytesIO(raw))
130130- text = pytesseract.image_to_string(image)
131131- self.markov.learn(text)
132132- await ctx.respond("> Bingus learned something from image!", ephemeral=True)
133133- await self.update_words()
134134- except PIL.UnidentifiedImageError:
135135- await ctx.respond("> Bingus only understands image files!", ephemeral=True)
136136-137137- @require_owner
138138- @commands.slash_command()
139139- async def study(
140140- self, ctx: discord.ApplicationContext, file: discord.Option(discord.Attachment)
141141- ):
142142- raw = await file.read()
143143- try:
144144- text = raw.decode()
145145- self.markov.learn(text)
146146- await ctx.respond("> Bingus learned from file!", ephemeral=True)
147147- await self.update_words()
148148- except UnicodeDecodeError:
149149- await ctx.respond(
150150- "> Bingus only understands UTF-8 text files!", ephemeral=True
151151- )
152152-153153- @require_owner
154154- @commands.slash_command()
155155- async def forget(self, ctx: discord.ApplicationContext):
156156- self.markov.forget()
157157- await ctx.respond("> Bingus forgot everything!", ephemeral=True)
158158- await self.update_words()
159159-160160- @commands.Cog.listener()
161161- async def on_ready(self):
162162- await self.update_words()
163163-164164- @commands.Cog.listener()
165165- async def on_message(self, msg: Message):
166166- if msg.is_system() or msg.flags.ephemeral or msg.channel.type == discord.ChannelType.private:
167167- return
168168-169169- if msg.author.id != self.bot.application_id:
170170- print("Bingus is learning!")
171171- self.markov.learn(msg.content)
172172- await self.update_words()
173173-174174- chance = 80 if msg.author.id != self.bot.application_id else 45
175175-176176- if msg.channel.id in self.reply_channels and random.randint(1, 100) <= chance:
177177- print("Bingus is responding!")
178178- response = self.markov.respond(msg.content)
179179- if response is not None and len(response) != 0:
180180- await msg.channel.trigger_typing()
181181- if msg.author.id == self.bot.application_id:
182182- await msg.channel.send(response[:1000])
183183- else:
184184- await msg.reply(response[:1000], mention_author=False)
185185- print("Bingus Responded!")
186186-187187-188188-def setup(bot):
189189- bot.add_cog(Markov(bot))
-31
src/bingus/dev/cog.template.py
···11-# import discord
22-from discord.ext import commands
33-# from discord.message import Message
44-55-66-class __CName__(commands.Cog):
77- # Setup any state for the cog here, this will
88- # exist for the run of the program
99- # If you need to change the parameters for __init__
1010- # make sure the [setup] function below this class.
1111- def __init__(self, bot) -> None:
1212- self.bot = bot
1313-1414- # Example of a slash command that will be loaded
1515- # with this cog
1616- # @commands.slash_command()
1717- # async def ping(self, ctx: discord.ApplicationContext):
1818- # await ctx.respond("pong!")
1919-2020- # Example of listening to all messages
2121- # ever sent in any server the bot is in
2222- # while active
2323- # @commands.Cog.listener()
2424- # async def on_message(self, msg: Message):
2525- # pass
2626-2727- # See the PyCord docs for more info and guides
2828-2929-3030-def setup(bot):
3131- bot.add_cog(__CName__(bot))
-12
src/bingus/dev/format.py
···11-import os
22-import pathlib
33-44-55-def main():
66- src = pathlib.Path(__file__).parent.parent.parent
77- print(f"Running black on src: {src}")
88- os.system(f"black {src.absolute()}")
99-1010-1111-if __name__ == "__main__":
1212- main()
-27
src/bingus/dev/new_cog.py
···11-import sys
22-import json
33-from pathlib import Path
44-55-66-def main():
77- name = input("Enter new cog's name (PascalCase!): ")
88- cogs_folder = Path(__file__).parent.parent.joinpath("cogs")
99- path = cogs_folder.joinpath(f"{name.lower()}.py")
1010- if path.exists():
1111- print("This cog seems to already exist! Refusing to overwrite.")
1212- sys.exit(1)
1313- else:
1414- print("Creating new Python file...")
1515- template_content = Path(__file__).parent.joinpath("cog.template.py").read_text()
1616- new_cog_file = template_content.replace("__CName__", name)
1717- path.write_text(new_cog_file)
1818- print('Adding to "cogs.json"...')
1919- cogs_json_path = cogs_folder.parent.joinpath("cogs.json")
2020- current: list[str] = json.loads(cogs_json_path.read_text())
2121- current.append(f"bingus.cogs.{name.lower()}")
2222- cogs_json_path.write_text(json.dumps(current))
2323- print(f"Cog Created at {path}!")
2424-2525-2626-if __name__ == "__main__":
2727- main()
src/bingus/lib/__init__.py
This is a binary file and will not be displayed.
-233
src/bingus/lib/markov.py
···11-from dataclasses import dataclass
22-import random
33-from typing import Optional
44-from pathlib import Path
55-from msgpack import packb, unpackb
66-from brotli import compress, decompress
77-88-99-@dataclass
1010-class Word:
1111- text: str
1212-1313- def __str__(self):
1414- return self.text
1515-1616- def __eq__(self, value: object):
1717- return self.text == value.text
1818-1919- def __hash__(self):
2020- return hash("WORD:" + self.text)
2121-2222-2323-@dataclass
2424-class End:
2525- def __str__(self):
2626- return "~END"
2727-2828- def __hash__(self):
2929- return hash("END")
3030-3131-3232-Token = Word | End
3333-3434-3535-def token_ser(t: Token) -> str:
3636- match t:
3737- case Word(w):
3838- return f"W-{w}"
3939- case _:
4040- return "E--"
4141-4242-4343-def token_de(s: str) -> Token:
4444- if s.startswith("W-"):
4545- return Word(s[2:])
4646- else:
4747- return End()
4848-4949-5050-def token_is_word(t: Token) -> bool:
5151- match t:
5252- case Word(_):
5353- return True
5454- case _:
5555- return False
5656-5757-5858-def token_is_end(t: Token) -> bool:
5959- match t:
6060- case End():
6161- return True
6262- case _:
6363- return False
6464-6565-6666-@dataclass
6767-class StateTransitions:
6868- to_tokens: dict[Token, int]
6969-7070- def merge(self, other):
7171- for key, val in other.to_tokens.items():
7272- if key in self.to_tokens.keys():
7373- self.to_tokens[key] += val
7474- else:
7575- self.to_tokens[key] = val
7676-7777- def register_transition(self, to_token: Token):
7878- if self.to_tokens.get(to_token) is None:
7979- self.to_tokens[to_token] = 0
8080- self.to_tokens[to_token] += 1
8181-8282- def pick_token(self, allow_end: bool = False) -> Optional[Token]:
8383- entries = [
8484- e for e in self.to_tokens.items() if allow_end or not token_is_end(e[0])
8585- ]
8686- if len(entries) == 0:
8787- return None
8888- else:
8989- return random.choices(
9090- [k for (k, _) in entries], weights=[v for (_, v) in entries]
9191- )[0]
9292-9393-9494-@dataclass
9595-class MarkovChain:
9696- edges: dict[Token, StateTransitions]
9797-9898- def _update(self, from_token: Token, to_token: Token):
9999- if self.edges.get(from_token) is None:
100100- new = StateTransitions({})
101101- new.register_transition(to_token)
102102- self.edges[from_token] = new
103103- else:
104104- self.edges[from_token].register_transition(to_token)
105105-106106- def _learn_from_tokens(self, tokens: list[Token]):
107107- for i, token in enumerate(tokens):
108108- if i == len(tokens) - 1:
109109- self._update(token, End())
110110- else:
111111- self._update(token, tokens[i + 1])
112112-113113- def _parse_source(self, source: str) -> list[Token]:
114114- return [
115115- Word(
116116- w if w.startswith("http://") or w.startswith("https://") else w.lower()
117117- )
118118- for w in source.split()
119119- if not (w.startswith("<@") and w.endswith(">"))
120120- ]
121121-122122- def get_edges(self, token: str) -> Optional[dict[str, int]]:
123123- edges = self.edges.get(Word(token))
124124- if edges is None:
125125- return None
126126- else:
127127- return edges.to_tokens
128128-129129- def learn(self, source: str):
130130- tokens = self._parse_source(source)
131131- self._learn_from_tokens(tokens)
132132-133133- def forget(self):
134134- self.edges = {}
135135-136136- def _pick_next(self, current_token: Token, allow_end: bool) -> Token:
137137- transitions = self.edges.get(current_token)
138138- if transitions is None:
139139- return End()
140140- else:
141141- next = transitions.pick_token(allow_end)
142142- if next is None:
143143- return End()
144144- else:
145145- return next
146146-147147- def _join_tokens(self, tokens: list[Token]) -> str:
148148- buf = []
149149- for i, c in enumerate(tokens):
150150- match c:
151151- case End():
152152- pass
153153- case Word(text):
154154- buf.append(text + " " if i < len(tokens) - 1 else text)
155155- return "".join(buf)
156156-157157- def _chain_tokens(
158158- self, starting_token: Optional[Token] = None, max_length: int = 20
159159- ) -> list[Token]:
160160- tokens = []
161161-162162- if starting_token is None:
163163- keys = self.edges.keys()
164164- if len(keys) == 0:
165165- return []
166166- else:
167167- starting_token = random.choice(list(keys))
168168- tokens.append(starting_token)
169169-170170- current_token = starting_token
171171-172172- while len(tokens) < max_length:
173173- next_token = self._pick_next(current_token, len(tokens) > 2)
174174- match next_token:
175175- case End():
176176- break
177177- case token:
178178- tokens.append(token)
179179- current_token = next_token
180180-181181- return tokens
182182-183183- def _chain(
184184- self, starting_token: Optional[Token] = None, max_length: int = 20
185185- ) -> str:
186186- tokens = self._chain_tokens(starting_token, max_length)
187187- joined = self._join_tokens(tokens)
188188- return joined
189189-190190- def respond(self, message: str, max_length: int = 20) -> str:
191191- tokens = self._parse_source(message)
192192- tt = [x for x in filter(token_is_word, tokens)]
193193- if len(tt) != 0 and tt[-1] in self.edges.keys():
194194- return self._chain(tt[-1], max_length=max_length)
195195- else:
196196- return self._chain(None, max_length=max_length)
197197-198198- def save_to_file(self, path: Path):
199199- if not path.parent.exists():
200200- path.parent.mkdir(parents=True)
201201- path.write_bytes(self.dumpb())
202202-203203- def load_from_file(path: Path):
204204- return MarkovChain.loadb(path.read_bytes())
205205-206206- def dumpb(self):
207207- return compress(packb(self.ser()))
208208-209209- def loadb(dat):
210210- return MarkovChain.deser(unpackb(decompress(dat)))
211211-212212- def ser(self):
213213- return {
214214- token_ser(e): {token_ser(k): v for k, v in w.to_tokens.items()}
215215- for e, w in self.edges.items()
216216- }
217217-218218- def deser(dat):
219219- edges = {
220220- token_de(e): StateTransitions({token_de(k): v for k, v in w.items()})
221221- for e, w in dat.items()
222222- }
223223- return MarkovChain(edges)
224224-225225- def merge(self, other):
226226- for key, val in other.edges.items():
227227- if key in self.edges.keys():
228228- self.edges[key].merge(val)
229229- else:
230230- self.edges[key] = val
231231-232232-233233-__all__ = (MarkovChain,)
-21
src/bingus/lib/permissions.py
···11-import discord
22-33-44-class NotOwnerError(discord.ApplicationCommandError):
55- def __init__(self) -> None:
66- super().__init__("You are not allowed to run this command")
77-88-99-def _check_owners(ctx: discord.ApplicationContext):
1010- if ctx.author.id not in ctx.bot.bingus_owners:
1111- raise NotOwnerError()
1212- else:
1313- return True
1414-1515-1616-def require_owner(cmd: discord.SlashCommand):
1717- cmd.checks.append(_check_owners)
1818- return cmd
1919-2020-2121-__all__ = (require_owner, NotOwnerError)
-68
src/bingus/main.py
···11-#!/usr/bin/env python
22-33-from dotenv import load_dotenv
44-from discord import Bot
55-import discord
66-import json
77-from pathlib import Path
88-from os import getenv
99-from .lib.permissions import NotOwnerError
1010-1111-1212-class BingusBot(Bot):
1313- def __init__(self, *args, **kwargs):
1414- self.bingus_owners = []
1515- super().__init__(*args, **kwargs)
1616-1717- async def get_owners(self):
1818- app = await self.application_info()
1919- if app.team is not None:
2020- return [m.id for m in app.team.members]
2121- else:
2222- return [app.owner.id]
2323-2424- async def on_application_command_error(
2525- self, context: discord.ApplicationContext, exception: discord.DiscordException
2626- ) -> None:
2727- if isinstance(exception, NotOwnerError):
2828- await context.respond(
2929- "You are not allowed to run this command!", ephemeral=True
3030- )
3131- else:
3232- return await super().on_application_command_error(context, exception)
3333-3434- async def on_ready(self):
3535- self.bingus_owners = await self.get_owners()
3636- print(f"Initialized Gateway as {self.user.name} ({self.user.id})")
3737-3838-3939-def main():
4040- load_dotenv(".env")
4141-4242- print("Initializing Base Bot...")
4343-4444- intents = discord.Intents.default()
4545- intents.message_content = True
4646-4747- bot = BingusBot(intents=intents)
4848-4949- EXTENSIONS: list[str] = json.loads(
5050- Path(__file__).parent.joinpath("cogs.json").read_text()
5151- )
5252-5353- for ext in EXTENSIONS:
5454- print(f'Initializing "{ext}"...')
5555- bot.load_extension(ext)
5656-5757- print("Connecting to Discord...")
5858-5959- token = getenv("TOKEN")
6060-6161- if token is None:
6262- print("Error: Env var TOKEN not found, are you missing it / don't have a .env?")
6363- else:
6464- bot.run(token)
6565-6666-6767-if __name__ == "__main__":
6868- main()