diff --git a/setup.py b/setup.py index b44ccfc7cc..3a32c619a7 100644 --- a/setup.py +++ b/setup.py @@ -66,6 +66,7 @@ def get_version() -> str: extras["mcp"] = [ "mcp>=1.8.0", + "typer", ] + extras["inference"] extras["testing"] = ( @@ -130,7 +131,10 @@ def get_version() -> str: packages=find_packages("src"), extras_require=extras, entry_points={ - "console_scripts": ["huggingface-cli=huggingface_hub.commands.huggingface_cli:main"], + "console_scripts": [ + "huggingface-cli=huggingface_hub.commands.huggingface_cli:main", + "tiny-agents=huggingface_hub.inference._mcp.cli:app", + ], "fsspec.specs": "hf=huggingface_hub.HfFileSystem", }, python_requires=">=3.8.0", diff --git a/src/huggingface_hub/inference/_mcp/agent.py b/src/huggingface_hub/inference/_mcp/agent.py new file mode 100644 index 0000000000..387ca178f1 --- /dev/null +++ b/src/huggingface_hub/inference/_mcp/agent.py @@ -0,0 +1,71 @@ +from __future__ import annotations + +import asyncio +from typing import AsyncGenerator, Dict, Iterable, List, Optional, Union + +from huggingface_hub import ChatCompletionInputMessage, ChatCompletionStreamOutput, MCPClient + +from .._providers import PROVIDER_OR_POLICY_T +from .constants import DEFAULT_SYSTEM_PROMPT, EXIT_LOOP_TOOLS, MAX_NUM_TURNS + + +class Agent(MCPClient): + """ + Python implementation of a Simple Agent + i.e. just a basic while loop on top of an Inference Client with MCP-powered tools + """ + + def __init__( + self, + *, + model: str, + servers: Iterable[Dict], + provider: Optional[PROVIDER_OR_POLICY_T] = None, + api_key: Optional[str] = None, + prompt: Optional[str] = None, + ): + super().__init__(model=model, provider=provider, api_key=api_key) + self._servers_cfg = list(servers) + self.messages: List[Union[Dict, ChatCompletionInputMessage]] = [ + {"role": "system", "content": prompt or DEFAULT_SYSTEM_PROMPT} + ] + + async def load_tools(self) -> None: + for cfg in self._servers_cfg: + await self.add_mcp_server(cfg["type"], **cfg["config"]) + + async def run( + self, + user_input: str, + *, + abort_event: Optional[asyncio.Event] = None, + ) -> AsyncGenerator[Union[ChatCompletionStreamOutput, ChatCompletionInputMessage], None]: + self.messages.append({"role": "user", "content": user_input}) + + num_turns: int = 0 + next_turn_should_call_tools = True + + while True: + if abort_event and abort_event.is_set(): + return + + async for item in self.process_single_turn_with_tools( + self.messages, + exit_loop_tools=EXIT_LOOP_TOOLS, + exit_if_first_chunk_no_tool=(num_turns > 0 and next_turn_should_call_tools), + ): + yield item + + num_turns += 1 + last = self.messages[-1] + + if last.get("role") == "tool" and last.get("name") in {t.function.name for t in EXIT_LOOP_TOOLS}: + return + + if last.get("role") != "tool" and num_turns > MAX_NUM_TURNS: + return + + if last.get("role") != "tool" and next_turn_should_call_tools: + return + + next_turn_should_call_tools = last.get("role") != "tool" diff --git a/src/huggingface_hub/inference/_mcp/cli.py b/src/huggingface_hub/inference/_mcp/cli.py new file mode 100644 index 0000000000..fd9d5aad6c --- /dev/null +++ b/src/huggingface_hub/inference/_mcp/cli.py @@ -0,0 +1,153 @@ +import asyncio +import os +import signal +from functools import partial +from typing import Any, Dict, List, Optional + +import typer +from rich import print + +from .agent import Agent +from .utils import _load_agent_config + + +app = typer.Typer( + rich_markup_mode="rich", + help="A squad of lightweight composable AI applications built on Hugging Face's Inference Client and MCP stack.", +) + +run_cli = typer.Typer( + name="run", + help="Run the Agent in the CLI", + invoke_without_command=True, +) +app.add_typer(run_cli, name="run") + + +async def _ainput(prompt: str = "» ") -> str: + loop = asyncio.get_running_loop() + return await loop.run_in_executor(None, partial(typer.prompt, prompt, prompt_suffix=" ")) + + +async def run_agent( + agent_path: Optional[str], +) -> None: + """ + Tiny Agent loop. + + Args: + agent_path (`str`, *optional*): + Path to a local folder containing an `agent.json` and optionally a custom `PROMPT.md` file or a built-in agent stored in a Hugging Face dataset. + + """ + config, prompt = _load_agent_config(agent_path) + + servers: List[Dict[str, Any]] = config.get("servers", []) + + abort_event = asyncio.Event() + first_sigint = True + + loop = asyncio.get_running_loop() + original_sigint_handler = signal.getsignal(signal.SIGINT) + + def _sigint_handler() -> None: + nonlocal first_sigint + if first_sigint: + first_sigint = False + abort_event.set() + print("\n[red]Interrupted. Press Ctrl+C again to quit.[/red]", flush=True) + return + + print("\n[red]Exiting...[/red]", flush=True) + + os._exit(130) + + try: + loop.add_signal_handler(signal.SIGINT, _sigint_handler) + + async with Agent( + provider=config["provider"], + model=config["model"], + servers=servers, + prompt=prompt, + ) as agent: + await agent.load_tools() + print(f"[bold blue]Agent loaded with {len(agent.available_tools)} tools:[/bold blue]") + for t in agent.available_tools: + print(f"[blue] • {t.function.name}[/blue]") + + while True: + abort_event.clear() + + try: + user_input = await _ainput() + first_sigint = True + except EOFError: + print("\n[red]EOF received, exiting.[/red]", flush=True) + break + except KeyboardInterrupt: + if not first_sigint and abort_event.is_set(): + continue + else: + print("\n[red]Keyboard interrupt during input processing.[/red]", flush=True) + break + + try: + async for chunk in agent.run(user_input, abort_event=abort_event): + if abort_event.is_set() and not first_sigint: + break + + if hasattr(chunk, "choices"): + delta = chunk.choices[0].delta + if delta.content: + print(delta.content, end="", flush=True) + if delta.tool_calls: + for call in delta.tool_calls: + if call.id: + print(f"", end="") + if call.function.name: + print(f"{call.function.name}", end=" ") + if call.function.arguments: + print(f"{call.function.arguments}", end="") + else: + print( + f"\n\n[green]Tool[{chunk.name}] {chunk.tool_call_id}\n{chunk.content}[/green]\n", + flush=True, + ) + + print() + + except Exception as e: + print(f"\n[bold red]Error during agent run: {e}[/bold red]", flush=True) + first_sigint = True # Allow graceful interrupt for the next command + + finally: + if loop and not loop.is_closed(): + loop.remove_signal_handler(signal.SIGINT) + elif original_sigint_handler: + signal.signal(signal.SIGINT, original_sigint_handler) + + +@run_cli.callback() +def run( + path: Optional[str] = typer.Argument( + None, + help=( + "Path to a local folder containing an agent.json file or a built-in agent " + "stored in a Hugging Face dataset (default: " + "https://huggingface.co/datasets/tiny-agents/tiny-agents)" + ), + ), +): + try: + asyncio.run(run_agent(path)) + except KeyboardInterrupt: + print("\n[red]Application terminated by KeyboardInterrupt.[/red]", flush=True) + raise typer.Exit(code=130) + except Exception as e: + print(f"\n[bold red]An unexpected error occurred: {e}[/bold red]", flush=True) + raise e + + +if __name__ == "__main__": + app() diff --git a/src/huggingface_hub/inference/_mcp/constants.py b/src/huggingface_hub/inference/_mcp/constants.py new file mode 100644 index 0000000000..dc5eb38ebf --- /dev/null +++ b/src/huggingface_hub/inference/_mcp/constants.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +import sys +from pathlib import Path +from typing import List + +from huggingface_hub import ChatCompletionInputTool + + +FILENAME_CONFIG = "agent.json" +FILENAME_PROMPT = "PROMPT.md" + +DEFAULT_AGENT = { + "model": "Qwen/Qwen2.5-72B-Instruct", + "provider": "nebius", + "servers": [ + { + "type": "stdio", + "config": { + "command": "npx", + "args": [ + "-y", + "@modelcontextprotocol/server-filesystem", + str(Path.home() / ("Desktop" if sys.platform == "darwin" else "")), + ], + }, + }, + { + "type": "stdio", + "config": { + "command": "npx", + "args": ["@playwright/mcp@latest"], + }, + }, + ], +} + + +DEFAULT_SYSTEM_PROMPT = """ +You are an agent - please keep going until the user’s query is completely +resolved, before ending your turn and yielding back to the user. Only terminate +your turn when you are sure that the problem is solved, or if you need more +info from the user to solve the problem. +If you are not sure about anything pertaining to the user’s request, use your +tools to read files and gather the relevant information: do NOT guess or make +up an answer. +You MUST plan extensively before each function call, and reflect extensively +on the outcomes of the previous function calls. DO NOT do this entire process +by making function calls only, as this can impair your ability to solve the +problem and think insightfully. +""".strip() + +MAX_NUM_TURNS = 10 + +TASK_COMPLETE_TOOL: ChatCompletionInputTool = ChatCompletionInputTool.parse_obj( # type: ignore[assignment] + { + "type": "function", + "function": { + "name": "task_complete", + "description": "Call this tool when the task given by the user is complete", + "parameters": {"type": "object", "properties": {}}, + }, + } +) + +ASK_QUESTION_TOOL: ChatCompletionInputTool = ChatCompletionInputTool.parse_obj( # type: ignore[assignment] + { + "type": "function", + "function": { + "name": "ask_question", + "description": "Ask the user for more info required to solve or clarify their problem.", + "parameters": {"type": "object", "properties": {}}, + }, + } +) + +EXIT_LOOP_TOOLS: List[ChatCompletionInputTool] = [TASK_COMPLETE_TOOL, ASK_QUESTION_TOOL] + + +DEFAULT_REPO_ID = "tiny-agents/tiny-agents" diff --git a/src/huggingface_hub/inference/_mcp/utils.py b/src/huggingface_hub/inference/_mcp/utils.py index 20a39f6897..0490eb600d 100644 --- a/src/huggingface_hub/inference/_mcp/utils.py +++ b/src/huggingface_hub/inference/_mcp/utils.py @@ -1,10 +1,17 @@ """ -Utility functions for formatting results from mcp.CallToolResult. +Utility functions for MCPClient and Tiny Agents. -Taken from the JS SDK: https://github.com/huggingface/huggingface.js/blob/main/packages/mcp-client/src/ResultFormatter.ts. +Formatting utilities taken from the JS SDK: https://github.com/huggingface/huggingface.js/blob/main/packages/mcp-client/src/ResultFormatter.ts. """ -from typing import TYPE_CHECKING, List +import json +from pathlib import Path +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple + +from huggingface_hub import snapshot_download +from huggingface_hub.errors import EntryNotFoundError + +from .constants import DEFAULT_AGENT, DEFAULT_REPO_ID, FILENAME_CONFIG, FILENAME_PROMPT if TYPE_CHECKING: @@ -74,3 +81,43 @@ def _get_base64_size(base64_str: str) -> int: padding = 1 return (len(base64_str) * 3) // 4 - padding + + +def _load_agent_config(agent_path: Optional[str]) -> Tuple[Dict[str, Any], Optional[str]]: + """Load server config and prompt.""" + + def _read_dir(directory: Path) -> Tuple[Dict[str, Any], Optional[str]]: + cfg_file = directory / FILENAME_CONFIG + if not cfg_file.exists(): + raise FileNotFoundError(f" Config file not found in {directory}! Please make sure it exists locally") + + config: Dict[str, Any] = json.loads(cfg_file.read_text(encoding="utf-8")) + prompt_file = directory / FILENAME_PROMPT + prompt: Optional[str] = prompt_file.read_text(encoding="utf-8") if prompt_file.exists() else None + return config, prompt + + if agent_path is None: + return DEFAULT_AGENT, None + + path = Path(agent_path).expanduser() + + if path.is_file(): + return json.loads(path.read_text(encoding="utf-8")), None + + if path.is_dir(): + return _read_dir(path) + + # fetch from the Hub + try: + repo_dir = Path( + snapshot_download( + repo_id=DEFAULT_REPO_ID, + allow_patterns=f"{agent_path}/*", + repo_type="dataset", + ) + ) + return _read_dir(repo_dir / agent_path) + except Exception as err: + raise EntryNotFoundError( + f" Agent {agent_path} not found in tiny-agents/tiny-agents! Please make sure it exists in https://huggingface.co/datasets/tiny-agents/tiny-agents." + ) from err