lollms-webui/update_script.py

import os
import sys
import git
import subprocess
import argparse
from pathlib import Path
from ascii_colors import ASCIIColors, trace_exception

import pipmaster as pm
if not pm.is_installed("PyQt5"):
    pm.install("PyQt5")
    
import sys
from PyQt5.QtWidgets import QApplication, QMessageBox

def show_error_dialog(message):
    try:
        app = QApplication(sys.argv)
        msg = QMessageBox()
        msg.setIcon(QMessageBox.Critical)
        msg.setText(message)
        msg.setWindowTitle("Error")
        msg.exec_()
    except:
        ASCIIColors.error(message)

def run_git_pull():
    try:
        ASCIIColors.info("----------------> Updating the code <-----------------------")
        
        repo = git.Repo(Path(__file__).parent)
        origin = repo.remotes.origin
        
        # Fetch the latest changes
        origin.fetch()
        
        # Check if there are any changes to pull
        if repo.head.commit == origin.refs.main.commit:
            ASCIIColors.success("Already up-to-date.")
        
        # Discard local changes and force update
        try:
            repo.git.reset('--hard', 'origin/main')
            repo.git.clean('-fd')
            origin.pull()
            ASCIIColors.success("Successfully updated the code.")
        except git.GitCommandError as e:
            error_message = f"Failed to update the code: {str(e)}"
            ASCIIColors.error(error_message)
            # show_error_dialog(error_message)

        print("Updating submodules")
        try:
            repo.git.submodule('update', '--init', '--recursive', '--force')
        except Exception as ex:
            error_message = f"Couldn't update submodules: {str(ex)}"
            ASCIIColors.error(error_message)
            # show_error_dialog(error_message)
        try:
            # Checkout the main branch on each submodule
            for submodule in repo.submodules:
                try:
                    submodule_repo = submodule.module()
                    submodule_repo.git.fetch('origin')
                    submodule_repo.git.reset('--hard', 'origin/main')
                    submodule_repo.git.clean('-fd')
                    ASCIIColors.success(f"Updated submodule: {submodule}.")
                except Exception as ex:
                    print(f"Couldn't update submodule {submodule}: {str(ex)}\nPlease report the error to ParisNeo either on Discord or on github.")
            
            execution_path = Path(os.getcwd())
        except Exception as ex:
            error_message = f"Couldn't update submodules: {str(ex)}\nPlease report the error to ParisNeo either on Discord or on github."
            ASCIIColors.error(error_message)
            # show_error_dialog(error_message)
        try:
            # Update lollms_core
            ASCIIColors.info("Updating lollms_core")
            lollms_core_path = execution_path / "lollms_core"
            if lollms_core_path.exists():
                subprocess.run(["git", "-C", str(lollms_core_path), "fetch", "origin"], check=True)
                subprocess.run(["git", "-C", str(lollms_core_path), "reset", "--hard", "origin/main"], check=True)
                subprocess.run(["git", "-C", str(lollms_core_path), "clean", "-fd"], check=True)
                ASCIIColors.success("Successfully updated lollms_core")
            else:
                ASCIIColors.warning("lollms_core directory not found")
                    
        except Exception as ex:
            error_message = f"Couldn't update submodules: {str(ex)}"
            ASCIIColors.error(error_message)
            # show_error_dialog(error_message)

        return True
    except Exception as e:
        error_message = f"Error during git operations: {str(e)}"
        ASCIIColors.error(error_message)
        # show_error_dialog(error_message)
        return False
    
def get_valid_input():
    while True:
        update_prompt = "New code is updated. Do you want to update the requirements? (y/n): "
        user_response = input(update_prompt).strip().lower()
        
        if user_response in ['y', 'yes']:
            return 'yes'
        elif user_response in ['n', 'no']:
            return 'no'
        else:
            print("Invalid input. Please respond with 'y' or 'n'.")
            
def install_requirements():
    try:
        # Get valid input from the user
        user_choice = get_valid_input()

        # Enhance the text based on the user's response
        if user_choice == 'yes':
            enhanced_text = "Great choice! Updating requirements ensures your project stays up-to-date with the latest changes."
            print(enhanced_text)
            subprocess.check_call([sys.executable, "-m", "pip", "install", "--upgrade", "-r", "requirements.txt"])
            subprocess.check_call([sys.executable, "-m", "pip", "install", "--upgrade", "-e", "lollms_core"])
            ASCIIColors.success("Successfully installed requirements")
        else:  # user_choice == 'no'
            enhanced_text = "Understood. Skipping the requirements update. Make sure to update them later if needed."
            print(enhanced_text)
    except subprocess.CalledProcessError as e:
        error_message = f"Error during pip install: {str(e)}"
        ASCIIColors.error(error_message)
        show_error_dialog(error_message)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--repo", type=str, default="https://github.com/ParisNeo/lollms-webui.git", help="Path to the Git repository")
    args = parser.parse_args()

    repo_path = args.repo

    # Perform git pull to update the repository
    if run_git_pull():
        # Install the new requirements
        install_requirements()

        # Reload the main script with the original arguments
        temp_file = "temp_args.txt"
        if os.path.exists(temp_file):
            with open(temp_file, "r") as file:
                args = file.read().split()
            main_script = "app.py"  # Replace with the actual name of your main script
            os.system(f"{sys.executable} {main_script} {' '.join(args)}")
            try:
                os.remove(temp_file)
            except Exception as e:
                error_message = f"Couldn't remove temp file. Try to remove it manually.\nThe file is located here: {temp_file}\nError: {str(e)}"
                ASCIIColors.warning(error_message)
        else:
            error_message = "Error: Temporary arguments file not found."
            ASCIIColors.error(error_message)
            show_error_dialog(error_message)
            sys.exit(1)
    else:
        error_message = "Update process failed. Please check the console for more details."
        ASCIIColors.error(error_message)
        show_error_dialog(error_message)
        sys.exit(1)

if __name__ == "__main__":
    main()