Add support formats for llamaindex

Author: yjg30737

pyproject.toml

@@ -21,9 +21,12 @@ dependencies = [
     "openai",
     "anthropic",
     "google-generativeai",
-    "llama-index",
     "replicate",
 
+    "llama-index",
+    "docx2txt",
+    "openpyxl",
+
     "g4f",
     "curl_cffi",
     "litellm",

pyqt_openai/__init__.py

@@ -487,7 +487,7 @@ def move_bin(filename, dst_dir):
         "env_var_name": "ANTHROPIC_API_KEY",
         "api_key": "",
         "manual_url": HOW_TO_GET_CLAUDE_API_KEY_URL,
-        "model_list": ["claude-3-5-sonnet-20240620"]
+        "model_list": ["claude-3-haiku-20240307", "claude-3-5-sonnet-20240620"]
     },
     # AWS Sagemaker
     {
@@ -903,6 +903,10 @@ def move_bin(filename, dst_dir):
 MAXIMUM_MESSAGES_IN_PARAMETER = 40
 MAXIMUM_MESSAGES_IN_PARAMETER_RANGE = 2, 1000
 
+# llamaIndex
+LLAMA_INDEX_DEFAULT_SUPPORTED_FORMATS_LIST = ['.txt']
+LLAMA_INDEX_DEFAULT_ALL_SUPPORTED_FORMATS_LIST = ['.txt', '.docx', '.hwp', '.ipynb', '.csv', '.jpeg', '.jpg', '.mbox', '.md', '.mp3', '.mp4', '.pdf', '.png', '.ppt', '.pptx', '.pptm']
+
 # PROMPT
 ## DEFAULT JSON FILENAME FOR PROMPT
 AWESOME_CHATGPT_PROMPTS_FILENAME = "prompt_res/awesome_chatgpt_prompts.json"
@@ -1142,6 +1146,7 @@ def move_bin(filename, dst_dir):
         # Llama Index
         "use_llama_index": False,
         "llama_index_directory": "",
+        "llama_index_supported_formats": LLAMA_INDEX_DEFAULT_SUPPORTED_FORMATS_LIST,
         # Customize
         "background_image": "",
         "user_image": DEFAULT_USER_IMAGE_PATH,

pyqt_openai/aboutDialog.py

@@ -51,10 +51,7 @@ def __initUi(self):
         descWidget3.setText(
             f"""
         <br/><br/>Contact: {CONTACT}<br/>
-        <p>Powered by PySide6</p>
-        <p>Powered by GPT4Free</p>
-        <p>Powered by LiteLLM</p>
-        <p>Powered by LlamaIndex</p>
+        <p>Powered by<br><br> PySide6, GPT4Free, LiteLLM,<br>LlamaIndex</p>
         """
         )
 

pyqt_openai/chat_widget/chatMainWidget.py

@@ -100,13 +100,6 @@ def __initUi(self):
 
         self.__chatWidget.setG4F(self.__chatRightSideBarWidget.currentTabIdx())
 
-        try:
-            self.__chatRightSideBarWidget.onDirectorySelected.connect(
-                LLAMAINDEX_WRAPPER.set_directory
-            )
-        except Exception as e:
-            QMessageBox.critical(self, LangClass.TRANSLATIONS["Error"], str(e))
-
         self.__promptGeneratorWidget = PromptGeneratorWidget()
 
         self.__sideBarBtn = Button()

pyqt_openai/chat_widget/right_sidebar/chatRightSideBarWidget.py

@@ -1,17 +1,18 @@
 from functools import partial
 
 from PySide6.QtCore import Signal
-from PySide6.QtWidgets import QScrollArea, QWidget, QTabWidget, QGridLayout
+from PySide6.QtWidgets import QScrollArea, QWidget, QTabWidget, QGridLayout, QMessageBox
 
 from pyqt_openai.chat_widget.right_sidebar.usingAPIPage import UsingAPIPage
 from pyqt_openai.chat_widget.right_sidebar.llama_widget.llamaPage import LlamaPage
 from pyqt_openai.chat_widget.right_sidebar.usingG4FPage import UsingG4FPage
 from pyqt_openai.config_loader import CONFIG_MANAGER
+from pyqt_openai.globals import LLAMAINDEX_WRAPPER
+from pyqt_openai.lang.translations import LangClass
 
 
 class ChatRightSideBarWidget(QScrollArea):
     onTabChanged = Signal(int)
-    onDirectorySelected = Signal(str)
     onToggleJSON = Signal(bool)
 
     def __init__(self, parent=None):
@@ -65,7 +66,11 @@ def __tabChanged(self, idx):
     def __onDirectorySelected(self, selected_dirname):
         self.__llama_index_directory = selected_dirname
         CONFIG_MANAGER.set_general_property("llama_index_directory", selected_dirname)
-        self.onDirectorySelected.emit(selected_dirname)
+        try:
+            LLAMAINDEX_WRAPPER.set_directory(selected_dirname)
+        except Exception as e:
+            QMessageBox.critical(self, LangClass.TRANSLATIONS["Error"], str(e))
+
 
     def currentTabIdx(self):
         return self.__cur_idx

pyqt_openai/chat_widget/right_sidebar/llama_widget/filesWidget.py

@@ -14,6 +14,7 @@
 )
 
 from pyqt_openai import TEXT_FILE_EXT_LIST, QFILEDIALOG_DEFAULT_DIRECTORY
