Black reformat

This commit is contained in:
Nexus 2024-04-16 00:46:26 +01:00
parent a634c64cd6
commit d5466a2b09
Signed by: nex
GPG key ID: 0FA334385D0B689F
9 changed files with 313 additions and 462 deletions

View file

@ -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", [])
@ -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)

View file

@ -83,18 +83,15 @@ class FFMeta(commands.Cog):
description="The quality of the resulting image from 1%-100%",
default=50,
min_value=1,
max_value=100
)
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"
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:
@ -135,17 +132,12 @@ class FFMeta(commands.Cog):
description="The bitrate in kilobits of the resulting audio from 1-512",
default=96,
min_value=0,
max_value=512
)
max_value=512,
),
] = 96,
mono: typing.Annotated[
bool,
discord.Option(
bool,
description="Whether to convert the audio to mono",
default=False
)
] = False
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,

View file

@ -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:

View file

@ -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:
@ -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)),
)
)
@ -359,7 +336,7 @@ class ChatHistory:
thread: str,
role: typing.Literal["user", "assistant", "system"],
content: str,
images: typing.Optional[list[str]] = None
images: typing.Optional[list[str]] = None,
) -> None:
"""
Appends a message to the given thread.
@ -503,49 +480,32 @@ class Ollama(commands.Cog):
discord.Option(
str,
"The query to feed into ollama. Not the system prompt.",
)
),
],
model: typing.Annotated[
str,
discord.Option(
str,
"The model to use for ollama. Defaults to 'llama2-uncensored:latest'.",
default="llama2-uncensored:7b-chat"
)
default="llama2-uncensored:7b-chat",
),
],
server: typing.Annotated[
str,
discord.Option(
str,
"The server to use for ollama.",
default="next",
choices=SERVER_KEYS
)
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
)
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
)
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
)
]
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"
@ -614,18 +573,13 @@ class Ollama(commands.Cog):
async with aiohttp.ClientSession(
base_url=server_config["base_url"],
timeout=aiohttp.ClientTimeout(
connect=30,
sock_read=10800,
sock_connect=30,
total=10830
)
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)
@ -871,7 +815,7 @@ class Ollama(commands.Cog):
description=f"Total: {total_duration}\nLoad: {load_duration}\n"
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)
@ -886,8 +830,8 @@ class Ollama(commands.Cog):
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)

View file

@ -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,7 +77,7 @@ 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")
@ -97,8 +93,8 @@ class QuoteQuota(commands.Cog):
description="How many days to look back on. Defaults to 7.",
default=7,
min_value=1,
max_value=365
)
max_value=365,
),
],
merge_other: Annotated[
bool,
@ -106,9 +102,9 @@ class QuoteQuota(commands.Cog):
bool,
name="merge_other",
description="Whether to merge authors with less than 5% of the total count into 'Other'.",
default=True
)
]
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
@ -164,33 +156,23 @@ class QuoteQuota(commands.Cog):
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)
' (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)

View file

@ -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:
@ -100,11 +98,8 @@ class ScreenshotCog(commands.Cog):
load_timeout: int = 10,
render_timeout: int = None,
eager: bool = None,
resolution: typing.Annotated[
str,
_RES_OPTION
] = "1440p",
use_proxy: bool = False
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)
@ -228,7 +219,7 @@ class ScreenshotCog(commands.Cog):
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))

View file

@ -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):
@ -90,7 +90,7 @@ class YTDLCog(commands.Cog):
format_id: str,
attachment_index: int = 0,
*,
snip: typing.Optional[str] = None
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
@ -217,14 +203,7 @@ class YTDLCog(commands.Cog):
async def yt_dl_command(
self,
ctx: discord.ApplicationContext,
url: typing.Annotated[
str,
discord.Option(
str,
description="The URL to download from.",
required=True
)
],
url: typing.Annotated[str, discord.Option(str, description="The URL to download from.", required=True)],
user_format: typing.Annotated[
typing.Optional[str],
discord.Option(
@ -232,8 +211,8 @@ class YTDLCog(commands.Cog):
name="format",
description="The name of the format to download. Can also specify resolutions for youtube.",
required=False,
default=None
)
default=None,
),
],
audio_only: typing.Annotated[
bool,
@ -243,15 +222,12 @@ class YTDLCog(commands.Cog):
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
)
]
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()),
)
previous = await self.get_saved(webpage_url, extracted_info["format_id"], snip or '*')
.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 "*")
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,7 +381,7 @@ 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(
@ -418,7 +389,7 @@ class YTDLCog(commands.Cog):
description="Failed to locate downloaded video file.\n"
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 <t:{int(eta_seconds + discord.utils.utcnow().timestamp()) + 2}:R>",
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,
)
)

View file

@ -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)

View file

@ -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)