wuyang/zk-data-agent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Improve Jupyter workspace bootstrap

Browse Source
This commit is contained in:
wuyang6
2026-05-12 22:33:15 +08:00
parent 295d5d1ea1
commit 86f49134a4
5 changed files with 858 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/agent_tools.py
+54 -3
View File
@@ -2313,7 +2313,7 @@ def _run_python_exec(arguments: dict[str, Any], context: ToolExecutionContext) -
)
payload = [
f'exit_code={result.exit_code}',
'interpreter=remote:python3',
f'interpreter={context.jupyter_runtime.python_interpreter}',
f'remote_workspace={context.jupyter_runtime.binding.workspace_cwd}',
'[stdout]',
result.stdout.rstrip(),
@@ -2327,7 +2327,7 @@ def _run_python_exec(arguments: dict[str, Any], context: ToolExecutionContext) -
{
'action': 'remote_python_exec',
'mode': 'code' if code else 'script',
'interpreter': 'remote:python3',
'interpreter': context.jupyter_runtime.python_interpreter,
'remote_workspace': context.jupyter_runtime.binding.workspace_cwd,
'timed_out': result.timed_out,
'timeout_seconds': timeout_seconds,
@@ -2537,13 +2537,64 @@ def _communicate_after_stop(process: subprocess.Popen[str]) -> tuple[str, str]:
def _run_python_package(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
_ensure_process_execution_allowed(context, 'Python package management')
action = _require_string(arguments, 'action')
interpreter = _resolve_python_interpreter(context)
timeout_seconds = _coerce_float(
arguments,
'timeout_seconds',
min(context.command_timeout_seconds * 4, 120.0),
)
max_output_chars = _coerce_int(arguments, 'max_output_chars', context.max_output_chars)
if context.jupyter_runtime is not None:
interpreter = context.jupyter_runtime.python_interpreter
if action == 'show':
result = context.jupyter_runtime.run_command(
f'{shlex.quote(interpreter)} -m pip --version',
timeout_seconds=timeout_seconds,
max_output_chars=max_output_chars,
cancel_event=context.cancel_event,
)
elif action == 'install':
packages = arguments.get('packages')
if not isinstance(packages, list) or not packages:
raise ToolExecutionError('packages must be a non-empty array for action=install')
package_args = [str(package).strip() for package in packages if str(package).strip()]
if not package_args:
raise ToolExecutionError('packages must contain at least one non-empty package name')
result = context.jupyter_runtime.install_python_packages(
package_args,
timeout_seconds=timeout_seconds,
max_output_chars=max_output_chars,
cancel_event=context.cancel_event,
)
else:
raise ToolExecutionError('action must be "show" or "install"')
payload = [
f'exit_code={result.exit_code}',
f'interpreter={interpreter}',
f'python_env={context.jupyter_runtime.python_env_path}',
'[stdout]',
result.stdout.rstrip(),
'[stderr]',
result.stderr.rstrip(),
]
if result.timed_out:
payload.insert(0, f'timed_out=true\ntimeout_seconds={timeout_seconds:g}')
return (
_truncate_output('\n'.join(payload).strip(), max_output_chars),
{
'action': 'remote_python_package',
'package_action': action,
'interpreter': interpreter,
'python_env_dir': context.jupyter_runtime.python_env_path,
'remote_workspace': context.jupyter_runtime.binding.workspace_cwd,
'timed_out': result.timed_out,
'timeout_seconds': timeout_seconds,
'exit_code': result.exit_code,
'stdout_preview': _snapshot_text(result.stdout),
'stderr_preview': _snapshot_text(result.stderr),
'output_preview': _snapshot_text('\n'.join(payload).strip()),
},
)
interpreter = _resolve_python_interpreter(context)
if action == 'show':
command = [interpreter, '-m', 'pip', '--version']
elif action == 'install':