Merge pull request #124 from yjg30737/Dev

Dev

pyqt_openai/chatGPTImportDialog.py

@@ -0,0 +1,119 @@
+from qtpy.QtCore import Qt
+from qtpy.QtWidgets import QMessageBox, QPushButton, QGroupBox, QTableWidgetItem, \
+    QLabel, QDialogButtonBox, QCheckBox, QDialog, QVBoxLayout, QSpinBox, QAbstractItemView
+
+from pyqt_openai.constants import THREAD_ORDERBY
+from pyqt_openai.util.script import get_conversation_from_chatgpt, get_chatgpt_data
+from pyqt_openai.widgets.checkBoxTableWidget import CheckBoxTableWidget
+from pyqt_openai.widgets.findPathWidget import FindPathWidget
+
+
+class ChatGPTImportDialog(QDialog):
+    def __init__(self):
+        super(ChatGPTImportDialog, self).__init__()
+        self.__initVal()
+        self.__initUi()
+
+    def __initVal(self):
+        # Get the most recent n conversation threads
+        self.__most_recent_n = 10
+        # Data to be imported
+        self.__data = []
+
+    def __initUi(self):
+        self.setWindowTitle("Import ChatGPT Data")
+        self.setWindowFlags(Qt.WindowType.Window | Qt.WindowType.WindowCloseButtonHint)
+
+        findPathWidget = FindPathWidget()
+        findPathWidget.added.connect(self.__setPath)
+        findPathWidget.getLineEdit().setPlaceholderText('Select a ChatGPT JSON file')
+        findPathWidget.setExtOfFiles('JSON Files (*.json)')
+
+        self.__chkBoxMostRecent = QCheckBox('Get most recent')
+
+        self.__mostRecentNSpinBox = QSpinBox()
+        self.__mostRecentNSpinBox.setRange(1, 10000)
+        self.__mostRecentNSpinBox.setValue(self.__most_recent_n)
+        self.__mostRecentNSpinBox.setEnabled(False)
+
+        self.__chkBoxMostRecent.stateChanged.connect(lambda state: self.__mostRecentNSpinBox.setEnabled(state == Qt.CheckState.Checked))
+
+        chatGPTImportGrpBox = QGroupBox('ChatGPT Import Options')
+        lay = QVBoxLayout()
+        lay.addWidget(self.__chkBoxMostRecent)
+        lay.addWidget(self.__mostRecentNSpinBox)
+        chatGPTImportGrpBox.setLayout(lay)
+
+        self.__checkBoxTableWidget = CheckBoxTableWidget()
+        self.__checkBoxTableWidget.setColumnCount(0)
+        # self.__checkBoxTableWidget.checkedSignal.connect(self.getData)
+
+        self.__allCheckBox = QCheckBox('Select All')
+        self.__allCheckBox.stateChanged.connect(self.__checkBoxTableWidget.toggleState) # if allChkBox is checked, tablewidget checkboxes will also be checked
+
+        lay = QVBoxLayout()
+        lay.addWidget(QLabel('Select the threads you want to import.'))
+        lay.addWidget(self.__allCheckBox)
+        lay.addWidget(self.__checkBoxTableWidget)
+
+        self.__chatGPTDataGroupBox = QGroupBox('ChatGPT Data')
+        self.__chatGPTDataGroupBox.setLayout(lay)
+        self.__chatGPTDataGroupBox.setEnabled(False)
+
+        self.__buttonBox = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel)
+        self.__buttonBox.accepted.connect(self.__accept)
+        self.__buttonBox.rejected.connect(self.reject)
+        self.__buttonBox.button(QDialogButtonBox.Ok).setEnabled(False)
+
+        lay = QVBoxLayout()
+        lay.addWidget(findPathWidget)
+        lay.addWidget(chatGPTImportGrpBox)
+        lay.addWidget(self.__chatGPTDataGroupBox)
+        lay.addWidget(self.__buttonBox)
+
+        self.setLayout(lay)
+
+        self.resize(800, 600)
+
+    def __setData(self):
+        checked_rows = self.__checkBoxTableWidget.getCheckedRows()
+        self.__data = get_chatgpt_data([self.__data[r] for r in checked_rows])
+
+    def __toggleOkButton(self):
+        self.__buttonBox.button(QDialogButtonBox.Ok).setEnabled(len(self.__checkBoxTableWidget.getCheckedRows()) > 0)
+
+    def __setPath(self, path):
+        try:
+            most_recent_n = self.__mostRecentNSpinBox.value() if self.__chkBoxMostRecent.isChecked() else None
+            result_dict = get_conversation_from_chatgpt(path, most_recent_n)
+            columns = result_dict['columns']
+            self.__data = result_dict['data']
+            self.__checkBoxTableWidget.setHorizontalHeaderLabels(columns)
+            self.__checkBoxTableWidget.setRowCount(len(self.__data))
+
+            for r_idx, r in enumerate(self.__data):
+                for c_idx, c in enumerate(columns):
+                    v = r[c]
+                    self.__checkBoxTableWidget.setItem(r_idx, c_idx+1, QTableWidgetItem(str(v)))
+
+            self.__checkBoxTableWidget.resizeColumnsToContents()
+            self.__checkBoxTableWidget.setEditTriggers(QAbstractItemView.EditTrigger.NoEditTriggers)
+            if THREAD_ORDERBY in columns:
+                self.__checkBoxTableWidget.sortByColumn(columns.index(THREAD_ORDERBY)+1, Qt.SortOrder.DescendingOrder)
+            self.__chatGPTDataGroupBox.setEnabled(True)
+            self.__allCheckBox.setChecked(True)
+            self.__toggleOkButton()
+
+            self.__checkBoxTableWidget.hideColumn(1)
+        except Exception as e:
+            QMessageBox.critical(self, "Error", 'Check whether the file is a valid JSON file for importing.')
+
+    def __accept(self):
+        if len(self.__checkBoxTableWidget.getCheckedRows()) > 0:
+            self.__setData()
+            self.accept()
+        else:
+            QMessageBox.critical(self, "Error", 'Select at least one thread to import.')
+
+    def getData(self):
+        return self.__data