+from pyqt_openai.config_loader import CONFIG_MANAGER
 from pyqt_openai.lang.translations import LangClass
 from pyqt_openai.util.common import getSeparator
 
@@ -29,14 +30,15 @@ def __init__(self, parent=None):
         self.__initUi()
 
     def __initVal(self):
-        self.__dirLblPrefix = LangClass.TRANSLATIONS["Directory"]
-        self.__curDirName = ""
+        self.__directory_label_prefix = LangClass.TRANSLATIONS["Directory"]
+        self.__current_directory_name = ""
+        self.__extension = CONFIG_MANAGER.get_general_property("llama_index_supported_formats")
 
     def __initUi(self):
         lbl = QLabel(LangClass.TRANSLATIONS["Files"])
         setDirBtn = QPushButton(LangClass.TRANSLATIONS["Set Directory"])
         setDirBtn.clicked.connect(self.setDirectory)
-        self.__dirLbl = QLabel(self.__dirLblPrefix)
+        self.__dirLbl = QLabel(self.__directory_label_prefix)
 
         lay = QHBoxLayout()
         lay.addWidget(lbl)
@@ -60,6 +62,15 @@ def __initUi(self):
         lay.setContentsMargins(0, 0, 0, 0)
         self.setLayout(lay)
 
+    def setExtension(self, ext):
+        # Set extension
+        self.__extension = ext
+        CONFIG_MANAGER.set_general_property("llama_index_supported_formats", ext)
+
+        # Refresh list based on new extension
+        self.__listWidget.clear()
+        self.setDirectory(self.__current_directory_name, called_from_btn=False)
+
     def setDirectory(self, directory=None, called_from_btn=True):
         try:
             if called_from_btn:
@@ -74,28 +85,28 @@ def setDirectory(self, directory=None, called_from_btn=True):
                 self.__listWidget.clear()
                 filenames = list(
                     filter(
-                        lambda x: os.path.splitext(x)[-1] in TEXT_FILE_EXT_LIST,
+                        lambda x: os.path.splitext(x)[-1] in self.__extension,
                         os.listdir(directory),
                     )
                 )
                 self.__listWidget.addItems(filenames)
                 self.itemUpdate.emit(len(filenames) > 0)
-                self.__curDirName = directory
-                self.__dirLbl.setText(self.__curDirName.split("/")[-1])
+                self.__current_directory_name = directory
+                self.__dirLbl.setText(self.__current_directory_name.split("/")[-1])
 
                 self.__listWidget.setCurrentRow(0)
                 # activate event as clicking first item (because this selects the first item anyway)
                 self.clicked.emit(
                     os.path.join(
-                        self.__curDirName, self.__listWidget.currentItem().text()
+                        self.__current_directory_name, self.__listWidget.currentItem().text()
                     )
                 )
                 self.onDirectorySelected.emit()
         except Exception as e:
             print(e)
 
     def getDirectory(self):
-        return self.__curDirName
+        return self.__current_directory_name
 
     def __sendDirectory(self, item):
-        self.clicked.emit(os.path.join(self.__curDirName, item.text()))
+        self.clicked.emit(os.path.join(self.__current_directory_name, item.text()))

pyqt_openai/chat_widget/right_sidebar/llama_widget/llamaPage.py

@@ -1,12 +1,13 @@
 from PySide6.QtCore import Qt, Signal
 from PySide6.QtGui import QFont
-from PySide6.QtWidgets import QTextBrowser
+from PySide6.QtWidgets import QTextBrowser, QMessageBox
 from PySide6.QtWidgets import QWidget, QLabel, QVBoxLayout
 
 from pyqt_openai import SMALL_LABEL_PARAM
 from pyqt_openai.config_loader import CONFIG_MANAGER
 from pyqt_openai.chat_widget.right_sidebar.llama_widget.filesWidget import FilesWidget
 from pyqt_openai.chat_widget.right_sidebar.llama_widget.supportedFileFormatsWidget import SupportedFileFormatsWidget
+from pyqt_openai.globals import LLAMAINDEX_WRAPPER
 from pyqt_openai.lang.translations import LangClass
 
 
@@ -26,6 +27,7 @@ def __initUi(self):
         self.__filesWidget.onDirectorySelected.connect(self.__onDirectorySelected)
 
         self.__supportedFileFormatsWidget = SupportedFileFormatsWidget()
