Added roi object names into system prompt and fixed a bug

Author: yeshaokai

amadeusgpt/analysis_objects/llm.py

@@ -229,11 +229,12 @@ def update_system_prompt(self, sandbox):
         task_program_docs = sandbox.get_task_program_docs()
         query_block = sandbox.get_query_block()
 
-        keypoint_names = sandbox.exec_namespace[
+        behavior_analysis = sandbox.exec_namespace[
             "behavior_analysis"
-        ].get_keypoint_names()
+        ]
+       
         self.system_prompt = _get_system_prompt(
-            query_block, core_api_docs, task_program_docs, keypoint_names
+            query_block, core_api_docs, task_program_docs, behavior_analysis
         )
 
         # update both history and context window

amadeusgpt/managers/relationship_manager.py

@@ -52,7 +52,7 @@ def get_animals_objects_relationships(
         for animal in animals:
             if animal_bodyparts_names is not None:
                 # the keypoints of animal get updated when we update the roi bodypart names
-                animal.update_roi_keypoint_names(animal_bodyparts_names)
+                animal.update_roi_keypoint_by_names(animal_bodyparts_names)
             for object in objs:
                 animal_object_relations = AnimalObjectRelationship(
                     animal, object, animal_bodyparts_names

amadeusgpt/system_prompts/code_generator.py

@@ -2,8 +2,10 @@ def _get_system_prompt(
     query,
     core_api_docs,
     task_program_docs,
-    keypoint_names,
+    behavior_analysis,
 ):
+    keypoint_names = behavior_analysis.get_keypoint_names()
+    roi_object_names = behavior_analysis.get_roi_object_names()
     system_prompt = f""" 
 You are helpful AI assistant. Your job is to answer user queries. 
 Importantly, before you write the code, you need to explain whether the question can be answered accurately by code. If not,  ask users to give more information.
@@ -66,6 +68,7 @@ def get_watching_events(config: Config):
 {query}\n{core_api_docs}\n{task_program_docs}\n
 
 The keypoint names for the animals are: {keypoint_names}
+Available ROI objects are: {roi_object_names}
 
 FORMATTING:
 1) If you are asked to provide plotting code, make sure you don't call plt.show() but return a tuple figure, axs