From d5cd5d6b3b47f61612fce4bd476500647fd7d649 Mon Sep 17 00:00:00 2001
From: Paul Palmieri <paulpalmieri@outlook.fr>
Date: Mon, 26 Aug 2024 13:36:16 +0200
Subject: [PATCH 01/11] Create profiling.py

Holds the decorators and chart logic
---
 .../lavague/core/utilities/profiling.py       | 184 ++++++++++++++++++
 1 file changed, 184 insertions(+)
 create mode 100644 lavague-core/lavague/core/utilities/profiling.py

diff --git a/lavague-core/lavague/core/utilities/profiling.py b/lavague-core/lavague/core/utilities/profiling.py
new file mode 100644
index 00000000..344b266b
--- /dev/null
+++ b/lavague-core/lavague/core/utilities/profiling.py
@@ -0,0 +1,184 @@
+import time
+from datetime import datetime
+import pandas as pd
+import matplotlib.pyplot as plt
+from functools import wraps
+import io
+from IPython.display import display, Image
+from itertools import cycle
+
+# stores llm and retriever calls
+agent_events = []
+
+# stores total runtime of each step
+agent_steps = []
+
+# call before each agent step to group events by steps
+def start_new_step():
+    global agent_events
+    agent_events.append([])
+
+# track the runtime of the run_step function
+def track_total_runtime():
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            start_time = time.time()
+            result = func(*args, **kwargs)
+            end_time = time.time()
+            duration = end_time - start_time
+            agent_steps.append({
+                'start_time': start_time,
+                # 'end_time': end_time,
+                'duration': duration
+            })
+            return result
+        return wrapper
+    return decorator
+
+# track runtime, prompt size, and completion size of a llm call (used on world model + action engine)
+def track_llm_call(event_name):
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            start_time = time.time()
+            result = func(*args, **kwargs)
+            end_time = time.time()
+            duration = end_time - start_time
+            
+            agent_events[-1].append({
+                'event_name': event_name,
+                'start_time': start_time,
+                # 'end_time': end_time,
+                'duration': duration,
+                'prompt_size': len(args[0]),
+                'completion_size': len(result.text)
+            })
+            return result
+        return wrapper
+    return decorator
+
+# track runtime and html size of a retriever call
+def track_retriever(event_name):
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            start_time = time.time()
+            result = func(*args, **kwargs)
+            end_time = time.time()
+            duration = end_time - start_time
+            
+            agent_events[-1].append({
+                'event_name': event_name,
+                'start_time': start_time,
+                # 'end_time': end_time,
+                'duration': duration,
+                'html_size': len(args[0].driver.get_html()),
+            })
+            return result
+        return wrapper
+    return decorator
+
+
+class ChartGenerator:
+    def __init__(self, agent_events, agent_steps):
+        self.agent_events = agent_events
+        self.total_step_runtime = agent_steps
+        self.step_color = "lightgrey"
+        self.event_color_scheme = ["lightblue", "bisque", "thistle", "lightgreen", "pink"]
+
+    def plot_waterfall(self):
+        # Calculate the earliest start time to align the x-axis to 0
+        base_start_time = self.total_step_runtime[0]['start_time']
+        
+        plt.figure(figsize=(20, 8))
+        ax = plt.gca()
+    
+        color_cycle = cycle(self.event_color_scheme)
+        event_colors = {}
+        
+        # Plot total step runtime (from run_step)
+        for i, step in enumerate(self.total_step_runtime):
+            duration = step['duration']
+            start_time = step['start_time'] - base_start_time  # Normalize to 0
+            
+            ax.barh(i, duration, left=start_time, color=self.step_color)
+            ax.text(start_time + duration / 2, i - 0.45, f"{duration:.2f}s",
+                    ha='center', va='center')
+        
+        # Plot each individual event on top of the step runtime
+        for step_index, step in enumerate(self.agent_events):
+            for event in step:
+                duration = event['duration']
+                event_name = event['event_name']
+                start_time = event['start_time'] - base_start_time  # Normalize to 0
+                
+                if event_name not in event_colors:
+                    event_colors[event_name] = next(color_cycle)
+                
+                color = event_colors[event_name]
+                ax.barh(step_index, duration, left=start_time, color=color)
+                ax.text(start_time + duration / 2, step_index, f"{duration:.2f}s",
+                        ha='center', va='center', fontsize=9, color='black', rotation=90)
+        
+        ax.invert_yaxis()
+        ax.set_yticks(range(len(self.total_step_runtime)))
+        ax.set_yticklabels([f'Step {i+1}' for i in range(len(self.total_step_runtime))])
+        ax.set_xlabel('Time (seconds)')
+        ax.set_title('Waterfall Chart with Events Overlay')
+
+        
+        # Add legend for event colors
+        legend_labels = [plt.Line2D([0], [0], color=color, lw=4) for color in event_colors.values()]
+        ax.legend(legend_labels, event_colors.keys(), title="Event Name")
+        
+        # Save to buffer
+        buf = io.BytesIO()
+        plt.savefig(buf, format='png')
+        buf.seek(0)
+        
+        plt.close()
+        
+        return Image(buf.read())
+
+    def get_summary_df(self):
+        summary_data = {}
+        
+        # Iterate over each step and each event in the step
+        for step_index, step_events in enumerate(self.agent_events):
+            # Step row key
+            step_key = f"Step {step_index + 1}"
+            summary_data[step_key] = {}
+            
+            # Count the number of each event type in the step
+            event_counts = {}
+            
+            # Iterate over each event in the step
+            for event in step_events:
+                event_name = event['event_name']
+                
+                # Increment the count for the event
+                if event_name not in event_counts:
+                    event_counts[event_name] = 1
+                else:
+                    event_counts[event_name] += 1
+                
+                # for each key in the event, excluding 'event_name', 'start_time', and 'end_time', add the value to the summary
+                for key, value in event.items():
+                    if key not in ['event_name', 'start_time', 'end_time']:
+                        metric_key = f"{event_name} {key}"
+                        
+                        if metric_key not in summary_data[step_key]:
+                            summary_data[step_key][metric_key] = value
+                        else:
+                            summary_data[step_key][metric_key] += value
+            
+            # add the event counts
+            for event_name, count in event_counts.items():
+                count_key = f"{event_name} count"
+                summary_data[step_key][count_key] = count
+
+        # Convert the dictionary to a DataFrame
+        df = pd.DataFrame(summary_data).T
+
+        return df
\ No newline at end of file

