diff --git a/model_tools.py b/model_tools.py index 835867424..538599890 100644 --- a/model_tools.py +++ b/model_tools.py @@ -518,8 +518,9 @@ def handle_terminal_function_call(function_name: str, function_args: Dict[str, A background = function_args.get("background", False) idle_threshold = function_args.get("idle_threshold", 5.0) timeout = function_args.get("timeout") + snapshot_id = function_args.get("snapshot_id") # Session management is handled internally - don't pass session_id from model - return terminal_tool(command, input_keys, None, background, idle_threshold, timeout) + return terminal_tool(command, input_keys, None, background, idle_threshold, timeout, snapshot_id=snapshot_id) else: return json.dumps({"error": f"Unknown terminal function: {function_name}"}) diff --git a/run_agent.py b/run_agent.py index eec9d63c7..86c7c5e0a 100644 --- a/run_agent.py +++ b/run_agent.py @@ -51,7 +51,8 @@ class AIAgent: disabled_tools: List[str] = None, enabled_toolsets: List[str] = None, disabled_toolsets: List[str] = None, - save_trajectories: bool = False + save_trajectories: bool = False, + morph_snapshot_id: str | None = None, ): """ Initialize the AI Agent. @@ -67,11 +68,13 @@ class AIAgent: enabled_toolsets (List[str]): Only enable tools from these toolsets (optional) disabled_toolsets (List[str]): Disable tools from these toolsets (optional) save_trajectories (bool): Whether to save conversation trajectories to JSONL files (default: False) + morph_snapshot_id (str | None): Morph Cloud snapshot id from which to start terminal tool """ self.model = model self.max_iterations = max_iterations self.tool_delay = tool_delay self.save_trajectories = save_trajectories + self.morph_snapshot_id = morph_snapshot_id # Store tool filtering options self.enabled_tools = enabled_tools @@ -388,6 +391,9 @@ class AIAgent: function_args = {} print(f" 📞 Tool {i}: {function_name}({list(function_args.keys())})") + + if function_name == "terminal" and self.morph_snapshot_id is not None: + function_args["snapshot_id"] = self.morph_snapshot_id # Execute the tool function_result = handle_function_call(function_name, function_args) @@ -480,7 +486,8 @@ def main( enabled_toolsets: str = None, disabled_toolsets: str = None, list_tools: bool = False, - save_trajectories: bool = False + save_trajectories: bool = False, + morph_snapshot_id: str | None = None, ): """ Main function for running the agent directly. @@ -497,6 +504,7 @@ def main( disabled_toolsets (str): Comma-separated list of toolsets to disable (e.g., "terminal_tools") list_tools (bool): Just list available tools and exit save_trajectories (bool): Save conversation trajectories to JSONL files. Defaults to False. + morph_snapshot_id (str | None): Morph Cloud snapshot id to start terminal tool from """ print("🤖 AI Agent with Tool Calling") print("=" * 50) @@ -573,7 +581,8 @@ def main( disabled_tools=disabled_tools_list, enabled_toolsets=enabled_toolsets_list, disabled_toolsets=disabled_toolsets_list, - save_trajectories=save_trajectories + save_trajectories=save_trajectories, + morph_snapshot_id=morph_snapshot_id ) except RuntimeError as e: print(f"❌ Failed to initialize agent: {e}") diff --git a/terminal_tool.py b/terminal_tool.py index e01d7a617..b40bd7cf7 100644 --- a/terminal_tool.py +++ b/terminal_tool.py @@ -78,7 +78,8 @@ def terminal_tool( session_id: Optional[str] = None, background: bool = False, idle_threshold: float = 5.0, - timeout: Optional[int] = None + timeout: Optional[int] = None, + snapshot_id: str | None = None, ) -> str: """ Execute a command on a Morph VM with optional interactive session support. @@ -136,7 +137,7 @@ def terminal_tool( ) # Execute with lifecycle management - result = run_tool_with_lifecycle_management(tool_call) + result = run_tool_with_lifecycle_management(tool_call, snapshot_id=snapshot_id) # Format the result with all possible fields # Map hecate's "stdout" to "output" for compatibility @@ -231,4 +232,4 @@ if __name__ == "__main__": print(f" MORPH_API_KEY: {'Set' if os.getenv('MORPH_API_KEY') else 'Not set'}") print(f" OPENAI_API_KEY: {'Set' if os.getenv('OPENAI_API_KEY') else 'Not set (optional)'}") print(f" HECATE_VM_LIFETIME_SECONDS: {os.getenv('HECATE_VM_LIFETIME_SECONDS', '300')} (default: 300)") - print(f" HECATE_DEFAULT_SNAPSHOT_ID: {os.getenv('HECATE_DEFAULT_SNAPSHOT_ID', 'snapshot_p5294qxt')} (default: snapshot_p5294qxt)") \ No newline at end of file + print(f" HECATE_DEFAULT_SNAPSHOT_ID: {os.getenv('HECATE_DEFAULT_SNAPSHOT_ID', 'snapshot_p5294qxt')} (default: snapshot_p5294qxt)")