Merge pull request #6 from JinyuanSun/sjy_dev

- OpenAI Python package: To install, enter `pip install openai` in the PyMOL command line.

## Installation

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

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/chatgpt.py` (replace `/path/to` with the actual path to the script).

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.

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

load https://raw.githubusercontent.com/JinyuanSun/ChatMol/main/chatmol.py

```

If you want a permentally installation, click `Plugin`, go to the `Plugin Manager`, navigate to the `Install New Plugin`, choose the local file or fetch from the url: `https://raw.githubusercontent.com/JinyuanSun/ChatMol/main/chatmol.py`

## 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 `chatgpt` command followed by your question or message in the PyMOL command line, e.g., `chatgpt "How do I align two proteins?”`. You will receive a helpful response such as:

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:

```text

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

```python

# Load the two protein structures

``

# Load two proteins

fetch 1ake

fetch 4ake

fetch 1tim

# Align the two structures

align 1ake, 4ake

# Align 1tim onto 1ake

align 1tim, 1ake

``

In this example, we first load two protein structures using the `fetch` command. Then, we use the `align` command to align the two structures. The first argument to the `align` command is the reference structure (the one that will not be moved), and the second argument is the mobile structure (the one that will be moved to align with the reference structure). 

In this example, we first load two proteins using the `fetch` command. If you already have the proteins loaded, you can skip this step.

By default, the `align` command aligns the two structures based on their backbone atoms. You can also specify which atoms to use for the alignment by adding the `atommask` option to the `align` command. For example:

Next, we use the `align` command to align `1tim` onto `1ake`. The first argument is the object to be aligned (`1tim`), and the second argument is the reference object (`1ake`). PyMOL will align the two proteins based on their structural similarity, and create a new object with the aligned structure.

```python

# Align the two structures using only the alpha-carbon atoms

align 1ake and name ca, 4ake and name ca

You can also specify which atoms to use for the alignment by adding the `atommask` option. For example:

``

# Align 1tim onto 1ake using only the backbone atoms

align 1tim and name CA+C+N+O, 1ake and name CA+C+N+O

``

In this example, we use the `name` selection operator to select only the alpha-carbon atoms (`ca`) of both structures for the alignment.

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.

```

3. Execute PyMOL commands automatically with the ChatGPT Plugin by adding a command execution flag “True” at the end of the prompt, like `chatgpt How do I color a protein by secondary structure?, True`.

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?`.

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

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

## Features

- Seamless integration with PyMOL.

import os

conversation_history = " "    

stashed_commands = []

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

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

def set_api_key(api_key):

    api_key = api_key.strip()

    openai.api_key = api_key

    try:

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

            api_key_file.write(api_key)

        print("API key set and saved to file successfully.")

    except:

        print("API key set successfully but could not be saved to file. You may need to reset the API key next time.")

def load_api_key():

    api_key = os.getenv(OPENAI_KEY_ENV)

    try:

        messages = [

    {"role": "system", "content": "You are an AI language model specialized in providing command line code solutions related to PyMOL. Generate clear and effective solutions in a continuous manner. Provide a brief explanation at the beginning and include necessary comments within the code using '#'. When providing demos or examples, try to use 'fetch' if object name is not provided. Prefer academic style visulizations. Code within triple backticks."}

    {"role": "system", "content": "You are an AI language model specialized in providing command line code solutions related to PyMOL. Generate clear and effective solutions in a continuous manner. When providing demos or examples, try to use 'fetch' if object name is not provided. Prefer academic style visulizations. Code within triple backticks, comment and code should not in the same line."}

]

        message_parts = conversation_history.strip().split("\n")

        for i, part in enumerate(message_parts):

        response = openai.ChatCompletion.create(

            model="gpt-3.5-turbo",

            messages=messages,

            max_tokens=300,

            max_tokens=1024,

            n=1,

            temperature=0.1,

        )

        print(f"Error: {e}")

        return ""

def start_chatgpt_cmd(message):

    response = chat_with_gpt(message)

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

def start_chatgpt_cmd(message, execute_code=False):

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

    global stashed_commands

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

        for command in stashed_commands:

            cmd.do(command)

        return 0

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

        execute = False

    response = chat_with_gpt(message)

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

    if execute_code:

    try:

            # print(response.split("```")[1])

            commands = response.split("```")[1].split("\n")[1:]

            for command in commands:

        command_blocks = []

        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}")

                    # cmd.

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

                    if "#" in command:

                        index_ = command.index("#")

                        if execute:

                            print(command[:index_])

                            cmd.do(command[:index_])

                        else:

                            stashed_commands.append(command[:index_])

                    else:

                        if execute:

                            print(command)

                            cmd.do(command)

                        else:

                            stashed_commands.append(command)

            print("Executed code successfully.")

    except Exception as e:

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

cmd.extend("set_api_key", set_api_key)

cmd.extend("chatgpt", start_chatgpt_cmd)

cmd.extend("chat", start_chatgpt_cmd)