ui.py

# ui.py
import os
import json
from PySide6.QtWidgets import (
    QApplication,
    QMainWindow,
    QVBoxLayout,
    QWidget,
    QListWidget,
    QLineEdit,
    QFormLayout,
    QSplitter,
    QPushButton,
    QLabel,
    QListWidgetItem,
    QHBoxLayout,
    QRadioButton,
    QButtonGroup,
    QMessageBox,
    QFileDialog,
    QTextEdit,
)
from PySide6.QtCore import Qt
from cryptography.fernet import Fernet
from datetime import datetime
import pandas as pd
from langchain_core.messages import AIMessage, HumanMessage
import chromadb
from chromadb.config import DEFAULT_TENANT, DEFAULT_DATABASE, Settings

from workers import (
    DatabaseConnectionWorker,
    AIProcessingWorkerSQL,
    AIProcessingWorkerExplain,
    AIProcessingWorkerUserExpresstion,
    AIProcessingWorkerGraph,
    ChromaDBClient,
)
from config_utils import CONFIG_FILE, CHROMA_DB_FILE, load_key, encrypt, decrypt


class MainWindow(QMainWindow):
    def __init__(self):
        super(MainWindow, self).__init__()
        self.setWindowTitle("GPT-like Chat with Database Connection")
        self.sql_model = None
        self.explain_model = None
        self.mainLayout = QVBoxLayout()
        self.message_used_for_sql = None
        self.popup = None
        self.db_schema = None
        self.db_connection = None
        self.visual_model = None
        self.sql_db_type = "postgresql"
        self.text_ehancer = None
        self.db_tables_details = None
        self.db_tables_names = None
        self.key = load_key()
        self.cipher_suite = Fernet(self.key)
        self.db_document = None
        self.upload_button = QPushButton("Upload Document")
        self.upload_button.clicked.connect(self.open_file_dialog)
        self.closeConnectionButton = QPushButton("Delete VecDB")
        self.closeConnectionButton.clicked.connect(self.close_connection)
        self.createConnectionButton = QPushButton("Connect VecDB")
        self.createConnectionButton.clicked.connect(self.create_connection)
        # List widget for chat
        self.listWidget = QListWidget()
        self.listWidget.setSelectionMode(QListWidget.NoSelection)
        self.chatHistory = []
        self.chroma_client = None
        self.screen_geometry = QApplication.primaryScreen().geometry()
        # LineEdit for user input
        self.lineEdit = QLineEdit()
        self.lineEdit.returnPressed.connect(self.update_chat)

        self.mainLayout.addWidget(self.listWidget)
        self.mainLayout.addWidget(self.lineEdit)

        # Sidebar layout for database connection details
        self.sidebarLayout = QVBoxLayout()

        # Create a form layout for the database connection details
        formLayout = QFormLayout()
        self.dbTypeGroup = QButtonGroup(self)
        self.postgresRadio = QRadioButton("PostgreSQL")
        self.mysqlRadio = QRadioButton("MySQL")
        self.sqliteRadio = QRadioButton("SQLite")
        self.oracleRadio = QRadioButton("Oracle")
        self.mssqlRadio = QRadioButton("MSSQL")
        self.dbTypeGroup.addButton(self.postgresRadio)
        self.dbTypeGroup.addButton(self.mysqlRadio)
        self.dbTypeGroup.addButton(self.sqliteRadio)
        self.dbTypeGroup.addButton(self.oracleRadio)
        self.dbTypeGroup.addButton(self.mssqlRadio)
        self.postgresRadio.setChecked(True)

        # Add radio buttons to form layout
        db_first_layout = QHBoxLayout()
        db_first_layout.addWidget(self.postgresRadio)
        db_first_layout.addWidget(self.mysqlRadio)
        db_first_layout.addWidget(self.sqliteRadio)
        formLayout.addRow(QLabel("Database Type:"))
        formLayout.addRow(db_first_layout)

        db_second_layout = QHBoxLayout()
        db_second_layout.addWidget(self.oracleRadio)
        db_second_layout.addWidget(self.mssqlRadio)
        formLayout.addRow(db_second_layout)

        # Connect radio button signals
        self.postgresRadio.toggled.connect(
            lambda: self.on_db_type_changed("postgresql")
        )
        self.mysqlRadio.toggled.connect(lambda: self.on_db_type_changed("mysql"))
        self.sqliteRadio.toggled.connect(lambda: self.on_db_type_changed("sqlite"))
        self.oracleRadio.toggled.connect(lambda: self.on_db_type_changed("oracle"))

        # Creating inputs for the connection details with empty default values
        self.hostInput = QLineEdit()
        self.dbnameInput = QLineEdit()
        self.userInput = QLineEdit()
        self.passwordInput = QLineEdit()
        self.portInput = QLineEdit("5432")
        self.apiKeyInput = QLineEdit()

        # Adding input fields to form layout
        formLayout.addRow(QLabel("Host:"), self.hostInput)
        formLayout.addRow(QLabel("Database Name:"), self.dbnameInput)
        formLayout.addRow(QLabel("User:"), self.userInput)
        formLayout.addRow(QLabel("Password:"), self.passwordInput)
        formLayout.addRow(QLabel("Port:"), self.portInput)
        formLayout.addRow(QLabel("API Key:"), self.apiKeyInput)
        formLayout.addRow(self.upload_button)

        # "Connect" button
        self.connectButton = QPushButton("Connect")
        self.connectButton.clicked.connect(self.connect_to_database)
        formLayout.addRow(self.connectButton)

        # "Load Credentials" button
        self.loadButton = QPushButton("Load Credentials")
        self.loadButton.clicked.connect(self.load_credentials)
        formLayout.addRow(self.loadButton)

        # Status label to show connection result (success or not connected)
        self.statusLabel = QLabel("")
        formLayout.addRow(self.statusLabel)

        # Add form layout to sidebar
        self.sidebarLayout.addLayout(formLayout)

        # Spacer to push the button to the bottom
        self.sidebarLayout.addStretch()

        # Create Chat Button at the bottom
        self.clearChatButton = QPushButton("Clear Chat")
        self.clearChatButton.clicked.connect(self.show_clear_chat_confirmation)
        self.sidebarLayout.addWidget(self.clearChatButton)
        self.sidebarLayout.addWidget(self.closeConnectionButton)
        self.closeConnectionButton.setDisabled(True)
        self.sidebarLayout.addWidget(self.createConnectionButton)
        self.createConnectionButton.setDisabled(True)
        if os.path.exists(CHROMA_DB_FILE) and self.chroma_client is not None:
            self.closeConnectionButton.setEnabled(True)
        if os.path.exists(CHROMA_DB_FILE):
            self.createConnectionButton.setEnabled(True)
        # Sidebar widget
        self.sidebarWidget = QWidget()
        self.sidebarWidget.setLayout(self.sidebarLayout)

        # Chat widget
        self.chatWidget = QWidget()
        self.chatWidget.setLayout(self.mainLayout)

        # Using QSplitter to divide the sidebar and the main chat area
        self.splitter = QSplitter(Qt.Horizontal)

        # Add sidebar (with a minimum width) and chat widget
        self.splitter.addWidget(self.sidebarWidget)
        self.splitter.addWidget(self.chatWidget)

        # Setting sidebar size limitations
        self.sidebarWidget.setMinimumWidth(200)
        self.sidebarWidget.setMaximumWidth(350)

        # Control the relative sizing - 1:4 means sidebar takes less space than chat
        self.splitter.setStretchFactor(0, 1)
        self.splitter.setStretchFactor(1, 4)

        # Create Toggle Button (Fixed on main layout, outside of sidebar)
        self.toggleSidebarButton = QPushButton("Close")
        self.toggleSidebarButton.clicked.connect(self.toggle_sidebar)

        # Create a layout to position the toggle button at the top-left corner
        self.toggleLayout = QVBoxLayout()
        self.toggleLayout.addWidget(self.toggleSidebarButton, alignment=Qt.AlignLeft)

        # Create a wrapper widget to hold the toggle button and the splitter
        self.wrapperWidget = QWidget()
        self.wrapperLayout = QVBoxLayout()
        self.wrapperLayout.addLayout(self.toggleLayout)  # Add toggle button layout
        self.wrapperLayout.addWidget(self.splitter)  # Add the main splitter
        self.wrapperWidget.setLayout(self.wrapperLayout)

        # Set the wrapper widget as the central widget
        self.setCentralWidget(self.wrapperWidget)
    def create_connection(self):
        if os.path.exists(CHROMA_DB_FILE):
            try:
                self.chroma_client = chromadb.PersistentClient(
                    settings=Settings(),
                    tenant=DEFAULT_TENANT,
                    database=DEFAULT_DATABASE,
                )
                self.statusLabel.setText("VecDB connected")
                self.statusLabel.setStyleSheet("color: green")
                self.createConnectionButton.setDisabled(True)
                self.closeConnectionButton.setEnabled(True)
            except Exception as e:
                self.statusLabel.setText("Error while Connecting")
                self.statusLabel.setStyleSheet("color: red")
                print(f"Error: {e}")
        else:
            self.statusLabel.setText("No VecDB to connect")
            self.statusLabel.setStyleSheet("color: red")
    def close_connection(self):
        def check_if_collection_exist(temp):
            try:
                temp.chroma_db.delete_collection("sql_data")
                return True
            except Exception as e:
                return False
        if self.chroma_client is not None:
            try:
                check_if_collection_exist(self.chroma_client)
                self.statusLabel.setText("DB deleted")
                self.statusLabel.setStyleSheet("color: green")
                self.createConnectionButton.setEnabled(True)
                self.closeConnectionButton.setDisabled(True)
                self.chroma_client = None
            except Exception as e:
                self.statusLabel.setText("Error while deleting")
                self.statusLabel.setStyleSheet("color: red")
                print(f"Error: {e}")
        else:
            self.statusLabel.setText("No DB to delete")
            self.statusLabel.setStyleSheet("color: red")

    def open_file_dialog(self):
        # Set filters to include .txt, .pdf, and .docx files
        file_filters = "All Files (*);;Text Files (*.txt);;PDF Files (*.pdf);;Word Documents (*.docx)"
        file_name, _ = QFileDialog.getOpenFileName(
            self, "Open Document", "", file_filters
        )
        self.temp = ChromaDBClient(file_name)
        self.temp.responce_generated.connect(self.chroma_db_client)
        self.upload_button.setEnabled(False)
        self.statusLabel.setText("Uploading...")
        self.statusLabel.setStyleSheet("color: orange;")
        self.temp.start()

    def chroma_db_client(self, client):
        if client:
            self.chroma_client = client
            self.upload_button.setEnabled(True)
            self.createConnectionButton.setDisabled(True)
            self.closeConnectionButton.setEnabled(True)
            self.statusLabel.setText("Uploaded")
            self.statusLabel.setStyleSheet("color: green;")
        else:
            self.statusLabel.setText("Not able to Upload")
            self.statusLabel.setStyleSheet("color: red;")

    def toggle_sidebar(self):
        # Check if sidebar is visible
        if self.sidebarWidget.isVisible():
            self.sidebarWidget.hide()
            self.toggleSidebarButton.setText("Open")  # Change button text when hidden
            self.splitter.setSizes(
                [0, self.width()]
            )  # Sidebar hidden, chat takes full width
        else:
            self.sidebarWidget.show()
            self.toggleSidebarButton.setText("Close")  # Reset button text when visible
            self.splitter.setSizes([400, self.width() - 200])  # Reset the sidebar width

    def show_clear_chat_confirmation(self):
        # Create a confirmation popup
        confirmation = QMessageBox()
        confirmation.setWindowTitle("Create Chat")
        confirmation.setText("Are you sure you want to clear?")
        confirmation.setStandardButtons(QMessageBox.Yes | QMessageBox.No)
        confirmation.setIcon(QMessageBox.Question)

        # Check the user's response
        result = confirmation.exec_()

        if result == QMessageBox.Yes:
            self.listWidget.clear()
            self.chatHistory = []

    def on_db_type_changed(self, db_type):
        self.sql_db_type = db_type

    def update_chat(self):
        if not self.sql_model:
            self.add_message_user("Not connected to the database.")
            return

        # Get the user's input
        user_text = self.lineEdit.text()
        self.lineEdit.clear()

        self.add_message_user(user_text)
        self.message_used_for_sql = user_text
        self.ehnace_user_text = AIProcessingWorkerUserExpresstion(
            user_text,
            self.text_ehancer,
            self.chatHistory,
            self.db_tables_names,
            self.chroma_client,
        )
        self.ehnace_user_text.responce_ehancer_.connect(
            self.user_expresstion_after_ehnance
        )

        self.ehnace_user_text.start()

    def user_expresstion_after_ehnance(self, user_text):
        if user_text == "User API key is invalid":
            self.statusLabel.setText("ERROR: While Connecting to model")
            self.statusLabel.setStyleSheet("color: red")
            return
        self.db_schema = ""
        for i in user_text["useful_tables"]:
            if i in self.db_tables_names:
                self.db_schema += self.db_tables_details[i] + "\n"

        # Start a separate thread to process AI response
        self.ai_sql_worker = AIProcessingWorkerSQL(
            self.sql_model,
            user_text,
            self.chatHistory,
            self.db_connection,
            self.db_schema,
            self.sql_db_type,
        )
        self.ai_sql_worker.response_generated.connect(self.add_ai_response)
        self.ai_sql_worker.start()

    def add_ai_response(self, ai_response, data, user_en):
        if "User API key is invalid" in ai_response:
            self.statusLabel.setText(
                "User API key is invalid",
            )
            self.statusLabel.setStyleSheet("color: red")
            return
        self.add_ai_sql_message(ai_response, data, user_en)

    def add_message_user(self, text):
        message_widget = QWidget()
        message_layout = QHBoxLayout()
        message_layout.setAlignment(Qt.AlignRight)
        message_widget.setLayout(message_layout)
        message_layout.addWidget(QLabel(text))
        list_item = QListWidgetItem()
        list_item.setSizeHint(message_widget.sizeHint())
        self.listWidget.addItem(list_item)
        self.listWidget.setItemWidget(list_item, message_widget)
        self.listWidget.scrollToBottom()

    def add_ai_sql_message(self, text, dataframe, user_en):
        message_widget = QWidget()
        message_layout = QHBoxLayout()
        message_layout.setAlignment(Qt.AlignLeft)
        message_widget.setLayout(message_layout)
        temp_var = QTextEdit()
        temp_var.setMarkdown("```sql\n" + text + "```")
        temp_var.setReadOnly(True)
        temp_var.setMaximumSize(
            int(self.screen_geometry.size().width() * 0.5),
            self.screen_geometry.size().height(),
        )
        message_layout.addWidget(temp_var)

        action_button = None
        action_button_graph = None

        if dataframe is not None and dataframe != "the sql query made a error":
            action_button = self.create_button(
                "Download", lambda: self.handle_action(dataframe)
            )
            message_layout.addWidget(action_button)
            action_button_graph = self.create_button(
                "Visualize",
                lambda: self.handle_action_graph(
                    pd.DataFrame(dataframe), user_en, text
                ),
            )
            message_layout.addWidget(action_button_graph)

        # Event handlers for button visibility
        def show_buttons(event):
            if action_button:
                action_button.setVisible(True)
            if action_button_graph:
                action_button_graph.setVisible(True)

        def hide_buttons(event):
            if action_button:
                action_button.setVisible(False)
            if action_button_graph:
                action_button_graph.setVisible(False)

        message_widget.enterEvent = show_buttons
        message_widget.leaveEvent = hide_buttons
        list_item = QListWidgetItem()
        list_item.setSizeHint(message_widget.sizeHint())
        self.listWidget.addItem(list_item)
        self.listWidget.setItemWidget(list_item, message_widget)
        self.listWidget.scrollToBottom()
        if self.message_used_for_sql:
            self.ai_explain_worker = AIProcessingWorkerExplain(
                self.explain_model,
                user_en,
                text,
                dataframe,
                self.chatHistory,
                self.db_schema,
            )
            self.chatHistory.append(HumanMessage(str(user_en)))
            self.chatHistory.append(AIMessage(text))
            self.ai_explain_worker.responce_generated.connect(
                self.add_ai_explain_response
            )

            self.ai_explain_worker.start()
            self.message_used_for_sql = None

    def create_button(self, text, action):
        button = QPushButton(text)
        button.setVisible(False)
        button.clicked.connect(action)
        return button

    def add_ai_explain_response(self, text):
        message_widget = QWidget()
        message_layout = (
            QVBoxLayout()
        )  # Use QVBoxLayout for vertical stacking of content
        message_layout.setAlignment(Qt.AlignLeft)
        self.chatHistory.append(AIMessage(text))

        message_label = QTextEdit()
        # Display the HTML content in the QTextEdit
        message_label.setMarkdown(text)
        message_label.setReadOnly(True)
        message_label.adjustSize()

        message_label.setMaximumSize(
            int(self.screen_geometry.size().width() * 0.5),
            self.screen_geometry.size().height(),
        )
        # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        message_widget.setLayout(message_layout)
        message_layout.addWidget(message_label)

        list_item = QListWidgetItem()
        list_item.setSizeHint(message_widget.sizeHint())
        self.listWidget.addItem(list_item)
        self.listWidget.setItemWidget(list_item, message_widget)
        self.listWidget.scrollToBottom()

    def handle_action_graph(self, df, user_en, sql):
        self.graph_df = pd.DataFrame(df)
        self.graph_user_en = user_en
        self.graph_sql = sql
        self.run_graph_generation()

    def run_graph_generation(self, complate=""):
        self.graphgen = AIProcessingWorkerGraph(
            self.visual_model,
            self.graph_user_en,
            self.graph_df,
            self.graph_sql,
            complate,
        )
        self.graphgen.responce_generated.connect(self.create_math_plot)
        self.graphgen.start()

    def create_math_plot(self, code):
        try:
            # Execute the code
            print(code)
            local_scope = {"df": self.graph_df}
            exec(code, {}, local_scope)
        except Exception as e:
            print(f"Error executing code: {e}")
            self.run_graph_generation(str(e))

    def handle_action(self, df):
        print("Downloading data...")
        try:
            if df is not pd.DataFrame:
                df = pd.DataFrame(list(df))
            # Convert timezone-aware datetime objects to timezone-unaware
            for col in df.select_dtypes(include=["datetimetz"]).columns:
                df[col] = df[col].apply(
                    lambda x: x.strftime("%Y-%m-%d %H:%M:%S.%f %Z")
                    if pd.notnull(x)
                    else ""
                )
            for col in df.select_dtypes(include=["timedelta"]).columns:
                df[col] = df[col].apply(lambda x: str(x) if pd.notnull(x) else "")

            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            excel_filename = f"downloaded_data_{timestamp}.xlsx"
            with pd.ExcelWriter(excel_filename, engine="openpyxl") as writer:
                df.to_excel(writer, index=False)

            self.statusLabel.setText(
                f"Data downloaded successfully as '{excel_filename}'."
            )
            self.statusLabel.setStyleSheet("color: green;")

            # Open the downloaded Excel file (Windows-specific)
            if os.name == "nt":
                os.startfile(excel_filename)
            else:
                print(f"Please open the file manually: {excel_filename}")

        except Exception as e:
            print(f"Error while downloading data: {e}")
            self.statusLabel.setText("Error downloading data.")
            self.statusLabel.setStyleSheet("color: red")

    def connect_to_database(self):
        self.statusLabel.setText("Connecting...")
        self.statusLabel.setStyleSheet("color: orange;")
        self.connectButton.setEnabled(False)
        self.db_worker = DatabaseConnectionWorker(
            self.hostInput.text(),
            self.dbnameInput.text(),
            self.userInput.text(),
            self.passwordInput.text(),
            self.portInput.text(),
            self.apiKeyInput.text(),
            self.sql_db_type,
        )
        self.db_worker.connection_result.connect(self.on_connection_result)
        self.db_worker.start()

    def on_connection_result(
        self,
        success,
        db_connection,
        sql_model,
        explain_model,
        table_info,
        visual_model,
        ehancer_model,
        table_names,
    ):
        if success:
            self.db_tables_details = table_info
            self.db_tables_names = table_names
            self.db_connection = db_connection
            self.sql_model = sql_model
            self.explain_model = explain_model
            self.visual_model = visual_model
            self.text_ehancer = ehancer_model
            self.statusLabel.setText("Success")
            self.statusLabel.setStyleSheet("color: green;")
            # Save credentials after successful connection
            self.save_credentials()
        else:
            self.statusLabel.setText("Not Connected")
            self.statusLabel.setStyleSheet("color: red")
        self.connectButton.setEnabled(True)

    def save_credentials(self):
        credentials = {
            "host": self.hostInput.text(),
            "dbname": self.dbnameInput.text(),
            "user": self.userInput.text(),
            "password": self.passwordInput.text(),
            "port": self.portInput.text(),
            "api_key": self.apiKeyInput.text(),
            "db_type": self.sql_db_type,
        }
        encrypted_credentials = {k: encrypt(v, self.cipher_suite) for k, v in credentials.items()}
        with open(CONFIG_FILE, "w") as f:
            json.dump(encrypted_credentials, f)

    def load_credentials(self):
        if os.path.exists(CONFIG_FILE):
            with open(CONFIG_FILE, "r") as f:
                encrypted_credentials = json.load(f)
                credentials = {
                    k: decrypt(v, self.cipher_suite) for k, v in encrypted_credentials.items()
                }
                self.db_type = credentials.get("db_type", "postgresql")
                if self.db_type == "postgresql":
                    self.postgresRadio.setChecked(True)
                elif self.db_type == "mysql":
                    self.mysqlRadio.setChecked(True)
                elif self.db_type == "sqlite":
                    self.sqliteRadio.setChecked(True)
                elif self.db_type == "oracle":
                    self.oracleRadio.setChecked(True)
                elif self.db_type == "mssql":
                    self.mssqlRadio.setChecked(True)
                self.hostInput.setText(credentials.get("host", ""))
                self.dbnameInput.setText(credentials.get("dbname", ""))
                self.userInput.setText(credentials.get("user", ""))
                self.passwordInput.setText(credentials.get("password", ""))
                self.portInput.setText(credentials.get("port", "5432"))
                self.apiKeyInput.setText(credentials.get("api_key", ""))
        else:
            self.statusLabel.setText("No saved credentials found.")
            self.statusLabel.setStyleSheet("color: red")

    def save_chat_history(self):
        with open("chat_history.json", "w") as f:
            json.dump([str(msg) for msg in self.chatHistory], f)

    def closeEvent(self, event):
        self.save_chat_history()
        event.accept()