add chatbox. add examples.

.gitignore

@@ -127,6 +127,7 @@ venv/
 ENV/
 env.bak/
 venv.bak/
+.python-version
 
 # Spyder project settings
 .spyderproject
@@ -159,4 +160,4 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
 
-.DS_Store
+.DS_Store

app.py

@@ -1,3 +1,4 @@
+import os
 import openai
 import numpy as np
 from tempfile import NamedTemporaryFile
@@ -8,20 +9,23 @@ from shapely.affinity import *
 from omegaconf import OmegaConf
 from moviepy.editor import ImageSequenceClip
 import gradio as gr
+from gradio import processing_utils
 
 from lmp import LMP, LMPFGen
 from sim import PickPlaceEnv, LMP_wrapper
 from consts import ALL_BLOCKS, ALL_BOWLS
 from md_logger import MarkdownLogger
 
+default_open_ai_key = os.getenv('OPEN_AI_SECRET')
 
 class DemoRunner:
-        
+
     def __init__(self):
         self._cfg = OmegaConf.to_container(OmegaConf.load('cfg.yaml'), resolve=True)
         self._env = None
         self._model_name = ''
         self._md_logger = MarkdownLogger()
+        self._temp_file_manager = processing_utils.TempFileManager()
 
     def make_LMP(self, env):
         # LMP env wrapper
@@ -50,7 +54,9 @@ class DemoRunner:
                 'get_corner_name', 'get_side_name',
             ]
         }
-        variable_vars['say'] = lambda msg: self._md_logger.log_text(f'Robot says: "{msg}"')
+        # variable_vars['say'] = lambda msg: self._md_logger.log_text(f'Robot says: "{msg}"')
+        variable_vars['say'] = lambda msg: self._md_logger.log_message(
+            f'{msg}')
 
         # creating the function-generating LMP
         lmp_fgen = LMPFGen(cfg['lmps']['fgen'], fixed_vars, variable_vars, self._md_logger)
@@ -86,9 +92,9 @@ class DemoRunner:
 
         return info, img
 
-    def run(self, instruction):
+    def run(self, instruction, history):
         if self._env is None:
-            return 'Please run setup first!', None, None
+            return 'Please run setup first!', None, None, history
 
         self._env.cache_video = []
         self._md_logger.clear()
@@ -96,7 +102,7 @@ class DemoRunner:
         try:
             self._lmp_tabletop_ui(instruction, f'objects = {self._env.object_list}')
         except Exception as e:
-            return f'Error: {e}', None, None
+            return f'Error: {e}', None, None, history
 
         video_file_name = None
         if self._env.cache_video:
@@ -104,27 +110,39 @@ class DemoRunner:
             video_file_name = NamedTemporaryFile(suffix='.mp4').name
             rendered_clip.write_videofile(video_file_name, fps=25)
 
-        return self._md_logger.get_log(), self._env.get_camera_image(), video_file_name
+        # Update chat messages
+        for message in self._md_logger.get_messages():
+            history.append((None, message))
+        if self._env.cache_video:
+            temp_name = self._temp_file_manager.make_temp_copy_if_needed(video_file_name)
+            history.append((None, (temp_name, )))
+
+        return self._md_logger.get_log(), self._env.get_camera_image(), video_file_name, history
 
 
 def setup(api_key, model_name, n_blocks, n_bowls):
     if not api_key:
         return 'Please enter your OpenAI API key!', None, None
-    
+
     if n_blocks + n_bowls == 0:
         return 'Please select at least one object!', None, None
 
     demo_runner = DemoRunner()
-    
+
     info, img = demo_runner.setup(api_key, model_name, n_blocks, n_bowls)
-    return info, img, demo_runner
+    welcome_message = 'How can I help you?'
+    return info, img, demo_runner, [(None, welcome_message)]
 
 
-def run(instruction, demo_runner):
+def run(demo_runner, chat_history):
     if demo_runner is None:
-        return 'Please run setup first!', None, None
-    return demo_runner.run(instruction)
+        return 'Please run setup first!', None, None, chat_history, None
+    instruction = chat_history[-1][0]
+    return *demo_runner.run(instruction, chat_history), ''
 
+def submit_chat(chat_message, history):
+    history += [[chat_message, None]]
+    return '', history
 
 if __name__ == '__main__':
     with open('README.md', 'r') as f:
@@ -140,35 +158,60 @@ if __name__ == '__main__':
         with gr.Row():
             with gr.Column():
                 with gr.Row():
-                    inp_api_key = gr.Textbox(label='OpenAI API Key (this is not stored anywhere)', lines=1)
+                    inp_api_key = gr.Textbox(value=default_open_ai_key,
+                                             label='OpenAI API Key (this is not stored anywhere)', lines=1)
                     inp_model_name = gr.Dropdown(label='Model Name', choices=['code-davinci-002', 'text-davinci-002'], value='code-davinci-002')
                 with gr.Row():
-                    inp_n_blocks = gr.Slider(label='Number of Blocks', minimum=0, maximum=4, value=3, step=1)
-                    inp_n_bowls = gr.Slider(label='Number of Bowls', minimum=0, maximum=4, value=3, step=1)
-            
+                    inp_n_blocks = gr.Slider(label='Number of Blocks', minimum=0, maximum=5, value=3, step=1)
+                    inp_n_bowls = gr.Slider(label='Number of Bowls', minimum=0, maximum=5, value=3, step=1)
+
                 btn_setup = gr.Button("Setup/Reset Simulation")
                 info_setup = gr.Markdown(label='Setup Info')
             with gr.Column():
