feat(alias-ds): refactor code execution logic and plan organization#111
feat(alias-ds): refactor code execution logic and plan organization#111SSSuperDan wants to merge 7 commits intoagentscope-ai:mainfrom
Conversation
Increased timeout settings for proxy connections.
…package versions.
Summary of ChangesHello @SSSuperDan, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a substantial refactoring of the Data Science Agent's operational framework. The primary goal is to enhance the agent's planning capabilities and standardize its code execution workflow by integrating a more structured planning notebook system. This change aims to improve the agent's ability to manage complex tasks, ensure consistent code execution practices, and provide more detailed reporting, ultimately leading to a more reliable and efficient data science agent. Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a significant refactoring of the data science agent's internal logic, transitioning from a simple 'todo list' mechanism to a more structured 'plan notebook' approach. This change enhances task management and provides clearer guidelines for the agent's operation. Key updates include the removal of the todo_write tool and its associated prompts, replaced by new planning tools like create_plan, update_subtask_state, finish_subtask, revise_current_plan, and finish_plan. Additionally, explicit 'Code Execution Protocols' have been added to the agent's system prompt, enforcing a safer and more controlled method for executing Python code by requiring writing to file and executing via shell. The run_ipython_cell tool has been removed from the data science agent's direct toolkit, aligning with these new protocols. Dependency versions for pydantic and mcp have been updated, and agentscope[full] and qdrant-client are now explicitly included in the sandbox requirements. Nginx configuration has also been updated with proxy timeout settings for improved stability. Overall, these changes improve the agent's robustness, maintainability, and adherence to best practices for code execution and task management.
| "read_file", | ||
| "write_file", | ||
| "run_ipython_cell", | ||
| "run_shell_command", |
There was a problem hiding this comment.
The ds_tool_list in init_ds_toolkit now includes read_file and removes run_ipython_cell. The removal of run_ipython_cell is a critical change that aligns with the new 'Code Execution Protocols' in the system prompt, promoting safer execution by requiring code to be written to a file and executed via run_shell_command.
| ## Code Execution Protocols | ||
|
|
||
| **For all Python code execution:** | ||
|
|
||
| 1. **Write code to file first** using `write_file` or similar tools | ||
| 2. **Execute via shell** using `run_shell_command` with `python <filename>.py` | ||
| 3. **For large and key results**: Save to file (CSV/JSON/pickle), load in next step. Avoid printing large outputs or passing data via stdout | ||
| 4. **Never execute code inline or outside this workflow** | ||
|
|
||
| **Example:** | ||
|
|
||
| Tool: `write_file` | ||
| Arguments: { | ||
| "file_path": "/workspace/code/analysis.py", | ||
| "content": "import pandas as pd\ndf = pd.read_csv('data.csv')\nprint(df.head())" | ||
| } | ||
|
|
||
| Tool: `run_shell_command` | ||
| Arguments: { | ||
| "command": "python /workspace/code/analysis.py" | ||
| } | ||
|
|
||
| Tool: `write_file` | ||
| Arguments: { | ||
| "file_path": "/workspace/code/step1_aggregate.py", | ||
| "content": "import pandas as pd\ndf = pd.read_csv('raw_data.csv')\nprocessed = df.groupby('category').sum()\nprocessed.to_csv('/workspace/data/intermediate_result.csv', index=False)\nprint('✓ Results saved')" | ||
| } |
There was a problem hiding this comment.
A new 'Code Execution Protocols' section has been added, explicitly outlining the recommended workflow for executing Python code: write to file, then execute via shell. This is a crucial improvement for security, reproducibility, and maintainability, preventing arbitrary inline code execution and promoting structured development.
| - **You must use `create_plan` to create a task plan**, especially for multi-step tasks. | ||
| - Use `update_subtask_state` to mark subtasks as 'in_progress' when starting them. | ||
| - Use `finish_subtask` to mark subtasks as 'done' with specific outcomes upon completion. | ||
| - Use `revise_current_plan` to adapt the plan accordingly if new insights or results emerge. | ||
| - Use `finish_plan` to finalize the entire task when all subtasks are complete. | ||
| - Skipping planning risks missing critical steps—this is unacceptable. |
There was a problem hiding this comment.
The 'Task Management Rules' section has been completely revised to detail the usage of the new PlanNotebook tools (create_plan, update_subtask_state, finish_subtask, revise_current_plan, finish_plan). This provides clear instructions for the agent on how to manage its tasks effectively.
| proxy_connect_timeout 120s; | ||
| proxy_send_timeout 120s; | ||
| proxy_read_timeout 120s; |
There was a problem hiding this comment.
Adding proxy_connect_timeout, proxy_send_timeout, and proxy_read_timeout to the Nginx configuration is a good operational improvement. These settings help prevent timeouts for long-running requests or slow network conditions, enhancing the stability and reliability of the service.
| return ( | ||
| f"{base_prompt}{self._selected_scenario_prompts}\n\n{todo_prompt}" | ||
| ) | ||
| return f"{base_prompt}{self._selected_scenario_prompts}" |
There was a problem hiding this comment.
The sys_prompt property has been simplified by removing the todo_prompt logic. This is consistent with the removal of the old todo list management system and the introduction of PlanNotebook.
| global_toolkit = AliasToolkit(sandbox, add_all=True) | ||
| worker_toolkit = init_ds_toolkit(global_toolkit) |
There was a problem hiding this comment.
The worker_toolkit is now initialized directly by calling init_ds_toolkit(global_toolkit). This refactoring improves modularity by centralizing the data science toolkit setup logic within init_ds_toolkit, making arun_datascience_agent cleaner and more focused.
| return response.content[0]["text"] | ||
|
|
||
| async def generate_report(self) -> Tuple[str, str]: | ||
| async def generate_report(self) -> Tuple[str, str, str]: |
There was a problem hiding this comment.
The return type hint for generate_report has been updated to Tuple[str, str, str], reflecting the new functionality of returning both markdown and HTML report content.
| } | ||
| } | ||
| } No newline at end of file | ||
| } |
There was a problem hiding this comment.
Adding a newline at the end of the file is a minor stylistic improvement, adhering to common file formatting conventions.
| pydantic==2.12.5 | ||
| requests==2.32.3 | ||
| mcp==1.9.0 | ||
| mcp==1.25.0 |
There was a problem hiding this comment.
Updating pydantic to 2.12.5 and mcp to 1.25.0 ensures that the sandbox environment uses the latest versions of these dependencies. This can bring bug fixes, performance improvements, and new features, contributing to overall system stability and security.
| agentscope[full]==1.0.11 | ||
| qdrant-client==1.15.1 |
There was a problem hiding this comment.
Explicitly adding agentscope[full]==1.0.11 and qdrant-client==1.15.1 to the sandbox requirements ensures that these critical packages are available within the sandbox environment. This is essential for the agent's functionality, especially after the removal of the install_package utility.
2 participants
📝 PR Type
📚 Description
This PR refactors the internal code execution logic and plan organization in the data science agent.
✅ Checklist
Please complete the following checks before submitting the PR:
pre-commit run --all-filesagentscope-sampleshas been updated (e.g.,README.md)