From 02e3038baa6757bbadb37f1d9ec461c942af4be7 Mon Sep 17 00:00:00 2001
From: nexy7574 <me@nexy7574.co.uk>
Date: Tue, 9 Jan 2024 09:34:55 +0000
Subject: [PATCH] Allow re-use of previous downloads

---
 .idea/workspace.xml | 20 ++++++++---
 docker-compose.yml  |  4 ++-
 requirements.txt    |  1 +
 src/cogs/ollama.py  |  2 +-
 src/cogs/ytdl.py    | 87 +++++++++++++++++++++++++++++++++++++++++++++
 5 files changed, 108 insertions(+), 6 deletions(-)

diff --git a/.idea/workspace.xml b/.idea/workspace.xml
index 39105f6..ed385cd 100644
--- a/.idea/workspace.xml
+++ b/.idea/workspace.xml
@@ -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>
\ No newline at end of file
diff --git a/docker-compose.yml b/docker-compose.yml
index cd40b79..bc87dbd 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -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:
diff --git a/requirements.txt b/requirements.txt
index ce1f46d..3634cc9 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -10,3 +10,4 @@ pillow==10.2
 selenium==4.16
 rich==13.7
 dnspython==2.4.2
+aiosqlite==0.19.0
diff --git a/src/cogs/ollama.py b/src/cogs/ollama.py
index a5f5f82..79b9dc7 100644
--- a/src/cogs/ollama.py
+++ b/src/cogs/ollama.py
@@ -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)
diff --git a/src/cogs/ytdl.py b/src/cogs/ytdl.py
index 28732ae..bd20097 100644
--- a/src/cogs/ytdl.py
+++ b/src/cogs/ytdl.py
@@ -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: