From 89bd0ae95fb7db36f9a32ed42fc8bd0e8c44b761 Mon Sep 17 00:00:00 2001
From: JoFrost <mebessaa@student.42.fr>
Date: Tue, 21 May 2024 16:47:55 +0300
Subject: [PATCH] feat: telemetry log can be exported

---
 lavague-core/lavague/core/agents.py           | 25 +++++++++++++++++--
 .../lavague/core/utilities/telemetry.py       | 22 ++++++++++++++++
 2 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/lavague-core/lavague/core/agents.py b/lavague-core/lavague/core/agents.py
index 367dd8d3..f6b9f1ac 100644
--- a/lavague-core/lavague/core/agents.py
+++ b/lavague-core/lavague/core/agents.py
@@ -31,9 +31,15 @@ def __init__(self, action_engine: ActionEngine, world_model: WorldModel):
     def get(self, url):
         self.driver.goto(url)
 
-    def run(self, objective, display=True):
+    def run(self, objective, display=True, log: bool = False):
         from selenium.webdriver.remote.webdriver import WebDriver
+        import os
 
+        log_lines = []
+
+        if log:
+            if os.path.isdir("./logs") == False:
+                os.mkdir("./logs")
         driver: WebDriver = self.driver.get_driver()
         action_engine: ActionEngine = self.action_engine
         world_model: WorldModel = self.world_model
@@ -74,6 +80,7 @@ def run(self, objective, display=True):
                 context = "\n".join(nodes)
                 for _ in range(N_ATTEMPTS):
                     try:
+                        code = ""
                         image = None
                         screenshot_after_action = None
                         error = ""
@@ -118,7 +125,7 @@ def run(self, objective, display=True):
                         pass
                     finally:
                         action_id = str(uuid.uuid4())
-                        send_telemetry(
+                        line = send_telemetry(
                             model_name=action_engine.llm.metadata.model_name,
                             code=action,
                             instruction=instruction,
@@ -136,6 +143,7 @@ def run(self, objective, display=True):
                             multi_modal_model=world_model.mm_llm.metadata.model_name,
                             step_id=step_id,
                             run_id=run_id,
+                            log=log
                         )
                         send_telemetry_scr(
                             action_id,
@@ -143,6 +151,19 @@ def run(self, objective, display=True):
                             image,
                             screenshot_after_action,
                         )
+                        if log:
+                            log_lines.append(line)
+
             else:
                 print("Objective reached")
                 break
+        
+        if log:
+            import csv
+
+            with open(f'./logs/{run_id}.csv', 'w', newline='') as output_file:
+                keys = log_lines[0].keys()
+                dict_writer = csv.DictWriter(output_file, keys)
+                dict_writer.writeheader()
+                dict_writer.writerows(log_lines)
+                print(f"Logs exported to logs/{run_id}.csv")
\ No newline at end of file
diff --git a/lavague-core/lavague/core/utilities/telemetry.py b/lavague-core/lavague/core/utilities/telemetry.py
index 1ffa610c..0b759688 100644
--- a/lavague-core/lavague/core/utilities/telemetry.py
+++ b/lavague-core/lavague/core/utilities/telemetry.py
@@ -73,11 +73,13 @@ def send_telemetry(
     multi_modal_model: str = "",
     step_id: str = "",
     run_id: str = "",
+    log: bool = False,
 ):
     """
     Telemetry to help performance.
     Mandatory telemetry variables - DO NOT DELETE ANY, else telemetry will fail: model_name, code, screenshot, html, source_nodes, instruction, url, origin, success
     """
+    line = None
     success_str = str(success)
     try:
         if TELEMETRY_VAR is None:
@@ -116,3 +118,23 @@ def send_telemetry(
             print("Telemetry failed with ", e)
         else:
             raise ValueError("Telemetry failed with ", e)
+    if log:
+        line = {
+            "run_id": run_id,
+            "step_id": step_id,
+            "action_id": action_id,
+            "lavague_version": get_installed_version("lavague"),
+            "code_produced": code,
+            "llm": model_name,
+            "multi_modal_model": multi_modal_model,
+            "url": url,
+            "objective": main_objective,
+            "instruction": instruction,
+            "chain_of_thoughts": objectives,
+            "source_nodes": source_nodes,
+            "bounding_box": bounding_box,
+            "viewport_size": viewport_size,
+            "success": str(success),
+            "error_msg": error,
+        }
+    return line
\ No newline at end of file