Merge pull request #3 from huggingface/job-tool-improvements

agent/core/agent_loop.py

@@ -15,6 +15,16 @@ from agent.core.tools import ToolRouter
 ToolCall = ChatCompletionMessageToolCall
 
 
+def _needs_approval(tool_name: str, tool_args: dict) -> bool:
+    """Check if a tool call requires user approval before execution"""
+    if tool_name != "hf_jobs":
+        return False
+
+    # Check if it's a run or uv operation
+    operation = tool_args.get("operation", "")
+    return operation in ["run", "uv"]
+
+
 class Handlers:
     """Handler functions for each operation type"""
 
@@ -33,9 +43,10 @@ class Handlers:
 
             Laminar.set_trace_session_id(session_id=session.session_id)
 
-        # Add user message to history
-        user_msg = Message(role="user", content=text)
-        session.context_manager.add_message(user_msg)
+        # Add user message to history only if there's actual content
+        if text:
+            user_msg = Message(role="user", content=text)
+            session.context_manager.add_message(user_msg)
 
         # Send event that we're processing
         await session.send_event(
@@ -97,6 +108,28 @@ class Handlers:
                     tool_name = tc.function.name
                     tool_args = json.loads(tc.function.arguments)
 
+                    # Check if this tool requires user approval
+                    if _needs_approval(tool_name, tool_args):
+                        await session.send_event(
+                            Event(
+                                event_type="approval_required",
+                                data={
+                                    "tool": tool_name,
+                                    "arguments": tool_args,
+                                    "tool_call_id": tc.id,
+                                },
+                            )
+                        )
+
+                        # Store pending approval and return early
+                        session.pending_approval = {
+                            "tool_call": tc,
+                            "arguments": tool_args,
+                        }
+
+                        # Return early - wait for EXEC_APPROVAL operation
+                        return None
+
                     await session.send_event(
                         Event(
                             event_type="tool_call",
@@ -191,6 +224,85 @@ class Handlers:
 
         await session.send_event(Event(event_type="undo_complete"))
 
+    @staticmethod
+    async def exec_approval(
+        session: Session, approved: bool, feedback: str | None = None
+    ) -> None:
+        """Handle job execution approval"""
+        if not session.pending_approval:
+            await session.send_event(
+                Event(
+                    event_type="error",
+                    data={"error": "No pending approval to process"},
+                )
+            )
+            return
+
+        tc = session.pending_approval["tool_call"]
+        tool_args = session.pending_approval["arguments"]
+        tool_name = tc.function.name
+
+        if approved:
+            # Execute the pending tool
+            await session.send_event(
+                Event(
+                    event_type="tool_call",
+                    data={"tool": tool_name, "arguments": tool_args},
+                )
+            )
+
+            output, success = await session.tool_router.call_tool(tool_name, tool_args)
+
+            # Add tool result to context
+            tool_msg = Message(
+                role="tool",
+                content=output,
+                tool_call_id=tc.id,
+                name=tool_name,
+            )
+            session.context_manager.add_message(tool_msg)
+
+            await session.send_event(
+                Event(
+                    event_type="tool_output",
+                    data={
+                        "tool": tool_name,
+                        "output": output,
+                        "success": success,
+                    },
+                )
+            )
+        else:
+            # User rejected - add cancellation message to context
+            cancellation_msg = "Job execution cancelled by user"
+            if feedback:
+                cancellation_msg += f". User feedback: {feedback}"
+
+            tool_msg = Message(
+                role="tool",
+                content=cancellation_msg,
+                tool_call_id=tc.id,
+                name=tool_name,
+            )
+            session.context_manager.add_message(tool_msg)
+
+            await session.send_event(
+                Event(
+                    event_type="tool_output",
+                    data={
+                        "tool": tool_name,
+                        "output": cancellation_msg,
+                        "success": False,
+                    },
+                )
+            )
+
+        # Clear pending approval
+        session.pending_approval = None
+
+        # Continue agent loop with empty input to process the tool result
+        await Handlers.run_agent(session, "")
+
     @staticmethod
     async def shutdown(session: Session) -> bool:
         """Handle shutdown (like shutdown in codex.rs:1329)"""
@@ -226,6 +338,12 @@ async def process_submission(session: Session, submission) -> bool:
         await Handlers.undo(session)
         return True
 
+    if op.op_type == OpType.EXEC_APPROVAL:
+        approved = op.data.get("approved", False) if op.data else False
+        feedback = op.data.get("feedback") if op.data else None
+        await Handlers.exec_approval(session, approved, feedback)
+        return True
+
     if op.op_type == OpType.SHUTDOWN:
         return not await Handlers.shutdown(session)
 

agent/core/session.py

@@ -53,6 +53,7 @@ class Session:
         )
         self.is_running = True
         self.current_task: asyncio.Task | None = None
+        self.pending_approval: Optional[dict[str, Any]] = None
 
     async def send_event(self, event: Event) -> None:
         """Send event back to client"""

agent/core/tools.py

@@ -8,6 +8,7 @@ from dataclasses import dataclass
 from typing import Any, Awaitable, Callable, Optional
 
 from fastmcp import Client
+from fastmcp.exceptions import ToolError
 from lmnr import observe
 from mcp.types import EmbeddedResource, ImageContent, TextContent
 
@@ -166,10 +167,14 @@ class ToolRouter:
 
         # Otherwise, use MCP client
         if self._mcp_initialized:
-            result = await self.mcp_client.call_tool(tool_name, arguments)
-            # Convert MCP content blocks to string
-            output = convert_mcp_content_to_string(result.content)
-            return output, not result.is_error
+            try:
+                result = await self.mcp_client.call_tool(tool_name, arguments)
+                output = convert_mcp_content_to_string(result.content)
+                return output, not result.is_error
+            except ToolError as e:
+                # Catch MCP tool errors and return them to the agent
+                error_msg = f"Tool error: {str(e)}"
+                return error_msg, False
 
         return "MCP client not initialized", False
 

agent/main.py

@@ -47,10 +47,13 @@ class Submission:
 
 async def event_listener(
     event_queue: asyncio.Queue,
+    submission_queue: asyncio.Queue,
     turn_complete_event: asyncio.Event,
     ready_event: asyncio.Event,
 ) -> None:
     """Background task that listens for events and displays them"""
+    submission_id = [1000]  # Use list to make it mutable in closure
+
     while True:
         try:
             event = await event_queue.get()
@@ -96,6 +99,69 @@ async def event_listener(
                 old_tokens = event.data.get("old_tokens", 0) if event.data else 0
                 new_tokens = event.data.get("new_tokens", 0) if event.data else 0
                 print(f"📦 Compacted context: {old_tokens} → {new_tokens} tokens")
+            elif event.event_type == "approval_required":
+                # Display job details and prompt for approval
+                tool_name = event.data.get("tool", "") if event.data else ""
+                arguments = event.data.get("arguments", {}) if event.data else {}
+
+                print("\n" + "=" * 60)
+                print("⚠️  JOB EXECUTION APPROVAL REQUIRED")
+                print("=" * 60)
+
+                operation = arguments.get("operation", "")
+                args = arguments.get("args", {})
+
+                print(f"Operation: {operation}")
+
+                if operation == "uv":
+                    script = args.get("script", "")
+                    dependencies = args.get("dependencies", [])
+                    print(f"Script to run:\n{script}")
+                    if dependencies:
+                        print(f"Dependencies: {', '.join(dependencies)}")
+                elif operation == "run":
+                    image = args.get("image", "")
+                    command = args.get("command", "")
+                    print(f"Docker image: {image}")
+                    print(f"Command: {command}")
+
+                # Common parameters
+                flavor = args.get("flavor", "cpu-basic")
+                detached = args.get("detached", False)
+                print(f"Hardware: {flavor}")
+                print(f"Detached mode: {detached}")
+
+                secrets = args.get("secrets", [])
+                if secrets:
+                    print(f"Secrets: {', '.join(secrets)}")
+
+                print("=" * 60)
+
+                # Get user decision
+                loop = asyncio.get_event_loop()
+                response = await loop.run_in_executor(
+                    None,
+                    input,
+                    "Approve? (y=yes, n=no, or provide feedback to reject): ",
+                )
+
+                response = response.strip()
+                approved = response.lower() in ["y", "yes"]
+                feedback = (
+                    None if approved or response.lower() in ["n", "no"] else response
+                )
+
+                # Submit approval
+                submission_id[0] += 1
+                approval_submission = Submission(
+                    id=f"approval_{submission_id[0]}",
+                    operation=Operation(
+                        op_type=OpType.EXEC_APPROVAL,
+                        data={"approved": approved, "feedback": feedback},
+                    ),
+                )
+                await submission_queue.put(approval_submission)
+                print("=" * 60 + "\n")
             # Silently ignore other events
 
         except asyncio.CancelledError:
@@ -145,7 +211,7 @@ async def main():
 
     # Start event listener in background
     listener_task = asyncio.create_task(
-        event_listener(event_queue, turn_complete_event, ready_event)
+        event_listener(event_queue, submission_queue, turn_complete_event, ready_event)
     )
 
     # Wait for agent to initialize

agent/prompts/system_prompt.yaml

@@ -85,6 +85,7 @@ system_prompt: |
   - Always search Hugging Face Hub for existing resources before suggesting custom implementations
   - When referencing models, datasets, or papers, include direct links from search results
   - Never assume a library is available - check documentation first
+  - Before processing any dataset: inspect its actual structure first using the mcp__hf-mcp-server__hub_repo_details tool. Never assume column names: verify them beforehand.
   - Follow ML best practices: proper train/val/test splits, reproducibility, evaluation metrics
   - For training tasks, consider compute requirements and suggest appropriate hardware
   - Never expose or log API keys, tokens, or secrets

agent/tools/jobs_tool.py

@@ -6,6 +6,7 @@ Refactored to use official huggingface-hub library instead of custom HTTP client
 
 import asyncio
 import base64
+import os
 from typing import Any, Dict, Literal, Optional
 
 from huggingface_hub import HfApi
@@ -60,10 +61,32 @@ OperationType = Literal[
 ]
 
 # Constants
-DEFAULT_LOG_WAIT_SECONDS = 10
 UV_DEFAULT_IMAGE = "ghcr.io/astral-sh/uv:python3.12-bookworm"
 
 
+def _substitute_hf_token(params: Dict[str, Any] | None) -> Dict[str, Any] | None:
+    """
+    Substitute $HF_TOKEN with actual token value from environment.
+
+    Args:
+        params: Dictionary that may contain "$HF_TOKEN" in values
+
+    Returns:
+        Dictionary with $HF_TOKEN substituted
+    """
+    if params is None:
+        return None
+
+    result = {}
+    for key, value in params.items():
+        if value == "$HF_TOKEN":
+            result[key] = os.environ.get("HF_TOKEN", "")
+        else:
+            result[key] = value
+
+    return result
+
+
 def _build_uv_command(
     script: str,
     with_deps: list[str] | None = None,
@@ -103,6 +126,20 @@ def _wrap_inline_script(
     return f'echo "{encoded}" | base64 -d | {uv_command_str}'
 
 
+def _ensure_hf_transfer_dependency(deps: list[str] | None) -> list[str]:
+    """Ensure hf-transfer is included in the dependencies list"""
+    if deps is None:
+        return ["hf-transfer"]
+
+    if isinstance(deps, list):
+        deps_copy = deps.copy()  # Don't modify the original
+        if "hf-transfer" not in deps_copy:
+            deps_copy.append("hf-transfer")
+        return deps_copy
+
+    return ["hf-transfer"]
+
+
 def _resolve_uv_command(
     script: str,
     with_deps: list[str] | None = None,
@@ -316,7 +353,8 @@ Call this tool with:
 {{
   "operation": "uv",
   "args": {{
-    "script": "import random\\nprint(42 + random.randint(1, 5))"
+    "script": "import random\\nprint(42 + random.randint(1, 5))",
+    "dependencies" : ["torch", "huggingface_hub"]
   }}
 }}
 ```
@@ -335,7 +373,6 @@ Call this tool with:
 **String format (simple cases only):**
 - Still accepted for backwards compatibility, parsed with POSIX shell semantics
 - Rejects shell operators and can mis-handle characters such as `&`; switch to arrays when things turn complex
-- `$HF_TOKEN` stays literal—forward it via `secrets: {{ "HF_TOKEN": "$HF_TOKEN" }}`
 
 ### Show command-specific help
 Call this tool with:
@@ -345,9 +382,10 @@ Call this tool with:
 
 ## Tips
 
-- Jobs default to non-detached mode (tail logs for up to {DEFAULT_LOG_WAIT_SECONDS}s or until completion). Set `detach: true` to return immediately.
+- Jobs default to non-detached mode (stream logs until completion). Set `detach: true` to return immediately.
 - Prefer array commands to avoid shell parsing surprises
-- To access private Hub assets, include `secrets: {{ "HF_TOKEN": "$HF_TOKEN" }}` to inject your auth token.
+- To access private Hub assets (spaces, private models, datasets, collections), pass `secrets: {{ "HF_TOKEN": "$HF_TOKEN" }}`
+- Before calling a job, think about dependencies (they must be specified), which hardware flavor to run on (choose simplest for task), and whether to include secrets.
 """
         return {"formatted": usage_text, "totalResults": 1, "resultsShared": 1}
 
@@ -356,6 +394,33 @@ Call this tool with:
         help_text = f"Help for operation: {operation}\n\nCall with appropriate arguments. Use the main help for examples."
         return {"formatted": help_text, "totalResults": 1, "resultsShared": 1}
 
+    async def _wait_for_job_completion(
+        self, job_id: str, namespace: Optional[str] = None
+    ) -> tuple[str, list[str]]:
+        """
+        Stream job logs until completion, printing them in real-time.
+
+        Returns:
+            tuple: (final_status, all_logs)
+        """
+        all_logs = []
+
+        # Fetch logs - generator streams logs as they arrive and ends when job completes
+        logs_gen = self.api.fetch_job_logs(job_id=job_id, namespace=namespace)
+
+        # Stream logs in real-time
+        for log_line in logs_gen:
+            print("\t" + log_line)
+            all_logs.append(log_line)
+
+        # After logs complete, fetch final job status
+        job_info = await _async_call(
+            self.api.inspect_job, job_id=job_id, namespace=namespace
+        )
+        final_status = job_info.status.stage
+
+        return final_status, all_logs
+
     async def _run_job(self, args: Dict[str, Any]) -> ToolResult:
         """Run a job using HfApi.run_job()"""
         try:
@@ -363,8 +428,8 @@ Call this tool with:
                 self.api.run_job,
                 image=args.get("image", "python:3.12"),
                 command=args.get("command"),
-                env=args.get("env"),
-                secrets=args.get("secrets"),
+                env=_substitute_hf_token(args.get("env")),
+                secrets=_substitute_hf_token(args.get("secrets")),
                 flavor=args.get("flavor", "cpu-basic"),
                 timeout=args.get("timeout", "30m"),
                 namespace=args.get("namespace") or self.namespace,
@@ -382,14 +447,28 @@ To check logs, call this tool with `{{"operation": "logs", "args": {{"job_id": "
 To inspect, call this tool with `{{"operation": "inspect", "args": {{"job_id": "{job.id}"}}}}`"""
                 return {"formatted": response, "totalResults": 1, "resultsShared": 1}
 
-            # Not detached - return job info
-            response = f"""Job started: {job.id}
+            # Not detached - wait for completion and stream logs
+            print(f"Job started: {job.id}")
+            print("Streaming logs...\n---\n")
 
-**Status:** {job.status.stage}
-**View logs at:** {job.url}
+            final_status, all_logs = await self._wait_for_job_completion(
+                job_id=job.id,
+                namespace=args.get("namespace") or self.namespace,
+            )
 
-Note: Logs are being collected. Check the job page for real-time logs.
-"""
+            # Format all logs for the agent
+            log_text = "\n".join(all_logs) if all_logs else "(no logs)"
+
+            response = f"""Job completed!
+
+**Job ID:** {job.id}
+**Final Status:** {final_status}
+**View at:** {job.url}
+
+**Logs:**
+```
+{log_text}
+```"""
             return {"formatted": response, "totalResults": 1, "resultsShared": 1}
 
         except Exception as e:
@@ -402,10 +481,18 @@ Note: Logs are being collected. Check the job page for real-time logs.
             if not script:
                 raise ValueError("script is required")
 
+            # Get dependencies and ensure hf-transfer is included
+            deps = (
+                args.get("with_deps")
+                or args.get("dependencies")
+                or args.get("packages")
+            )
+            deps = _ensure_hf_transfer_dependency(deps)
+
             # Resolve the command based on script type (URL, inline, or file)
             command = _resolve_uv_command(
                 script=script,
-                with_deps=args.get("with_deps") or args.get("dependencies"),
+                with_deps=deps,
                 python=args.get("python"),
                 script_args=args.get("script_args"),
             )
@@ -415,20 +502,46 @@ Note: Logs are being collected. Check the job page for real-time logs.
                 self.api.run_job,
                 image=UV_DEFAULT_IMAGE,
                 command=command,
-                env=args.get("env"),
-                secrets=args.get("secrets"),
-                flavor=args.get("flavor", "cpu-basic"),
+                env=_substitute_hf_token(args.get("env")),
+                secrets=_substitute_hf_token(args.get("secrets")),
+                flavor=args.get("flavor") or args.get("hardware") or "cpu-basic",
                 timeout=args.get("timeout", "30m"),
                 namespace=args.get("namespace") or self.namespace,
             )
 
-            response = f"""UV Job started: {job.id}
+            # If detached, return immediately
+            if args.get("detach", False):
+                response = f"""UV Job started successfully!
 
+**Job ID:** {job.id}
 **Status:** {job.status.stage}
 **View at:** {job.url}
 
-To check logs, call this tool with `{{"operation": "logs", "args": {{"job_id": "{job.id}"}}}}`
-"""
+To check logs, call this tool with `{{"operation": "logs", "args": {{"job_id": "{job.id}"}}}}`"""
+                return {"formatted": response, "totalResults": 1, "resultsShared": 1}
+
+            # Not detached - wait for completion and stream logs
+            print(f"UV Job started: {job.id}")
+            print("Streaming logs...\n---\n")
+
+            final_status, all_logs = await self._wait_for_job_completion(
+                job_id=job.id,
+                namespace=args.get("namespace") or self.namespace,
+            )
+
+            # Format all logs for the agent
+            log_text = "\n".join(all_logs) if all_logs else "(no logs)"
+
+            response = f"""UV Job completed!
+
+**Job ID:** {job.id}
+**Final Status:** {final_status}
+**View at:** {job.url}
+
+**Logs:**
+```
+{log_text}
+```"""
             return {"formatted": response, "totalResults": 1, "resultsShared": 1}
 
         except Exception as e:
@@ -578,8 +691,8 @@ To verify, call this tool with `{{"operation": "inspect", "args": {{"job_id": "{
                 image=args.get("image", "python:3.12"),
                 command=args.get("command"),
                 schedule=args.get("schedule"),
-                env=args.get("env"),
-                secrets=args.get("secrets"),
+                env=_substitute_hf_token(args.get("env")),
+                secrets=_substitute_hf_token(args.get("secrets")),
                 flavor=args.get("flavor", "cpu-basic"),
                 timeout=args.get("timeout", "30m"),
                 namespace=args.get("namespace") or self.namespace,
@@ -613,10 +726,18 @@ To list all, call this tool with `{{"operation": "scheduled ps"}}`"""
             if not schedule:
                 raise ValueError("schedule is required")
 
+            # Get dependencies and ensure hf-transfer is included
+            deps = (
+                args.get("with_deps")
+                or args.get("dependencies")
+                or args.get("packages")
+            )
+            deps = _ensure_hf_transfer_dependency(deps)
+
             # Resolve the command based on script type
             command = _resolve_uv_command(
                 script=script,
-                with_deps=args.get("with_deps") or args.get("dependencies"),
+                with_deps=deps,
                 python=args.get("python"),
                 script_args=args.get("script_args"),
             )
@@ -627,9 +748,9 @@ To list all, call this tool with `{{"operation": "scheduled ps"}}`"""
                 image=UV_DEFAULT_IMAGE,
                 command=command,
                 schedule=schedule,
-                env=args.get("env"),
-                secrets=args.get("secrets"),
-                flavor=args.get("flavor", "cpu-basic"),
+                env=_substitute_hf_token(args.get("env")),
+                secrets=_substitute_hf_token(args.get("secrets")),
+                flavor=args.get("flavor") or args.get("hardware") or "cpu-basic",
                 timeout=args.get("timeout", "30m"),
                 namespace=args.get("namespace") or self.namespace,
             )
@@ -788,6 +909,7 @@ HF_JOBS_TOOL_SPEC = {
     "description": (
         "Manage Hugging Face CPU/GPU compute jobs. Run commands in Docker containers, "
         "execute Python scripts with UV. List, schedule and monitor jobs/logs. "
+        "Example hardware/flavor: cpu-basic, cpu-performance, t4-medium. "
         "Call this tool with no operation for full usage instructions and examples."
     ),
     "parameters": {
@@ -818,7 +940,12 @@ HF_JOBS_TOOL_SPEC = {
             },
             "args": {
                 "type": "object",
-                "description": "Operation-specific arguments as a JSON object",
+                "description": (
+                    "Operation-specific arguments as a JSON object. "
+                    "Common args: script (for uv), packages/dependencies (array), "
+                    "flavor/hardware (e.g., a10g-large, cpu-basic), command (array), "
+                    "image (string), env (object), secrets (object)."
+                ),
                 "additionalProperties": True,
             },
         },