Merge pull request #4 from huggingface/planning-and-design

agent/core/agent_loop.py

@@ -365,7 +365,7 @@ async def submission_loop(
 
     # Create session with tool router
     session = Session(event_queue, config=config, tool_router=tool_router)
-    print("🤖 Agent loop started")
+    print("Agent loop started")
 
     # Main processing loop
     async with tool_router:

agent/core/tools.py

@@ -14,6 +14,7 @@ from mcp.types import EmbeddedResource, ImageContent, TextContent
 
 from agent.config import MCPServerConfig
 from agent.tools.jobs_tool import HF_JOBS_TOOL_SPEC, hf_jobs_handler
+from agent.tools.plan_tool import PLAN_TOOL_SPEC, plan_tool_handler
 
 # Suppress aiohttp deprecation warning
 warnings.filterwarnings(
@@ -186,7 +187,7 @@ class ToolRouter:
 
 def create_builtin_tools() -> list[ToolSpec]:
     """Create built-in tool specifications"""
-    print(f"Creating built-in tools: {HF_JOBS_TOOL_SPEC['name']}")
+    print(f"Creating built-in tools: {HF_JOBS_TOOL_SPEC['name']}, {PLAN_TOOL_SPEC['name']}")
     return [
         ToolSpec(
             name=HF_JOBS_TOOL_SPEC["name"],
@@ -194,4 +195,10 @@ def create_builtin_tools() -> list[ToolSpec]:
             parameters=HF_JOBS_TOOL_SPEC["parameters"],
             handler=hf_jobs_handler,
         ),
+        ToolSpec(
+            name=PLAN_TOOL_SPEC["name"],
+            description=PLAN_TOOL_SPEC["description"],
+            parameters=PLAN_TOOL_SPEC["parameters"],
+            handler=plan_tool_handler,
+        ),
     ]

agent/main.py

@@ -16,6 +16,16 @@ from agent.config import load_config
 from agent.core.agent_loop import submission_loop
 from agent.core.session import OpType
 from agent.core.tools import ToolRouter
+from agent.utils.terminal_display import (
+    format_error,
+    format_header,
+    format_plan_display,
+    format_separator,
+    format_success,
+    format_tool_call,
+    format_tool_output,
+    format_turn_complete,
+)
 
 litellm.drop_params = True
 
@@ -24,9 +34,9 @@ if lmnr_api_key:
     try:
         Laminar.initialize(project_api_key=lmnr_api_key)
         litellm.callbacks = [LaminarLiteLLMCallback()]
-        print("✅ Laminar initialized")
+        print("Laminar initialized")
     except Exception as e:
-        print(f"⚠️ Failed to initialize Laminar: {e}")
+        print(f"Failed to initialize Laminar: {e}")
 
 
 @dataclass
@@ -53,6 +63,7 @@ async def event_listener(
 ) -> None:
     """Background task that listens for events and displays them"""
     submission_id = [1000]  # Use list to make it mutable in closure
+    last_tool_name = [None]  # Track last tool called
 
     while True:
         try:
@@ -60,27 +71,32 @@ async def event_listener(
 
             # Display event
             if event.event_type == "ready":
-                print("✅ Agent ready")
+                print(format_success("\U0001F917 Agent ready"))
                 ready_event.set()
             elif event.event_type == "assistant_message":
                 content = event.data.get("content", "") if event.data else ""
                 if content:
-                    print(f"\n🤖 Assistant: {content}")
+                    print(f"\nAssistant: {content}")
             elif event.event_type == "tool_call":
                 tool_name = event.data.get("tool", "") if event.data else ""
                 arguments = event.data.get("arguments", {}) if event.data else {}
                 if tool_name:
-                    print(
-                        f"🔧 Calling tool: {tool_name} with arguments: {json.dumps(arguments)[:100]}..."
-                    )
+                    last_tool_name[0] = tool_name  # Store for tool_output event
+                    args_str = json.dumps(arguments)[:100] + "..."
+                    print(format_tool_call(tool_name, args_str))
             elif event.event_type == "tool_output":
                 output = event.data.get("output", "") if event.data else ""
                 success = event.data.get("success", False) if event.data else False
-                status = "✅" if success else "❌"
                 if output:
-                    print(f"{status} Tool output: {output}")
+                    # Don't truncate plan_tool output, truncate everything else
+                    should_truncate = last_tool_name[0] != "plan_tool"
+                    print(format_tool_output(output, success, truncate=should_truncate))
             elif event.event_type == "turn_complete":
-                print("✅ Turn complete\n")
+                print(format_turn_complete())
+                # Display plan after turn complete
+                plan_display = format_plan_display()
+                if plan_display:
+                    print(plan_display)
                 turn_complete_event.set()
             elif event.event_type == "error":
                 error = (
@@ -88,30 +104,26 @@ async def event_listener(
                     if event.data
                     else "Unknown error"
                 )
-                print(f"❌ Error: {error}")
+                print(format_error(error))
                 turn_complete_event.set()
             elif event.event_type == "shutdown":
-                print("🛑 Agent shutdown")
+                print("Agent shutdown")
                 break
             elif event.event_type == "processing":
-                print("⏳ Processing...", flush=True)
+                print("Processing...", flush=True)
             elif event.event_type == "compacted":
                 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")
+                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}")
+                print(f"\nOperation: {operation}")
 
                 if operation == "uv":
                     script = args.get("script", "")
@@ -135,9 +147,10 @@ async def event_listener(
                 if secrets:
                     print(f"Secrets: {', '.join(secrets)}")
 
-                print("=" * 60)
-
                 # Get user decision
+                print("\n" + format_separator())
+                print(format_header("JOB EXECUTION APPROVAL REQUIRED"))
+                print(format_separator())
                 loop = asyncio.get_event_loop()
                 response = await loop.run_in_executor(
                     None,
@@ -161,13 +174,13 @@ async def event_listener(
                     ),
                 )
                 await submission_queue.put(approval_submission)
-                print("=" * 60 + "\n")
+                print(format_separator() + "\n")
             # Silently ignore other events
 
         except asyncio.CancelledError:
             break
         except Exception as e:
-            print(f"⚠️ Event listener error: {e}")
+            print(f"Event listener error: {e}")
 
 
 async def get_user_input() -> str:
@@ -178,10 +191,26 @@ async def get_user_input() -> str:
 
 async def main():
     """Interactive chat with the agent"""
-    print("=" * 60)
-    print("🤖 Interactive Agent Chat")
-    print("=" * 60)
+    from agent.utils.terminal_display import Colors
+
+    # Clear screen
+    os.system('clear' if os.name != 'nt' else 'cls')
+
+
+    banner = r"""
+  _   _                   _               _____                   _                    _   
+ | | | |_   _  __ _  __ _(_)_ __   __ _  |  ___|_ _  ___ ___     / \   __ _  ___ _ __ | |_ 
+ | |_| | | | |/ _` |/ _` | | '_ \ / _` | | |_ / _` |/ __/ _ \   / _ \ / _` |/ _ \ '_ \| __|
+ |  _  | |_| | (_| | (_| | | | | | (_| | |  _| (_| | (_|  __/  / ___ \ (_| |  __/ | | | |_ 
+ |_| |_|\__,_|\__, |\__, |_|_| |_|\__, | |_|  \__,_|\___\___| /_/   \_\__, |\___|_| |_|\__|
+              |___/ |___/         |___/                               |___/
+    """
+
+
+    print(format_separator())
+    print(f"{Colors.YELLOW} {banner}{Colors.RESET}")
     print("Type your messages below. Type 'exit', 'quit', or '/quit' to end.\n")
+    print(format_separator())
 
     # Create queues for communication
     submission_queue = asyncio.Queue()
@@ -215,7 +244,7 @@ async def main():
     )
 
     # Wait for agent to initialize
-    print("⏳ Initializing agent...")
+    print("Initializing agent...")
     await ready_event.wait()
 
     submission_id = 0
@@ -253,7 +282,7 @@ async def main():
             await submission_queue.put(submission)
 
     except KeyboardInterrupt:
-        print("\n\n⚠️ Interrupted by user")
+        print("\n\nInterrupted by user")
 
     # Shutdown
     print("\n🛑 Shutting down agent...")

agent/prompts/system_prompt.yaml

@@ -1,5 +1,22 @@
 system_prompt: |
   You are HF Agent, a powerful AI assistant for Machine Learning Engineering, particularly training Large Language Models. You have access to {{ num_tools }} tools for interacting with Hugging Face Hub and performing ML tasks.
+  
+  # Task Approach
+
+  1. Always formulate a plan. Pass the todos to the PlanTool. Update the plan as progress is made.
+  2. Search for relevant models, datasets, and documentation on Hugging Face Hub.
+  3. Use all available tools to complete the task. Leverage existing resources before creating new ones.
+  4. Invoke multiple independent tools simultaneously for efficiency
+
+  # Autonomy / Subordinate trade-off.
+
+  Your main goal is to achieve what the user asked. For this:
+  1. Take action, follow-up, launch jobs. Ask for as little action from the user as possible. Do not ask them to do things you could do via a script.
+
+  However !! :
+  1. Don't surprise the user with costly, irreversible, or strange actions without asking.
+  2. Don't be shy to ask questions if needed.
+  3. Don't be overly talkative, explaining everything after a task ended.
 
   # Available Tools
 
@@ -9,23 +26,8 @@ system_prompt: |
   - Spaces: Use and discover ML applications
   - Jobs: Manage compute jobs for training and inference
   - Image Generation: Generate and transform images
-
-  # Agency
-
-  You take initiative when the user asks you to do something, maintaining an appropriate balance between:
-
-  1. Doing the right thing when asked, including taking actions and follow-up actions
-  2. Not surprising the user with actions you take without asking
-  3. Not adding unnecessary explanations after completing tasks
-
-  # Task Approach
-
-  For ML engineering tasks:
-  1. Use all available tools to complete the task
-  2. Search for relevant models, datasets, and documentation on Hugging Face Hub
-  3. Leverage existing resources before creating new ones
-  4. Invoke multiple independent tools simultaneously for efficiency
-
+  - Planning : a planning/to-do tool.
+  
   # Examples
 
   <example>
@@ -83,12 +85,16 @@ system_prompt: |
   # Conventions
 
   - Always search Hugging Face Hub for existing resources before suggesting custom implementations
+  - Keep in mind that a space is a repo, so you can create a space directly by uploading files that way. Repos should also be used to store files permanently : post-execution, files from jobs are not available.
+  - To run jobs, you must always pass the whole content of the file to execute. No files are available on server. Your local files and distant files are entirely seperate scopes.
+  - To access, create, or modify private Hub assets (spaces, private models, datasets, collections), pass `secrets: {% raw %}{{ "HF_TOKEN": "$HF_TOKEN" }}{% endraw %}` along with the jobs parameters. This is important. Without it, you will encounter authentification issues. Do not assume the user is connected on the jobs' server.
   - 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
+  - Unless absolutely necessary, don't ask user for action. This does not apply to follow-up questions you have.
+  - For training tasks, consider compute requirements and choose appropriate hardware.
+  - Never expose or log API keys, tokens, or secrets. Do not assume keys or secrets are available. Only Hugging Face private resources are available.
 
   # Communication Style
 

agent/tools/jobs_tool.py

@@ -66,20 +66,21 @@ 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.
+    Substitute HF_TOKEN key with actual token value from environment.
 
     Args:
-        params: Dictionary that may contain "$HF_TOKEN" in values
+        params: Dictionary that may contain "HF_TOKEN" as a key
 
     Returns:
-        Dictionary with $HF_TOKEN substituted
+        Dictionary with HF_TOKEN value substituted from environment
     """
+    print("DEBUG !! : ", params)
     if params is None:
         return None
 
     result = {}
     for key, value in params.items():
-        if value == "$HF_TOKEN":
+        if key == "HF_TOKEN":
             result[key] = os.environ.get("HF_TOKEN", "")
         else:
             result[key] = value
@@ -354,7 +355,8 @@ Call this tool with:
   "operation": "uv",
   "args": {{
     "script": "import random\\nprint(42 + random.randint(1, 5))",
-    "dependencies" : ["torch", "huggingface_hub"]
+    "dependencies": ["torch", "huggingface_hub"],
+    "secrets": {{"HF_TOKEN": "$HF_TOKEN"}}
   }}
 }}
 ```
@@ -384,7 +386,7 @@ Call this tool with:
 
 - 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 (spaces, private models, datasets, collections), pass `secrets: {{ "HF_TOKEN": "$HF_TOKEN" }}`
+- To access, create, or modify private Hub assets (spaces, private models, datasets, collections), pass `secrets: {{ "HF_TOKEN": "$HF_TOKEN" }}`. This is important. Without it, you will encounter authentification issues. Do not assume the user is connected on the jobs' server.
 - 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}
@@ -448,7 +450,7 @@ To inspect, call this tool with `{{"operation": "inspect", "args": {{"job_id": "
                 return {"formatted": response, "totalResults": 1, "resultsShared": 1}
 
             # Not detached - wait for completion and stream logs
-            print(f"Job started: {job.id}")
+            print(f"Job started: {job.url}")
             print("Streaming logs...\n---\n")
 
             final_status, all_logs = await self._wait_for_job_completion(
@@ -521,7 +523,7 @@ To check logs, call this tool with `{{"operation": "logs", "args": {{"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(f"UV Job started: {job.url}")
             print("Streaming logs...\n---\n")
 
             final_status, all_logs = await self._wait_for_job_completion(

agent/tools/plan_tool.py

@@ -0,0 +1,113 @@
+import json
+from typing import Any, Dict, List
+from .types import ToolResult
+from agent.utils.terminal_display import format_plan_tool_output
+
+
+# In-memory storage for the current plan (raw structure from agent)
+_current_plan: List[Dict[str, str]] = []
+
+
+class PlanTool:
+    """Tool for managing a list of todos with status tracking."""
+
+    def __init__(self):
+        pass
+
+    async def execute(self, params: Dict[str, Any]) -> ToolResult:
+        """
+        Execute the WritePlan operation.
+
+        Args:
+            params: Dictionary containing:
+                - todos: List of todo items, each with id, content, and status
+
+        Returns:
+            ToolResult with formatted output
+        """
+        global _current_plan
+
+        todos = params.get("todos", [])
+
+        # Validate todos structure
+        for todo in todos:
+            if not isinstance(todo, dict):
+                return {
+                    "formatted": "Error: Each todo must be an object. Re call the tool with correct format (mandatory).",
+                    "isError": True,
+                }
+
+            required_fields = ["id", "content", "status"]
+            for field in required_fields:
+                if field not in todo:
+                    return {
+                        "formatted": f"Error: Todo missing required field '{field}'. Re call the tool with correct format (mandatory).",
+                        "isError": True,
+                    }
+
+            # Validate status
+            valid_statuses = ["pending", "in_progress", "completed"]
+            if todo["status"] not in valid_statuses:
+                return {
+                    "formatted": f"Error: Invalid status '{todo['status']}'. Must be one of: {', '.join(valid_statuses)}. Re call the tool with correct format (mandatory).",
+                    "isError": True,
+                }
+
+        # Store the raw todos structure in memory
+        _current_plan = todos
+
+        # Format only for display using terminal_display utility
+        formatted_output = format_plan_tool_output(todos)
+
+        return {
+            "formatted": formatted_output,
+            "totalResults": len(todos),
+            "isError": False,
+        }
+
+
+def get_current_plan() -> List[Dict[str, str]]:
+    """Get the current plan (raw structure)."""
+    return _current_plan
+
+
+# Tool specification
+PLAN_TOOL_SPEC = {
+    "name": "plan_tool",
+    "description": "Manage a plan with a list of todos. Each call replaces the entire plan with the provided todos list.",
+    "parameters": {
+        "type": "object",
+        "properties": {
+            "todos": {
+                "type": "array",
+                "description": "List of todo items",
+                "items": {
+                    "type": "object",
+                    "properties": {
+                        "id": {
+                            "type": "string",
+                            "description": "Unique identifier for the todo"
+                        },
+                        "content": {
+                            "type": "string",
+                            "description": "Description of the todo task"
+                        },
+                        "status": {
+                            "type": "string",
+                            "enum": ["pending", "in_progress", "completed"],
+                            "description": "Current status of the todo"
+                        }
+                    },
+                    "required": ["id", "content", "status"]
+                }
+            }
+        },
+        "required": ["todos"]
+    }
+}
+
+
+async def plan_tool_handler(arguments: Dict[str, Any]) -> tuple[str, bool]:
+    tool = PlanTool()
+    result = await tool.execute(arguments)
+    return result["formatted"], not result.get("isError", False)

agent/utils/terminal_display.py

@@ -0,0 +1,143 @@
+"""
+Terminal display utilities with colors and formatting
+"""
+
+
+# ANSI color codes
+class Colors:
+    RED = "\033[91m"
+    GREEN = "\033[92m"
+    YELLOW = "\033[93m"
+    BLUE = "\033[94m"
+    MAGENTA = "\033[95m"
+    CYAN = "\033[96m"
+    BOLD = "\033[1m"
+    UNDERLINE = "\033[4m"
+    RESET = "\033[0m"
+
+
+def truncate_to_lines(text: str, max_lines: int = 6) -> str:
+    """Truncate text to max_lines, adding '...' if truncated"""
+    lines = text.split('\n')
+    if len(lines) <= max_lines:
+        return text
+    return '\n'.join(lines[:max_lines]) + f'\n{Colors.CYAN}... ({len(lines) - max_lines} more lines){Colors.RESET}'
+
+
+def format_header(text: str, emoji: str = "") -> str:
+    """Format a header with bold"""
+    full_text = f"{emoji} {text}" if emoji else text
+    return f"{Colors.BOLD}{full_text}{Colors.RESET}"
+
+
+def format_plan_display() -> str:
+    """Format the current plan for display (no colors, full visibility)"""
+    from agent.tools.plan_tool import get_current_plan
+
+    plan = get_current_plan()
+    if not plan:
+        return ""
+
+    lines = ["\n" + "=" * 60]
+    lines.append("CURRENT PLAN")
+    lines.append("=" * 60 + "\n")
+
+    # Group by status
+    completed = [t for t in plan if t["status"] == "completed"]
+    in_progress = [t for t in plan if t["status"] == "in_progress"]
+    pending = [t for t in plan if t["status"] == "pending"]
+
+    if completed:
+        lines.append("Completed:")
+        for todo in completed:
+            lines.append(f"  [x] {todo['id']}. {todo['content']}")
+        lines.append("")
+
+    if in_progress:
+        lines.append("In Progress:")
+        for todo in in_progress:
+            lines.append(f"  [~] {todo['id']}. {todo['content']}")
+        lines.append("")
+
+    if pending:
+        lines.append("Pending:")
+        for todo in pending:
+            lines.append(f"  [ ] {todo['id']}. {todo['content']}")
+        lines.append("")
+
+    lines.append(f"Total: {len(plan)} todos ({len(completed)} completed, {len(in_progress)} in progress, {len(pending)} pending)")
+    lines.append("=" * 60 + "\n")
+
+    return '\n'.join(lines)
+
+
+def format_error(message: str) -> str:
+    """Format an error message in red"""
+    return f"{Colors.RED}ERROR: {message}{Colors.RESET}"
+
+
+def format_success(message: str, emoji: str = "") -> str:
+    """Format a success message in green"""
+    prefix = f"{emoji} " if emoji else ""
+    return f"{Colors.GREEN}{prefix}{message}{Colors.RESET}"
+
+
+def format_tool_call(tool_name: str, arguments: str) -> str:
+    """Format a tool call message"""
+    return f"{Colors.YELLOW}Calling tool: {Colors.BOLD}{tool_name}{Colors.RESET}{Colors.YELLOW} with arguments: {arguments}{Colors.RESET}"
+
+
+def format_tool_output(output: str, success: bool, truncate: bool = True) -> str:
+    """Format tool output with color and optional truncation"""
+    if truncate:
+        output = truncate_to_lines(output, max_lines=6)
+
+    if success:
+        return f"{Colors.YELLOW}Tool output:{Colors.RESET}\n{output}"
+    else:
+        return f"{Colors.RED}Tool output:{Colors.RESET}\n{output}"
+
+
+def format_turn_complete() -> str:
+    """Format turn complete message in green with hugging face emoji"""
+    return f"{Colors.GREEN}{Colors.BOLD}\U0001F917 Turn complete{Colors.RESET}\n"
+
+
+def format_separator(char: str = "=", length: int = 60) -> str:
+    """Format a separator line"""
+    return char * length
+
+
+def format_plan_tool_output(todos: list) -> str:
+    """Format the plan tool output (no colors, full visibility)"""
+    if not todos:
+        return "Plan is empty."
+
+    lines = ["Plan updated successfully", ""]
+
+    # Group by status
+    completed = [t for t in todos if t["status"] == "completed"]
+    in_progress = [t for t in todos if t["status"] == "in_progress"]
+    pending = [t for t in todos if t["status"] == "pending"]
+
+    if completed:
+        lines.append("Completed:")
+        for todo in completed:
+            lines.append(f"  [x] {todo['id']}. {todo['content']}")
+        lines.append("")
+
+    if in_progress:
+        lines.append("In Progress:")
+        for todo in in_progress:
+            lines.append(f"  [~] {todo['id']}. {todo['content']}")
+        lines.append("")
+
+    if pending:
+        lines.append("Pending:")
+        for todo in pending:
+            lines.append(f"  [ ] {todo['id']}. {todo['content']}")
+        lines.append("")
+
+    lines.append(f"Total: {len(todos)} todos ({len(completed)} completed, {len(in_progress)} in progress, {len(pending)} pending)")
+
+    return "\n".join(lines)