pyqt_openai/chatNavWidget.py

@@ -1,13 +1,15 @@
-from PyQt5.QtWidgets import QDialog
 from qtpy.QtCore import Signal, QSortFilterProxyModel, Qt
 from qtpy.QtSql import QSqlTableModel, QSqlQuery
-from qtpy.QtWidgets import QApplication, QWidget, QVBoxLayout, QMessageBox, QStyledItemDelegate, QTableView, \
+from qtpy.QtWidgets import QWidget, QVBoxLayout, QMessageBox, QPushButton, QStyledItemDelegate, QTableView, \
     QAbstractItemView, \
     QHBoxLayout, \
-    QLabel, QSpacerItem, QSizePolicy, QFileDialog, QComboBox
+    QLabel, QSpacerItem, QSizePolicy, QFileDialog, QComboBox, QDialog
 
 # for search feature
+from pyqt_openai.chatGPTImportDialog import ChatGPTImportDialog
+from pyqt_openai.constants import THREAD_ORDERBY
 from pyqt_openai.exportDialog import ExportDialog
+from pyqt_openai.importDialog import ImportDialog
 from pyqt_openai.models import ChatThreadContainer
 from pyqt_openai.pyqt_openai_data import DB
 from pyqt_openai.widgets.button import Button
@@ -58,6 +60,8 @@ class ChatNavWidget(QWidget):
     cleared = Signal()
     onImport = Signal(str)
     onExport = Signal(list)
+    onChatGPTImport = Signal(list)
+    onFavoriteClicked = Signal(bool)
 
     def __init__(self, columns, table_nm):
         super().__init__()
@@ -143,7 +147,7 @@ def __initUi(self):
             self.__model.setHeaderData(i, Qt.Orientation.Horizontal, self.__columns[i])
         self.__model.select()
         # descending order by insert date
-        idx = self.__columns.index('insert_dt')
+        idx = self.__columns.index(THREAD_ORDERBY)
         self.__model.sort(idx, Qt.SortOrder.DescendingOrder)
 
         # init the proxy model
@@ -171,9 +175,14 @@ def __initUi(self):
         self.__tableView.clicked.connect(self.__clicked)
         self.__tableView.activated.connect(self.__clicked)
 
+        self.__favoriteBtn = QPushButton('Favorite List')
+        self.__favoriteBtn.setCheckable(True)
+        self.__favoriteBtn.toggled.connect(self.__onFavoriteClicked)
+
         lay = QVBoxLayout()
         lay.addWidget(menuWidget)
         lay.addWidget(self.__tableView)
+        lay.addWidget(self.__favoriteBtn)
         self.setLayout(lay)
 
         self.refreshData()
@@ -186,19 +195,34 @@ def add(self, called_from_parent=False):
         self.__model.select()
 
     def __import(self):
