utils.py 4.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128
  1. """
  2. Utility functions for MCPClient and Tiny Agents.
  3. Formatting utilities taken from the JS SDK: https://github.com/huggingface/huggingface.js/blob/main/packages/mcp-client/src/ResultFormatter.ts.
  4. """
  5. import json
  6. from pathlib import Path
  7. from typing import TYPE_CHECKING, List, Optional, Tuple
  8. from huggingface_hub import snapshot_download
  9. from huggingface_hub.errors import EntryNotFoundError
  10. from .constants import DEFAULT_AGENT, DEFAULT_REPO_ID, FILENAME_CONFIG, PROMPT_FILENAMES
  11. from .types import AgentConfig
  12. if TYPE_CHECKING:
  13. from mcp import types as mcp_types
  14. def format_result(result: "mcp_types.CallToolResult") -> str:
  15. """
  16. Formats a mcp.types.CallToolResult content into a human-readable string.
  17. Args:
  18. result (CallToolResult)
  19. Object returned by mcp.ClientSession.call_tool.
  20. Returns:
  21. str
  22. A formatted string representing the content of the result.
  23. """
  24. content = result.content
  25. if len(content) == 0:
  26. return "[No content]"
  27. formatted_parts: List[str] = []
  28. for item in content:
  29. if item.type == "text":
  30. formatted_parts.append(item.text)
  31. elif item.type == "image":
  32. formatted_parts.append(
  33. f"[Binary Content: Image {item.mimeType}, {_get_base64_size(item.data)} bytes]\n"
  34. f"The task is complete and the content accessible to the User"
  35. )
  36. elif item.type == "audio":
  37. formatted_parts.append(
  38. f"[Binary Content: Audio {item.mimeType}, {_get_base64_size(item.data)} bytes]\n"
  39. f"The task is complete and the content accessible to the User"
  40. )
  41. elif item.type == "resource":
  42. resource = item.resource
  43. if hasattr(resource, "text"):
  44. formatted_parts.append(resource.text)
  45. elif hasattr(resource, "blob"):
  46. formatted_parts.append(
  47. f"[Binary Content ({resource.uri}): {resource.mimeType}, {_get_base64_size(resource.blob)} bytes]\n"
  48. f"The task is complete and the content accessible to the User"
  49. )
  50. return "\n".join(formatted_parts)
  51. def _get_base64_size(base64_str: str) -> int:
  52. """Estimate the byte size of a base64-encoded string."""
  53. # Remove any prefix like "data:image/png;base64,"
  54. if "," in base64_str:
  55. base64_str = base64_str.split(",")[1]
  56. padding = 0
  57. if base64_str.endswith("=="):
  58. padding = 2
  59. elif base64_str.endswith("="):
  60. padding = 1
  61. return (len(base64_str) * 3) // 4 - padding
  62. def _load_agent_config(agent_path: Optional[str]) -> Tuple[AgentConfig, Optional[str]]:
  63. """Load server config and prompt."""
  64. def _read_dir(directory: Path) -> Tuple[AgentConfig, Optional[str]]:
  65. cfg_file = directory / FILENAME_CONFIG
  66. if not cfg_file.exists():
  67. raise FileNotFoundError(f" Config file not found in {directory}! Please make sure it exists locally")
  68. config: AgentConfig = json.loads(cfg_file.read_text(encoding="utf-8"))
  69. prompt: Optional[str] = None
  70. for filename in PROMPT_FILENAMES:
  71. prompt_file = directory / filename
  72. if prompt_file.exists():
  73. prompt = prompt_file.read_text(encoding="utf-8")
  74. break
  75. return config, prompt
  76. if agent_path is None:
  77. return DEFAULT_AGENT, None # type: ignore[return-value]
  78. path = Path(agent_path).expanduser()
  79. if path.is_file():
  80. return json.loads(path.read_text(encoding="utf-8")), None
  81. if path.is_dir():
  82. return _read_dir(path)
  83. # fetch from the Hub
  84. try:
  85. repo_dir = Path(
  86. snapshot_download(
  87. repo_id=DEFAULT_REPO_ID,
  88. allow_patterns=f"{agent_path}/*",
  89. repo_type="dataset",
  90. )
  91. )
  92. return _read_dir(repo_dir / agent_path)
  93. except Exception as err:
  94. raise EntryNotFoundError(
  95. 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."
  96. ) from err