-                img_setup = gr.Image(label='Current Simulation')
+                img_setup = gr.Image(label='Current Simulation State')
 
         with gr.Row():
             with gr.Column():
-                
-                inp_instruction = gr.Textbox(label='Instruction', lines=1) 
+                chat_box = gr.Chatbot()
+                inp_instruction = gr.Textbox(label='Instruction', lines=1)
+                examples = gr.Examples(
+                    [
+                        'stack two of the blocks',
+                        'what color is the rightmost block?',
+                        'arrange the blocks into figure 3',
+                        'put blocks into non-matching bowls'
+                    ],
+                    inp_instruction,
+                )
                 btn_run = gr.Button("Run (this may take 30+ seconds)")
                 info_run = gr.Markdown(label='Generated Code')
             with gr.Column():
                 video_run = gr.Video(label='Video of Last Instruction')
-        
+
         btn_setup.click(
-            setup, 
-            inputs=[inp_api_key, inp_model_name, inp_n_blocks, inp_n_bowls], 
-            outputs=[info_setup, img_setup, state]
+            setup,
+            inputs=[inp_api_key, inp_model_name, inp_n_blocks, inp_n_bowls],
+            outputs=[info_setup, img_setup, state, chat_box],
         )
         btn_run.click(
-            run, 
-            inputs=[inp_instruction, state], 
-            outputs=[info_run, img_setup, video_run]
+            submit_chat,
+            [inp_instruction, chat_box],
+            [inp_instruction, chat_box],
+        ).then(
+            run,
+            inputs=[state, chat_box],
+            outputs=[info_run, img_setup, video_run, chat_box, inp_instruction],
         )
-        
+        inp_instruction.submit(
+            submit_chat,
+            [inp_instruction, chat_box],
+            [inp_instruction, chat_box],
+        ).then(
+            run,
+            inputs=[state, chat_box],
+            outputs=[info_run, img_setup, video_run,
+                     chat_box, inp_instruction],
+        )
+
+    demo.queue()
     demo.launch()

lmp.py

@@ -62,8 +62,8 @@ class LMP:
                 break
             except (RateLimitError, APIConnectionError) as e:
                 print(f'OpenAI API got err {e}')
-                print('Retrying after 10s.')
-                sleep(10)
+                print('Retrying after 2s.')
+                sleep(2)
 
         if self._cfg['include_context'] and context != '':
             to_exec = f'{context}\n{code_str}'
@@ -117,7 +117,7 @@ class LMPFGen:
         while True:
             try:
                 f_src = openai.Completion.create(
-                    prompt=prompt, 
+                    prompt=prompt,
                     stop=self._stop_tokens,
                     temperature=self._cfg['temperature'],
                     engine=self._cfg['engine'],
@@ -126,8 +126,8 @@ class LMPFGen:
                 break
             except (RateLimitError, APIConnectionError) as e:
                 print(f'OpenAI API got err {e}')
-                print('Retrying after 10s.')
-                sleep(10)
+                print('Retrying after 2s.')
+                sleep(2)
 
         if fix_bugs:
             f_src = openai.Edit.create(
@@ -141,7 +141,7 @@ class LMPFGen:
             other_vars = {}
         gvars = merge_dicts([self._fixed_vars, self._variable_vars, other_vars])
         lvars = {}
-        
+
         exec_safe(f_src, gvars, lvars)
 
         f = lvars[f_name]
@@ -187,9 +187,9 @@ class LMPFGen:
                     # redefine parent f so newly created child_fs are in scope
                     gvars = merge_dicts([self._fixed_vars, self._variable_vars, new_fs, other_vars])
                     lvars = {}
-                    
+
                     exec_safe(f_src, gvars, lvars)
-                    
+
                     f = lvars[f_name]
 
                 new_fs[f_name], srcs[f_name] = f, f_src
@@ -235,17 +235,17 @@ def var_exists(name, all_vars):
 
 def merge_dicts(dicts):
     return {
-        k : v 
+        k : v
         for d in dicts
         for k, v in d.items()
     }
-    
+
 
 def exec_safe(code_str, gvars=None, lvars=None):
     banned_phrases = ['import', '__']
     for phrase in banned_phrases:
         assert phrase not in code_str
-  
+
     if gvars is None:
         gvars = {}
     if lvars is None:

md_logger.py

@@ -2,15 +2,24 @@ class MarkdownLogger:
 
     def __init__(self):
         self._log = ''
+        self._messages = []
 
     def log_text(self, text):
         self._log += '\n' + text + '\n'
-    
+
     def log_code(self, code):
         self._log += f'\n```python\n{code}\n```\n'
 
+    def log_message(self, text):
+        self._messages.append(text)
+
     def clear(self):
         self._log = ''
 
     def get_log(self):
-        return self._log
+        return self._log
+
+    def get_messages(self):
+        m = self._messages
+        self._messages = []
+        return m

requirements.txt

@@ -8,4 +8,5 @@ pybullet
 imageio==2.4.1
 imageio-ffmpeg
 moviepy
-omegaconf
+omegaconf
+gradio