Yoan Di Cosmo commited on
Commit
3fa386c
·
1 Parent(s): bbb86f1

Add auto_file_upload and confirm_cpu_jobs config settings

Browse files
agent/config.py CHANGED
@@ -23,6 +23,22 @@ class Config(BaseModel):
23
  session_dataset_repo: str = "smolagents/hf-agent-sessions"
24
  auto_save_interval: int = 3 # Save every N user turns (0 = disabled)
25
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
26
 
27
  def substitute_env_vars(obj: Any) -> Any:
28
  """
 
23
  session_dataset_repo: str = "smolagents/hf-agent-sessions"
24
  auto_save_interval: int = 3 # Save every N user turns (0 = disabled)
25
 
26
+ # Permission control parameters
27
+ confirm_cpu_jobs: bool = True
28
+ """
29
+ If True, agent must ask for user confirmation before launching CPU-only jobs.
30
+ If False, agent can run CPU-only jobs automatically without confirmation.
31
+ Default: True (require confirmation for safety)
32
+ """
33
+
34
+ auto_file_upload: bool = False
35
+ """
36
+ If True, agent may upload files automatically without additional prompts
37
+ (user has given prior permission via this config).
38
+ If False, agent must ask for user permission before uploading any file.
39
+ Default: False (ask permission for safety)
40
+ """
41
+
42
 
43
  def substitute_env_vars(obj: Any) -> Any:
44
  """
agent/config_mcp_example.json ADDED
@@ -0,0 +1,15 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "model_name": "anthropic/claude-sonnet-4-5-20250929",
3
+ "system_prompt_path": "",
4
+ "confirm_cpu_jobs": false,
5
+ "auto_file_upload": true,
6
+ "mcpServers": {
7
+ "hf-mcp-server": {
8
+ "transport": "http",
9
+ "url": "https://huggingface.co/mcp?login",
10
+ "headers": {
11
+ "Authorization": "Bearer ${HF_TOKEN}"
12
+ }
13
+ }
14
+ }
15
+ }
agent/core/agent_loop.py CHANGED
@@ -11,6 +11,7 @@ from lmnr import observe
11
  from agent.config import Config
12
  from agent.core.session import Event, OpType, Session
13
  from agent.core.tools import ToolRouter
 
14
 
15
  ToolCall = ChatCompletionMessageToolCall
16
 
@@ -37,23 +38,47 @@ def _validate_tool_args(tool_args: dict) -> tuple[bool, str | None]:
37
  return True, None
38
 
39
 
40
- def _needs_approval(tool_name: str, tool_args: dict) -> bool:
41
- """Check if a tool call requires user approval before execution"""
42
  # If args are malformed, skip approval (validation error will be shown later)
43
  args_valid, _ = _validate_tool_args(tool_args)
44
  if not args_valid:
45
  return False
46
 
47
  if tool_name == "hf_jobs":
48
- # Check if it's a run or uv operation
49
  operation = tool_args.get("operation", "")
50
- return operation in ["run", "uv"]
51
-
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
52
  if tool_name == "hf_private_repos":
53
- # Repo creation and file uploads require approval
54
  operation = tool_args.get("operation", "")
55
- return operation in ["create_repo", "upload_file"]
56
-
 
 
 
 
 
 
 
 
57
  return False
58
 
59
 
@@ -143,7 +168,7 @@ class Handlers:
143
  tool_name = tc.function.name
144
  tool_args = json.loads(tc.function.arguments)
145
 
146
- if _needs_approval(tool_name, tool_args):
147
  approval_required_tools.append(tc)
148
  else:
149
  non_approval_tools.append(tc)
 
11
  from agent.config import Config
12
  from agent.core.session import Event, OpType, Session
13
  from agent.core.tools import ToolRouter
14
+ from agent.tools.jobs_tool import CPU_FLAVORS
15
 
16
  ToolCall = ChatCompletionMessageToolCall
17
 
 
38
  return True, None
39
 
40
 
41
+ def _needs_approval(tool_name: str, tool_args: dict, config: Config | None = None) -> bool:
42
+ """Check if a tool call requires user approval before execution."""
43
  # If args are malformed, skip approval (validation error will be shown later)
44
  args_valid, _ = _validate_tool_args(tool_args)
45
  if not args_valid:
46
  return False
47
 
48
  if tool_name == "hf_jobs":
49
+ # Check if it's a run or scheduled run operation
50
  operation = tool_args.get("operation", "")
51
+ if operation not in ["run", "uv", "scheduled run", "scheduled uv"]:
52
+ return False
53
+
54
+ # Check if this is a CPU-only job
55
+ args = tool_args.get("args", {})
56
+ hardware_flavor = args.get("flavor") or args.get("hardware") or args.get("hardware_flavor") or "cpu-basic"
57
+ is_cpu_job = hardware_flavor in CPU_FLAVORS
58
+
59
+ if is_cpu_job:
60
+ # Check config: if confirm_cpu_jobs is False, skip approval
61
+ if config and not config.confirm_cpu_jobs:
62
+ return False
63
+ # Otherwise, require approval for CPU jobs
64
+ return True
65
+
66
+ # For GPU jobs, always require approval (existing behavior)
67
+ return True
68
+
69
+ # Check for file upload operations (hf_private_repos or other tools)
70
  if tool_name == "hf_private_repos":
 
71
  operation = tool_args.get("operation", "")
72
+ if operation == "upload_file":
73
+ # Check config: if auto_file_upload is True, skip approval
74
+ if config and config.auto_file_upload:
75
+ return False
76
+ # Otherwise, require approval for file uploads
77
+ return True
78
+ # Other operations (create_repo, etc.) always require approval
79
+ if operation in ["create_repo"]:
80
+ return True
81
+
82
  return False
83
 
84
 
 
168
  tool_name = tc.function.name
169
  tool_args = json.loads(tc.function.arguments)
170
 
171
+ if _needs_approval(tool_name, tool_args, session.config):
172
  approval_required_tools.append(tc)
173
  else:
174
  non_approval_tools.append(tc)
configs/main_agent_config.json CHANGED
@@ -2,6 +2,8 @@
2
  "model_name": "anthropic/claude-opus-4-5-20251101",
3
  "save_sessions": true,
4
  "session_dataset_repo": "smolagents/hf-agent-sessions",
 
 
5
  "mcpServers": {
6
  "hf-mcp-server": {
7
  "transport": "http",
 
2
  "model_name": "anthropic/claude-opus-4-5-20251101",
3
  "save_sessions": true,
4
  "session_dataset_repo": "smolagents/hf-agent-sessions",
5
+ "confirm_cpu_jobs": false,
6
+ "auto_file_upload": false,
7
  "mcpServers": {
8
  "hf-mcp-server": {
9
  "transport": "http",