diff --git a/src/cogs/auto_responder.py b/src/cogs/auto_responder.py index 46c7e81..8162ad8 100644 --- a/src/cogs/auto_responder.py +++ b/src/cogs/auto_responder.py @@ -41,7 +41,9 @@ class AutoResponder(commands.Cog): links.append(url.geturl()) return links - async def _transcode_hevc_to_h264(self, uri: str | pathlib.Path) -> tuple[discord.File, pathlib.Path] | None: + async def _transcode_hevc_to_h264( + self, uri: str | pathlib.Path, *, update: discord.Message = None + ) -> tuple[discord.File, pathlib.Path] | None: """ Transcodes the given URL or file to H264 from HEVC, trying to preserve quality and file size. @@ -64,7 +66,7 @@ class AutoResponder(commands.Cog): self.log.warning( "Video %r is %.2f seconds long (more than 10 minutes). Refusing to process further.", uri, - float(info["format"].get("duration", 600.1)) + float(info["format"].get("duration", 600.1)), ) return streams = info.get("streams", []) @@ -82,7 +84,7 @@ class AutoResponder(commands.Cog): async with aiohttp.ClientSession( headers={ "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp," - "*/*;q=0.8", + "*/*;q=0.8", "Accept-Language": "en-US,en;q=0.5", "Sec-Fetch-Dest": "document", "Sec-Fetch-Mode": "navigate", @@ -105,21 +107,35 @@ class AutoResponder(commands.Cog): self.log.info("Transcoding %r to %r", uri, tmp_path) args = [ "-hide_banner", - "-hwaccel", "auto", - "-i", tmp_dl.name, - "-c:v", "libx264", - "-crf", "25", - "-maxrate", "5M", - "-minrate", "100K", - "-bufsize", "5M", - "-c:a", "libopus", - "-b:a", "64k", - "-preset", "faster", - "-vsync", "2", - "-pix_fmt", "yuv420p", - "-movflags", "faststart", - "-profile:v", "main", - "-y" + "-hwaccel", + "auto", + "-i", + tmp_dl.name, + "-c:v", + "libx264", + "-crf", + "25", + "-maxrate", + "5M", + "-minrate", + "100K", + "-bufsize", + "5M", + "-c:a", + "libopus", + "-b:a", + "64k", + "-preset", + "faster", + "-vsync", + "2", + "-pix_fmt", + "yuv420p", + "-movflags", + "faststart", + "-profile:v", + "main", + "-y", ] process = await asyncio.create_subprocess_exec( "ffmpeg", @@ -150,7 +166,7 @@ class AutoResponder(commands.Cog): if message.channel.name == "spam" and message.author.id in { 1101439218334576742, 1229496078726860921, - 421698654189912064 + 421698654189912064, }: # links = self.extract_links(message.content, "static-assets-1.truthsocial.com") links = self.extract_links(message.content) @@ -167,7 +183,7 @@ class AutoResponder(commands.Cog): ".ps", ".ts", ".3gp", - ".3g2" + ".3g2", } # ^ All containers allowed to contain HEVC # per https://en.wikipedia.org/wiki/Comparison_of_video_container_formats @@ -175,7 +191,7 @@ class AutoResponder(commands.Cog): self.log.info("Found link to transcode: %r", link) try: async with message.channel.typing(): - _r = await self._transcode_hevc_to_h264(link) + _r = await self._transcode_hevc_to_h264(link, update=message) if not _r: continue file, _p = _r @@ -186,7 +202,7 @@ class AutoResponder(commands.Cog): self.log.warning( "Transcoded file too large: %r (%.2f)MB", _p, - _p.stat().st_size / 1024 / 1024 + _p.stat().st_size / 1024 / 1024, ) if _p.stat().st_size <= 510 * 1024 * 1024: file.fp.seek(0) @@ -194,17 +210,11 @@ class AutoResponder(commands.Cog): async with httpx.AsyncClient() as client: response = await client.post( "https://0x0.st", - files={ - "file": (_p.name, file.fp, "video/mp4") - }, - headers={ - "User-Agent": "CollegeBot (matrix: @nex:nexy7574.co.uk)" - } + files={"file": (_p.name, file.fp, "video/mp4")}, + headers={"User-Agent": "CollegeBot (matrix: @nex:nexy7574.co.uk)"}, ) if response.status_code == 200: - await message.reply( - "https://embeds.video/" + response.text.strip() - ) + await message.reply("https://embeds.video/" + response.text.strip()) _p.unlink() except Exception as e: self.log.error("Failed to transcode %r: %r", link, e) diff --git a/src/cogs/ffmeta.py b/src/cogs/ffmeta.py index 9adae24..76ea2c2 100644 --- a/src/cogs/ffmeta.py +++ b/src/cogs/ffmeta.py @@ -72,29 +72,26 @@ class FFMeta(commands.Cog): @commands.slash_command() async def jpegify( - self, - ctx: discord.ApplicationContext, - url: str = None, - attachment: discord.Attachment = None, - quality: typing.Annotated[ + self, + ctx: discord.ApplicationContext, + url: str = None, + attachment: discord.Attachment = None, + quality: typing.Annotated[ + int, + discord.Option( int, - discord.Option( - int, - description="The quality of the resulting image from 1%-100%", - default=50, - min_value=1, - max_value=100 - ) - ] = 50, - image_format: typing.Annotated[ - str, - discord.Option( - str, - description="The format of the resulting image", - choices=["jpeg", "webp"], - default="jpeg" - ) - ] = "jpeg" + description="The quality of the resulting image from 1%-100%", + default=50, + min_value=1, + max_value=100, + ), + ] = 50, + image_format: typing.Annotated[ + str, + discord.Option( + str, description="The format of the resulting image", choices=["jpeg", "webp"], default="jpeg" + ), + ] = "jpeg", ): """Converts a given URL or attachment to a JPEG""" if url is None: @@ -124,28 +121,23 @@ class FFMeta(commands.Cog): @commands.slash_command() async def opusinate( - self, - ctx: discord.ApplicationContext, - url: str = None, - attachment: discord.Attachment = None, - bitrate: typing.Annotated[ + self, + ctx: discord.ApplicationContext, + url: str = None, + attachment: discord.Attachment = None, + bitrate: typing.Annotated[ + int, + discord.Option( int, - discord.Option( - int, - description="The bitrate in kilobits of the resulting audio from 1-512", - default=96, - min_value=0, - max_value=512 - ) - ] = 96, - mono: typing.Annotated[ - bool, - discord.Option( - bool, - description="Whether to convert the audio to mono", - default=False - ) - ] = False + description="The bitrate in kilobits of the resulting audio from 1-512", + default=96, + min_value=0, + max_value=512, + ), + ] = 96, + mono: typing.Annotated[ + bool, discord.Option(bool, description="Whether to convert the audio to mono", default=False) + ] = False, ): """Converts a given URL or attachment to an Opus file""" if bitrate == 0: @@ -175,8 +167,10 @@ class FFMeta(commands.Cog): probe_process = await asyncio.create_subprocess_exec( "ffprobe", - "-v", "quiet", - "-print_format", "json", + "-v", + "quiet", + "-print_format", + "json", "-show_format", "-i", temp.name, @@ -207,12 +201,16 @@ class FFMeta(commands.Cog): "-stats", "-i", temp.name, - "-c:a", "libopus", - "-b:a", f"{bitrate}k", + "-c:a", + "libopus", + "-b:a", + f"{bitrate}k", "-vn", "-sn", - "-ac", str(channels), - "-f", "opus", + "-ac", + str(channels), + "-f", + "opus", "-y", "pipe:1", stdout=asyncio.subprocess.PIPE, diff --git a/src/cogs/net.py b/src/cogs/net.py index 7716193..5cf7b6d 100644 --- a/src/cogs/net.py +++ b/src/cogs/net.py @@ -103,8 +103,7 @@ class NetworkCog(commands.Cog): file.write(stderr) file.seek(0) return await ctx.respond( - "Seemingly all output was filtered. Returning raw command output.", - file=discord.File(file, "whois.txt") + "Seemingly all output was filtered. Returning raw command output.", file=discord.File(file, "whois.txt") ) for page in paginator.pages: @@ -243,7 +242,7 @@ class NetworkCog(commands.Cog): paginator.add_line(f"Error: {e}") for page in paginator.pages: await ctx.respond(page) - + @commands.slash_command(name="what-are-matthews-bank-details") async def matthew_bank(self, ctx: discord.ApplicationContext): """For the 80th time""" diff --git a/src/cogs/ollama.py b/src/cogs/ollama.py index 1b6d92e..2e3ec5c 100644 --- a/src/cogs/ollama.py +++ b/src/cogs/ollama.py @@ -166,12 +166,7 @@ class OllamaChatHandler: async def __aiter__(self): async with aiohttp.ClientSession(base_url=self.base_url) as client: async with client.post( - "/api/chat", - json={ - "model": self.model, - "stream": True, - "messages": self.messages - } + "/api/chat", json={"model": self.model, "stream": True, "messages": self.messages} ) as response: response.raise_for_status() async for line in ollama_stream(response.content): @@ -191,10 +186,7 @@ class OllamaClient: self.base_url = base_url self.authorisation = authorisation - def with_client( - self, - timeout: aiohttp.ClientTimeout | float | int | None = None - ) -> aiohttp.ClientSession: + def with_client(self, timeout: aiohttp.ClientTimeout | float | int | None = None) -> aiohttp.ClientSession: """ Creates an instance for a request, with properly populated values. :param timeout: @@ -240,9 +232,9 @@ class OllamaClient: return handler def new_chat( - self, - model: str, - messages: list[dict[str, str]], + self, + model: str, + messages: list[dict[str, str]], ) -> OllamaChatHandler: """ Starts a chat with the given messages. @@ -291,9 +283,7 @@ class ChatHistory: def save_thread(self, thread_id: str): self.log.info("Saving thread:%s - %r", thread_id, self._internal[thread_id]) - self.redis.set( - "threads:" + thread_id, json.dumps(self._internal[thread_id]) - ) + self.redis.set("threads:" + thread_id, json.dumps(self._internal[thread_id])) def create_thread(self, member: discord.Member, default: str | None = None) -> str: """ @@ -304,26 +294,15 @@ class ChatHistory: :return: The thread's ID. """ key = os.urandom(3).hex() - self._internal[key] = { - "member": member.id, - "seed": round(time.time()), - "messages": [] - } + self._internal[key] = {"member": member.id, "seed": round(time.time()), "messages": []} with open("./assets/ollama-prompt.txt") as file: system_prompt = default or file.read() - self.add_message( - key, - "system", - system_prompt - ) + self.add_message(key, "system", system_prompt) return key @staticmethod def _construct_message(role: str, content: str, images: typing.Optional[list[str]]) -> dict[str, str]: - x = { - "role": role, - "content": content - } + x = {"role": role, "content": content} if images: x["images"] = images return x @@ -335,10 +314,8 @@ class ChatHistory: instance = cog.history return list( filter( - lambda v: (ctx.value or v) in v, map( - lambda d: list(d.keys()), - instance.threads_for(ctx.interaction.user) - ) + lambda v: (ctx.value or v) in v, + map(lambda d: list(d.keys()), instance.threads_for(ctx.interaction.user)), ) ) @@ -355,11 +332,11 @@ class ChatHistory: return t def add_message( - self, - thread: str, - role: typing.Literal["user", "assistant", "system"], - content: str, - images: typing.Optional[list[str]] = None + self, + thread: str, + role: typing.Literal["user", "assistant", "system"], + content: str, + images: typing.Optional[list[str]] = None, ) -> None: """ Appends a message to the given thread. @@ -496,56 +473,39 @@ class Ollama(commands.Cog): @commands.slash_command() async def ollama( - self, - ctx: discord.ApplicationContext, - query: typing.Annotated[ + self, + ctx: discord.ApplicationContext, + query: typing.Annotated[ + str, + discord.Option( str, - discord.Option( - str, - "The query to feed into ollama. Not the system prompt.", - ) - ], - model: typing.Annotated[ + "The query to feed into ollama. Not the system prompt.", + ), + ], + model: typing.Annotated[ + str, + discord.Option( str, - discord.Option( - str, - "The model to use for ollama. Defaults to 'llama2-uncensored:latest'.", - default="llama2-uncensored:7b-chat" - ) - ], - server: typing.Annotated[ - str, - discord.Option( - str, - "The server to use for ollama.", - default="next", - choices=SERVER_KEYS - ) - ], - context: typing.Annotated[ - str, - discord.Option( - str, - "The context key of a previous ollama response to use as context.", - default=None - ) - ], - give_acid: typing.Annotated[ - bool, - discord.Option( - bool, - "Whether to give the AI acid, LSD, and other hallucinogens before responding.", - default=False - ) - ], - image: typing.Annotated[ - discord.Attachment, - discord.Option( - discord.Attachment, - "An image to feed into ollama. Only works with llava.", - default=None - ) - ] + "The model to use for ollama. Defaults to 'llama2-uncensored:latest'.", + default="llama2-uncensored:7b-chat", + ), + ], + server: typing.Annotated[ + str, discord.Option(str, "The server to use for ollama.", default="next", choices=SERVER_KEYS) + ], + context: typing.Annotated[ + str, discord.Option(str, "The context key of a previous ollama response to use as context.", default=None) + ], + give_acid: typing.Annotated[ + bool, + discord.Option( + bool, "Whether to give the AI acid, LSD, and other hallucinogens before responding.", default=False + ), + ], + image: typing.Annotated[ + discord.Attachment, + discord.Option(discord.Attachment, "An image to feed into ollama. Only works with llava.", default=None), + ], ): system_query = None if context is not None: @@ -578,8 +538,7 @@ class Ollama(commands.Cog): if image: if fnmatch(model, "llava:*") is False: await ctx.respond( - "You can only use images with llava. Switching model to `llava:latest`.", - delete_after=5 + "You can only use images with llava. Switching model to `llava:latest`.", delete_after=5 ) model = "llava:latest" @@ -613,19 +572,14 @@ class Ollama(commands.Cog): return async with aiohttp.ClientSession( - base_url=server_config["base_url"], - timeout=aiohttp.ClientTimeout( - connect=30, - sock_read=10800, - sock_connect=30, - total=10830 - ) + base_url=server_config["base_url"], + timeout=aiohttp.ClientTimeout(connect=30, sock_read=10800, sock_connect=30, total=10830), ) as session: embed = discord.Embed( title="Checking server...", description=f"Checking that specified model and tag ({model}) are available on the server.", color=discord.Color.blurple(), - timestamp=discord.utils.utcnow() + timestamp=discord.utils.utcnow(), ) embed.set_footer(text="Using server %r" % server, icon_url=server_config.get("icon_url")) await ctx.respond(embed=embed) @@ -636,7 +590,7 @@ class Ollama(commands.Cog): title="Server was offline. Trying next server.", description=f"Trying server {server}...", color=discord.Color.gold(), - timestamp=discord.utils.utcnow() + timestamp=discord.utils.utcnow(), ) embed.set_footer(text="Using server %r" % server, icon_url=server_config.get("icon_url")) await ctx.edit(embed=embed) @@ -650,7 +604,7 @@ class Ollama(commands.Cog): title="All servers are offline.", description="Please try again later.", color=discord.Color.red(), - timestamp=discord.utils.utcnow() + timestamp=discord.utils.utcnow(), ) embed.set_footer(text="Unable to continue.") return await ctx.edit(embed=embed) @@ -665,7 +619,7 @@ class Ollama(commands.Cog): title=f"HTTP {resp.status} {resp.reason!r} while checking for model.", description=f"```{await resp.text() or 'No response body'}```"[:4096], color=discord.Color.red(), - timestamp=discord.utils.utcnow() + timestamp=discord.utils.utcnow(), ) embed.set_footer(text="Unable to continue.") return await ctx.edit(embed=embed) @@ -674,7 +628,7 @@ class Ollama(commands.Cog): title="Connection error while checking for model.", description=f"```{e}```"[:4096], color=discord.Color.red(), - timestamp=discord.utils.utcnow() + timestamp=discord.utils.utcnow(), ) embed.set_footer(text="Unable to continue.") return await ctx.edit(embed=embed) @@ -694,7 +648,7 @@ class Ollama(commands.Cog): title=f"Downloading {model!r}", description=f"Downloading {model!r} from {server_config['base_url']}", color=discord.Color.blurple(), - timestamp=discord.utils.utcnow() + timestamp=discord.utils.utcnow(), ) embed.add_field(name="Progress", value=progress_bar(0)) await ctx.edit(embed=embed) @@ -708,7 +662,7 @@ class Ollama(commands.Cog): title=f"HTTP {response.status} {response.reason!r} while downloading model.", description=f"```{await response.text() or 'No response body'}```"[:4096], color=discord.Color.red(), - timestamp=discord.utils.utcnow() + timestamp=discord.utils.utcnow(), ) embed.set_footer(text="Unable to continue.") return await ctx.edit(embed=embed) @@ -718,7 +672,7 @@ class Ollama(commands.Cog): embed = discord.Embed( title="Download cancelled.", colour=discord.Colour.red(), - timestamp=discord.utils.utcnow() + timestamp=discord.utils.utcnow(), ) return await ctx.edit(embed=embed, view=None) if time.time() >= (last_update + 5.1): @@ -739,17 +693,15 @@ class Ollama(commands.Cog): title="Generating response...", description=">>> ", color=discord.Color.blurple(), - timestamp=discord.utils.utcnow() + timestamp=discord.utils.utcnow(), ) embed.set_author( name=model, url="https://ollama.ai/library/" + model.split(":")[0], - icon_url="https://ollama.ai/public/ollama.png" + icon_url="https://ollama.ai/public/ollama.png", ) embed.add_field( - name="Prompt", - value=">>> " + textwrap.shorten(query, width=1020, placeholder="..."), - inline=False + name="Prompt", value=">>> " + textwrap.shorten(query, width=1020, placeholder="..."), inline=False ) embed.set_footer(text="Using server %r" % server, icon_url=server_config.get("icon_url")) if image_data: @@ -774,10 +726,7 @@ class Ollama(commands.Cog): context = list(__thread.keys())[0] messages = self.history.get_history(context) - user_message = { - "role": "user", - "content": query - } + user_message = {"role": "user", "content": query} if image_data: user_message["images"] = [image_data] messages.append(user_message) @@ -789,12 +738,7 @@ class Ollama(commands.Cog): params["top_p"] = 2 params["repeat_penalty"] = 2 - payload = { - "model": model, - "stream": True, - "options": params, - "messages": messages - } + payload = {"model": model, "stream": True, "options": params, "messages": messages} async with session.post( "/api/chat", json=payload, @@ -805,7 +749,7 @@ class Ollama(commands.Cog): title=f"HTTP {response.status} {response.reason!r} while generating response.", description=f"```{await response.text() or 'No response body'}```"[:4096], color=discord.Color.red(), - timestamp=discord.utils.utcnow() + timestamp=discord.utils.utcnow(), ) embed.set_footer(text="Unable to continue.") return await ctx.edit(embed=embed) @@ -846,7 +790,7 @@ class Ollama(commands.Cog): value = buffer.getvalue() if len(value) >= 4096: embeds = [discord.Embed(title="Done!", colour=discord.Color.green())] - + current_page = "" for word in value.split(): if len(current_page) + len(word) >= 4096: @@ -855,7 +799,7 @@ class Ollama(commands.Cog): current_page += word + " " else: embeds.append(discord.Embed(description=current_page)) - + await ctx.edit(embeds=embeds, view=None) else: await ctx.edit(embed=embed, view=None) @@ -869,25 +813,25 @@ class Ollama(commands.Cog): embed = discord.Embed( title="Timings", description=f"Total: {total_duration}\nLoad: {load_duration}\n" - f"Prompt Eval: {prompt_eval_duration}\nEval: {eval_duration}", + f"Prompt Eval: {prompt_eval_duration}\nEval: {eval_duration}", color=discord.Color.blurple(), - timestamp=discord.utils.utcnow() + timestamp=discord.utils.utcnow(), ) return await ctx.respond(embed=embed, ephemeral=True) @commands.slash_command(name="ollama-history") async def ollama_history( - self, - ctx: discord.ApplicationContext, - thread_id: typing.Annotated[ - str, - discord.Option( - name="thread_id", - description="Thread/Context ID", - type=str, - autocomplete=ChatHistory.autocomplete, - ) - ] + self, + ctx: discord.ApplicationContext, + thread_id: typing.Annotated[ + str, + discord.Option( + name="thread_id", + description="Thread/Context ID", + type=str, + autocomplete=ChatHistory.autocomplete, + ), + ], ): """Shows the history for a thread.""" # await ctx.defer(ephemeral=True) @@ -904,17 +848,11 @@ class Ollama(commands.Cog): if message["role"] == "system": continue max_length = 4000 - len("> **%s**: " % message["role"]) - paginator.add_line( - "> **{}**: {}".format(message["role"], textwrap.shorten(message["content"], max_length)) - ) + paginator.add_line("> **{}**: {}".format(message["role"], textwrap.shorten(message["content"], max_length))) embeds = [] for page in paginator.pages: - embeds.append( - discord.Embed( - description=page - ) - ) + embeds.append(discord.Embed(description=page)) ephemeral = len(embeds) > 1 for chunk in discord.utils.as_chunks(iter(embeds or [discord.Embed(title="No Content.")]), 10): await ctx.respond(embeds=chunk, ephemeral=ephemeral) @@ -929,10 +867,7 @@ class Ollama(commands.Cog): if not content: return await ctx.respond("No content to send to AI.", ephemeral=True) await ctx.defer() - user_message = { - "role": "user", - "content": message.content - } + user_message = {"role": "user", "content": message.content} self.history.add_message(thread, "user", user_message["content"]) for _ in range(10): @@ -947,10 +882,7 @@ class Ollama(commands.Cog): with client.download_model("orca-mini", "3b") as handler: async for _ in handler: self.log.info( - "Downloading orca-mini:3b on server %r - %s (%.2f%%)", - server, - handler.status, - handler.percent + "Downloading orca-mini:3b on server %r - %s (%.2f%%)", server, handler.status, handler.percent ) messages = self.history.get_history(thread) diff --git a/src/cogs/quote_quota.py b/src/cogs/quote_quota.py index 156a47b..dd56333 100644 --- a/src/cogs/quote_quota.py +++ b/src/cogs/quote_quota.py @@ -28,11 +28,7 @@ class QuoteQuota(commands.Cog): return c @staticmethod - def generate_pie_chart( - usernames: list[str], - counts: list[int], - no_other: bool = False - ) -> discord.File: + def generate_pie_chart(usernames: list[str], counts: list[int], no_other: bool = False) -> discord.File: """ Converts the given username and count tuples into a nice pretty pie chart. @@ -81,34 +77,34 @@ class QuoteQuota(commands.Cog): ) fig.subplots_adjust(left=0.1, bottom=0.1, right=0.9, top=0.9, wspace=0.3, hspace=0.4) fio = io.BytesIO() - fig.savefig(fio, format='png') + fig.savefig(fio, format="png") fio.seek(0) return discord.File(fio, filename="pie.png") @commands.slash_command() async def quota( - self, - ctx: discord.ApplicationContext, - days: Annotated[ + self, + ctx: discord.ApplicationContext, + days: Annotated[ + int, + discord.Option( int, - discord.Option( - int, - name="lookback", - description="How many days to look back on. Defaults to 7.", - default=7, - min_value=1, - max_value=365 - ) - ], - merge_other: Annotated[ + name="lookback", + description="How many days to look back on. Defaults to 7.", + default=7, + min_value=1, + max_value=365, + ), + ], + merge_other: Annotated[ + bool, + discord.Option( bool, - discord.Option( - bool, - name="merge_other", - description="Whether to merge authors with less than 5% of the total count into 'Other'.", - default=True - ) - ] + name="merge_other", + description="Whether to merge authors with less than 5% of the total count into 'Other'.", + default=True, + ), + ], ): """Checks the quote quota for the quotes channel.""" now = discord.utils.utcnow() @@ -123,11 +119,7 @@ class QuoteQuota(commands.Cog): authors = {} filtered_messages = 0 total = 0 - async for message in channel.history( - limit=None, - after=oldest, - oldest_first=False - ): + async for message in channel.history(limit=None, after=oldest, oldest_first=False): total += 1 if not message.content: filtered_messages += 1 @@ -163,34 +155,24 @@ class QuoteQuota(commands.Cog): if total: return await ctx.edit( content="No valid messages found in the last {!s} days. " - "Make sure quotes are formatted properly ending with ` - AuthorName`" - " (e.g. `\"This is my quote\" - Jimmy`)".format(days) + "Make sure quotes are formatted properly ending with ` - AuthorName`" + ' (e.g. `"This is my quote" - Jimmy`)'.format(days) ) else: - return await ctx.edit( - content="No messages found in the last {!s} days.".format(days) - ) + return await ctx.edit(content="No messages found in the last {!s} days.".format(days)) file = await asyncio.to_thread( - self.generate_pie_chart, - list(authors.keys()), - list(authors.values()), - merge_other + self.generate_pie_chart, list(authors.keys()), list(authors.values()), merge_other ) return await ctx.edit( content="{:,} messages (out of {:,}) were filtered (didn't follow format?)".format( - filtered_messages, - total + filtered_messages, total ), - file=file + file=file, ) def _metacounter( - self, - messages: list[discord.Message], - filter_func: Callable[[discord.Message], bool], - *, - now: datetime = None + self, messages: list[discord.Message], filter_func: Callable[[discord.Message], bool], *, now: datetime = None ) -> dict[str, float | int]: now = now or discord.utils.utcnow().replace(minute=0, second=0, microsecond=0) counts = { @@ -227,6 +209,7 @@ class QuoteQuota(commands.Cog): :param messages: The messages to process :returns: The stats """ + def is_truth(msg: discord.Message) -> bool: if msg.author.id == 1101439218334576742: # if msg.created_at.timestamp() <= 1713202855.80234: @@ -250,6 +233,7 @@ class QuoteQuota(commands.Cog): :param messages: The messages to process :returns: The stats """ + def is_truth(msg: discord.Message) -> bool: if msg.author.id == 1229496078726860921: # All the tate truths are already tagged. @@ -268,14 +252,9 @@ class QuoteQuota(commands.Cog): :param channel: The channel to process :returns: The stats """ - embed = discord.Embed( - title="Truth Counts", - color=discord.Color.blurple(), - timestamp=discord.utils.utcnow() - ) + embed = discord.Embed(title="Truth Counts", color=discord.Color.blurple(), timestamp=discord.utils.utcnow()) messages: list[discord.Message] = await channel.history( - limit=None, - after=discord.Object(1229487065117233203) + limit=None, after=discord.Object(1229487065117233203) ).flatten() trump_stats = await self._process_trump_truths(messages) tate_stats = await self._process_tate_truths(messages) @@ -296,7 +275,7 @@ class QuoteQuota(commands.Cog): f"**Last Week:** {tate_stats['week']:,} ({tate_stats['per_day']:.1f}/day)\n" f"**Last Day:** {tate_stats['day']:,} ({tate_stats['per_hour']:.1f}/hour)\n" f"**Last Hour:** {tate_stats['hour']:,} ({tate_stats['per_minute']:.1f}/min)" - ) + ), ) return embed @@ -314,7 +293,7 @@ class QuoteQuota(commands.Cog): title="Counting truths, please wait.", description="This may take a minute depending on how insane Trump and Tate are feeling.", color=discord.Color.blurple(), - timestamp=now + timestamp=now, ) await ctx.respond(embed=embed) embed = await self._process_all_messages(channel) diff --git a/src/cogs/screenshot.py b/src/cogs/screenshot.py index 204da13..d3434f6 100644 --- a/src/cogs/screenshot.py +++ b/src/cogs/screenshot.py @@ -29,12 +29,12 @@ RESOLUTIONS = { "480p": "854x480", "360p": "640x360", "240p": "426x240", - "144p": "256x144" + "144p": "256x144", } _RES_OPTION = discord.Option( name="resolution", description="The resolution of the browser, can be WxH, or a preset (e.g. 1080p).", - autocomplete=discord.utils.basic_autocomplete(tuple(RESOLUTIONS.keys())) + autocomplete=discord.utils.basic_autocomplete(tuple(RESOLUTIONS.keys())), ) @@ -64,9 +64,7 @@ class ScreenshotCog(commands.Cog): "plugins.always_open_pdf_externally": False, "download_restrictions": 3, } - _chrome_options.add_experimental_option( - "prefs", prefs - ) + _chrome_options.add_experimental_option("prefs", prefs) return _chrome_options def compress_png(self, input_file: io.BytesIO) -> io.BytesIO: @@ -94,17 +92,14 @@ class ScreenshotCog(commands.Cog): # noinspection PyTypeChecker @commands.slash_command() async def screenshot( - self, - ctx: discord.ApplicationContext, - url: str, - load_timeout: int = 10, - render_timeout: int = None, - eager: bool = None, - resolution: typing.Annotated[ - str, - _RES_OPTION - ] = "1440p", - use_proxy: bool = False + self, + ctx: discord.ApplicationContext, + url: str, + load_timeout: int = 10, + render_timeout: int = None, + eager: bool = None, + resolution: typing.Annotated[str, _RES_OPTION] = "1440p", + use_proxy: bool = False, ): """Screenshots a webpage.""" await ctx.defer() @@ -125,7 +120,7 @@ class ScreenshotCog(commands.Cog): load_timeout, render_timeout, "eager" if eager else "lazy", - resolution + resolution, ) parsed = urlparse(url) await ctx.respond("Initialising...") @@ -139,11 +134,7 @@ class ScreenshotCog(commands.Cog): else: use_proxy = False service = await asyncio.to_thread(ChromeService) - driver: webdriver.Chrome = await asyncio.to_thread( - webdriver.Chrome, - service=service, - options=options - ) + driver: webdriver.Chrome = await asyncio.to_thread(webdriver.Chrome, service=service, options=options) driver.set_page_load_timeout(load_timeout) if resolution: resolution = RESOLUTIONS.get(resolution.lower(), resolution) @@ -218,17 +209,17 @@ class ScreenshotCog(commands.Cog): embed = discord.Embed( title=f"Screenshot of {parsed.hostname}", description=f"Init time: {seconds(start_init, end_init)}s\n" - f"Request time: {seconds(start_request, end_request)}s\n" - f"Wait time: {seconds(start_wait, end_wait)}s\n" - f"Save time: {seconds(start_save, end_save)}s\n" - f"Compress time: {seconds(start_compress, end_compress)}s\n" - f"Cleanup time: {seconds(start_cleanup, end_cleanup)}s\n" - f"Total time: {seconds(start_init, end_cleanup)}s\n" - f"Screenshot size: {screenshot_size_mb}MB\n" - f"Resolution: {resolution}" - f"Used proxy: {use_proxy}", + f"Request time: {seconds(start_request, end_request)}s\n" + f"Wait time: {seconds(start_wait, end_wait)}s\n" + f"Save time: {seconds(start_save, end_save)}s\n" + f"Compress time: {seconds(start_compress, end_compress)}s\n" + f"Cleanup time: {seconds(start_cleanup, end_cleanup)}s\n" + f"Total time: {seconds(start_init, end_cleanup)}s\n" + f"Screenshot size: {screenshot_size_mb}MB\n" + f"Resolution: {resolution}" + f"Used proxy: {use_proxy}", colour=discord.Colour.dark_theme(), - timestamp=discord.utils.utcnow() + timestamp=discord.utils.utcnow(), ) embed.set_image(url="attachment://" + fn) return await ctx.edit(content=None, embed=embed, file=discord.File(file, filename=fn)) diff --git a/src/cogs/ytdl.py b/src/cogs/ytdl.py index c0d7fe6..03cb2ae 100644 --- a/src/cogs/ytdl.py +++ b/src/cogs/ytdl.py @@ -59,11 +59,11 @@ class YTDLCog(commands.Cog): # "max_filesize": (25 * 1024 * 1024) - 256 } self.colours = { - "youtube.com": 0xff0000, - "youtu.be": 0xff0000, + "youtube.com": 0xFF0000, + "youtu.be": 0xFF0000, "tiktok.com": 0x25F5EF, - "instagram.com": 0xe1306c, - "shronk.net": 0xFFF952 + "instagram.com": 0xE1306C, + "shronk.net": 0xFFF952, } async def _init_db(self): @@ -84,13 +84,13 @@ class YTDLCog(commands.Cog): return async def save_link( - self, - message: discord.Message, - webpage_url: str, - format_id: str, - attachment_index: int = 0, - *, - snip: typing.Optional[str] = None + self, + message: discord.Message, + webpage_url: str, + format_id: str, + attachment_index: int = 0, + *, + snip: typing.Optional[str] = None, ): """ Saves a link to discord to prevent having to re-download it. @@ -101,7 +101,7 @@ class YTDLCog(commands.Cog): :param snip: The start and end time to snip the video. e.g. 00:00:00-00:10:00 :return: The created hash key """ - snip = snip or '*' + snip = snip or "*" await self._init_db() async with aiosqlite.connect("./data/ytdl.db") as db: _hash = hashlib.md5(f"{webpage_url}:{format_id}:{snip}".encode()).hexdigest() @@ -113,7 +113,7 @@ class YTDLCog(commands.Cog): snip, message.channel.id, message.id, - attachment_index + attachment_index, ) await db.execute( """ @@ -124,17 +124,12 @@ class YTDLCog(commands.Cog): channel_id=excluded.channel_id, attachment_index=excluded.attachment_index """, - (_hash, message.id, message.channel.id, webpage_url, format_id, attachment_index) + (_hash, message.id, message.channel.id, webpage_url, format_id, attachment_index), ) await db.commit() return _hash - async def get_saved( - self, - webpage_url: str, - format_id: str, - snip: str - ) -> typing.Optional[str]: + async def get_saved(self, webpage_url: str, format_id: str, snip: str) -> typing.Optional[str]: """ Attempts to retrieve the attachment URL of a previously saved download. :param webpage_url: The webpage url @@ -146,15 +141,10 @@ class YTDLCog(commands.Cog): async with aiosqlite.connect("./data/ytdl.db") as db: _hash = hashlib.md5(f"{webpage_url}:{format_id}:{snip}".encode()).hexdigest() self.log.debug( - "Attempting to find a saved download for '%s:%s:%s' (%r).", - webpage_url, - format_id, - snip, - _hash + "Attempting to find a saved download for '%s:%s:%s' (%r).", webpage_url, format_id, snip, _hash ) cursor = await db.execute( - "SELECT message_id, channel_id, attachment_index FROM downloads WHERE key=?", - (_hash,) + "SELECT message_id, channel_id, attachment_index FROM downloads WHERE key=?", (_hash,) ) entry = await cursor.fetchone() if not entry: @@ -199,14 +189,10 @@ class YTDLCog(commands.Cog): "-movflags", "faststart", "-y", - str(new_file) + str(new_file), ] self.log.debug("Running command: ffmpeg %s", " ".join(args)) - process = subprocess.run( - ["ffmpeg", *args], - stdout=subprocess.PIPE, - stderr=subprocess.PIPE - ) + process = subprocess.run(["ffmpeg", *args], stdout=subprocess.PIPE, stderr=subprocess.PIPE) if process.returncode != 0: raise RuntimeError(process.stderr.decode()) return new_file @@ -215,43 +201,33 @@ class YTDLCog(commands.Cog): @commands.max_concurrency(1, wait=False) # @commands.bot_has_permissions(send_messages=True, embed_links=True, attach_files=True) async def yt_dl_command( - self, - ctx: discord.ApplicationContext, - url: typing.Annotated[ + self, + ctx: discord.ApplicationContext, + url: typing.Annotated[str, discord.Option(str, description="The URL to download from.", required=True)], + user_format: typing.Annotated[ + typing.Optional[str], + discord.Option( str, - discord.Option( - str, - description="The URL to download from.", - required=True - ) - ], - user_format: typing.Annotated[ - typing.Optional[str], - discord.Option( - str, - name="format", - description="The name of the format to download. Can also specify resolutions for youtube.", - required=False, - default=None - ) - ], - audio_only: typing.Annotated[ + name="format", + description="The name of the format to download. Can also specify resolutions for youtube.", + required=False, + default=None, + ), + ], + audio_only: typing.Annotated[ + bool, + discord.Option( bool, - discord.Option( - bool, - name="audio-only", - description="Whether to convert result into an m4a file. Overwrites `format` if True.", - required=False, - default=False, - ) - ], - snip: typing.Annotated[ - typing.Optional[str], - discord.Option( - description="A start and end position to trim. e.g. 00:00:00-00:10:00.", - required=False - ) - ] + name="audio-only", + description="Whether to convert result into an m4a file. Overwrites `format` if True.", + required=False, + default=False, + ), + ], + snip: typing.Annotated[ + typing.Optional[str], + discord.Option(description="A start and end position to trim. e.g. 00:00:00-00:10:00.", required=False), + ], ): """Runs yt-dlp and outputs into discord.""" await ctx.defer() @@ -279,10 +255,7 @@ class YTDLCog(commands.Cog): # Overwrite format here to be best audio under 25 megabytes. chosen_format = "ba[filesize<20M]" # Also force sorting by the best audio bitrate first. - options["format_sort"] = [ - "abr", - "br" - ] + options["format_sort"] = ["abr", "br"] options["postprocessors"] = [ {"key": "FFmpegExtractAudio", "preferredquality": "96", "preferredcodec": "best"} ] @@ -290,9 +263,7 @@ class YTDLCog(commands.Cog): options["paths"] = paths with yt_dlp.YoutubeDL(options) as downloader: - await ctx.respond( - embed=discord.Embed().set_footer(text="Downloading (step 1/10)") - ) + await ctx.respond(embed=discord.Embed().set_footer(text="Downloading (step 1/10)")) try: # noinspection PyTypeChecker extracted_info = await asyncio.to_thread(downloader.extract_info, url, download=False) @@ -309,7 +280,7 @@ class YTDLCog(commands.Cog): "fps": "1", "vcodec": "error", "acodec": "error", - "filesize": 0 + "filesize": 0, } title = "error" description = str(e) @@ -362,10 +333,12 @@ class YTDLCog(commands.Cog): title=title, description=description, url=webpage_url, - colour=self.colours.get(domain, discord.Colour.og_blurple()) - ).set_footer(text="Downloading (step 2/10)").set_thumbnail(url=thumbnail_url) + colour=self.colours.get(domain, discord.Colour.og_blurple()), + ) + .set_footer(text="Downloading (step 2/10)") + .set_thumbnail(url=thumbnail_url) ) - previous = await self.get_saved(webpage_url, extracted_info["format_id"], snip or '*') + previous = await self.get_saved(webpage_url, extracted_info["format_id"], snip or "*") if previous: await ctx.edit( content=previous, @@ -375,10 +348,8 @@ class YTDLCog(commands.Cog): colour=discord.Colour.green(), timestamp=discord.utils.utcnow(), url=previous, - fields=[ - discord.EmbedField(name="URL", value=previous, inline=False) - ] - ).set_image(url=previous) + fields=[discord.EmbedField(name="URL", value=previous, inline=False)], + ).set_image(url=previous), ) return try: @@ -410,15 +381,15 @@ class YTDLCog(commands.Cog): "Failed to locate downloaded file. Was supposed to be looking for a file extension of " "%r amongst files %r, however none were found.", extracted_info["ext"], - list(map(str, temp_dir.iterdir())) + list(map(str, temp_dir.iterdir())), ) return await ctx.edit( embed=discord.Embed( title="Error", description="Failed to locate downloaded video file.\n" - f"Files: {', '.join(list(map(str, temp_dir.iterdir())))}", + f"Files: {', '.join(list(map(str, temp_dir.iterdir())))}", colour=discord.Colour.red(), - url=webpage_url + url=webpage_url, ) ) @@ -454,7 +425,7 @@ class YTDLCog(commands.Cog): "-y", "-strict", "2", - str(new_file) + str(new_file), ] async with ctx.channel.typing(): await ctx.edit( @@ -462,15 +433,12 @@ class YTDLCog(commands.Cog): title=f"Trimming from {trim_start} to {trim_end}.", description="Please wait, this may take a couple of minutes.", colour=discord.Colour.og_blurple(), - timestamp=discord.utils.utcnow() + timestamp=discord.utils.utcnow(), ) ) self.log.debug("Running command: 'ffmpeg %s'", " ".join(args)) process = await asyncio.create_subprocess_exec( - "ffmpeg", - *args, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE + "ffmpeg", *args, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE ) stdout, stderr = await process.communicate() self.log.debug("STDOUT:\n%r", stdout.decode()) @@ -481,7 +449,7 @@ class YTDLCog(commands.Cog): title="Error", description=f"Trimming failed:\n```\n{stderr.decode()}\n```", colour=discord.Colour.red(), - url=webpage_url + url=webpage_url, ) ) file = new_file @@ -498,7 +466,7 @@ class YTDLCog(commands.Cog): title="Error", description=f"File is too large to upload ({round(size_bytes / 1024 / 1024)}MB).", colour=discord.Colour.red(), - url=webpage_url + url=webpage_url, ) ) size_megabits = (size_bytes * 8) / 1024 / 1024 @@ -509,7 +477,7 @@ class YTDLCog(commands.Cog): title="Uploading...", description=f"ETA ", colour=discord.Colour.og_blurple(), - timestamp=discord.utils.utcnow() + timestamp=discord.utils.utcnow(), ) ) try: @@ -520,10 +488,10 @@ class YTDLCog(commands.Cog): description="Views: {:,} | Likes: {:,}".format(views or 0, likes or 0), colour=discord.Colour.green(), timestamp=discord.utils.utcnow(), - url=webpage_url - ) + url=webpage_url, + ), ) - await self.save_link(msg, webpage_url, chosen_format_id, snip=snip or '*') + await self.save_link(msg, webpage_url, chosen_format_id, snip=snip or "*") except discord.HTTPException as e: self.log.error(e, exc_info=True) return await ctx.edit( @@ -531,7 +499,7 @@ class YTDLCog(commands.Cog): title="Error", description=f"Upload failed:\n```\n{e}\n```", colour=discord.Colour.red(), - url=webpage_url + url=webpage_url, ) ) diff --git a/src/conf.py b/src/conf.py index e260f8d..0a81373 100644 --- a/src/conf.py +++ b/src/conf.py @@ -13,7 +13,7 @@ if (Path.cwd() / ".git").exists(): stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True, - check=True + check=True, ).stdout.strip() except subprocess.CalledProcessError: log.debug("Unable to auto-detect running version using git.", exc_info=True) @@ -23,29 +23,15 @@ else: VERSION = "unknown" try: - CONFIG = toml.load('config.toml') + CONFIG = toml.load("config.toml") CONFIG.setdefault("logging", {}) CONFIG.setdefault("jimmy", {}) CONFIG.setdefault("ollama", {}) CONFIG.setdefault("rss", {"meta": {"channel": None}}) CONFIG.setdefault("screenshot", {}) CONFIG.setdefault("quote_a", {"channel": None}) - CONFIG.setdefault( - "server", - { - "host": "0.0.0.0", - "port": 8080, - "channel": 1032974266527907901 - } - ) - CONFIG.setdefault( - "redis", - { - "host": "redis", - "port": 6379, - "decode_responses": True - } - ) + CONFIG.setdefault("server", {"host": "0.0.0.0", "port": 8080, "channel": 1032974266527907901}) + CONFIG.setdefault("redis", {"host": "redis", "port": 6379, "decode_responses": True}) except FileNotFoundError: cwd = Path.cwd() log.critical("Unable to locate config.toml in %s.", cwd, exc_info=True) diff --git a/src/main.py b/src/main.py index c2231c7..80b41b1 100644 --- a/src/main.py +++ b/src/main.py @@ -35,15 +35,15 @@ class KumaThread(KillableThread): self.interval = interval self.kill = Event() self.retries = 0 - + def calculate_backoff(self) -> float: rnd = random.uniform(0, 1) retries = min(self.retries, 1000) - t = (2 * 2 ** retries) + rnd + t = (2 * 2**retries) + rnd self.log.debug("Backoff: 2 * (2 ** %d) + %f = %f", retries, rnd, t) # T can never exceed self.interval return max(0, min(self.interval, t)) - + def run(self) -> None: with httpx.Client(http2=True) as client: while not self.kill.is_set(): @@ -79,18 +79,10 @@ logging.basicConfig( show_time=False, show_path=False, markup=True, - console=Console( - width=cols, - height=lns - ) + console=Console(width=cols, height=lns), ), - FileHandler( - filename=CONFIG["logging"].get("file", "jimmy.log"), - mode="a", - encoding="utf-8", - errors="replace" - ) - ] + FileHandler(filename=CONFIG["logging"].get("file", "jimmy.log"), mode="a", encoding="utf-8", errors="replace"), + ], ) for logger in CONFIG["logging"].get("suppress", []): logging.getLogger(logger).setLevel(logging.WARNING) @@ -106,8 +98,7 @@ class Client(commands.Bot): async def start(self, token: str, *, reconnect: bool = True) -> None: if CONFIG["jimmy"].get("uptime_kuma_url"): self.uptime_thread = KumaThread( - CONFIG["jimmy"]["uptime_kuma_url"], - CONFIG["jimmy"].get("uptime_kuma_interval", 60.0) + CONFIG["jimmy"]["uptime_kuma_url"], CONFIG["jimmy"].get("uptime_kuma_interval", 60.0) ) self.uptime_thread.start() await super().start(token, reconnect=reconnect) @@ -124,7 +115,7 @@ bot = Client( case_insensitive=True, strip_after_prefix=True, debug_guilds=CONFIG["jimmy"].get("debug_guilds"), - intents=discord.Intents.all() + intents=discord.Intents.all(), ) for ext in ("ytdl", "net", "screenshot", "ollama", "ffmeta", "quote_quota", "auto_responder"): @@ -163,8 +154,7 @@ async def on_application_command_error(ctx: discord.ApplicationContext, exc: Exc for page in paginator.pages: await ctx.respond(page) else: - await ctx.respond(f"An error occurred while processing your command. Please try again later.\n" - f"{exc}") + await ctx.respond(f"An error occurred while processing your command. Please try again later.\n" f"{exc}") @bot.listen() @@ -183,9 +173,7 @@ async def delete_message(ctx: discord.ApplicationContext, message: discord.Messa if message.author != bot.user: return await ctx.respond("I don't have permission to delete messages in this channel.", delete_after=30) - log.info( - "%s deleted message %s>%s: %r", ctx.author, ctx.channel.name, message.id, message.content - ) + log.info("%s deleted message %s>%s: %r", ctx.author, ctx.channel.name, message.id, message.content) await message.delete(delay=3) await ctx.respond(f"\N{white heavy check mark} Deleted message by {message.author.display_name}.") await ctx.delete(delay=15)