kshitijthakkar commited on
Commit
16e44ce
·
1 Parent(s): 5cc4a9e

fix: Make HF token optional in push_dataset_to_hub tool

Browse files

- Changed hf_token parameter from required to optional (default: None)
- Auto-detect HF_TOKEN from environment variables if not provided
- Updated docstring to clarify token can come from Settings/env
- Return clear error message if no token available from either source
- Fixed all token references in function to use the resolved token variable

Users can now:
1. Provide token directly via parameter
2. Set token in Settings tab (env var)
3. Set HF_TOKEN environment variable

Resolves UX issue where token had to be provided every time despite being saved in Settings.

Files changed (1) hide show
  1. mcp_tools.py +16 -5
mcp_tools.py CHANGED
@@ -1769,7 +1769,7 @@ dataset.push_to_hub("your-username/smoltrace-{domain.lower()}-tasks")
1769
  async def push_dataset_to_hub(
1770
  dataset_json: str,
1771
  repo_name: str,
1772
- hf_token: str,
1773
  private: bool = False,
1774
  prompt_template: str = None
1775
  ) -> str:
@@ -1784,12 +1784,13 @@ async def push_dataset_to_hub(
1784
  - Format: {username}/smoltrace-{domain}-tasks or {username}/smoltrace-{domain}-tasks-v{version}
1785
  - Examples: "mycompany/smoltrace-finance-tasks", "alice/smoltrace-healthcare-tasks-v2"
1786
 
1787
- **Security**: Requires valid HuggingFace token with write permissions.
 
1788
 
1789
  Args:
1790
  dataset_json (str): JSON string containing the tasks array (from generate_synthetic_dataset output, use the "tasks" field)
1791
  repo_name (str): HuggingFace repository name following SMOLTRACE naming: {username}/smoltrace-{domain}-tasks
1792
- hf_token (str): HuggingFace API token with write permissions (get from https://huggingface.co/settings/tokens)
1793
  private (bool): Whether to create a private dataset. Default: False (public)
1794
  prompt_template (str): Optional YAML prompt template to include in dataset card (from generate_prompt_template)
1795
 
@@ -1797,8 +1798,18 @@ async def push_dataset_to_hub(
1797
  str: JSON response with upload status, dataset URL, and next steps
1798
  """
1799
  try:
 
1800
  from huggingface_hub import HfApi
1801
 
 
 
 
 
 
 
 
 
 
1802
  # Validate repo name follows SMOLTRACE convention
1803
  if "smoltrace-" not in repo_name and "-tasks" not in repo_name:
1804
  return json.dumps({
@@ -1843,7 +1854,7 @@ async def push_dataset_to_hub(
1843
  # Push to hub
1844
  dataset.push_to_hub(
1845
  repo_name,
1846
- token=hf_token,
1847
  private=private
1848
  )
1849
 
@@ -1931,7 +1942,7 @@ Part of the MCP's 1st Birthday Hackathon project.
1931
  path_in_repo="README.md",
1932
  repo_id=repo_name,
1933
  repo_type="dataset",
1934
- token=hf_token
1935
  )
1936
 
1937
  print(f"[PUSH_DATASET_TO_HUB] Prompt template added to dataset card successfully")
 
1769
  async def push_dataset_to_hub(
1770
  dataset_json: str,
1771
  repo_name: str,
1772
+ hf_token: str = None,
1773
  private: bool = False,
1774
  prompt_template: str = None
1775
  ) -> str:
 
1784
  - Format: {username}/smoltrace-{domain}-tasks or {username}/smoltrace-{domain}-tasks-v{version}
1785
  - Examples: "mycompany/smoltrace-finance-tasks", "alice/smoltrace-healthcare-tasks-v2"
1786
 
1787
+ **Security**: Requires valid HuggingFace token with write permissions. If not provided,
1788
+ will use HF_TOKEN from environment variables or Settings.
1789
 
1790
  Args:
1791
  dataset_json (str): JSON string containing the tasks array (from generate_synthetic_dataset output, use the "tasks" field)
1792
  repo_name (str): HuggingFace repository name following SMOLTRACE naming: {username}/smoltrace-{domain}-tasks
1793
+ hf_token (str): HuggingFace API token with write permissions (optional - uses HF_TOKEN env var if not provided)
1794
  private (bool): Whether to create a private dataset. Default: False (public)
1795
  prompt_template (str): Optional YAML prompt template to include in dataset card (from generate_prompt_template)
1796
 
 
1798
  str: JSON response with upload status, dataset URL, and next steps
1799
  """
1800
  try:
1801
+ import os
1802
  from huggingface_hub import HfApi
1803
 
1804
+ # Use provided token or fallback to environment variable
1805
+ token = hf_token or os.environ.get("HF_TOKEN")
1806
+ if not token:
1807
+ return json.dumps({
1808
+ "error": "HuggingFace token required",
1809
+ "message": "Please provide hf_token parameter or set HF_TOKEN environment variable in Settings",
1810
+ "get_token": "https://huggingface.co/settings/tokens"
1811
+ }, indent=2)
1812
+
1813
  # Validate repo name follows SMOLTRACE convention
1814
  if "smoltrace-" not in repo_name and "-tasks" not in repo_name:
1815
  return json.dumps({
 
1854
  # Push to hub
1855
  dataset.push_to_hub(
1856
  repo_name,
1857
+ token=token,
1858
  private=private
1859
  )
1860
 
 
1942
  path_in_repo="README.md",
1943
  repo_id=repo_name,
1944
  repo_type="dataset",
1945
+ token=token
1946
  )
1947
 
1948
  print(f"[PUSH_DATASET_TO_HUB] Prompt template added to dataset card successfully")