Added compatibility for docx and pdf

2024-06-21 17:16:59 -06:00
parent b7f11f17d2
commit 3db0a374e6
4 changed files with 86 additions and 39 deletions
--- a/src/dialogs.py
+++ b/src/dialogs.py
@@ -270,15 +270,24 @@ def create_model_from_file(self):

 # FILE CHOOSER | WORKS

-def attach_file_response(self, file_dialog, result, file_type):
+def attach_file_response(self, file_dialog, result):
+    file_types = {
+        "plain_text": ["txt", "md", "html", "css", "js", "py", "java", "json", "xml"],
+        "image": ["png", "jpeg", "jpg", "webp", "gif"],
+        "pdf": ["pdf"],
+        "docx": ["docx"]
+    }
    try: file = file_dialog.open_finish(result)
    except: return
-    self.attach_file(file.get_path(), file_type)
-
-
-def attach_file(self, filter, file_type):
+    extension = file.get_path().split(".")[-1]
+    file_type = next(key for key, value in file_types.items() if extension in value)
+    if not file_type: return
    if file_type == 'image' and not self.verify_if_image_can_be_used():
        self.show_toast('error', 8, self.main_overlay)
        return
+    self.attach_file(file.get_path(), file_type)
+
+
+def attach_file(self, filter):
    file_dialog = Gtk.FileDialog(default_filter=filter)
-    file_dialog.open(self, None, lambda file_dialog, result, file_type=file_type: attach_file_response(self, file_dialog, result, file_type))
+    file_dialog.open(self, None, lambda file_dialog, result: attach_file_response(self, file_dialog, result))
--- a/src/window.py
+++ b/src/window.py
@@ -21,10 +21,11 @@ import gi
 gi.require_version('GtkSource', '5')
 gi.require_version('GdkPixbuf', '2.0')
 from gi.repository import Adw, Gtk, Gdk, GLib, GtkSource, Gio, GdkPixbuf
-import json, requests, threading, os, re, base64, sys, gettext, locale, webbrowser, subprocess, uuid, shutil, tarfile, tempfile
+import json, requests, threading, os, re, base64, sys, gettext, locale, webbrowser, subprocess, uuid, shutil, tarfile, tempfile, docx
 from time import sleep
 from io import BytesIO
 from PIL import Image
+from pypdf import PdfReader
 from datetime import datetime
 from .available_models import available_models
 from . import dialogs, local_instance, connection_handler, update_history
@@ -90,10 +91,10 @@ class AlpacaWindow(Adw.ApplicationWindow):
    chats_menu_button = Gtk.Template.Child()
    attachment_container = Gtk.Template.Child()
    attachment_box = Gtk.Template.Child()
-    file_filter_image = Gtk.Template.Child()
    file_filter_tar = Gtk.Template.Child()
    file_filter_gguf = Gtk.Template.Child()
-    file_filter_text = Gtk.Template.Child()
+    file_filter_attachments = Gtk.Template.Child()
+    attachment_button = Gtk.Template.Child()
    model_drop_down = Gtk.Template.Child()
    model_string_list = Gtk.Template.Child()

@@ -1129,6 +1130,20 @@ class AlpacaWindow(Adw.ApplicationWindow):
        elif file_type == 'plain_text':
            with open(file_path, 'r') as f:
                return f.read()
+        elif file_type == 'pdf':
+            reader = PdfReader(file_path)
+            if len(reader.pages) == 0: return None
+            text = ""
+            for i, page in enumerate(reader.pages):
+                text += f"\n- Page {i}\n{page.extract_text()}\n"
+            return text
+        elif file_type == 'docx':
+            document = docx.Document(file_path)
+            if len(document.paragraphs) == 0: return None
+            text = ""
+            for paragraph in document.paragraphs:
+                text += f"{paragraph.text}\n"
+            return text

    def remove_attached_file(self, button):
        del self.attachments[button.get_name()]
