llm/gptPlotCreator.py

import re
import random
import linecache
import subprocess
from langchain.prompts import PromptTemplate
from langchain.chat_models import ChatOpenAI
from langchain.chains import LLMChain
from langchain.llms import OpenAI
from langchain.chains import ConversationChain
from langchain.memory import ConversationBufferMemory
from langchain.prompts.chat import (
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
)
from langchain.vectorstores import Chroma
from langchain.embeddings.openai import OpenAIEmbeddings
import os
from dotenv import load_dotenv
from PIL import Image

# Import required modules
from pymavlink import mavutil
import json

class PlotCreator:
    """
    PlotCreator is a class that generates Python scripts to plot MAVLink data
    provided by a user, leveraging OpenAI's models for conversational agents.
    """

    last_code = ""  # stores the last code generated
    logfile_name = ""
    script_path = ""
    plot_path = ""

    def __init__(self):
        """
        Initialize an instance of PlotCreator.
        """

        load_dotenv()  # load environment variables from a .env file

        self.model = os.getenv("OPENAI_MODEL")  # get the name of the OpenAI model to use

        # create an instance of ChatOpenAI with the specified model, maximum tokens, and temperature
        llm = ChatOpenAI(model_name=self.model, max_tokens=2000, temperature=0)

        # define the input variables and template for the prompt to generate Python scripts
        mavlink_data_prompt = PromptTemplate(
            input_variables=["data_types", "history", "human_input", "file", "output_file"],
            template="You are an AI conversation agent that will be used for generating python scripts to plot mavlink data provided by the user. Please create a python script using matplotlib and pymavlink's mavutil to plot the data provided by the user. Please do not explain the code just return the script. Please plot each independent variable over time in seconds. Please save the plot to file named {output_file} with at least 400 dpi and do not call plt.show(). please use blocking=false in your call to recv_match and be sure to break the loop if a msg in None. here are the relevant data types in the log:\n\n{data_types} \n\nChat History:\n{history} \n\nHUMAN: {human_input} \n\nplease read this data from the file {file}.",
        )

        # create an instance of LLMChain with the defined prompt and verbosity
        self.chain = LLMChain(verbose=True, llm=llm, prompt=mavlink_data_prompt)

    @staticmethod
    def extract_code_snippets(text):
        """
        Extracts code snippets from a text.

        This function searches the text for substrings enclosed in '```', which are assumed to be code snippets.

        Args:
            text (str): The text to search for code snippets.

        Returns:
            list: A list of code snippets found in the text. If no snippets are found, returns a list containing the original text.
        """

        pattern = r'```.*?\n(.*?)```'  # pattern to match code snippets enclosed in '```'
        # use regex to find all matches of the pattern in the text
        snippets = re.findall(pattern, text, re.DOTALL | re.MULTILINE)
        if len(snippets) == 0:  # if no snippets were found
            snippets = [text]  # treat the entire text as a single snippet
        return snippets

    @staticmethod
    def write_plot_script(filename, text):
        """
        Writes a script to a file.

        Args:
            filename (str): The name of the file to write the script to.
            text (str): The script to write to the file.
        """

        with open(filename, 'w') as file:  # open the file for writing
            file.write(text)  # write the script to the file

    def attempt_to_fix_sctript(self, filename, error_message):
        """
        Attempts to fix a script that caused an error.

        Args:
            filename (str): The name of the file containing the script.
            error_message (str): The error message produced by the script.

        Returns:
            list: A list containing the fixed script, or the original script with an error message if it couldn't be fixed.
        """

        # create an instance of ChatOpenAI with the specified model, maximum tokens, and temperature
        llm = ChatOpenAI(model_name=self.model , max_tokens=8000, temperature=0)

        # define the input variables and template for the prompt to generate Python scripts
        fix_plot_script_template = PromptTemplate(
            input_variables=["data_types", "error", "script"],
            template="You are an AI agent that is designed to debug scripts created to plot mavlink data using matplotlib and pymavlink's mavutil. the following script produced this error: \n\n{script}\n\nThe error is: \n\n{error}\n\n Here are message definitions that are possibly relevant for the script:\n\n {data_types}\n\n. Please fix the script so that it produces the correct plot. please return the fixed script in a markdown code block.",
        )

        # read script from file
        with open(filename, 'r') as file:
            script = file.read()

        # create an instance of LLMChain with the defined prompt and verbosity
        chain = LLMChain(verbose=True, llm=llm, prompt=fix_plot_script_template)
        try:
            response = chain.run({"data_types" : self.message_types, "error": error_message, "script": script})  # run the LLMChain with the error and script as input
        except:
            return "Sorry I couldn't fix the script. Here is the original script I tried:\n\n" + script
        print(response)
        code = PlotCreator.extract_code_snippets(response)  # extract the fixed script from the response
        PlotCreator.write_plot_script("plot.py", code[0])  # write the fixed script to a file

        # run the fixed script 
        try:
            subprocess.check_output(["python", self.script_path], stderr=subprocess.STDOUT)
        except:
            code[0] = "Sorry I was unable to fix the script.\nThis is my attempt to fix it:\n\n" + code[0]
        return code

    def set_logfile_name(self, filename):
        """
        Set the name of the log file.
        
        :param filename: The name of the log file.
        :type filename: str
        """
        # extract the path to the log file

        path = os.path.dirname(filename)
        self.logfile_name = filename
        self.script_path = os.path.join(path, "plot.py")
        self.plot_path = os.path.join(path, "plot.png")


    def find_relevant_data_types(self, human_input):
        # Search the database for documents that are similar to the human input
        docs = self.db.similarity_search(human_input)

        # Concatenate the content of the documents into a string
        data_type_info_text = ""
        for doc in docs:
            data_type_info_text += doc.page_content + "\n\n"
        
        return data_type_info_text

    def run_script(self):
        # Run the script and if it doesn't work, capture the output and call attempt_to_fix_script
        try:
            subprocess.check_output(["python", self.script_path], stderr=subprocess.STDOUT)
        except subprocess.CalledProcessError as e:
            print(e.output.decode())
            code = self.attempt_to_fix_sctript(self.script_path, e.output.decode())
            self.last_code = code[0]

        except Exception as e:
            print(e)
            code = self.attempt_to_fix_sctript(self.script_path, str(e))
            self.last_code = code[0]


        # Return a list containing the filename of the plot and the code used to generate it
        return [[(None, (self.plot_path,))], self.last_code]

    def create_plot(self, human_input, data_type_info_text):
        """
        Create a plot based on the input provided by the human.
        
        :param human_input: Input provided by the human.
        :type human_input: str
        """

        # Create a history of generated scripts if one exists
        if self.last_code != "":
            history = "\n\nLast script generated:\n\n" + self.last_code
        else:
            history = ""


        # Generate a response by running the chain with the relevant data types, history, file name and human input
        response = self.chain.run({"data_types" : data_type_info_text, "history" : history, "file": self.logfile_name, "human_input": human_input, "output_file": self.plot_path})
        print(response)

        # Parse the code from the response 
        code = self.extract_code_snippets(response)

        # Write the code to a file named "plot.py"
        self.write_plot_script(self.script_path, code[0])

        # Store the code for the next iteration
        self.last_code = code[0]


        return code[0]


    def parse_mavlink_log(self):
        """
        Parse the MAVLink log to extract unique message types and their fields.
        
        :return: A JSON string representation of the unique message types and their fields.
        :rtype: str
        """
        # Initialize a dictionary to store unique message types and their fields
        self.message_types = {}

        # Establish a MAVLink connection
        mav_log = mavutil.mavlink_connection(self.logfile_name)

        # Loop through the log file and extract all unique message types
        while True:
            try:
                # Receive a message
                msg = mav_log.recv_match(blocking=False, type=None)
                # Check if we received a message
                if msg is None:
                    break

                # Store the unique message types and their fields in the dictionary
                if msg.get_type() not in self.message_types:
                    # Add the message type and its fields to the dictionary
                    self.message_types[msg.get_type()] = {
                        "count": 1,
                        "fields": {field: type(getattr(msg, field)).__name__ for field in msg.get_fieldnames()}
                    }
                else:
                    # Increment the count for this message type
                    self.message_types[msg.get_type()]["count"] += 1

            except KeyboardInterrupt:
                break
            except:
                print("Unknown error")
                break

        # Create embeddings for the message types
        self.create_embeddings(self.message_types)

        # Return a JSON string of the message types
        return json.dumps(self.message_types, indent=4)

    def create_embedding(self, texts):
        """
        Create OpenAI embeddings for a list of texts.
        
        :param texts: A list of texts to create embeddings for.
        :type texts: list of str
        """
        # Initialize a dictionary to store the embeddings
        embeddings = OpenAIEmbeddings(model="text-embedding-ada-002")
        self.db = Chroma.from_texts(texts, embeddings)

    def create_embeddings(self, message_types):
        """
        Create OpenAI embeddings for a dictionary of message types.
        
        :param message_types: A dictionary of message types to create embeddings for.
        :type message_types: dict
        """
        print(message_types)

        # Convert the message types to a list of JSON strings
        texts = []
        for message_type in message_types:
            texts.append(json.dumps({ message_type : message_types[message_type]}))

        print(f"Texts: {texts}")
        # Create the embeddings
        self.create_embedding(texts)