Start replay

app.py

@@ -12,7 +12,7 @@ from e2b_desktop import Sandbox
 from smolagents import CodeAgent
 from smolagents.monitoring import LogLevel
 from smolagents.gradio_ui import GradioUI, stream_to_gradio
-
+from model_replay import FakeModelClass
 
 from e2bqwen import QwenVLAPIModel, E2BVisionAgent
 
@@ -29,7 +29,6 @@ if not os.path.exists(TMP_DIR):
 hf_token = os.getenv("HUGGINGFACE_API_KEY")
 login(token=hf_token)
 
-
 custom_css = """
 .sandbox-container {
     position: relative;
@@ -471,6 +470,7 @@ def create_agent(data_dir, desktop):
         planning_interval=10,
     )
 
+
 class EnrichedGradioUI(GradioUI):
     def log_user_message(self, text_input):
         import gradio as gr
@@ -480,7 +480,7 @@ class EnrichedGradioUI(GradioUI):
             gr.Button(interactive=False),
         )
 
-    def interact_with_agent(self, task_input, stored_messages, session_state, session_hash, request: gr.Request):
+    def interact_with_agent(self, task_input, stored_messages, session_state, session_hash, replay_log, request: gr.Request):
         import gradio as gr
 
         interaction_id = generate_interaction_id(request)
@@ -491,11 +491,14 @@ class EnrichedGradioUI(GradioUI):
         if not os.path.exists(data_dir):
             os.makedirs(data_dir)
 
-        
         if "agent" in session_state:
             session_state["agent"].data_dir = data_dir # Update data dir to new interaction
         else:
             session_state["agent"] = create_agent(data_dir=data_dir, desktop=desktop)
+
+        if replay_log is not None:
+            original_model = session_state["agent"].model
+            session_state["agent"].model = FakeModelReplayLog(replay_log)
         
         try:
             stored_messages.append(gr.ChatMessage(role="user", content=task_input))
@@ -528,6 +531,8 @@ class EnrichedGradioUI(GradioUI):
             save_final_status(data_dir, "failed", summary=[], error_message=error_message)
 
         finally:
+            if replay_log: # Replace the model with original model
+                session_state["agent"].model = original_model
             upload_to_hf_and_remove(data_dir)
 
 theme = gr.themes.Default(font=["Oxanium", "sans-serif"], primary_hue="amber", secondary_hue="blue")
@@ -571,13 +576,8 @@ with gr.Blocks(theme=theme, css=custom_css, js=custom_js) as demo:
             session_state = gr.State({})
             stored_messages = gr.State([])
 
-            with gr.Group(visible=False) as results_container:
-                results_output = gr.Textbox(
-                    label="Results",
-                    interactive=False,
-                    elem_id="results-output"
-                )
 
+            replay_btn = gr.Button("Replay an agent run")
 
             minimalist_toggle = gr.Checkbox(label="Innie/Outie", value=False)
 
@@ -664,8 +664,8 @@ with gr.Blocks(theme=theme, css=custom_css, js=custom_js) as demo:
 
     # Function to set view-only mode
     def clear_and_set_view_only(task_input, request: gr.Request):
-        # First clear the results, then set view-only mode
-        return "", update_html(False, request), gr.update(visible=False)
+        # set view-only mode
+        return update_html(False, request)
 
     def set_interactive(request: gr.Request):
         return update_html(True, request)
@@ -676,11 +676,29 @@ with gr.Blocks(theme=theme, css=custom_css, js=custom_js) as demo:
     view_only_event = update_btn.click(
         fn=clear_and_set_view_only,
         inputs=[task_input], 
-        outputs=[results_output, sandbox_html, results_container]
-    ).then(agent_ui.interact_with_agent, [task_input, stored_messages, session_state, session_hash_state], [chatbot_display]).then(
+        outputs=[sandbox_html]
+    ).then(
+        agent_ui.interact_with_agent,
+        inputs=[task_input, stored_messages, session_state, session_hash_state, None],
+        outputs=[chatbot_display]
+    ).then(
+        fn=set_interactive,
+        inputs=[],
+        outputs=[sandbox_html]
+    )
+
+    replay_btn.click(
+        fn=clear_and_set_view_only,
+        inputs=[task_input], 
+        outputs=[sandbox_html]
+    ).then(
+        agent_ui.interact_with_agent,
+        inputs=[task_input, stored_messages, session_state, session_hash_state, "udupp2fyavq_1743170323"],
+        outputs=[chatbot_display]
+    ).then(
         fn=set_interactive,
         inputs=[],
-        outputs=sandbox_html
+        outputs=[sandbox_html]
     )
 
     demo.load(

model_replay.py

@@ -0,0 +1,90 @@
+from smolagents.models import Model, ChatMessage, Tool, MessageRole
+from time import time
+
+class FakeModelClass(Model):
+    """A model class that returns pre-recorded responses from a log file.
+    
+    This class is useful for testing and debugging purposes, as it doesn't make
+    actual API calls but instead returns responses from a pre-recorded log file.
+    
+    Parameters:
+        log_url (str, optional): 
+            URL to the log file. Defaults to the smolagents example log.
+        **kwargs: Additional keyword arguments passed to the Model base class.
+    """
+    
+    def __init__(
+        self,
+        log_folder: str,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.dataset_name = "smolagents/computer-agent-logs",
+        self.log_folder = log_folder
+        self.call_counter = 0
+        self.model_outputs = self._load_model_outputs()
+        
+    def _load_model_outputs(self) -> List[str]:
+        """Load model outputs from the log file using HuggingFace datasets library."""
+        # Download the file from Hugging Face Hub
+        file_path = hf_hub_download(
+            repo_id=self.dataset_name,
+            filename=self.log_folder + "/metadata.json",
+            repo_type="dataset"
+        )
+        
+        # Load and parse the JSON data
+        with open(file_path, 'r') as f:
+            log_data = json.load(f)
+        
+        # Extract only the model_output from each step in tool_calls
+        model_outputs = []
+        
+        for step in log_data.get("tool_calls", []):
+            if "model_output_message" in step:
+                model_outputs.append(step["model_output_message"])
+
+        print(f"Loaded {len(model_outputs)} model outputs from log file")
+        return model_outputs
+
+    def __call__(
+        self,
+        messages: List[Dict[str, str]],
+        stop_sequences: Optional[List[str]] = None,
+        grammar: Optional[str] = None,
+        tools_to_call_from: Optional[List[Tool]] = None,
+        **kwargs
+    ) -> ChatMessage:
+        """Return the next pre-recorded response from the log file.
+        
+        Parameters:
+            messages: List of input messages (ignored).
+            stop_sequences: Optional list of stop sequences (ignored).
+            grammar: Optional grammar specification (ignored).
+            tools_to_call_from: Optional list of tools (ignored).
+            **kwargs: Additional keyword arguments (ignored).
+            
+        Returns:
+            ChatMessage: The next pre-recorded response.
+        """
+        time.sleep(1.0)
+
+        # Get the next model output
+        if self.call_counter < len(self.model_outputs):
+            content = self.model_outputs[self.call_counter]
+            self.call_counter += 1
+        else:
+            content = "No more pre-recorded responses available."
+
+        # Token counts are simulated
+        self.last_input_token_count = len(str(messages)) // 4  # Rough approximation
+        self.last_output_token_count = len(content) // 4  # Rough approximation
+        
+        # Create and return a ChatMessage
+        return ChatMessage(
+            role=MessageRole.ASSISTANT,
+            content=content,
+            tool_calls=None,
+            raw={"source": "pre-recorded log", "call_number": self.call_counter}
+        )
+