From b8ccb19b7d77c8d3ac7db0e8fa34a1f5bc947374 Mon Sep 17 00:00:00 2001
From: Paul Palmieri <paulpalmieri@outlook.fr>
Date: Mon, 26 Aug 2024 13:38:41 +0200
Subject: [PATCH 02/11] Use decorators

---
 lavague-core/lavague/core/agents.py      | 4 ++++
 lavague-core/lavague/core/navigation.py  | 6 +++++-
 lavague-core/lavague/core/world_model.py | 9 ++++++++-
 3 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/lavague-core/lavague/core/agents.py b/lavague-core/lavague/core/agents.py
index 2541323f..f7f951e1 100644
--- a/lavague-core/lavague/core/agents.py
+++ b/lavague-core/lavague/core/agents.py
@@ -22,6 +22,8 @@
 from lavague.core.token_counter import TokenCounter
 from lavague.core.utilities.config import is_flag_true
 
+from lavague.core.utilities.profiling import ChartGenerator, track_total_runtime, start_new_step
+
 logging_print = logging.getLogger(__name__)
 logging_print.setLevel(logging.INFO)
 format = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
@@ -446,6 +448,7 @@ def _run_step_gradio(
         self.process_token_usage()
         self.logger.end_step()
 
+    @track_total_runtime()
     def run_step(self, objective: str) -> Optional[ActionResult]:
         obs = self.driver.get_obs()
         current_state, past = self.st_memory.get_state()
@@ -503,6 +506,7 @@ def run(
 
         try:
             for _ in range(self.n_steps):
+                start_new_step()
                 result = self.run_step(objective)
 
                 if result is not None:
diff --git a/lavague-core/lavague/core/navigation.py b/lavague-core/lavague/core/navigation.py
index cc15f0ae..bede4907 100644
--- a/lavague-core/lavague/core/navigation.py
+++ b/lavague-core/lavague/core/navigation.py
@@ -24,6 +24,7 @@
 from PIL import Image
 from llama_index.core.base.llms.base import BaseLLM
 from llama_index.core.embeddings import BaseEmbedding
+from lavague.core.utilities.profiling import track_retriever, track_llm_call
 
 NAVIGATION_ENGINE_PROMPT_TEMPLATE = ActionTemplate(
     """
@@ -143,6 +144,7 @@ def from_context(
             extractor,
         )
 
+    @track_retriever()
     def get_nodes(self, query: str) -> List[str]:
         """
         Get the nodes from the html page
@@ -452,7 +454,9 @@ def execute_instruction(self, instruction: str) -> ActionResult:
                 query_str=instruction,
                 authorized_xpaths=authorized_xpaths,
             )
-            response = self.llm.complete(prompt).text
+
+            # response = self.llm.complete(prompt).text
+            response = track_llm_call("Navigation")(self.llm.complete)(prompt).text
             end = time.time()
             action_generation_time = end - start
             action_outcome = {
diff --git a/lavague-core/lavague/core/world_model.py b/lavague-core/lavague/core/world_model.py
index 65ac7e7c..0e3be5f2 100644
--- a/lavague-core/lavague/core/world_model.py
+++ b/lavague-core/lavague/core/world_model.py
@@ -11,6 +11,7 @@
 from lavague.core.utilities.model_utils import get_model_name
 import time
 import yaml
+from lavague.core.utilities.profiling import track_llm_call
 
 WORLD_MODEL_GENERAL_EXAMPLES = """
 Objective:  Go to the first issue you can find
@@ -430,7 +431,13 @@ def get_instruction(
         )
 
         start = time.time()
-        mm_llm_output = mm_llm.complete(prompt, image_documents=image_documents).text
+
+        # decorated llm call
+        mm_llm_output = track_llm_call("World Model")(mm_llm.complete)(
+            prompt, image_documents=image_documents
+        ).text
+        # mm_llm_output = mm_llm.complete(prompt, image_documents=image_documents).text
+
         end = time.time()
         world_model_inference_time = end - start
 

From 77a27b63faf61e1ef34ec233e0925a9c122f5eb2 Mon Sep 17 00:00:00 2001
From: Paul Palmieri <paulpalmieri@outlook.fr>
Date: Mon, 26 Aug 2024 13:39:09 +0200
Subject: [PATCH 03/11] Add get_summary to agents.py

---
 lavague-core/lavague/core/agents.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/lavague-core/lavague/core/agents.py b/lavague-core/lavague/core/agents.py
index f7f951e1..04df3b67 100644
--- a/lavague-core/lavague/core/agents.py
+++ b/lavague-core/lavague/core/agents.py
@@ -602,3 +602,11 @@ def display_all_nodes(self) -> None:
 
     def set_origin(self, origin: str):
         self.origin = origin
+
+    def get_summary(self):
+        from lavague.core.utilities.profiling import agent_events, agent_steps
+        chart_generator = ChartGenerator(agent_events=agent_events, agent_steps=agent_steps)
+        plot = chart_generator.plot_waterfall()
+        table = chart_generator.get_summary_df()
+        
+        return plot, table

From 32f4cbe8e64177b9700ab4d709c20de90093c7c4 Mon Sep 17 00:00:00 2001
From: Paul Palmieri <paulpalmieri@outlook.fr>
Date: Mon, 26 Aug 2024 14:52:17 +0200
Subject: [PATCH 04/11] Merge all decorators into a single one

---
 .../lavague/core/utilities/profiling.py       | 92 ++++++++-----------
 1 file changed, 39 insertions(+), 53 deletions(-)

diff --git a/lavague-core/lavague/core/utilities/profiling.py b/lavague-core/lavague/core/utilities/profiling.py
index 344b266b..8bd66676 100644
--- a/lavague-core/lavague/core/utilities/profiling.py
+++ b/lavague-core/lavague/core/utilities/profiling.py
@@ -1,10 +1,9 @@
 import time
-from datetime import datetime
 import pandas as pd
 import matplotlib.pyplot as plt
 from functools import wraps
 import io
-from IPython.display import display, Image
+from IPython.display import Image
 from itertools import cycle
 
 # stores llm and retriever calls
@@ -17,9 +16,14 @@
 def start_new_step():
     global agent_events
     agent_events.append([])
+    
+def clear_profiling_data():
+    global agent_events, agent_steps
+    agent_events = []
+    agent_steps = []
 
-# track the runtime of the run_step function
-def track_total_runtime():
+# The profile decorator that handles different event types
+def profile_agent(event_type: str, event_name: str = None):
     def decorator(func):
         @wraps(func)
         def wrapper(*args, **kwargs):
@@ -27,54 +31,36 @@ def wrapper(*args, **kwargs):
             result = func(*args, **kwargs)
             end_time = time.time()
             duration = end_time - start_time
-            agent_steps.append({
-                'start_time': start_time,
-                # 'end_time': end_time,
-                'duration': duration
-            })
-            return result
-        return wrapper
-    return decorator
 
-# track runtime, prompt size, and completion size of a llm call (used on world model + action engine)
-def track_llm_call(event_name):
-    def decorator(func):
-        @wraps(func)
-        def wrapper(*args, **kwargs):
-            start_time = time.time()
-            result = func(*args, **kwargs)
-            end_time = time.time()
-            duration = end_time - start_time
-            
-            agent_events[-1].append({
-                'event_name': event_name,
-                'start_time': start_time,
-                # 'end_time': end_time,
-                'duration': duration,
-                'prompt_size': len(args[0]),
-                'completion_size': len(result.text)
-            })
-            return result
-        return wrapper
-    return decorator
+            if event_type == "RUN_STEP":
+                agent_steps.append({
+                    'start_time': start_time,
+                    'duration': duration
+                })
+            elif event_type == "LLM_CALL":
+                agent_events[-1].append({
+                    'event_name': event_name,
+                    'start_time': start_time,
+                    'duration': duration,
+                    'prompt_size': len(args[0]),  # Assuming args[0] is the prompt
+                    'completion_size': len(result.text)  # Assuming result.text is the completion
+                })
+            elif event_type == "RETRIEVER_CALL":
+                agent_events[-1].append({
+                    'event_name': 'Retriever',
+                    'start_time': start_time,
+                    'duration': duration,
+                    'html_size': len(args[0].driver.get_html()),  # Assuming args[0] is the retriever object
+                })
+            elif event_type == "DEFAULT":
+                agent_events[-1].append({
+                    'event_name': event_name,
+                    'start_time': start_time,
+                    'duration': duration
+                })
+            else:
+                raise ValueError(f"Unknown event type: {event_type}")
 
-# track runtime and html size of a retriever call
-def track_retriever(event_name):
-    def decorator(func):
-        @wraps(func)
-        def wrapper(*args, **kwargs):
-            start_time = time.time()
-            result = func(*args, **kwargs)
-            end_time = time.time()
-            duration = end_time - start_time
-            
-            agent_events[-1].append({
-                'event_name': event_name,
-                'start_time': start_time,
-                # 'end_time': end_time,
-                'duration': duration,
-                'html_size': len(args[0].driver.get_html()),
-            })
             return result
         return wrapper
     return decorator
@@ -84,7 +70,7 @@ class ChartGenerator:
     def __init__(self, agent_events, agent_steps):
         self.agent_events = agent_events
         self.total_step_runtime = agent_steps
-        self.step_color = "lightgrey"
+        self.step_color = "grey"
         self.event_color_scheme = ["lightblue", "bisque", "thistle", "lightgreen", "pink"]
 
     def plot_waterfall(self):
@@ -117,7 +103,7 @@ def plot_waterfall(self):
                     event_colors[event_name] = next(color_cycle)
                 
                 color = event_colors[event_name]
-                ax.barh(step_index, duration, left=start_time, color=color)
+                ax.barh(step_index, duration, left=start_time, color=color, alpha=1)
                 ax.text(start_time + duration / 2, step_index, f"{duration:.2f}s",
                         ha='center', va='center', fontsize=9, color='black', rotation=90)
         
@@ -125,7 +111,7 @@ def plot_waterfall(self):
         ax.set_yticks(range(len(self.total_step_runtime)))
         ax.set_yticklabels([f'Step {i+1}' for i in range(len(self.total_step_runtime))])
         ax.set_xlabel('Time (seconds)')
-        ax.set_title('Waterfall Chart with Events Overlay')
+        ax.set_title('Agent Event Waterfall')
 
         
         # Add legend for event colors

From 0bf4fbffdd084a727e794267eef29c5dfd02ea8b Mon Sep 17 00:00:00 2001
From: Paul Palmieri <paulpalmieri@outlook.fr>
Date: Mon, 26 Aug 2024 14:52:33 +0200
Subject: [PATCH 05/11] Use new profiler

---
 lavague-core/lavague/core/agents.py      | 4 ++--
 lavague-core/lavague/core/navigation.py  | 8 +++++---
 lavague-core/lavague/core/world_model.py | 4 ++--
 3 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/lavague-core/lavague/core/agents.py b/lavague-core/lavague/core/agents.py
index 04df3b67..0d3a4a83 100644
--- a/lavague-core/lavague/core/agents.py
+++ b/lavague-core/lavague/core/agents.py
@@ -22,7 +22,7 @@
 from lavague.core.token_counter import TokenCounter
 from lavague.core.utilities.config import is_flag_true
 
-from lavague.core.utilities.profiling import ChartGenerator, track_total_runtime, start_new_step
+from lavague.core.utilities.profiling import ChartGenerator, profile_agent, start_new_step, clear_profiling_data
 
 logging_print = logging.getLogger(__name__)
 logging_print.setLevel(logging.INFO)
@@ -448,7 +448,7 @@ def _run_step_gradio(
         self.process_token_usage()
         self.logger.end_step()
 
-    @track_total_runtime()
+    @profile_agent(event_type="RUN_STEP")
     def run_step(self, objective: str) -> Optional[ActionResult]:
         obs = self.driver.get_obs()
         current_state, past = self.st_memory.get_state()
diff --git a/lavague-core/lavague/core/navigation.py b/lavague-core/lavague/core/navigation.py
index bede4907..466253d9 100644
--- a/lavague-core/lavague/core/navigation.py
+++ b/lavague-core/lavague/core/navigation.py
@@ -24,7 +24,7 @@
 from PIL import Image
 from llama_index.core.base.llms.base import BaseLLM
 from llama_index.core.embeddings import BaseEmbedding
-from lavague.core.utilities.profiling import track_retriever, track_llm_call
+from lavague.core.utilities.profiling import profile_agent
 
 NAVIGATION_ENGINE_PROMPT_TEMPLATE = ActionTemplate(
     """
@@ -144,7 +144,7 @@ def from_context(
             extractor,
         )
 
-    @track_retriever()
+    @profile_agent(event_type="RETRIEVER_CALL")
     def get_nodes(self, query: str) -> List[str]:
         """
         Get the nodes from the html page
@@ -456,7 +456,7 @@ def execute_instruction(self, instruction: str) -> ActionResult:
             )
 
             # response = self.llm.complete(prompt).text
-            response = track_llm_call("Navigation")(self.llm.complete)(prompt).text
+            response = profile_agent(event_type="LLM_CALL", event_name="Navigation Engine")(self.llm.complete)(prompt).text
             end = time.time()
             action_generation_time = end - start
             action_outcome = {
@@ -480,6 +480,8 @@ def execute_instruction(self, instruction: str) -> ActionResult:
                     for item in vision_data:
                         display_screenshot(item["screenshot"])
                         time.sleep(0.2)
+                        
+                profile_agent(event_type="DEFAULT", event_name="Execute code")(self.driver.exec_code)(action)
                 self.driver.exec_code(action)
                 time.sleep(self.time_between_actions)
                 if self.display:
diff --git a/lavague-core/lavague/core/world_model.py b/lavague-core/lavague/core/world_model.py
index 0e3be5f2..c22d9729 100644
--- a/lavague-core/lavague/core/world_model.py
+++ b/lavague-core/lavague/core/world_model.py
@@ -11,7 +11,7 @@
 from lavague.core.utilities.model_utils import get_model_name
 import time
 import yaml
-from lavague.core.utilities.profiling import track_llm_call
+from lavague.core.utilities.profiling import profile_agent
 
 WORLD_MODEL_GENERAL_EXAMPLES = """
 Objective:  Go to the first issue you can find
@@ -433,7 +433,7 @@ def get_instruction(
         start = time.time()
 
         # decorated llm call
-        mm_llm_output = track_llm_call("World Model")(mm_llm.complete)(
+        mm_llm_output = profile_agent(event_type="LLM_CALL", event_name="World Model")(mm_llm.complete)(
             prompt, image_documents=image_documents
         ).text
         # mm_llm_output = mm_llm.complete(prompt, image_documents=image_documents).text

From bcbfcdbfdeaff7d035e2cd1d4a6fec9fc0123ca8 Mon Sep 17 00:00:00 2001
From: Paul Palmieri <paulpalmieri@outlook.fr>
Date: Mon, 26 Aug 2024 14:52:43 +0200
Subject: [PATCH 06/11] Clear data after plot is generated

---
 lavague-core/lavague/core/agents.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/lavague-core/lavague/core/agents.py b/lavague-core/lavague/core/agents.py
index 0d3a4a83..e5553208 100644
--- a/lavague-core/lavague/core/agents.py
+++ b/lavague-core/lavague/core/agents.py
@@ -609,4 +609,6 @@ def get_summary(self):
         plot = chart_generator.plot_waterfall()
         table = chart_generator.get_summary_df()
         
+        clear_profiling_data()
+        
         return plot, table

From 17dfaa4cebc91715cc11e4c36dd76d33eef83dd0 Mon Sep 17 00:00:00 2001
From: Paul Palmieri <paulpalmieri@outlook.fr>
Date: Mon, 26 Aug 2024 15:29:46 +0200
Subject: [PATCH 07/11] Remove duplicate code

---
 lavague-core/lavague/core/navigation.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lavague-core/lavague/core/navigation.py b/lavague-core/lavague/core/navigation.py
index 466253d9..5b99c983 100644
--- a/lavague-core/lavague/core/navigation.py
+++ b/lavague-core/lavague/core/navigation.py
@@ -482,7 +482,7 @@ def execute_instruction(self, instruction: str) -> ActionResult:
                         time.sleep(0.2)
                         
                 profile_agent(event_type="DEFAULT", event_name="Execute code")(self.driver.exec_code)(action)
-                self.driver.exec_code(action)
+                # self.driver.exec_code(action)
                 time.sleep(self.time_between_actions)
                 if self.display:
                     try:

From fd515f91ed8541f20be5444520f53267673f4580 Mon Sep 17 00:00:00 2001
From: Paul Palmieri <paulpalmieri@outlook.fr>
Date: Mon, 26 Aug 2024 15:30:14 +0200
Subject: [PATCH 08/11] Manually add Step to the legend

---
 lavague-core/lavague/core/utilities/profiling.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/lavague-core/lavague/core/utilities/profiling.py b/lavague-core/lavague/core/utilities/profiling.py
index 8bd66676..79b32fe0 100644
--- a/lavague-core/lavague/core/utilities/profiling.py
+++ b/lavague-core/lavague/core/utilities/profiling.py
@@ -115,8 +115,15 @@ def plot_waterfall(self):
 
         
         # Add legend for event colors
+        # Existing legend labels
         legend_labels = [plt.Line2D([0], [0], color=color, lw=4) for color in event_colors.values()]
-        ax.legend(legend_labels, event_colors.keys(), title="Event Name")
+
+        # Adding the "Step" label in grey
+        step_label = plt.Line2D([0], [0], color='grey', lw=4)
+        legend_labels.append(step_label)
+
+        # Update the legend to include "Step"
+        ax.legend(legend_labels, list(event_colors.keys()) + ["Step"], title="Event Name")
         
         # Save to buffer
         buf = io.BytesIO()

From 8b5edae99a226f266156f3c4a594aafefc03ffac Mon Sep 17 00:00:00 2001
From: Paul Palmieri <paulpalmieri@outlook.fr>
Date: Mon, 26 Aug 2024 15:31:28 +0200
Subject: [PATCH 09/11] chore: format

---
 lavague-core/lavague/core/agents.py           |  16 +-
 lavague-core/lavague/core/navigation.py       |  10 +-
 .../lavague/core/utilities/profiling.py       | 160 +++++++++++-------
 lavague-core/lavague/core/world_model.py      |   6 +-
 4 files changed, 119 insertions(+), 73 deletions(-)

diff --git a/lavague-core/lavague/core/agents.py b/lavague-core/lavague/core/agents.py
index e5553208..9d168a5b 100644
--- a/lavague-core/lavague/core/agents.py
+++ b/lavague-core/lavague/core/agents.py
@@ -22,7 +22,12 @@
 from lavague.core.token_counter import TokenCounter
 from lavague.core.utilities.config import is_flag_true
 
-from lavague.core.utilities.profiling import ChartGenerator, profile_agent, start_new_step, clear_profiling_data
+from lavague.core.utilities.profiling import (
+    ChartGenerator,
+    profile_agent,
+    start_new_step,
+    clear_profiling_data,
+)
 
 logging_print = logging.getLogger(__name__)
 logging_print.setLevel(logging.INFO)
@@ -605,10 +610,13 @@ def set_origin(self, origin: str):
 
     def get_summary(self):
         from lavague.core.utilities.profiling import agent_events, agent_steps
-        chart_generator = ChartGenerator(agent_events=agent_events, agent_steps=agent_steps)
+
+        chart_generator = ChartGenerator(
+            agent_events=agent_events, agent_steps=agent_steps
+        )
         plot = chart_generator.plot_waterfall()
         table = chart_generator.get_summary_df()
-        
+
         clear_profiling_data()
-        
+
         return plot, table
diff --git a/lavague-core/lavague/core/navigation.py b/lavague-core/lavague/core/navigation.py
index 5b99c983..1b1990f9 100644
--- a/lavague-core/lavague/core/navigation.py
+++ b/lavague-core/lavague/core/navigation.py
@@ -456,7 +456,9 @@ def execute_instruction(self, instruction: str) -> ActionResult:
             )
 
             # response = self.llm.complete(prompt).text
-            response = profile_agent(event_type="LLM_CALL", event_name="Navigation Engine")(self.llm.complete)(prompt).text
+            response = profile_agent(
+                event_type="LLM_CALL", event_name="Navigation Engine"
+            )(self.llm.complete)(prompt).text
             end = time.time()
             action_generation_time = end - start
             action_outcome = {
@@ -480,8 +482,10 @@ def execute_instruction(self, instruction: str) -> ActionResult:
                     for item in vision_data:
                         display_screenshot(item["screenshot"])
                         time.sleep(0.2)
-                        
-                profile_agent(event_type="DEFAULT", event_name="Execute code")(self.driver.exec_code)(action)
+
+                profile_agent(event_type="DEFAULT", event_name="Execute code")(
+                    self.driver.exec_code
+                )(action)
                 # self.driver.exec_code(action)
                 time.sleep(self.time_between_actions)
                 if self.display:
diff --git a/lavague-core/lavague/core/utilities/profiling.py b/lavague-core/lavague/core/utilities/profiling.py
index 79b32fe0..913048c6 100644
--- a/lavague-core/lavague/core/utilities/profiling.py
+++ b/lavague-core/lavague/core/utilities/profiling.py
@@ -12,16 +12,19 @@
 # stores total runtime of each step
 agent_steps = []
 
+
 # call before each agent step to group events by steps
 def start_new_step():
     global agent_events
     agent_events.append([])
-    
+
+
 def clear_profiling_data():
     global agent_events, agent_steps
     agent_events = []
     agent_steps = []
 
+
 # The profile decorator that handles different event types
 def profile_agent(event_type: str, event_name: str = None):
     def decorator(func):
@@ -33,36 +36,45 @@ def wrapper(*args, **kwargs):
             duration = end_time - start_time
 
             if event_type == "RUN_STEP":
-                agent_steps.append({
-                    'start_time': start_time,
-                    'duration': duration
-                })
+                agent_steps.append({"start_time": start_time, "duration": duration})
             elif event_type == "LLM_CALL":
-                agent_events[-1].append({
-                    'event_name': event_name,
-                    'start_time': start_time,
-                    'duration': duration,
-                    'prompt_size': len(args[0]),  # Assuming args[0] is the prompt
-                    'completion_size': len(result.text)  # Assuming result.text is the completion
-                })
+                agent_events[-1].append(
+                    {
+                        "event_name": event_name,
+                        "start_time": start_time,
+                        "duration": duration,
+                        "prompt_size": len(args[0]),  # Assuming args[0] is the prompt
+                        "completion_size": len(
+                            result.text
+                        ),  # Assuming result.text is the completion
+                    }
+                )
             elif event_type == "RETRIEVER_CALL":
-                agent_events[-1].append({
-                    'event_name': 'Retriever',
-                    'start_time': start_time,
-                    'duration': duration,
-                    'html_size': len(args[0].driver.get_html()),  # Assuming args[0] is the retriever object
-                })
+                agent_events[-1].append(
+                    {
+                        "event_name": "Retriever",
+                        "start_time": start_time,
+                        "duration": duration,
+                        "html_size": len(
+                            args[0].driver.get_html()
+                        ),  # Assuming args[0] is the retriever object
+                    }
+                )
             elif event_type == "DEFAULT":
-                agent_events[-1].append({
-                    'event_name': event_name,
-                    'start_time': start_time,
-                    'duration': duration
-                })
+                agent_events[-1].append(
+                    {
+                        "event_name": event_name,
+                        "start_time": start_time,
+                        "duration": duration,
+                    }
+                )
             else:
                 raise ValueError(f"Unknown event type: {event_type}")
 
             return result
+
         return wrapper
+
     return decorator
 
 
@@ -71,101 +83,123 @@ def __init__(self, agent_events, agent_steps):
         self.agent_events = agent_events
         self.total_step_runtime = agent_steps
         self.step_color = "grey"
-        self.event_color_scheme = ["lightblue", "bisque", "thistle", "lightgreen", "pink"]
+        self.event_color_scheme = [
+            "lightblue",
+            "bisque",
+            "thistle",
+            "lightgreen",
+            "pink",
+        ]
 
     def plot_waterfall(self):
         # Calculate the earliest start time to align the x-axis to 0
-        base_start_time = self.total_step_runtime[0]['start_time']
-        
+        base_start_time = self.total_step_runtime[0]["start_time"]
+
         plt.figure(figsize=(20, 8))
         ax = plt.gca()
-    
+
         color_cycle = cycle(self.event_color_scheme)
         event_colors = {}
-        
+
         # Plot total step runtime (from run_step)
         for i, step in enumerate(self.total_step_runtime):
-            duration = step['duration']
-            start_time = step['start_time'] - base_start_time  # Normalize to 0
-            
+            duration = step["duration"]
+            start_time = step["start_time"] - base_start_time  # Normalize to 0
+
             ax.barh(i, duration, left=start_time, color=self.step_color)
-            ax.text(start_time + duration / 2, i - 0.45, f"{duration:.2f}s",
-                    ha='center', va='center')
-        
+            ax.text(
+                start_time + duration / 2,
+                i - 0.45,
+                f"{duration:.2f}s",
+                ha="center",
+                va="center",
+            )
+
         # Plot each individual event on top of the step runtime
         for step_index, step in enumerate(self.agent_events):
             for event in step:
-                duration = event['duration']
-                event_name = event['event_name']
-                start_time = event['start_time'] - base_start_time  # Normalize to 0
-                
+                duration = event["duration"]
+                event_name = event["event_name"]
+                start_time = event["start_time"] - base_start_time  # Normalize to 0
+
                 if event_name not in event_colors:
                     event_colors[event_name] = next(color_cycle)
-                
+
                 color = event_colors[event_name]
                 ax.barh(step_index, duration, left=start_time, color=color, alpha=1)
-                ax.text(start_time + duration / 2, step_index, f"{duration:.2f}s",
-                        ha='center', va='center', fontsize=9, color='black', rotation=90)
-        
+                ax.text(
+                    start_time + duration / 2,
+                    step_index,
+                    f"{duration:.2f}s",
+                    ha="center",
+                    va="center",
+                    fontsize=9,
+                    color="black",
+                    rotation=90,
+                )
+
         ax.invert_yaxis()
         ax.set_yticks(range(len(self.total_step_runtime)))
-        ax.set_yticklabels([f'Step {i+1}' for i in range(len(self.total_step_runtime))])
-        ax.set_xlabel('Time (seconds)')
-        ax.set_title('Agent Event Waterfall')
+        ax.set_yticklabels([f"Step {i+1}" for i in range(len(self.total_step_runtime))])
+        ax.set_xlabel("Time (seconds)")
+        ax.set_title("Agent Event Waterfall")
 
-        
         # Add legend for event colors
         # Existing legend labels
-        legend_labels = [plt.Line2D([0], [0], color=color, lw=4) for color in event_colors.values()]
+        legend_labels = [
+            plt.Line2D([0], [0], color=color, lw=4) for color in event_colors.values()
+        ]
 
         # Adding the "Step" label in grey
-        step_label = plt.Line2D([0], [0], color='grey', lw=4)
+        step_label = plt.Line2D([0], [0], color="grey", lw=4)
         legend_labels.append(step_label)
 
         # Update the legend to include "Step"
-        ax.legend(legend_labels, list(event_colors.keys()) + ["Step"], title="Event Name")
-        
+        ax.legend(
+            legend_labels, list(event_colors.keys()) + ["Step"], title="Event Name"
+        )
+
         # Save to buffer
         buf = io.BytesIO()
-        plt.savefig(buf, format='png')
+        plt.savefig(buf, format="png")
         buf.seek(0)
-        
+
         plt.close()
-        
+
         return Image(buf.read())
 
     def get_summary_df(self):
         summary_data = {}
-        
+
         # Iterate over each step and each event in the step
         for step_index, step_events in enumerate(self.agent_events):
             # Step row key
             step_key = f"Step {step_index + 1}"
             summary_data[step_key] = {}
-            
+
             # Count the number of each event type in the step
             event_counts = {}
-            
+
             # Iterate over each event in the step
             for event in step_events:
-                event_name = event['event_name']
-                
+                event_name = event["event_name"]
+
                 # Increment the count for the event
                 if event_name not in event_counts:
                     event_counts[event_name] = 1
                 else:
                     event_counts[event_name] += 1
-                
+
                 # for each key in the event, excluding 'event_name', 'start_time', and 'end_time', add the value to the summary
                 for key, value in event.items():
-                    if key not in ['event_name', 'start_time', 'end_time']:
+                    if key not in ["event_name", "start_time", "end_time"]:
                         metric_key = f"{event_name} {key}"
-                        
+
                         if metric_key not in summary_data[step_key]:
                             summary_data[step_key][metric_key] = value
                         else:
                             summary_data[step_key][metric_key] += value
-            
+
             # add the event counts
             for event_name, count in event_counts.items():
                 count_key = f"{event_name} count"
@@ -174,4 +208,4 @@ def get_summary_df(self):
         # Convert the dictionary to a DataFrame
         df = pd.DataFrame(summary_data).T
 
-        return df
\ No newline at end of file
+        return df
diff --git a/lavague-core/lavague/core/world_model.py b/lavague-core/lavague/core/world_model.py
index c22d9729..9caeab55 100644
--- a/lavague-core/lavague/core/world_model.py
+++ b/lavague-core/lavague/core/world_model.py
@@ -433,9 +433,9 @@ def get_instruction(
         start = time.time()
 
         # decorated llm call
-        mm_llm_output = profile_agent(event_type="LLM_CALL", event_name="World Model")(mm_llm.complete)(
-            prompt, image_documents=image_documents
-        ).text
+        mm_llm_output = profile_agent(event_type="LLM_CALL", event_name="World Model")(
+            mm_llm.complete
+        )(prompt, image_documents=image_documents).text
         # mm_llm_output = mm_llm.complete(prompt, image_documents=image_documents).text
 
         end = time.time()

From c803e24bddc3fde195a5e97b079d97204f3769e9 Mon Sep 17 00:00:00 2001
From: Paul Palmieri <paulpalmieri@outlook.fr>
Date: Mon, 26 Aug 2024 17:51:00 +0200
Subject: [PATCH 10/11] Remove decorator, use contextmanager

---
 lavague-core/lavague/core/agents.py           |  6 +-
 lavague-core/lavague/core/navigation.py       | 26 +++---
 .../lavague/core/utilities/profiling.py       | 86 ++++++++-----------
 lavague-core/lavague/core/world_model.py      | 11 ++-
 4 files changed, 56 insertions(+), 73 deletions(-)

diff --git a/lavague-core/lavague/core/agents.py b/lavague-core/lavague/core/agents.py
index 9d168a5b..4111336c 100644
--- a/lavague-core/lavague/core/agents.py
+++ b/lavague-core/lavague/core/agents.py
@@ -24,7 +24,7 @@
 
 from lavague.core.utilities.profiling import (
     ChartGenerator,
-    profile_agent,
+    time_profiler,
     start_new_step,
     clear_profiling_data,
 )
@@ -453,7 +453,6 @@ def _run_step_gradio(
         self.process_token_usage()
         self.logger.end_step()
 
-    @profile_agent(event_type="RUN_STEP")
     def run_step(self, objective: str) -> Optional[ActionResult]:
         obs = self.driver.get_obs()
         current_state, past = self.st_memory.get_state()
@@ -512,7 +511,8 @@ def run(
         try:
             for _ in range(self.n_steps):
                 start_new_step()
-                result = self.run_step(objective)
+                with time_profiler("Run step", full_step_profiling=True):
+                    result = self.run_step(objective)
 
                 if result is not None:
                     break
diff --git a/lavague-core/lavague/core/navigation.py b/lavague-core/lavague/core/navigation.py
index 1b1990f9..fe7dd9b4 100644
--- a/lavague-core/lavague/core/navigation.py
+++ b/lavague-core/lavague/core/navigation.py
@@ -24,7 +24,7 @@
 from PIL import Image
 from llama_index.core.base.llms.base import BaseLLM
 from llama_index.core.embeddings import BaseEmbedding
-from lavague.core.utilities.profiling import profile_agent
+from lavague.core.utilities.profiling import time_profiler
 
 NAVIGATION_ENGINE_PROMPT_TEMPLATE = ActionTemplate(
     """
@@ -144,7 +144,6 @@ def from_context(
             extractor,
         )
 
-    @profile_agent(event_type="RETRIEVER_CALL")
     def get_nodes(self, query: str) -> List[str]:
         """
         Get the nodes from the html page
@@ -156,9 +155,13 @@ def get_nodes(self, query: str) -> List[str]:
             `List[str]`: The nodes
         """
         viewport_only = not self.driver.previously_scanned
-        source_nodes = self.retriever.retrieve(
-            QueryBundle(query_str=query), [self.driver.get_html()], viewport_only
-        )
+
+        html = self.driver.get_html()
+
+        with time_profiler("Retriever Inference", html_size=len(html)):
+            source_nodes = self.retriever.retrieve(
+                QueryBundle(query_str=query), [html], viewport_only
+            )
         return source_nodes
 
     def add_knowledge(self, knowledge: str):
@@ -455,10 +458,9 @@ def execute_instruction(self, instruction: str) -> ActionResult:
                 authorized_xpaths=authorized_xpaths,
             )
 
-            # response = self.llm.complete(prompt).text
-            response = profile_agent(
-                event_type="LLM_CALL", event_name="Navigation Engine"
-            )(self.llm.complete)(prompt).text
+            with time_profiler("Navigation Engine Inference", prompt_size=len(prompt)):
+                response = self.llm.complete(prompt).text
+
             end = time.time()
             action_generation_time = end - start
             action_outcome = {
@@ -483,10 +485,8 @@ def execute_instruction(self, instruction: str) -> ActionResult:
                         display_screenshot(item["screenshot"])
                         time.sleep(0.2)
 
-                profile_agent(event_type="DEFAULT", event_name="Execute code")(
-                    self.driver.exec_code
-                )(action)
-                # self.driver.exec_code(action)
+                with time_profiler("Execute Code"):
+                    self.driver.exec_code(action)
                 time.sleep(self.time_between_actions)
                 if self.display:
                     try:
diff --git a/lavague-core/lavague/core/utilities/profiling.py b/lavague-core/lavague/core/utilities/profiling.py
index 913048c6..cdfaeb51 100644
--- a/lavague-core/lavague/core/utilities/profiling.py
+++ b/lavague-core/lavague/core/utilities/profiling.py
@@ -5,6 +5,7 @@
 import io
 from IPython.display import Image
 from itertools import cycle
+from contextlib import contextmanager
 
 # stores llm and retriever calls
 agent_events = []
@@ -25,57 +26,40 @@ def clear_profiling_data():
     agent_steps = []
 
 
-# The profile decorator that handles different event types
-def profile_agent(event_type: str, event_name: str = None):
-    def decorator(func):
-        @wraps(func)
-        def wrapper(*args, **kwargs):
-            start_time = time.time()
-            result = func(*args, **kwargs)
-            end_time = time.time()
-            duration = end_time - start_time
-
-            if event_type == "RUN_STEP":
-                agent_steps.append({"start_time": start_time, "duration": duration})
-            elif event_type == "LLM_CALL":
-                agent_events[-1].append(
-                    {
-                        "event_name": event_name,
-                        "start_time": start_time,
-                        "duration": duration,
-                        "prompt_size": len(args[0]),  # Assuming args[0] is the prompt
-                        "completion_size": len(
-                            result.text
-                        ),  # Assuming result.text is the completion
-                    }
-                )
-            elif event_type == "RETRIEVER_CALL":
-                agent_events[-1].append(
-                    {
-                        "event_name": "Retriever",
-                        "start_time": start_time,
-                        "duration": duration,
-                        "html_size": len(
-                            args[0].driver.get_html()
-                        ),  # Assuming args[0] is the retriever object
-                    }
-                )
-            elif event_type == "DEFAULT":
-                agent_events[-1].append(
-                    {
-                        "event_name": event_name,
-                        "start_time": start_time,
-                        "duration": duration,
-                    }
-                )
-            else:
-                raise ValueError(f"Unknown event type: {event_type}")
-
-            return result
-
-        return wrapper
-
-    return decorator
+@contextmanager
+def time_profiler(
+    event_name, prompt_size=None, html_size=None, full_step_profiling=False
+):
+    """
+    A context manager to profile the execution time of code blocks.
+
+    Parameters:
+    - event_name: The name of the event being profiled.
+    - prompt_size: Optional size of the prompt, if applicable.
+    - html_size: Optional size of the HTML, if applicable.
+    - full_step_profiling: Boolean indicating whether to profile full steps or individual events.
+    """
+    start_time = time.perf_counter()
+    try:
+        yield
+    finally:
+        end_time = time.perf_counter()
+        duration = end_time - start_time
+
+        # create profiling record
+        record = {
+            "event_name": event_name,
+            "start_time": start_time,
+            "duration": duration,
+            **({"prompt_size": prompt_size} if prompt_size is not None else {}),
+            **({"html_size": html_size} if html_size is not None else {}),
+        }
+
+        # append the record to the appropriate list
+        if full_step_profiling:
+            agent_steps.append(record)
+        else:
+            agent_events[-1].append(record)
 
 
 class ChartGenerator:
diff --git a/lavague-core/lavague/core/world_model.py b/lavague-core/lavague/core/world_model.py
index 9caeab55..2d7395ab 100644
--- a/lavague-core/lavague/core/world_model.py
+++ b/lavague-core/lavague/core/world_model.py
@@ -11,7 +11,7 @@
 from lavague.core.utilities.model_utils import get_model_name
 import time
 import yaml
-from lavague.core.utilities.profiling import profile_agent
+from lavague.core.utilities.profiling import time_profiler
 
 WORLD_MODEL_GENERAL_EXAMPLES = """
 Objective:  Go to the first issue you can find
@@ -432,11 +432,10 @@ def get_instruction(
 
         start = time.time()
 
-        # decorated llm call
-        mm_llm_output = profile_agent(event_type="LLM_CALL", event_name="World Model")(
-            mm_llm.complete
-        )(prompt, image_documents=image_documents).text
-        # mm_llm_output = mm_llm.complete(prompt, image_documents=image_documents).text
+        with time_profiler("World Model Inference", prompt_size=len(prompt)):
+            mm_llm_output = mm_llm.complete(
+                prompt, image_documents=image_documents
+            ).text
 
         end = time.time()
         world_model_inference_time = end - start

From 007ed7b948df92834a3d6fc1341ee6b52defb9cd Mon Sep 17 00:00:00 2001
From: Paul Palmieri <paulpalmieri@outlook.fr>
Date: Mon, 26 Aug 2024 17:51:06 +0200
Subject: [PATCH 11/11] Update colors

---
 lavague-core/lavague/core/utilities/profiling.py | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/lavague-core/lavague/core/utilities/profiling.py b/lavague-core/lavague/core/utilities/profiling.py
index cdfaeb51..f3d84d8e 100644
--- a/lavague-core/lavague/core/utilities/profiling.py
+++ b/lavague-core/lavague/core/utilities/profiling.py
@@ -68,11 +68,16 @@ def __init__(self, agent_events, agent_steps):
         self.total_step_runtime = agent_steps
         self.step_color = "grey"
         self.event_color_scheme = [
-            "lightblue",
-            "bisque",
-            "thistle",
-            "lightgreen",
-            "pink",
+            "#FFB3B3",  # Pastel Red
+            "#ADD8E6",  # Pastel Blue
+            "#B2D8B2",  # Pastel Green
+            "#FFCC99",  # Pastel Orange
+            "#D1B3FF",  # Pastel Purple
+            "#FFB3DE",  # Pastel Pink
+            "#B3FFFF",  # Pastel Cyan
+            "#FFFFB3",  # Pastel Yellow
+            "#FFB3FF",  # Pastel Magenta
+            "#D2B48C",  # Pastel Brown
         ]
 
     def plot_waterfall(self):