-        filename = QFileDialog.getOpenFileName(self, 'Import', '', 'SQLite DB files (*.db)')
-        if filename:
-            filename = filename[0]
-            self.onImport.emit(filename)
+        dialog = ImportDialog()
+        reply = dialog.exec()
+        if reply == QDialog.Accepted:
+            import_type = dialog.getImportType()
+            if import_type == 'General':
+                filename = QFileDialog.getOpenFileName(self, 'Import', '', 'JSON files (*.json)')
+                if filename:
+                    filename = filename[0]
+                    if filename:
+                        self.onImport.emit(filename)
+            else:
+                chatgptDialog = ChatGPTImportDialog()
+                reply = chatgptDialog.exec()
+                if reply == QDialog.Accepted:
+                    data = chatgptDialog.getData()
+                    self.onChatGPTImport.emit(data)
 
     def __export(self):
         columns = ChatThreadContainer.get_keys()
-        data = DB.selectAllConv()
-        sort_by = 'update_dt'
-        dialog = ExportDialog(columns, data, sort_by=sort_by)
-        reply = dialog.exec()
-        if reply == QDialog.Accepted:
-            self.onExport.emit(dialog.getSelectedIds())
+        data = DB.selectAllThread()
+        sort_by = THREAD_ORDERBY
+        if len(data) > 0:
+            dialog = ExportDialog(columns, data, sort_by=sort_by)
+            reply = dialog.exec()
+            if reply == QDialog.Accepted:
+                self.onExport.emit(dialog.getSelectedIds())
+        else:
+            QMessageBox.information(self, 'Information', 'No data to export.')
 
     def __updated(self, i, r):
         # send updated signal
@@ -213,22 +237,28 @@ def refreshData(self, title=None):
         self.__proxyModel.setFilterRegularExpression(title)
 
     def __clicked(self, idx):
-        # get id of record
-        id = self.__model.data(self.__proxyModel.mapToSource(idx.siblingAtColumn(0)), role=Qt.ItemDataRole.DisplayRole)
-        title = self.__model.data(self.__proxyModel.mapToSource(idx.siblingAtColumn(1)), role=Qt.ItemDataRole.DisplayRole)
+        # get the source index
+        source_idx = self.__proxyModel.mapToSource(idx)
+        # get the primary key value of the row
+        cur_id = self.__model.record(source_idx.row()).value("id")
+        clicked_thread = DB.selectThread(cur_id)
+        # get the title
+        title = clicked_thread['name']
 
-        self.clicked.emit(id, title)
+        self.clicked.emit(cur_id, title)
 
     def __getSelectedIds(self):
-        idx_s = [idx.siblingAtColumn(0) for idx in self.__tableView.selectedIndexes()]
-        idx_s = list(set(idx_s))
-        ids = [self.__model.data(idx, role=Qt.ItemDataRole.DisplayRole) for idx in idx_s]
+        selected_idx_s = self.__tableView.selectedIndexes()
+        ids = []
+        for idx in selected_idx_s:
+            ids.append(self.__model.data(self.__proxyModel.mapToSource(idx.siblingAtColumn(0)), role=Qt.ItemDataRole.DisplayRole))
+        ids = list(set(ids))
         return ids
 
     def __delete(self):
         ids = self.__getSelectedIds()
         for _id in ids:
-            DB.deleteConv(_id)
+            DB.deleteThread(_id)
         self.__model.select()
         self.cleared.emit()
 
@@ -239,7 +269,7 @@ def __clear(self):
         # Before clearing, confirm the action
         reply = QMessageBox.question(self, 'Confirm', 'Are you sure to clear all data?', QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No)
         if reply == QMessageBox.StandardButton.Yes:
-            DB.deleteConv()
+            DB.deleteThread()
             self.__model.select()
             self.cleared.emit()
 
@@ -250,8 +280,8 @@ def __search(self, search_text):
         # content
         elif self.__searchOptionCmbBox.currentText() == 'Content':
             if search_text:
-                convs = DB.selectAllContentOfConv(content_to_select=search_text)
-                ids = [_[0] for _ in convs]
+                threads = DB.selectAllContentOfThread(content_to_select=search_text)
+                ids = [_[0] for _ in threads]
                 self.__model.setQuery(QSqlQuery(f"SELECT {','.join(self.__columns)} FROM {self.__table_nm} "
                                                 f"WHERE id IN ({','.join(map(str, ids))})"))
             else:
@@ -265,4 +295,10 @@ def setColumns(self, columns):
         self.__model.clear()
         self.__model.setTable(self.__table_nm)
         self.__model.setQuery(QSqlQuery(f"SELECT {','.join(self.__columns)} FROM {self.__table_nm}"))
