Merge pull request #17 from JinyuanSun/local_gui

import json

from pymol import cmd

conversation_history = ""

def query_qaserver(question):

    headers = {

        'Content-Type': 'application/x-www-form-urlencoded',

    return response.text

def chatlit(question):

    global conversation_history

    question = conversation_history + question

    answer = query_qaserver(question)

    data = json.loads(answer)

    commands = data['answer']

# A local GUI of ChatMol

## Requirements

- Python 3.6 or higher

- Pymol 2.3 or higher

## Installation

1. Clone this repository

2. Install Pymol

```

conda install -c conda-forge pymol-open-source

```

## Usage

```bash

python local_gui.py

```

import tkinter as tk

from tkinter import ttk

import requests

import json

import subprocess

import threading

lite_conversation_history = ""

def launch_pymol():

    """start a new process to launch PyMOL

    Only tested on MacOS, may be different on other OS

    """

    process = subprocess.Popen(["pymol", "pymol_server.py"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)

    stdout, stderr = process.communicate()

    if process.returncode != 0:

        print(f"Error occurred while launching PyMOL: {stderr}")

    else:

        print(stdout)

def query_qaserver(question):

    headers = {

        'Content-Type': 'application/x-www-form-urlencoded',

    }

    data = 'question=' + question.replace('"','')

    response = requests.post('https://chatmol.org/qa/lite/', headers=headers, data=data)

    return response.text

def chatlite(question):

    global lite_conversation_history

    question = lite_conversation_history + "Instructions: " + question

    answer = query_qaserver(question)

    data = json.loads(answer)

    lite_conversation_history = data['conversation_history']

    lite_conversation_history += "\nAnswer: "

    lite_conversation_history += data['answer']

    lite_conversation_history += "\n"

    commands = data['answer']

    commands = commands.split('\n')

    print("Answers from ChatMol-Lite: ")

    for command in commands:

        if command == '':

            continue

        else:

            print(command)

    return commands

def send_message():

    message = entry.get()

    chat.config(state='normal')

    chat.insert(tk.END, "You: " + message + "\n")

    response = chatlite(message)

    chat.insert(tk.END, "ChatMol-Lite: " + '\n'.join(response) + "\n")

    chat.config(state='disabled')

    script.insert(tk.END, '\n'.join(response) + "\n")

    entry.delete(0, tk.END)

def send_response_to_server():

    command = script.get("1.0", "end-1c")

    response = requests.post('http://localhost:8101/send_message', data=command)

    if response.status_code == 200:

        print('Command sent to server successfully.')

        script.delete('1.0', tk.END)  # Clear the script box after sending command to server

    else:

        print(f'Failed to send command to server. Status code: {response.status_code}')

# GUI code

window = tk.Tk()

window.title("ChatMol-Lite")

# Create the main container

frame = ttk.Frame(window, padding="10 10 10 10")

frame.grid(column=0, row=0, sticky=(tk.W, tk.E, tk.N, tk.S))

window.columnconfigure(0, weight=1)

window.rowconfigure(0, weight=1)

frame.columnconfigure(0, weight=1)

frame.columnconfigure(1, weight=3)  # Make the right panel take up more space

# left 

left_frame = ttk.Frame(frame, padding="10 10 10 10")

left_frame.grid(column=0, row=0, sticky=(tk.W, tk.E, tk.N, tk.S))

# right

right_frame = ttk.Frame(frame, padding="10 10 10 10")

right_frame.grid(column=1, row=0, sticky=(tk.W, tk.E, tk.N, tk.S))

# Create the chat box in the left container

chat = tk.Text(left_frame, state='disabled')

chat.pack(fill='both', expand=True)

# Create the entry box in the left container

entry = tk.Entry(left_frame)

entry.pack(fill='x')

# Create the send button in the left container

send_button = ttk.Button(left_frame, text="Send", command=send_message)

send_button.pack(fill='x')

# Create the script box in the right container

script = tk.Text(right_frame)

script.pack(fill='both', expand=True)

# Create the send to server button in the right container

send_to_server_button = ttk.Button(right_frame, text="Send to PyMOL", command=send_response_to_server)

send_to_server_button.pack(fill='x')

# Launch PyMOL in a separate thread

if __name__ == "__main__":

    pymol_thread = threading.Thread(target=launch_pymol)

    pymol_thread.start()

if __name__ == "pymol": # still cannot start from pymol interpreter

    from pymol_server import *

window.mainloop()

 No newline at end of file

import http.server

from http import HTTPStatus

from pymol import cmd

import urllib.parse

import threading

class PyMOLCommandHandler(http.server.BaseHTTPRequestHandler):

    def _send_cors_headers(self):

        """Sets headers required for CORS"""

        self.send_header("Access-Control-Allow-Origin", "*")

        self.send_header("Access-Control-Allow-Methods", "GET,POST,OPTIONS")

        self.send_header("Access-Control-Allow-Headers", "x-api-key,Content-Type")

    def do_OPTIONS(self):

        """Respond to a OPTIONS request."""

        self.send_response(HTTPStatus.NO_CONTENT)

        self._send_cors_headers()

        self.end_headers()

    def do_POST(self):

        if self.path != "/send_message":

            self.send_response(HTTPStatus.NOT_FOUND)

            self.end_headers()

            return

        content_length = int(self.headers['Content-Length'])

        post_data = self.rfile.read(content_length)

        post_data = urllib.parse.unquote(post_data.decode())

        try:

            cmd.do(post_data)

            self.send_response(HTTPStatus.OK)

            self._send_cors_headers()

            self.end_headers()

            self.wfile.write(b'Command executed')

        except Exception as e:

            self.send_response(HTTPStatus.INTERNAL_SERVER_ERROR)

            self.end_headers()

            self.wfile.write(str(e).encode())

def start_server():

    httpd = http.server.HTTPServer(('localhost', 8101), PyMOLCommandHandler)

    httpd.serve_forever()

server_thread = threading.Thread(target=start_server)

server_thread.start()

print("Server started")