some clean up and some docs

Author: yeshaokai

amadeusgpt/analysis_objects/analysis_factory.py

@@ -2,7 +2,6 @@
 
 analysis_fac = {}
 
-
 def create_analysis(config):
     if str(config) not in analysis_fac:
         analysis_fac[str(config)] = AnimalBehaviorAnalysis(config)

amadeusgpt/main.py

@@ -77,11 +77,17 @@ def step(self, user_query):
         return result
 
     def get_analysis(self):
-        sandbox = self.sandbox
-        analysis = sandbox.exec_namespace["behavior_analysis"]
+        """
+        Every sandbox stores a unique "behavior analysis" instance in its namespace
+        Therefore, get analysis gets the current sandbox's analysis.
+        """
+        analysis = self.sandbox.exec_namespace["behavior_analysis"]
         return analysis
 
     def run_task_program(self, task_program_name: str):
+        """
+        Execute the task program on the currently holding sandbox
+        """
         return self.sandbox.run_task_program(task_program_name)
 
     def save_results(self, out_folder: str| None = None):

amadeusgpt/managers/visual_manager.py

@@ -678,14 +678,15 @@ def write_video(self, out_folder, video_file_path, out_name, events):
         return out_videos
 
     def generate_video_clips_from_events(
-        self, out_folder, video_file, events: List[BaseEvent], behavior_name
+        self, out_folder, events: List[BaseEvent], behavior_name
     ):
         """
         This function takes a list of events and generates video clips from the events
         1) For the same events, we first group events based on the video
         2) For the same event on the same video, we plot the animal name and the "sender" of the event
         3) Then we write those videos to the disk
         """
+        video_file = self.config["video_info"]["video_file_path"]
 
         videoname = video_file.split("/")[-1].replace(".mp4", "").replace(".avi", "")
         video_name = f"{videoname}_{behavior_name}_video.mp4"

amadeusgpt/programs/sandbox.py

@@ -305,15 +305,6 @@ def update_config(self, config):
     def copy(self):
         return Sandbox(self.config, self.api_registry)
 
-    def visual_validate(self, video_file, events, behavior_name):
-        # change video and keypoint file
-        analysis = create_analysis(self.config)
-        out_folder = os.path.join(self.config["evo_info"]["data_folder"], "inspection")
-        discovered_behaviors = []
-        for name, task_program in self.task_program_library.items():
-            if task_program["creator"] != "human":
-                discovered_behaviors.append(name)
-
     def update_matched_integration_modules(self, matched_modules):
         self.matched_modules = matched_modules
 
@@ -358,7 +349,8 @@ def update_namespace(self):
         self.exec_namespace["task_programs"] = TaskProgramLibrary.get_task_programs()
 
     def code_execution(self, qa_message):
-        # add main function into the namespace
+        # update the namespace in the beginning of code execution makes sure that 
+        # if there is a change in the config, we always use the newest config
         self.update_namespace()
         code = qa_message["code"]
         # not need to do further if there was no code found
@@ -424,12 +416,16 @@ def events_to_videos(self, events, function_name):
         out_folder = str(self.result_folder)
         os.makedirs(out_folder, exist_ok=True)
         behavior_name = "_".join(function_name.split(" "))
-        video_file = self.config["video_info"]["video_file_path"]
         return visual_manager.generate_video_clips_from_events(
-            out_folder, video_file, events, behavior_name
+            out_folder, events, behavior_name
         )
 
     def render_qa_message(self, qa_message):
+        """
+        To be called after code execution.
+        If the function returns a list of events, we visualize those events to keypoint plot, ethogram plot and videos
+        if the function returns is a tuple of axe and figure, we put them into the plots filed 
+        """
         function_rets = qa_message["function_rets"]
         behavior_analysis = self.exec_namespace["behavior_analysis"]
         bodypart_names = behavior_analysis.animal_manager.get_keypoint_names()
@@ -487,6 +483,11 @@ def render_qa_message(self, qa_message):
         return qa_message
 
     def llm_step(self, user_query):
+        """
+        1) We first use gpt-4o to create meta_info describing the scene
+        2) We then ask LLM to generate code based on the query
+        3) We also cache the qa_message for future reference
+        """
         qa_message = create_message(user_query, self)
 
         # so that the frontend can display it too
@@ -502,7 +503,8 @@ def llm_step(self, user_query):
 
     def run_task_program(self, task_program_name):
         """
-        Sandbox is also responsible for running task program
+        1) sandbox is also responsible for running task program
+        2) self.task_program_library references to a singleton so a different sandbox still has reference to the task program
         """
         task_program = self.task_program_library[task_program_name]
         # there might be better way to set this
@@ -542,6 +544,9 @@ def step(self, user_query, number_of_debugs=1):
 
 
 def save_figure_to_tempfile(fig):
+    """
+    Only used for debug
+    """
     import tempfile
 
     # save the figure
@@ -566,6 +571,9 @@ def save_figure_to_tempfile(fig):
 
 
 def render_temp_message(query, sandbox):
+    """
+    Only used for debug
+    """
     import streamlit as st
 
     qa_message = create_message("random query", sandbox)

amadeusgpt/utils.py

@@ -1,21 +1,14 @@
 import ast
-import copy
 import inspect
-import re
 import sys
 import time
 import traceback
 from itertools import groupby
 from operator import itemgetter