@@ -1143,7 +1158,12 @@ class AlpacaWindow(Adw.ApplicationWindow):

            button_content = Adw.ButtonContent(
                label=shown_name,
-                icon_name={"image": "image-x-generic-symbolic", "plain_text": "document-text-symbolic"}[file_type]
+                icon_name={
+                    "image": "image-x-generic-symbolic",
+                    "plain_text": "document-text-symbolic",
+                    "pdf": "document-text-symbolic",
+                    "docx": "document-text-symbolic"
+                }[file_type]
            )
            button = Gtk.Button(
                vexpand=True,
@@ -1176,10 +1196,8 @@ class AlpacaWindow(Adw.ApplicationWindow):
        self.get_application().create_action('import_chat', lambda *_: self.import_chat())
        self.get_application().create_action('create_model_from_existing', lambda *_: dialogs.create_model_from_existing(self))
        self.get_application().create_action('create_model_from_file', lambda *_: dialogs.create_model_from_file(self))
-        self.get_application().create_action('attach_image', lambda *_: dialogs.attach_file(self, self.file_filter_image, "image"))
-        self.get_application().create_action('attach_plain_text', lambda *_: dialogs.attach_file(self, self.file_filter_text, "plain_text"))
        self.add_chat_button.connect("clicked", lambda button : self.new_chat())
-
+        self.attachment_button.connect("clicked", lambda button, file_filter=self.file_filter_attachments: dialogs.attach_file(self, file_filter))
        self.create_model_name.get_delegate().connect("insert-text", self.check_alphanumeric)
        self.remote_connection_entry.connect("entry-activated", lambda entry : entry.set_css_classes([]))
        self.remote_connection_switch.connect("notify", lambda pspec, user_data : self.connection_switched())
--- a/src/window.ui
+++ b/src/window.ui
@@ -189,9 +189,7 @@
                                <property name="orientation">0</property>
                                <property name="spacing">12</property>
                                <child>
-                                  <object class="GtkMenuButton" id="attachment_button">
-                                    <property name="menu-model">attachment_menu</property>
-                                    <property name="direction">0</property>
+                                  <object class="GtkButton" id="attachment_button">
                                    <property name="vexpand">false</property>
                                    <property name="valign">3</property>
                                    <property name="tooltip-text" translatable="yes">Attach file</property>
@@ -903,21 +901,8 @@
      </item>
    </section>
  </menu>
-  <menu id="attachment_menu">
-    <section>
-      <item>
-        <attribute name="label" translatable="yes">Plain text file</attribute>
-        <attribute name="action">app.attach_plain_text</attribute>
-      </item>
-      <item>
-        <attribute name="label" translatable="yes">Image</attribute>
-        <attribute name="action">app.attach_image</attribute>
-      </item>
-    </section>
-  </menu>
-  <object class="GtkFileFilter" id="file_filter_text">
+  <object class="GtkFileFilter" id="file_filter_attachments">
    <suffixes>
-      <suffix></suffix>
      <suffix>txt</suffix>
      <suffix>md</suffix>
      <suffix>html</suffix>
@@ -927,17 +912,14 @@
      <suffix>java</suffix>
      <suffix>json</suffix>
      <suffix>xml</suffix>
+      <suffix>pdf</suffix>
+      <suffix>docx</suffix>
+      <suffix>png</suffix>
+      <suffix>jpeg</suffix>
+      <suffix>webp</suffix>
+      <suffix>gif</suffix>
    </suffixes>
  </object>
-  <object class="GtkFileFilter" id="file_filter_image">
-    <mime-types>
-      <mime-type>image/svg+xml</mime-type>
-      <mime-type>image/png</mime-type>
-      <mime-type>image/jpeg</mime-type>
-      <mime-type>image/webp</mime-type>
-      <mime-type>image/gif</mime-type>
-    </mime-types>
-  </object>
  <object class="GtkFileFilter" id="file_filter_tar">
    <mime-types>
      <mime-type>application/x-tar</mime-type>