From 64ce627bb7e2c34c9e65c30b1c72014b11c58b66 Mon Sep 17 00:00:00 2001
From: Rene Teigen <rene9917@gmail.com>
Date: Mon, 30 Jan 2023 20:03:05 +0000
Subject: [PATCH] Update filetypes

---
 models/index_model.py | 9 ++++++++-
 pyproject.toml        | 1 +
 requirements.txt      | 3 ++-
 3 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/models/index_model.py b/models/index_model.py
index e53909d..2d6d2af 100644
--- a/models/index_model.py
+++ b/models/index_model.py
@@ -29,7 +29,14 @@ class Index_handler:
     
         try:
             temp_path = tempfile.TemporaryDirectory()
-            temp_file = tempfile.NamedTemporaryFile(suffix=".txt", dir=temp_path.name, delete=False)
+            if file.content_type.startswith("text/plain"):
+                suffix = ".txt"
+            elif file.content_type.startswith("application/pdf"):
+                suffix = ".pdf"
+            else:
+                await ctx.respond("Only accepts txt or pdf files")
+                return
+            temp_file = tempfile.NamedTemporaryFile(suffix=suffix, dir=temp_path.name, delete=False)
             await file.save(temp_file.name)
             index = await self.loop.run_in_executor(None, partial(self.index_file, temp_path.name))
             self.index_storage[ctx.user.id] = index
diff --git a/pyproject.toml b/pyproject.toml
index 9213da9..a9a20b0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -32,6 +32,7 @@ dependencies = [
   "flask",
   "beautifulsoup4",
   "gpt-index",
+  "PyPDF2",
 ]
 dynamic = ["version"]
 [project.scripts]
diff --git a/requirements.txt b/requirements.txt
index 03aea68..91c86b1 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -10,4 +10,5 @@ sqlitedict==2.1.0
 backoff==2.2.1
 flask==2.2.2
 beautifulsoup4==4.11.1
-gpt-index==0.2.16
\ No newline at end of file
+gpt-index==0.2.16
+PyPDF2==3.0.1
\ No newline at end of file