+        self.__supportedFileFormatsWidget.checkedSignal.connect(self.__formatCheckedSignal)
 
         self.__txtBrowser = QTextBrowser()
         self.__txtBrowser.setPlaceholderText(
@@ -49,10 +51,18 @@ def __onDirectorySelected(self):
         selected_dirname = self.__filesWidget.getDirectory()
         self.onDirectorySelected.emit(selected_dirname)
 
-    def __setTextInBrowser(self, txt_file):
-        with open(txt_file, "r", encoding="utf-8") as f:
-            self.__txtBrowser.setText(f.read())
+    def __setTextInBrowser(self, file):
+        try:
+            with open(file, "r", encoding="utf-8") as f:
+                self.__txtBrowser.setText(f.read())
+        except UnicodeDecodeError as e:
+            self.__txtBrowser.setText('Some files like Excel files cannot be previewed.')
+        except Exception as e:
+            print(e)
 
     def setDirectory(self):
         directory = CONFIG_MANAGER.get_general_property("llama_index_directory")
         self.__filesWidget.setDirectory(directory, called_from_btn=False)
+
+    def __formatCheckedSignal(self, ext):
+        self.__filesWidget.setExtension(ext)

pyqt_openai/chat_widget/right_sidebar/llama_widget/supportedFileFormatsWidget.py

@@ -1,9 +1,14 @@
+from PySide6.QtCore import Signal, Qt
 from PySide6.QtWidgets import QWidget, QVBoxLayout, QLabel
 
+from pyqt_openai import LLAMA_INDEX_DEFAULT_ALL_SUPPORTED_FORMATS_LIST
+from pyqt_openai.config_loader import CONFIG_MANAGER
 from pyqt_openai.widgets.checkBoxListWidget import CheckBoxListWidget
 
 
 class SupportedFileFormatsWidget(QWidget):
+    checkedSignal = Signal(list)
+
     def __init__(self, parent=None):
         super().__init__(parent)
         self.__initUi()
@@ -12,15 +17,21 @@ def __initUi(self):
         self.__listWidget = CheckBoxListWidget()
         self.__listWidget.checkedSignal.connect(self.__sendCheckedSignal)
 
-        ext_lst = ['.txt', '.docx', '.xlsx', '.md', '.pdf', '.html']
+        all_supported_format_lst = LLAMA_INDEX_DEFAULT_ALL_SUPPORTED_FORMATS_LIST
+        self.__listWidget.addItems(all_supported_format_lst)
 
-        self.__listWidget.addItems(ext_lst, checked=True)
+        current_supported_format_lst = CONFIG_MANAGER.get_general_property("llama_index_supported_formats")
+        for i in range(self.__listWidget.count()):
+            supported_format = self.__listWidget.item(i).text()
+            if supported_format in current_supported_format_lst:
+                self.__listWidget.item(i).setCheckState(Qt.CheckState.Checked)
 
         lay = QVBoxLayout()
         # TODO LANGUAGE
         lay.addWidget(QLabel('Supported File Formats'))
         lay.addWidget(self.__listWidget)
+        lay.setContentsMargins(0, 0, 0, 0)
         self.setLayout(lay)
 
     def __sendCheckedSignal(self, r_idx, state):
-        print(r_idx, state)
+        self.checkedSignal.emit(self.__listWidget.getCheckedItemsText())

pyqt_openai/chat_widget/right_sidebar/usingAPIPage.py

@@ -279,7 +279,7 @@ def __initUi(self):
         self.__llamaChkBox = QCheckBox()
         self.__llamaChkBox.setChecked(self.__use_llama_index)
         self.__llamaChkBox.toggled.connect(self.__use_llama_indexChecked)
-        self.__llamaChkBox.setText(LangClass.TRANSLATIONS["Use LlamaIndex"])
+        self.__llamaChkBox.setText(LangClass.TRANSLATIONS["Use LlamaIndex (You need OpenAI API key)"])
 
         lay = QVBoxLayout()
         lay.addWidget(manualBrowser)

pyqt_openai/util/common.py

@@ -578,7 +578,7 @@ def get_chat_model(is_g4f=False):
     else:
         all_models = []
         for obj in DEFAULT_API_CONFIGS:
-            all_models.extend(obj["model_list"])
+            all_models.extend(obj.get("model_list", []))
         return all_models
 
 def get_gemini_argument(model, system, messages, cur_text, stream, images):
@@ -678,7 +678,7 @@ def get_message_obj(role, content):
 # Check which provider a specific model belongs to
 def get_provider_from_model(model):
     for obj in DEFAULT_API_CONFIGS:
-        if model in obj["model_list"]:
+        if model in obj.get("model_list", []):
             return obj["display_name"]
     return None