Allow re-use of previous downloads
This commit is contained in:
parent
e2143e94e9
commit
02e3038baa
5 changed files with 108 additions and 6 deletions
|
@ -4,9 +4,12 @@
|
||||||
<option name="autoReloadType" value="SELECTIVE" />
|
<option name="autoReloadType" value="SELECTIVE" />
|
||||||
</component>
|
</component>
|
||||||
<component name="ChangeListManager">
|
<component name="ChangeListManager">
|
||||||
<list default="true" id="aa6d62a8-d64d-4a60-a85f-8d9fa52b6b49" name="Changes" comment="Fix the space strip and JSON output">
|
<list default="true" id="aa6d62a8-d64d-4a60-a85f-8d9fa52b6b49" name="Changes" comment="Fix the last few words not being added">
|
||||||
<change beforePath="$PROJECT_DIR$/.idea/workspace.xml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/workspace.xml" afterDir="false" />
|
<change beforePath="$PROJECT_DIR$/.idea/workspace.xml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/workspace.xml" afterDir="false" />
|
||||||
|
<change beforePath="$PROJECT_DIR$/docker-compose.yml" beforeDir="false" afterPath="$PROJECT_DIR$/docker-compose.yml" afterDir="false" />
|
||||||
|
<change beforePath="$PROJECT_DIR$/requirements.txt" beforeDir="false" afterPath="$PROJECT_DIR$/requirements.txt" afterDir="false" />
|
||||||
<change beforePath="$PROJECT_DIR$/src/cogs/ollama.py" beforeDir="false" afterPath="$PROJECT_DIR$/src/cogs/ollama.py" afterDir="false" />
|
<change beforePath="$PROJECT_DIR$/src/cogs/ollama.py" beforeDir="false" afterPath="$PROJECT_DIR$/src/cogs/ollama.py" afterDir="false" />
|
||||||
|
<change beforePath="$PROJECT_DIR$/src/cogs/ytdl.py" beforeDir="false" afterPath="$PROJECT_DIR$/src/cogs/ytdl.py" afterDir="false" />
|
||||||
</list>
|
</list>
|
||||||
<option name="SHOW_DIALOG" value="false" />
|
<option name="SHOW_DIALOG" value="false" />
|
||||||
<option name="HIGHLIGHT_CONFLICTS" value="true" />
|
<option name="HIGHLIGHT_CONFLICTS" value="true" />
|
||||||
|
@ -81,7 +84,7 @@
|
||||||
<workItem from="1704572901106" duration="125000" />
|
<workItem from="1704572901106" duration="125000" />
|
||||||
<workItem from="1704573038888" duration="7749000" />
|
<workItem from="1704573038888" duration="7749000" />
|
||||||
<workItem from="1704760229948" duration="946000" />
|
<workItem from="1704760229948" duration="946000" />
|
||||||
<workItem from="1704790508193" duration="66000" />
|
<workItem from="1704790508193" duration="2117000" />
|
||||||
</task>
|
</task>
|
||||||
<task id="LOCAL-00001" summary="Update gitignore">
|
<task id="LOCAL-00001" summary="Update gitignore">
|
||||||
<option name="closed" value="true" />
|
<option name="closed" value="true" />
|
||||||
|
@ -243,7 +246,15 @@
|
||||||
<option name="project" value="LOCAL" />
|
<option name="project" value="LOCAL" />
|
||||||
<updated>1704760944356</updated>
|
<updated>1704760944356</updated>
|
||||||
</task>
|
</task>
|
||||||
<option name="localTasksCounter" value="21" />
|
<task id="LOCAL-00021" summary="Fix the last few words not being added">
|
||||||
|
<option name="closed" value="true" />
|
||||||
|
<created>1704790607299</created>
|
||||||
|
<option name="number" value="00021" />
|
||||||
|
<option name="presentableId" value="LOCAL-00021" />
|
||||||
|
<option name="project" value="LOCAL" />
|
||||||
|
<updated>1704790607299</updated>
|
||||||
|
</task>
|
||||||
|
<option name="localTasksCounter" value="22" />
|
||||||
<servers />
|
<servers />
|
||||||
</component>
|
</component>
|
||||||
<component name="TypeScriptGeneratedFilesManager">
|
<component name="TypeScriptGeneratedFilesManager">
|
||||||
|
@ -271,6 +282,7 @@
|
||||||
<MESSAGE value="Fix ollama streaming" />
|
<MESSAGE value="Fix ollama streaming" />
|
||||||
<MESSAGE value="Set output to TTY because why not" />
|
<MESSAGE value="Set output to TTY because why not" />
|
||||||
<MESSAGE value="Fix the space strip and JSON output" />
|
<MESSAGE value="Fix the space strip and JSON output" />
|
||||||
<option name="LAST_COMMIT_MESSAGE" value="Fix the space strip and JSON output" />
|
<MESSAGE value="Fix the last few words not being added" />
|
||||||
|
<option name="LAST_COMMIT_MESSAGE" value="Fix the last few words not being added" />
|
||||||
</component>
|
</component>
|
||||||
</project>
|
</project>
|
|
@ -8,8 +8,10 @@ services:
|
||||||
- ./config.toml:/app/config.toml
|
- ./config.toml:/app/config.toml
|
||||||
- ./jimmy.log:/app/jimmy.log
|
- ./jimmy.log:/app/jimmy.log
|
||||||
- /dev/dri:/dev/dri
|
- /dev/dri:/dev/dri
|
||||||
|
- jimmy-data:/app/data
|
||||||
ollama:
|
ollama:
|
||||||
image: ollama/ollama:latest
|
image: ollama/ollama:latest
|
||||||
|
container_name: ollama
|
||||||
restart: unless-stopped
|
restart: unless-stopped
|
||||||
ports:
|
ports:
|
||||||
- 11434:11434
|
- 11434:11434
|
||||||
|
@ -18,4 +20,4 @@ services:
|
||||||
|
|
||||||
volumes:
|
volumes:
|
||||||
ollama-data:
|
ollama-data:
|
||||||
|
jimmy-data:
|
||||||
|
|
|
@ -10,3 +10,4 @@ pillow==10.2
|
||||||
selenium==4.16
|
selenium==4.16
|
||||||
rich==13.7
|
rich==13.7
|
||||||
dnspython==2.4.2
|
dnspython==2.4.2
|
||||||
|
aiosqlite==0.19.0
|
||||||
|
|
|
@ -226,7 +226,6 @@ class Ollama(commands.Cog):
|
||||||
json={
|
json={
|
||||||
"model": model,
|
"model": model,
|
||||||
"prompt": query,
|
"prompt": query,
|
||||||
# "format": "json",
|
|
||||||
"system": system_prompt,
|
"system": system_prompt,
|
||||||
"stream": True
|
"stream": True
|
||||||
}
|
}
|
||||||
|
@ -259,6 +258,7 @@ class Ollama(commands.Cog):
|
||||||
await ctx.edit(embed=embed)
|
await ctx.edit(embed=embed)
|
||||||
self.log.debug(f"Updating message ({last_update} -> {time.time()})")
|
self.log.debug(f"Updating message ({last_update} -> {time.time()})")
|
||||||
last_update = time.time()
|
last_update = time.time()
|
||||||
|
self.log.debug("Ollama finished consuming.")
|
||||||
embed.title = "Done!"
|
embed.title = "Done!"
|
||||||
embed.color = discord.Color.green()
|
embed.color = discord.Color.green()
|
||||||
await ctx.edit(embed=embed)
|
await ctx.edit(embed=embed)
|
||||||
|
|
|
@ -1,5 +1,6 @@
|
||||||
import asyncio
|
import asyncio
|
||||||
import functools
|
import functools
|
||||||
|
import hashlib
|
||||||
import logging
|
import logging
|
||||||
import tempfile
|
import tempfile
|
||||||
import textwrap
|
import textwrap
|
||||||
|
@ -8,6 +9,7 @@ from pathlib import Path
|
||||||
from urllib.parse import urlparse
|
from urllib.parse import urlparse
|
||||||
|
|
||||||
import discord
|
import discord
|
||||||
|
import aiosqlite
|
||||||
import yt_dlp
|
import yt_dlp
|
||||||
from discord.ext import commands
|
from discord.ext import commands
|
||||||
|
|
||||||
|
@ -62,6 +64,76 @@ class YTDLCog(commands.Cog):
|
||||||
"shronk.net": 0xFFF952
|
"shronk.net": 0xFFF952
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async def _init_db(self):
|
||||||
|
async with aiosqlite.connect("./data/ytdl.db") as db:
|
||||||
|
await db.execute(
|
||||||
|
"""
|
||||||
|
CREATE TABLE IF NOT EXISTS downloads (
|
||||||
|
key TEXT PRIMARY KEY,
|
||||||
|
message_id INTEGER NOT NULL UNIQUE,
|
||||||
|
channel_id INTEGER NOT NULL,
|
||||||
|
webpage_url TEXT NOT NULL,
|
||||||
|
format_id TEXT NOT NULL,
|
||||||
|
attachment_index INTEGER NOT NULL DEFAULT 0,
|
||||||
|
)
|
||||||
|
"""
|
||||||
|
)
|
||||||
|
await db.commit()
|
||||||
|
return aiosqlite.connect("./data/ytdl.db")
|
||||||
|
|
||||||
|
async def save_link(self, message: discord.Message, webpage_url: str, format_id: str, attachment_index: int = 0):
|
||||||
|
"""
|
||||||
|
Saves a link to discord to prevent having to re-download it.
|
||||||
|
:param message: The download message with the attachment.
|
||||||
|
:param webpage_url: The "webpage_url" key of the metadata
|
||||||
|
:param format_id: The "format_Id" key of the metadata
|
||||||
|
:param attachment_index: The index of the attachment. Defaults to 0
|
||||||
|
:return: The created hash key
|
||||||
|
"""
|
||||||
|
async with self._init_db() as db:
|
||||||
|
_hash = hashlib.md5(f"{webpage_url}:{format_id}".encode())
|
||||||
|
await db.execute(
|
||||||
|
"""
|
||||||
|
INSERT INTO downloads (key, message_id, channel_id, webpage_url, format_id, attachment_index)
|
||||||
|
VALUES (?, ?, ?, ?, ?, ?)
|
||||||
|
""",
|
||||||
|
(_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) -> typing.Optional[str]:
|
||||||
|
"""
|
||||||
|
Attempts to retrieve the attachment URL of a previously saved download.
|
||||||
|
:param webpage_url: The webpage url
|
||||||
|
:param format_id: The format ID
|
||||||
|
:return: the URL, if found and valid.
|
||||||
|
"""
|
||||||
|
async with self._init_db() as db:
|
||||||
|
_hash = hashlib.md5(f"{webpage_url}:{format_id}".encode())
|
||||||
|
cursor = await db.execute(
|
||||||
|
"SELECT (message_id, channel_id, attachment_index) FROM downloads WHERE key=?",
|
||||||
|
(_hash,)
|
||||||
|
)
|
||||||
|
entry = await cursor.fetchone()
|
||||||
|
if not entry:
|
||||||
|
return
|
||||||
|
message_id, channel_id, attachment_index = entry
|
||||||
|
channel = self.bot.get_channel(channel_id)
|
||||||
|
if not channel:
|
||||||
|
return
|
||||||
|
try:
|
||||||
|
message = await channel.fetch_message(message_id)
|
||||||
|
except discord.HTTPException:
|
||||||
|
await db.execute("DELETE FROM downloads WHERE key=?", (_hash,))
|
||||||
|
return
|
||||||
|
|
||||||
|
try:
|
||||||
|
return message.attachments[attachment_index].url
|
||||||
|
except IndexError:
|
||||||
|
return
|
||||||
|
|
||||||
|
|
||||||
@commands.slash_command(name="yt-dl")
|
@commands.slash_command(name="yt-dl")
|
||||||
@commands.max_concurrency(1, wait=False)
|
@commands.max_concurrency(1, wait=False)
|
||||||
# @commands.bot_has_permissions(send_messages=True, embed_links=True, attach_files=True)
|
# @commands.bot_has_permissions(send_messages=True, embed_links=True, attach_files=True)
|
||||||
|
@ -200,6 +272,21 @@ class YTDLCog(commands.Cog):
|
||||||
colour=self.colours.get(domain, discord.Colour.og_blurple())
|
colour=self.colours.get(domain, discord.Colour.og_blurple())
|
||||||
).set_footer(text="Downloading (step 2/10)").set_thumbnail(url=thumbnail_url)
|
).set_footer(text="Downloading (step 2/10)").set_thumbnail(url=thumbnail_url)
|
||||||
)
|
)
|
||||||
|
previous = await self.get_saved(webpage_url, extracted_info["format_id"])
|
||||||
|
if previous:
|
||||||
|
await ctx.edit(
|
||||||
|
embed=discord.Embed(
|
||||||
|
title=f"Downloaded {title}!",
|
||||||
|
description="Used previously downloaded attachment.",
|
||||||
|
colour=discord.Colour.green(),
|
||||||
|
timestamp=discord.utils.utcnow(),
|
||||||
|
url=previous,
|
||||||
|
fields=[
|
||||||
|
discord.EmbedField(name="URL", value=previous, inline=False)
|
||||||
|
]
|
||||||
|
).set_image(url=previous)
|
||||||
|
)
|
||||||
|
return
|
||||||
try:
|
try:
|
||||||
await asyncio.to_thread(functools.partial(downloader.download, [url]))
|
await asyncio.to_thread(functools.partial(downloader.download, [url]))
|
||||||
except yt_dlp.DownloadError as e:
|
except yt_dlp.DownloadError as e:
|
||||||
|
|
Loading…
Reference in a new issue