black and isort

Author: yeshaokai

amadeusgpt/__init__.py

@@ -5,10 +5,10 @@
 Apache-2.0 license
 """
 
-from amadeusgpt.integration_modules import *
 from matplotlib import pyplot as plt
 
 from amadeusgpt.implementation import AnimalBehaviorAnalysis
+from amadeusgpt.integration_modules import *
 from amadeusgpt.main import AMADEUS
 from amadeusgpt.version import VERSION, __version__
 

amadeusgpt/analysis_objects/event.py

@@ -172,11 +172,13 @@ def event_negate(cls, events: List[BaseEvent]) -> List[BaseEvent]:
             mask |= event.generate_mask()
         negate_mask = ~mask
 
-        
-
         negate_events = Event.mask2events(
-            negate_mask, video_file_path, sender_animal_name,
-            set(), set(), smooth_window_size = 1
+            negate_mask,
+            video_file_path,
+            sender_animal_name,
+            set(),
+            set(),
+            smooth_window_size=1,
         )
 
         return negate_events
@@ -540,7 +542,7 @@ def fuse_subgraph_by_kvs(
         For example, if there are two conditions to be met in the masks we look for locations that have overlap as 2
         """
         # retrieve all events that satisfy the conditions (k=v)
-        events = graph.traverse_by_kvs(merge_kvs)        
+        events = graph.traverse_by_kvs(merge_kvs)
         if not allow_more_than_2_overlap:
             assert (
                 Event.check_max_in_sum(events) <= number_of_overlap_for_fusion

amadeusgpt/analysis_objects/llm.py

@@ -1,15 +1,19 @@
+import base64
+import io
 import json
 import os
 import re
 import time
 import traceback
-from amadeusgpt.utils import AmadeusLogger
-from .base import AnalysisObject
+
+import cv2
 import openai
 from openai import OpenAI
-import base64
-import cv2
-import io
+
+from amadeusgpt.utils import AmadeusLogger
+
+from .base import AnalysisObject
+
 
 class LLM(AnalysisObject):
     total_tokens = 0
@@ -26,7 +30,6 @@ def __init__(self, config):
         self.context_window = []
         # only for logging and long-term memory usage.
         self.history = []
-    
 
     def encode_image(self, image_path):
         with open(image_path, "rb") as image_file:
@@ -42,10 +45,10 @@ def connect_gpt(self, messages, **kwargs):
         # if openai version is less than 1
         return self.connect_gpt_oai_1(messages, **kwargs)
 
-    def connect_gpt_oai_1(self, messages, **kwargs): 
+    def connect_gpt_oai_1(self, messages, **kwargs):
         """
         This is routed to openai > 1.0 interfaces
-        """       
+        """
 
         if self.config.get("use_streamlit", False):
             if "OPENAI_API_KEY" in os.environ:
@@ -82,10 +85,15 @@ def connect_gpt_oai_1(self, messages, **kwargs):
                 }
                 response = client.chat.completions.create(**json_data)
 
-                LLM.total_tokens =  LLM.total_tokens + response.usage.prompt_tokens + response.usage.completion_tokens
+                LLM.total_tokens = (
+                    LLM.total_tokens
+                    + response.usage.prompt_tokens
+                    + response.usage.completion_tokens
+                )
                 LLM.total_cost += (
                     LLM.prices[self.gpt_model]["input"] * response.usage.prompt_tokens
-                    + LLM.prices[self.gpt_model]["output"] * response.usage.completion_tokens
+                    + LLM.prices[self.gpt_model]["output"]
+                    * response.usage.completion_tokens
                 )
                 print("current total cost", round(LLM.total_cost, 2), "$")
                 print("current total tokens", LLM.total_tokens)
@@ -110,7 +118,7 @@ def connect_gpt_oai_1(self, messages, **kwargs):
 
         return response
 
-    def update_history(self, role, content, encoded_image = None, replace=False):
+    def update_history(self, role, content, encoded_image=None, replace=False):
         if role == "system":
             if len(self.history) > 0:
                 self.history[0]["content"] = content
@@ -124,7 +132,7 @@ def update_history(self, role, content, encoded_image = None, replace=False):
                 self.history.append({"role": role, "content": content})
                 num_AI_messages = (len(self.context_window) - 1) // 2
                 if num_AI_messages == self.keep_last_n_messages:
-                    print ("doing active forgetting")
+                    print("doing active forgetting")
                     # we forget the oldest AI message and corresponding answer
                     self.context_window.pop(1)
                     self.context_window.pop(1)
@@ -134,23 +142,25 @@ def update_history(self, role, content, encoded_image = None, replace=False):
                     self.history.append({"role": role, "content": content})
                     num_AI_messages = (len(self.context_window) - 1) // 2
                     if num_AI_messages == self.keep_last_n_messages:
-                        print ("doing active forgetting")
+                        print("doing active forgetting")
                         # we forget the oldest AI message and corresponding answer
                         self.context_window.pop(1)
                         self.context_window.pop(1)
                     self.context_window.append({"role": role, "content": content})
                 else:
                     message = {
-                        "role": "user", "content": [
-                        {"type": "text", "text": content},
-                        {"type": "image_url", "image_url": {
-                        "url": f"data:image/png;base64,{encoded_image}"}
-                        }]
-                    }                                            
-                    self.context_window.append(message) 
-    
-                    
-               
+                        "role": "user",
+                        "content": [
+                            {"type": "text", "text": content},
+                            {
+                                "type": "image_url",
+                                "image_url": {
+                                    "url": f"data:image/png;base64,{encoded_image}"
+                                },
+                            },
+                        ],
+                    }
+                    self.context_window.append(message)
 
     def clean_context_window(self):
         while len(self.context_window) > 1:
@@ -185,6 +195,7 @@ def parse_openai_response(cls, response):
 class VisualLLM(LLM):
     def __init__(self, config):
         super().__init__(config)
+
     def speak(self, sandbox):
         """
         Only to comment about one image
@@ -195,25 +206,29 @@ def speak(self, sandbox):
         """
 
         from amadeusgpt.system_prompts.visual_llm import _get_system_prompt
+
         self.system_prompt = _get_system_prompt()
         analysis = sandbox.exec_namespace["behavior_analysis"]
         scene_image = analysis.visual_manager.get_scene_image()
-        result, buffer = cv2.imencode('.jpeg', scene_image)     
+        result, buffer = cv2.imencode(".jpeg", scene_image)
         image_bytes = io.BytesIO(buffer)
-        base64_image = base64.b64encode(image_bytes.getvalue()).decode('utf-8')   
+        base64_image = base64.b64encode(image_bytes.getvalue()).decode("utf-8")
 
         self.update_history("system", self.system_prompt)
-        self.update_history("user", "here is the image", encoded_image = base64_image, replace = True)
-        response = self.connect_gpt(self.context_window, max_tokens=2000)        
+        self.update_history(
+            "user", "here is the image", encoded_image=base64_image, replace=True
+        )
+        response = self.connect_gpt(self.context_window, max_tokens=2000)
         text = response.choices[0].message.content.strip()
-        print (text)
+        print(text)
         pattern = r"```json(.*?)```"
         if len(re.findall(pattern, text, re.DOTALL)) == 0:
             raise ValueError("can't parse the json string correctly", text)
         else:
             json_string = re.findall(pattern, text, re.DOTALL)[0]
             json_obj = json.loads(json_string)
-            return json_obj      
+            return json_obj
+
 
 class CodeGenerationLLM(LLM):
     """
@@ -222,7 +237,6 @@ class CodeGenerationLLM(LLM):
 
     def __init__(self, config):
         super().__init__(config)
-  
 
     def speak(self, sandbox):
         """
@@ -265,10 +279,8 @@ def update_system_prompt(self, sandbox):
         task_program_docs = sandbox.get_task_program_docs()
         query_block = sandbox.get_query_block()
 
-        behavior_analysis = sandbox.exec_namespace[
-            "behavior_analysis"
-        ]
-       
+        behavior_analysis = sandbox.exec_namespace["behavior_analysis"]
+
         self.system_prompt = _get_system_prompt(
             query_block, core_api_docs, task_program_docs, behavior_analysis
         )
@@ -280,7 +292,6 @@ def update_system_prompt(self, sandbox):
 class MutationLLM(LLM):
     def __init__(self, config):
         super().__init__(config)
-    
 
     def update_system_prompt(self, sandbox):
         from amadeusgpt.system_prompts.mutation import _get_system_prompt
@@ -307,7 +318,6 @@ def speak(self, sandbox):
 class BreedLLM(LLM):
     def __init__(self, config):
         super().__init__(config)
-    
 
     def update_system_prompt(self, sandbox):
         from amadeusgpt.system_prompts.breed import _get_system_prompt
@@ -342,7 +352,6 @@ class DiagnosisLLM(LLM):
     """
     Resource management for testing and error handling
     """
-    
 
     @classmethod
     def get_system_prompt(
@@ -410,8 +419,9 @@ def speak(self, sandbox):
 
 
 if __name__ == "__main__":
-    from amadeusgpt.config import Config   
+    from amadeusgpt.config import Config
     from amadeusgpt.main import create_amadeus
+
     config = Config("amadeusgpt/configs/EPM_template.yaml")
 
     amadeus = create_amadeus(config)

amadeusgpt/analysis_objects/object.py

@@ -298,7 +298,7 @@ def set_body_orientation_keypoints(
 
     def set_head_orientation_keypoints(
         self, head_orientation_keypoints: Dict[str, Any]
-    ):        
+    ):
         self.nose_name = head_orientation_keypoints["nose"]
         self.neck_name = head_orientation_keypoints["neck"]
         self.support_head_orientation = True

amadeusgpt/analysis_objects/relationship.py

@@ -71,7 +71,7 @@ def calc_angle_in_egocentric_animal(mouse_cs_inv, p):
         p_in_mouse[:, 1], p_in_mouse[:, 0]
     )  # relative angle between the object and the mouse body axis
     theta = np.rad2deg(theta % (2 * np.pi))
- 
+
     return theta
 
 
@@ -253,9 +253,9 @@ def _animal_animal_relationship(
         sender_pos = sender_animal.get_center()
         receiver_pos = receiver_animal.get_center()
         direction_vector = receiver_pos - sender_pos
-        sender_velocity = np.nanmean(sender_animal.get_velocity(), axis = 1)
+        sender_velocity = np.nanmean(sender_animal.get_velocity(), axis=1)
         norm_direction_vector = direction_vector / np.linalg.norm(direction_vector)
-        relative_speed = np.einsum('ij,ij->i', sender_velocity, norm_direction_vector)
+        relative_speed = np.einsum("ij,ij->i", sender_velocity, norm_direction_vector)
 
         closest_distance = np.nanmin(
             get_pairwise_distance(sender_animal.keypoints, receiver_animal.keypoints),
@@ -274,6 +274,6 @@ def _animal_animal_relationship(
         if angles is not None:
             ret["relative_angle"] = angles
         if orientation is not None:
-            ret["orientation"] = orientation 
+            ret["orientation"] = orientation
 
         return ret

amadeusgpt/app.py

@@ -23,7 +23,6 @@ def main():
             st.session_state["exist_valid_openai_api_key"] = True
         else:
             st.session_state["exist_valid_openai_api_key"] = False
-    
 
     example_to_page = {}
 

amadeusgpt/app_utils.py

@@ -112,7 +112,7 @@ def __init__(self, amadeus_answer=None, json_entry=None):
             amadeus_answer = amadeus_answer.to_dict()
         self.data.update(amadeus_answer)
 
-    def render(self, debug = False):
+    def render(self, debug=False):
         """
         We use the getter for better encapsulation
         overall structure of what to be rendered
@@ -181,9 +181,9 @@ def render(self, debug = False):
                         # Remind users we are fixing the error by self debugging
                         st.markdown(f"Let me try to fix the error by self-debugging\n ")
                         if not debug:
-                            sandbox.llms["self_debug"].speak(sandbox)                        
+                            sandbox.llms["self_debug"].speak(sandbox)
                             qa_message = sandbox.code_execution(qa_message)
-                            self.render(debug = True)
+                            self.render(debug=True)
                     # do not need to execute the block one more time
                     if not self.rendered:
                         self.rendered = True
@@ -558,7 +558,9 @@ def render_page_by_example(example):
             if uploaded_keypoint_file is not None:
                 path = save_uploaded_file(uploaded_keypoint_file, save_dir)
                 st.session_state["uploaded_keypoint_file"] = path
-                config["keypoint_info"]["keypoint_file_path"] = st.session_state["uploaded_keypoint_file"]
+                config["keypoint_info"]["keypoint_file_path"] = st.session_state[
+                    "uploaded_keypoint_file"
+                ]
 
         if "uploaded_video_file" not in st.session_state:
             uploaded_video_file = st.file_uploader(
@@ -569,7 +571,9 @@ def render_page_by_example(example):
             if uploaded_video_file is not None:
                 path = save_uploaded_file(uploaded_video_file, save_dir)
                 st.session_state["uploaded_video_file"] = uploaded_video_file
-                config["video_info"]["video_file_path"] = st.session_state["uploaded_video_file"]      
+                config["video_info"]["video_file_path"] = st.session_state[
+                    "uploaded_video_file"
+                ]
 
         ###### USER INPUT PANEL ######
         # get user input once getting the uploaded files
@@ -675,12 +679,15 @@ def render_page_by_example(example):
             st.caption("Raw video from Horse-30")
         else:
             st.caption("DeepLabCut-SuperAnimal tracked video")
-        if config["video_info"]["video_file_path"] and config["video_info"]["video_file_path"] is not None: 
+        if (
+            config["video_info"]["video_file_path"]
+            and config["video_info"]["video_file_path"] is not None
+        ):
             st.video(config["video_info"]["video_file_path"])
 
         if "uploaded_video_file" in st.session_state:
             st.video(st.session_state["uploaded_video_file"])
-        
+
         # we only show objects for MausHaus for demo
         # if sam_image is not None:
         #     st.caption("SAM segmentation results")

amadeusgpt/main.py

@@ -56,7 +56,7 @@ def __init__(self, config: Dict[str, Any]):
 
         # can only do this after the register process
         self.sandbox.configure_using_vlm()
-            
+
     def match_integration_module(self, user_query: str):
         """
         Return a list of matched integration modules
@@ -87,17 +87,18 @@ def step(self, user_query):
 
     def get_analysis(self):
         sandbox = self.sandbox
-        analysis = sandbox.exec_namespace['behavior_analysis']
+        analysis = sandbox.exec_namespace["behavior_analysis"]
         return analysis
-    
+
     def run_task_program(self, task_program_name: str):
         return self.sandbox.run_task_program(task_program_name)
 
 
 if __name__ == "__main__":
-    from amadeusgpt.config import Config   
-    from amadeusgpt.main import create_amadeus
     from amadeusgpt.analysis_objects.llm import VisualLLM
+    from amadeusgpt.config import Config
+    from amadeusgpt.main import create_amadeus
+
     config = Config("amadeusgpt/configs/EPM_template.yaml")
 
     amadeus = create_amadeus(config)