-from pydoc import doc
 from typing import Any, Dict, Sequence
-
 import cv2
-import matplotlib.pyplot as plt
 import numpy as np
-from numpy import ndarray
 from scipy.ndimage.filters import uniform_filter1d
-
 from amadeusgpt.logger import AmadeusLogger
 
 
@@ -63,23 +56,6 @@ def moving_average(x: Sequence, window_size: int, pos: str = "centered"):
     return uniform_filter1d(x, window_size, mode="constant", origin=origin)
 
 
-def moving_variance(x, window_size):
-    """
-    Blazing fast implementation of a moving variance.
-    :param x: ndarray, 1D input
-    :param window_size: int, window length
-    :return: 1D ndarray of length len(x)-window+1
-    """
-    nrows = x.size - window_size + 1
-    n = x.strides[0]
-    mat = np.lib.stride_tricks.as_strided(
-        x,
-        shape=(nrows, window_size),
-        strides=(n, n),
-    )
-    return np.var(mat, axis=1)
-
-
 def smooth_boolean_mask(x: Sequence, window_size: int):
     # `window_size` should be at least twice as large as the
     # minimal number of consecutive frames to be smoothed out.
@@ -109,136 +85,6 @@ def get_video_length(video_path):
     return int(n_frames)
 
 
-def frame_number_to_minute_seconds(frame_number, video_path):
-    fps = get_fps(video_path)
-    temp = frame_number / fps
-    minutes = int(temp // 60)
-    seconds = int(temp % 60)
-    ret = f"{minutes:02d}:{seconds:02d}"
-    return ret
-
-
-def search_generated_func(text):
-    functions = []
-    lines = text.split("\n")
-    func_names = []
-    i = 0
-
-    while i < len(lines):
-        line = lines[i]
-        func_signature = "def task_program"
-        if line.startswith(func_signature):
-            start = line.index("def ") + 4
-            end = line.index("(")
-            func_name = line[start:end]
-            func_names.append(func_name)
-            function_lines = [line]
-            nesting_level = 0
-            i += 1
-
-            while i < len(lines):
-                line = lines[i]
-                # Check for nested function definitions
-                if line.lstrip().startswith("def "):
-                    nesting_level += 1
-                elif line.lstrip().startswith("return") and nesting_level > 0:
-                    nesting_level -= 1
-                elif line.lstrip().startswith("return") and nesting_level == 0:
-                    function_lines.append(line)
-                    break
-
-                function_lines.append(line)
-                i += 1
-
-            functions.append("\n".join(function_lines))
-        i += 1
-
-    return functions, func_names
-
-
-def search_external_module_for_context_window(text):
-    """
-    just include everything
-    """
-    functions = []
-    i = 0
-    lines = text.split("\n")
-    func_names = []
-    while i < len(lines):
-        line = lines[i].strip()
-        func_signature = "def "
-        if line.strip() == "":
-            i += 1
-            continue
-        if line.startswith(func_signature):
-            start = line.index("def ") + 4
-            end = line.index("(")
-            func_name = line[start:end]
-            func_names.append(func_name)
-            function_lines = [line]
-            in_function = True
-            while in_function:
-                i += 1
-                if i == len(lines):
-                    break
-                next_line = lines[i].rstrip()
-                if not next_line.startswith((" ", "\t")):
-                    in_function = False
-                    continue
-                function_lines.append(next_line)
-            functions.append("\n".join(function_lines))
-        else:
-            i += 1
-    return functions, func_names
-
-
-def search_external_module_for_task_program_table(text):
-    """
-    in this case, just include everything
-    """
-    functions = []
-    i = 0
-    lines = text.split("\n")
-    lines_copy = copy.deepcopy(lines)
-    func_names = []
-    example_indentation = " " * 4
-    while i < len(lines):
-        if lines[i].startswith("def "):
-            start = lines[i].index("def ") + 4
-            end = lines[i].index("(")
-            func_name = lines[i][start:end]
-            func_names.append(func_name)
-        if "" not in lines[i]:
-            i += 1
-            continue
-        else:
-            lines[i] = lines[i].replace("", "").strip()
-        line = lines[i].strip()
-        func_signature = "def "
-        if line.strip() == "":
-            i += 1
-            continue
-        if line.startswith("def "):
-            function_lines = [line]
-            in_function = True
-            while in_function:
-                i += 1
-                if i == len(lines):
-                    break
-                next_line = lines[i].rstrip()
-                if "" not in lines_copy[i]:
-                    in_function = False
-                    continue
-                function_lines.append(
-                    next_line.replace("", "").replace(example_indentation, "", 1)
-                )
-            functions.append("\n".join(function_lines))
-        else:
-            i += 1
-
-    return functions, func_names
-
-
 def filter_kwargs_for_function(func, kwargs):
     sig = inspect.signature(func)
     return {k: v for k, v in kwargs.items() if k in sig.parameters}
@@ -385,20 +231,3 @@ def func2json(func):
         }
         return json_obj
 
-
-def get_func_name_from_func_string(function_string: str):
-    import ast
-
-    # Parse the string into an AST
-    parsed_ast = ast.parse(function_string)
-
-    # Initialize a variable to hold the function name
-    function_name = None
-
-    # Traverse the AST
-    for node in ast.walk(parsed_ast):
-        if isinstance(node, ast.FunctionDef):
-            function_name = node.name
-            break
-
-    return function_name