From c4c34a2ec19f6c1cd7b4fe4f6e0cd07664d17b8e Mon Sep 17 00:00:00 2001
From: Locochoco <ivanrwpf@gmail.com>
Date: Tue, 4 Mar 2025 17:17:25 -0300
Subject: [PATCH] added learning from pdf (admin only)

---
 pyproject.toml            |  3 ++-
 src/bingus/cogs/markov.py | 23 +++++++++++++++++++++++
 uv.lock                   | 11 +++++++++++
 3 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index bc01a1d..6fded67 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -13,7 +13,8 @@ dependencies = [
     "py-cord>=2.6.1",
     "python-dotenv>=1.0.1",
     "pytesseract>=0.3.13",
-    "pillow>=11.1.0"
+    "pillow>=11.1.0",
+    "pypdf>=5.3.1"
 ]
 
 [project.scripts]
diff --git a/src/bingus/cogs/markov.py b/src/bingus/cogs/markov.py
index afc97e4..0843854 100644
--- a/src/bingus/cogs/markov.py
+++ b/src/bingus/cogs/markov.py
@@ -4,6 +4,7 @@ import io
 import discord
 import pytesseract
 import PIL
+import pypdf
 from discord.ext import commands
 from discord.message import Message
 from pathlib import Path
@@ -96,6 +97,28 @@ class Markov(commands.Cog):
             else:
                 await ctx.respond(f"{head}:\n{msg}")
 
+    @require_owner
+    @commands.slash_command()
+    async def pdf(
+        self, ctx: discord.ApplicationContext, file: discord.Option(discord.Attachment)
+    ):
+        await ctx.defer(ephemeral=True)
+        raw = await file.read()
+        try:
+            pdf = pypdf.PdfReader(io.BytesIO(raw))
+            # i = 0
+            for page in pdf.pages:
+                # i += 1
+                # printf("Bingus learned a page! {i}/{pdf.get_num_pages()}")
+                text = page.extract_text()
+                self.markov.learn(text)
+                await self.update_words()
+            await ctx.respond(
+                "> Bingus learned something from the pdf!", ephemeral=True
+            )
+        except pypdf.errors.PdfReadError:
+            await ctx.respond("> Bingus only understands pdf files!", ephemeral=True)
+
     @require_owner
     @commands.slash_command()
     async def ocr(
diff --git a/uv.lock b/uv.lock
index d5ceb59..9ebffe9 100644
--- a/uv.lock
+++ b/uv.lock
@@ -90,6 +90,7 @@ dependencies = [
     { name = "msgpack" },
     { name = "pillow" },
     { name = "py-cord" },
+    { name = "pypdf" },
     { name = "pytesseract" },
     { name = "python-dotenv" },
 ]
@@ -100,6 +101,7 @@ requires-dist = [
     { name = "msgpack", specifier = ">=1.1.0" },
     { name = "pillow", specifier = ">=11.1.0" },
     { name = "py-cord", specifier = ">=2.6.1" },
+    { name = "pypdf", specifier = ">=5.3.1" },
     { name = "pytesseract", specifier = ">=0.3.13" },
     { name = "python-dotenv", specifier = ">=1.0.1" },
 ]
@@ -362,6 +364,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/e7/90/2690ded84e34b15ca2619932a358c1b7dc6d28fe845dfbd01929fc33c9da/py_cord-2.6.1-py3-none-any.whl", hash = "sha256:e3d3b528c5e37b0e0825f5b884cbb9267860976c1e4878e28b55da8fd3af834b", size = 1089154 },
 ]
 
+[[package]]
+name = "pypdf"
+version = "5.3.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/da/5b/67df68ec4b934aae9ca89edfb43a869c5edb3bd504dd275be9e83001d3e9/pypdf-5.3.1.tar.gz", hash = "sha256:0b9b715252b3c60bacc052e6a780e8b742cee9b9a2135f6007bb018e22a5adad", size = 5011845 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/f4/0c/75da081f5948e07f373a92087e4808739a3248d308f01c78c9bd4a51defa/pypdf-5.3.1-py3-none-any.whl", hash = "sha256:20ea5b8686faad1b695fda054462b667d5e5f51e25fbbc092f12c5e0bb20d738", size = 302042 },
+]
+
 [[package]]
 name = "pytesseract"
 version = "0.3.13"
-- 
2.51.2