added Vector DB for message definitions. plotter now working reliably

README.md

@@ -2,6 +2,12 @@
 
 MAVPlot is a Python-based project which uses Gradio as an interface and GPT-X powered by OpenAI as a chatbot to generate and plot MAVLink data. It provides an easy-to-use, chatbot-like interface for users to describe the plot they would like to generate.
 
+![chat bot](docs/chat_bot_if.PNG)
+
+## Architecture
+
+![arch](docs/GPT_MAVPlot_Arch.png)
+
 ## Installation
 
 1. Clone the repository:
@@ -10,23 +16,26 @@ MAVPlot is a Python-based project which uses Gradio as an interface and GPT-X po
 git clone https://github.com/yourusername/mavplot.git
 ```
 
+TODO
 2. Install the requirements:
 
 ```shell
 pip install -r requirements.txt
 ```
 
+3. Setup .env File 
+
+Copy the `template.env` file to a file named `.env` in your root directory. Add your Openai API key to the file  
+
 ## Usage
 
 After installing all dependencies, run the main script using:
 
 ```shell
-python llm_plot.py
+python gpt_mavplot.py
 ```
 
-A web-based Gradio interface will launch. You can then input the description of the plot you would like to generate in the textbox, or upload a file.
-
-The chatbot will process your request and generate the corresponding plot, which will be displayed in the chat interface.
+A web-based Gradio interface will launch. You can upload a mavlink tlog then prompt the bot to generate plots from the log. The chatbot will process your request and generate the corresponding plot, which will be displayed in the chat interface. The script use to generate the log will also be posted to the chat interface. 
 
 ## Contributing
 

llm_plot.py → gpt_mavplot.py

@@ -27,9 +27,8 @@ def bot(history):
 
         print(history)
 
-        history_str = format_history(history)
-        response = plot_creator.create_plot(user_input, history_str)
-        print(response)
+        # history_str = format_history(history)
+        response = plot_creator.create_plot(user_input)
         history[-1][1] = response[0]
         history = history + [(None, f"Here is the code used to generate the plot:")]
         history = history + [(None, f"{response[1]}")]
@@ -41,7 +40,10 @@ def bot(history):
         filename, extension = os.path.splitext(os.path.basename(file_path))
         
         history[-1][0] = f"user uploaded file: {filename}{extension}"
-        history[-1][1] = "I will be using the file you uploaded to generate the plot. Please describe the plot you would like to generate."
+        history[-1][1] = "processing file..."
+        data_types = plot_creator.parse_mavlink_log()
+        history = history + [(None, f"here are the data types in the log {data_types}")]
+        history = history + [(None, f"I am done processing the file. Now you can ask me to generate a plot.")]
     
     return history
 

llm/gptPlotCreator.py

@@ -12,75 +12,115 @@ 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 = ""
+    last_code = ""  # stores the last code generated
 
     def __init__(self):
-        load_dotenv()
-        self.model = os.getenv("OPENAI_MODEL")
-        # llm = ChatOpenAI(model_name="gpt-3.5-turbo", max_tokens=2000, temperature=0)
+        """
+        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=["history", "human_input", "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 plot.png in the same directory as plot.py with at least 400 dpi. Also be careful not to write a script that gets stuck in an endless loop.\n\nChat History:\n{history} \n\nHUMAN: {human_input} \n\nplease read this data from the file {file}.",
+            input_variables=["data_types", "history", "human_input", "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 plot.png in the same directory as plot.py with at least 400 dpi. 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}.",
         )
-        self.chain = LLMChain(verbose=True, llm=llm, prompt=mavlink_data_prompt)
 
-    @staticmethod
-    def sample_lines(filename, num_lines=5):
-        with open(filename) as f:
-            total_lines = sum(1 for _ in f)
-        
-        if total_lines < num_lines:
-            raise ValueError("File has fewer lines than the number of lines requested.")
-        
-        line_numbers = random.sample(range(1, total_lines + 1), num_lines)
-        lines = [linecache.getline(filename, line_number).rstrip() for line_number in line_numbers]
-        
-        return '\n'.join(lines)
+        # 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):
-        pattern = r'```.*?\n(.*?)```'
+        """
+        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:
-            snippets = [text]
+        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):
-        with open(filename, 'w') as file:
-            file.write(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):
-        # llm = ChatOpenAI(model_name="gpt-3.5-turbo", max_tokens=2000, temperature=0)
-        llm = ChatOpenAI(model_name=self.model , max_tokens=2000, temperature=0)
+        """
+        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=["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\nPlease fix the script so that it produces the correct plot. please return the fixed script in a markdown code block.",
+            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)
-        response = chain.run({"error": error_message, "script": script})
+        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)
-        PlotCreator.write_plot_script("plot.py", code[0])
+        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 script 
+        # run the fixed script 
         try:
             subprocess.check_output(["python", "plot.py"], stderr=subprocess.STDOUT)
         except:
@@ -88,22 +128,48 @@ class PlotCreator:
         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
+        """
         self.logfile_name = filename
 
-    def create_plot(self, human_input, history):
-
+    def create_plot(self, human_input):
+        """
+        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 = history + "\n\nLast script generated:\n\n" + self.last_code
+            history = "\n\nLast script generated:\n\n" + self.last_code
+        else:
+            history = ""
+
+        # 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"
 
+        print(docs)
 
-        response = self.chain.run({"history" : history, "file": self.logfile_name, "human_input": human_input})
+        # 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})
         print(response)
 
-        # parse the code from the 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("plot.py", code[0])
 
-        # run the script if it doesn't work capture output and call attempt_to_fix_script
+        # Run the script and if it doesn't work, capture the output and call attempt_to_fix_script
         try:
             subprocess.check_output(["python", "plot.py"], stderr=subprocess.STDOUT)
         except subprocess.CalledProcessError as e:
@@ -113,6 +179,83 @@ class PlotCreator:
             print(e)
             code = self.attempt_to_fix_sctript("plot.py", str(e))
 
+        # Store the code for the next iteration
         self.last_code = code[0]
 
+        # Return a list containing the filename of the plot and the code used to generate it
         return [("plot.png", None), 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)