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" />
|
||||
</component>
|
||||
<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$/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/ytdl.py" beforeDir="false" afterPath="$PROJECT_DIR$/src/cogs/ytdl.py" afterDir="false" />
|
||||
</list>
|
||||
<option name="SHOW_DIALOG" value="false" />
|
||||
<option name="HIGHLIGHT_CONFLICTS" value="true" />
|
||||
|
|
@ -81,7 +84,7 @@
|
|||
<workItem from="1704572901106" duration="125000" />
|
||||
<workItem from="1704573038888" duration="7749000" />
|
||||
<workItem from="1704760229948" duration="946000" />
|
||||
<workItem from="1704790508193" duration="66000" />
|
||||
<workItem from="1704790508193" duration="2117000" />
|
||||
</task>
|
||||
<task id="LOCAL-00001" summary="Update gitignore">
|
||||
<option name="closed" value="true" />
|
||||
|
|
@ -243,7 +246,15 @@
|
|||
<option name="project" value="LOCAL" />
|
||||
<updated>1704760944356</updated>
|
||||
</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 />
|
||||
</component>
|
||||
<component name="TypeScriptGeneratedFilesManager">
|
||||
|
|
@ -271,6 +282,7 @@
|
|||
<MESSAGE value="Fix ollama streaming" />
|
||||
<MESSAGE value="Set output to TTY because why not" />
|
||||
<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>
|
||||
</project>
|
||||
|
|
@ -8,8 +8,10 @@ services:
|
|||
- ./config.toml:/app/config.toml
|
||||
- ./jimmy.log:/app/jimmy.log
|
||||
- /dev/dri:/dev/dri
|
||||
- jimmy-data:/app/data
|
||||
ollama:
|
||||
image: ollama/ollama:latest
|
||||
container_name: ollama
|
||||
restart: unless-stopped
|
||||
ports:
|
||||
- 11434:11434
|
||||
|
|
@ -18,4 +20,4 @@ services:
|
|||
|
||||
volumes:
|
||||
ollama-data:
|
||||
|
||||
jimmy-data:
|
||||
|
|
|
|||
|
|
@ -10,3 +10,4 @@ pillow==10.2
|
|||
selenium==4.16
|
||||
rich==13.7
|
||||
dnspython==2.4.2
|
||||
aiosqlite==0.19.0
|
||||
|
|
|
|||
|
|
@ -226,7 +226,6 @@ class Ollama(commands.Cog):
|
|||
json={
|
||||
"model": model,
|
||||
"prompt": query,
|
||||
# "format": "json",
|
||||
"system": system_prompt,
|
||||
"stream": True
|
||||
}
|
||||
|
|
@ -259,6 +258,7 @@ class Ollama(commands.Cog):
|
|||
await ctx.edit(embed=embed)
|
||||
self.log.debug(f"Updating message ({last_update} -> {time.time()})")
|
||||
last_update = time.time()
|
||||
self.log.debug("Ollama finished consuming.")
|
||||
embed.title = "Done!"
|
||||
embed.color = discord.Color.green()
|
||||
await ctx.edit(embed=embed)
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import asyncio
|
||||
import functools
|
||||
import hashlib
|
||||
import logging
|
||||
import tempfile
|
||||
import textwrap
|
||||
|
|
@ -8,6 +9,7 @@ from pathlib import Path
|
|||
from urllib.parse import urlparse
|
||||
|
||||
import discord
|
||||
import aiosqlite
|
||||
import yt_dlp
|
||||
from discord.ext import commands
|
||||
|
||||
|
|
@ -62,6 +64,76 @@ class YTDLCog(commands.Cog):
|
|||
"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.max_concurrency(1, wait=False)
|
||||
# @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())
|
||||
).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:
|
||||
await asyncio.to_thread(functools.partial(downloader.download, [url]))
|
||||
except yt_dlp.DownloadError as e:
|
||||
|
|
|
|||
Loading…
Reference in a new issue