From b2bcc9fc28a5906ee1e0ab7e0a77c618c512ad1c Mon Sep 17 00:00:00 2001
From: Dillon Laird <dillonalaird@gmail.com>
Date: Tue, 15 Oct 2024 16:53:35 -0700
Subject: [PATCH] add another prompt for editing code

---
 vision_agent/agent/vision_agent_prompts.py | 65 +++++++++++++++++++++-
 1 file changed, 62 insertions(+), 3 deletions(-)

diff --git a/vision_agent/agent/vision_agent_prompts.py b/vision_agent/agent/vision_agent_prompts.py
index 7e06bfb2..66d6614e 100644
--- a/vision_agent/agent/vision_agent_prompts.py
+++ b/vision_agent/agent/vision_agent_prompts.py
@@ -54,7 +54,7 @@
 <let_user_respond>false</let_user_respond>
 
 OBSERVATION:
-[Artifact dog_detector.py]
+[Artifact dog_detector.py (5 lines total)]
 0|from vision_agent.tools import load_image, owl_v2
 1|def detect_dogs(image_path: str):
 2|    image = load_image(image_path)
@@ -95,7 +95,7 @@
 <let_user_respond>false</let_user_respond>
 
 OBSERVATION:
-[Artifact dog_detector.py]
+[Artifact dog_detector.py (5 lines total)]
 0|from vision_agent.tools import load_image, owl_v2
 1|def detect_dogs(image_path: str):
 2|    image = load_image(image_path)
@@ -167,7 +167,7 @@
 <let_user_respond>false</let_user_respond>
 
 OBSERVATION:
-[Artifact code.py]
+[Artifact code.py (17 lines total)]
 0|from vision_agent.tools import load_image, florence2_phrase_grounding, closest_box_distance, overlay_bounding_boxes, save_image
 1|def count_workers_with_helmets(image_path: str, output_path: str):
 2|    image = load_image(image_path)
@@ -202,7 +202,9 @@
 AGENT: <thinking>Two workers with helmets are detected, I will show this to the user and ask them if the result looks good.</thinking>
 <response>The code to detect workers with helmets is saved in code.py and the visualization under 'workers_viz.png'.</response>
 <let_user_respond>true</let_user_respond>
+"""
 
+EXAMPLES_CODE3_EXTRA1 = """
 USER: The detections are slightly off. Can you fine tune florence2 using these labels? "[{'image_path': 'image1.png': 'labels': ['worker', 'helmet'], 'bboxes': [[235, 118, 294, 241], [232, 118, 294, 128]]}, {'image_path': 'image2.png': 'labels': ['worker', 'helmet'], 'bboxes': [[235, 118, 294, 241], [232, 118, 294, 128]]}]"
 
 AGENT: <thinking>Because the user has supplied me with labels I can call object_detection_fine_tuning on their behalf to fine tune the model</thinking>
@@ -250,3 +252,60 @@ def count_workers_with_helmets(image_path: str, output_path: str):
 <response>I have updated the code using the fine tuned florence2_phrase_grounding model to count the workers wearing helmets in code.py and saved the visualization under 'workers_viz.png'.</response>
 <let_user_respond>true</let_user_respond>
 """
+
+EXAMPLES_CODE3_EXTRA2 = """
+USER: Can you write a streamlit app that allows me to upload an image and count the workers with helmets?
+
+AGENT: <thinking>I will use edit_code_artifact to generate a streamlit app that allows the user to upload an image and count the workers with helmets.</thinking>
+<execute_python>
+CODE = '''
+import streamlit as st
+from code imoprt count_workers_with_helmets
+import os
+
+def main():
+    st.title("Worker with Helmet Counter")
+    uploaded_file = st.file_uploader("Choose image...", type=["jpeg", "jpg", "png"])
+    if uploaded_file is not None:
+        with open("temp_image.jpg", "wb") as f:
+            f.write(uploaded_file.read())
+
+        count = count_workers_with_helmets("temp_image.jpg", "temp_image_viz.jpg")
+        st.image("temp_image_viz.jpg")
+        st.write(f"Number of workers with helmets: {count}")
+        os.remove("temp_image.jpg")
+        os.remove("temp_image_viz.jpg")
+
+if __name__ == "__main__":
+    main()
+'''
+edit_code_artifact(artifacts, 'streamlit_app.py', CODE, 0, 0)
+</execute_python>
+
+OBSERVATION:
+[Artifact streamlit_app.py (19 lines total)]
+0|import streamlit as st
+1|from code imoprt count_workers_with_helmets
+2|import os
+3|
+4|def main():
+5|    st.title("Worker with Helmet Counter")
+6|    uploaded_file = st.file_uploader("Choose image...", type=["jpeg", "jpg", "png"])
+7|    if uploaded_file is not None:
+8|        with open("temp_image.jpg", "wb") as f:
+9|            f.write(uploaded_file.read())
+10|
+11|        count = count_workers_with_helmets("temp_image.jpg", "temp_image_viz.jpg")
+12|        st.image("temp_image_viz.jpg")
+13|        st.write(f"Number of workers with helmets: {count}")
+14|        os.remove("temp_image.jpg")
+15|        os.remove("temp_image_viz.jpg")
+16|
+17|if __name__ == "__main__":
+18|    main()
+[End of artifact]
+
+AGENT: <thinking>I have generated the streamlit app that allows the user to upload an image and count the workers</thinking>
+<response>The streamlit app is saved in streamlit_app.py, you can run the app by running `streamlit run streamlit_app.py`.</response>
+<let_user_respond>true</let_user_respond>
+"""