124 lines
4 KiB
Python
124 lines
4 KiB
Python
"""
|
|
Utility functions for MCPClient and Tiny Agents.
|
|
|
|
Formatting utilities taken from the JS SDK: https://github.com/huggingface/huggingface.js/blob/main/packages/mcp-client/src/ResultFormatter.ts.
|
|
"""
|
|
|
|
import json
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, 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
|
|
from .types import AgentConfig
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
from mcp import types as mcp_types
|
|
|
|
|
|
def format_result(result: "mcp_types.CallToolResult") -> str:
|
|
"""
|
|
Formats a mcp.types.CallToolResult content into a human-readable string.
|
|
|
|
Args:
|
|
result (CallToolResult)
|
|
Object returned by mcp.ClientSession.call_tool.
|
|
|
|
Returns:
|
|
str
|
|
A formatted string representing the content of the result.
|
|
"""
|
|
content = result.content
|
|
|
|
if len(content) == 0:
|
|
return "[No content]"
|
|
|
|
formatted_parts: List[str] = []
|
|
|
|
for item in content:
|
|
if item.type == "text":
|
|
formatted_parts.append(item.text)
|
|
|
|
elif item.type == "image":
|
|
formatted_parts.append(
|
|
f"[Binary Content: Image {item.mimeType}, {_get_base64_size(item.data)} bytes]\n"
|
|
f"The task is complete and the content accessible to the User"
|
|
)
|
|
|
|
elif item.type == "audio":
|
|
formatted_parts.append(
|
|
f"[Binary Content: Audio {item.mimeType}, {_get_base64_size(item.data)} bytes]\n"
|
|
f"The task is complete and the content accessible to the User"
|
|
)
|
|
|
|
elif item.type == "resource":
|
|
resource = item.resource
|
|
|
|
if hasattr(resource, "text"):
|
|
formatted_parts.append(resource.text)
|
|
|
|
elif hasattr(resource, "blob"):
|
|
formatted_parts.append(
|
|
f"[Binary Content ({resource.uri}): {resource.mimeType}, {_get_base64_size(resource.blob)} bytes]\n"
|
|
f"The task is complete and the content accessible to the User"
|
|
)
|
|
|
|
return "\n".join(formatted_parts)
|
|
|
|
|
|
def _get_base64_size(base64_str: str) -> int:
|
|
"""Estimate the byte size of a base64-encoded string."""
|
|
# Remove any prefix like "data:image/png;base64,"
|
|
if "," in base64_str:
|
|
base64_str = base64_str.split(",")[1]
|
|
|
|
padding = 0
|
|
if base64_str.endswith("=="):
|
|
padding = 2
|
|
elif base64_str.endswith("="):
|
|
padding = 1
|
|
|
|
return (len(base64_str) * 3) // 4 - padding
|
|
|
|
|
|
def _load_agent_config(agent_path: Optional[str]) -> Tuple[AgentConfig, Optional[str]]:
|
|
"""Load server config and prompt."""
|
|
|
|
def _read_dir(directory: Path) -> Tuple[AgentConfig, 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: AgentConfig = 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 # type: ignore[return-value]
|
|
|
|
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
|