-        self.__model.select()
+        self.__model.select()
+
+    def __onFavoriteClicked(self, f):
+        self.onFavoriteClicked.emit(f)
+
+    def activateFavoriteFromParent(self, f):
+        self.__favoriteBtn.setChecked(f)

pyqt_openai/chat_widget/aiChatUnit.py

@@ -1,12 +1,15 @@
 import pyperclip
 from qtpy.QtCore import Qt
-from qtpy.QtGui import QPalette, QColor
-from qtpy.QtWidgets import QLabel, QWidget, QVBoxLayout, QApplication, QHBoxLayout, QSpacerItem, QSizePolicy, \
+from qtpy.QtGui import QPalette
+from qtpy.QtWidgets import QLabel, QMessageBox, QWidget, QVBoxLayout, QApplication, QHBoxLayout, QSpacerItem, QSizePolicy, \
     QTextBrowser, QAbstractScrollArea
 
-from pyqt_openai.chat_widget.convUnitResultDialog import ConvUnitResultDialog
-from pyqt_openai.widgets.circleProfileImage import RoundedImage
+from pyqt_openai.chat_widget.messageResultDialog import MessageResultDialog
+from pyqt_openai.constants import DEFAULT_ICON_SIZE
+from pyqt_openai.models import ChatMessageContainer
+from pyqt_openai.pyqt_openai_data import DB
 from pyqt_openai.widgets.button import Button
+from pyqt_openai.widgets.circleProfileImage import RoundedImage
 
 
 class SourceBrowser(QWidget):
@@ -79,7 +82,12 @@ def __initUi(self):
         lay = QHBoxLayout()
 
         self.__icon = RoundedImage()
-        self.__icon.setMaximumSize(24, 24)
+        self.__icon.setMaximumSize(*DEFAULT_ICON_SIZE)
+
+        self.__favoriteBtn = Button()
+        self.__favoriteBtn.setStyleAndIcon('ico/favorite_no.svg')
+        self.__favoriteBtn.setCheckable(True)
+        self.__favoriteBtn.toggled.connect(self.__favorite)
 
         self.__infoBtn = Button()
         self.__infoBtn.setStyleAndIcon('ico/info.svg')
@@ -92,6 +100,7 @@ def __initUi(self):
 
         lay.addWidget(self.__icon)
         lay.addSpacerItem(QSpacerItem(10, 10, QSizePolicy.Policy.MinimumExpanding))
+        lay.addWidget(self.__favoriteBtn)
         lay.addWidget(self.__infoBtn)
         lay.addWidget(self.__copyBtn)
         lay.setContentsMargins(2, 2, 2, 2)
@@ -112,14 +121,25 @@ def __initUi(self):
 
         self.setLayout(lay)
 
-        self.setBackgroundRole(QPalette.ColorRole.Midlight)
+        self.setBackgroundRole(QPalette.ColorRole.AlternateBase)
         self.setAutoFillBackground(True)
 
     def __copy(self):
         pyperclip.copy(self.text())
 
+    def __favorite(self, f, insert_f=True):
+        favorite = 1 if f else 0
+        if favorite:
+            self.__favoriteBtn.setStyleAndIcon('ico/favorite_yes.svg')
+        else:
+            self.__favoriteBtn.setStyleAndIcon('ico/favorite_no.svg')
+        if insert_f:
+            current_date = DB.updateMessage(self.__result_info.id, favorite)
+            self.__result_info.favorite = favorite
+            self.__result_info.favorite_set_date = current_date
+
     def __showInfo(self):
-        dialog = ConvUnitResultDialog(self.__result_info)
+        dialog = MessageResultDialog(self.__result_info)
         dialog.exec()
 
     def text(self):
@@ -147,13 +167,16 @@ def setAlignment(self, a0):
                     widget.setAlignment(a0)
 
     def disableGUIDuringGenerateResponse(self):
+        self.__favoriteBtn.setEnabled(False)
         self.__copyBtn.setEnabled(False)
         self.__infoBtn.setEnabled(False)
 
-    def showConvResultInfo(self, info_dict):
+    def showConvResultInfo(self, arg: ChatMessageContainer):
+        self.__favoriteBtn.setEnabled(True)
         self.__copyBtn.setEnabled(True)
         self.__infoBtn.setEnabled(True)
-        self.__result_info = info_dict
+        self.__result_info = arg
+        self.__favorite(True if arg.favorite else False, insert_f=False)
 
     def setText(self, text: str):
         self.__lbl = QLabel(text)