Merge pull request #7 from JinyuanSun/jli_dev

1. Download the plugin script `chatmol.py` and save it to a convenient location on your computer.

2. Open PyMOL.

3. In the PyMOL command line, enter `run /path/to/chatmol.py` (replace `/path/to` with the actual path to the script).

4. The plugin is now installed and ready to use.

4. Create a .PyMOL folder in your home director for saving PyMOL related files, like ChatGPT API keys, PyMOL license file etc.

5. The plugin is now installed and ready to use.

Alternatively, you can use the following command to load the plugin directly:

## Usage

1. Set your OpenAI API key by entering the following command in the PyMOL command line: `set_api_key your_api_key_here` (replace `your_api_key_here` with your actual API key). The API key will be saved in the same directory as the plugin script for future use.

2. Ask ChatGPT about how to perform PyMOL tasks. Use the `chat` command followed by your question or message in the PyMOL command line, e.g., `chat "How do I align two proteins?", False`. You will receive a helpful response such as:

2. Execute PyMOL commands automatically with the ChatGPT Plugin. In PyMOL command line, use `chat` as the wake up word for ChatGPT plugin, followed by a message of your PyMOL tasks or even a question for how to use some of the PyMOL commands. Detailed instructions on how to use PyMOL commands will be displaed and the commands will be executed automatically by default. For example, `chat Show me example to color a protein by its secondary structures`. An example protein molecule will be displayed in PyMOL 3D window with molecule colored by its seconary structure. 

3. Ask ChatMol about how to perform PyMOL tasks without execution the PyMOL commands. You can disable the automatic execution by adding a question mark `?` at the end of ChatMol prompt, e.g., `chat How do I align two proteins?`. You will receive a helpful response such as:

```text

ChatGPT: To align two proteins in PyMOL, you can use the `align` command. Here's an example:

In this example, we use the `and` operator to select only the backbone atoms (`CA`, `C`, `N`, and `O`) for both proteins. This can be useful if you only want to align the backbone of the proteins, and ignore any side chain differences.

```

commands will be saved every time, if you believe this is what you want, run `chat e` will execute satshed commands.

3. If you trust ChatGPT, just remove the fasle from your command: `chat How do I color a protein by secondary structure?`.

  commands from each query will be saved internally. if you want to execute all saved commands, run `chat e` or `chat execute`. After execution, the stashed commands are cleared.

4. To start a new chat session, just enter the following in the PyMOL command line: `chat new`. This will let ChatMol clear the conversation history.

![img](./assets/img_ss.png)

conversation_history = " "    

stashed_commands = []

PLUGIN_DIR = os.path.dirname(os.path.abspath(__file__))

API_KEY_FILE = os.path.join(PLUGIN_DIR, "apikey.txt")

# Save API Key in ~/.PyMOL/apikey.txt

API_KEY_FILE = os.path.expanduser('~')+"/.PyMOL/apikey.txt"

OPENAI_KEY_ENV = "OPENAI_API_KEY"

def set_api_key(api_key):

    api_key = api_key.strip()

    openai.api_key = api_key

    print("APIKEYFILE = ",API_KEY_FILE)

    try:

        with open(API_KEY_FILE, "w") as api_key_file:

            api_key_file.write(api_key)

def load_api_key():

    api_key = os.getenv(OPENAI_KEY_ENV)

    print("APIKEYFILE = ",API_KEY_FILE)

    if not api_key:

        try:

            with open(API_KEY_FILE, "r") as api_key_file:

            messages=messages,

            max_tokens=1024,

            n=1,

            temperature=0.1,

            temperature=0,

        )

        answer = response.choices[0].message['content'].strip()

def start_chatgpt_cmd(message, execute:bool=True):

    global stashed_commands

    if message.strip() == "e":

    global conversation_history

    message = message.strip()

    execution = True

    if (message[-1] == '?'):

        execution = False

    if message.strip() == "e" or message.strip() == 'execute' :

        if (len(stashed_commands) == 0):

            print("There is no stashed commands")

        for command in stashed_commands:

            cmd.do(command)

        # clear stash

        stashed_commands.clear()

        return 0

    if execute.lower() == "false":

    if message.strip() == "new":

        # clear conversation history and stash

        conversation_history = " "

        stashed_commands.clear()

        return 0

    if (message[-1] == '?'):

        execute = False

    response = chat_with_gpt(message)

    print("ChatGPT: " + response.strip())

    try:

        command_blocks = []

        # I think it would be better to reset stashed_commands to empty for each chat.

        # Stash should only keep commands of last conversation, not all commands of the whole history

        stashed_commands.clear()

        for i, block in enumerate(response.split("```")):

            if i%2 == 1:

                command_blocks.append(block)

        for command_block in command_blocks:

            for command in command_block.split("\n"):

                if command.strip() != "" and not command.strip().startswith("#"):

                    # print(f"Executing code: {command}")

                    # Should skip "python", otherwise, not action will be displayed in 3D window

                    if (command.strip() == "python"):

                        continue

                    if "#" in command:

                        index_ = command.index("#")

                        if execute:

                            stashed_commands.append(command)

    except Exception as e:

        print(f"Error executing code: {e}")

        print(f"Error command execution code: {e}")

cmd.extend("set_api_key", set_api_key)

cmd.extend("chat", start_chatgpt_cmd)