akseljoonas HF Staff Claude Opus 4.6 commited on
Commit
7edb225
·
1 Parent(s): 085cd02

feat: prompt for HF token interactively when not found

Browse files

Instead of printing a warning and continuing without a token,
prompt the user to paste their token, validate it via the API,
and persist it via huggingface_hub.login().

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Files changed (1) hide show
  1. agent/main.py +48 -6
agent/main.py CHANGED
@@ -51,7 +51,7 @@ def _safe_get_args(arguments: dict) -> dict:
51
 
52
 
53
  def _get_hf_token() -> str | None:
54
- """Get HF token from environment or huggingface_hub login."""
55
  token = os.environ.get("HF_TOKEN")
56
  if token:
57
  return token
@@ -65,6 +65,48 @@ def _get_hf_token() -> str | None:
65
  pass
66
  return None
67
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
68
  @dataclass
69
  class Operation:
70
  """Operation to be executed by the agent"""
@@ -584,12 +626,15 @@ async def main():
584
  # Wait for agent to initialize
585
  print("Initializing agent...")
586
 
587
- # HF token
 
 
 
588
  hf_token = _get_hf_token()
589
  if hf_token:
590
  print("HF token loaded")
591
  else:
592
- print("Warning: No HF token found. Set HF_TOKEN or run `huggingface-cli login`.")
593
 
594
  # Create queues for communication
595
  submission_queue = asyncio.Queue()
@@ -608,9 +653,6 @@ async def main():
608
  print(f"Loading MCP servers: {', '.join(config.mcpServers.keys())}")
609
  tool_router = ToolRouter(config.mcpServers, hf_token=hf_token, local_mode=True)
610
 
611
- # Create prompt session for input
612
- prompt_session = PromptSession()
613
-
614
  # Session holder for interrupt/model/status access
615
  session_holder = [None]
616
 
 
51
 
52
 
53
  def _get_hf_token() -> str | None:
54
+ """Get HF token from environment or huggingface_hub cached login."""
55
  token = os.environ.get("HF_TOKEN")
56
  if token:
57
  return token
 
65
  pass
66
  return None
67
 
68
+
69
+ async def _prompt_and_save_hf_token(prompt_session: PromptSession) -> str:
70
+ """Prompt user for HF token, validate it, save via huggingface_hub.login(). Loops until valid."""
71
+ from prompt_toolkit.formatted_text import HTML
72
+ from huggingface_hub import HfApi, login
73
+
74
+ print("\nA Hugging Face token is required.")
75
+ print("Get one at: https://huggingface.co/settings/tokens\n")
76
+
77
+ while True:
78
+ try:
79
+ token = await prompt_session.prompt_async(
80
+ HTML("<b>Paste your HF token: </b>")
81
+ )
82
+ except (EOFError, KeyboardInterrupt):
83
+ print("\nToken is required to continue.")
84
+ continue
85
+
86
+ token = token.strip()
87
+ if not token:
88
+ print("Token cannot be empty.")
89
+ continue
90
+
91
+ # Validate token against the API
92
+ try:
93
+ api = HfApi(token=token)
94
+ user_info = api.whoami()
95
+ username = user_info.get("name", "unknown")
96
+ print(f"Token valid (user: {username})")
97
+ except Exception:
98
+ print("Invalid token. Please try again.")
99
+ continue
100
+
101
+ # Save for future sessions
102
+ try:
103
+ login(token=token, add_to_git_credential=False)
104
+ print("Token saved to ~/.cache/huggingface/token")
105
+ except Exception as e:
106
+ print(f"Warning: could not persist token ({e}), using for this session only.")
107
+
108
+ return token
109
+
110
  @dataclass
111
  class Operation:
112
  """Operation to be executed by the agent"""
 
626
  # Wait for agent to initialize
627
  print("Initializing agent...")
628
 
629
+ # Create prompt session for input (needed early for token prompt)
630
+ prompt_session = PromptSession()
631
+
632
+ # HF token — required, prompt if missing
633
  hf_token = _get_hf_token()
634
  if hf_token:
635
  print("HF token loaded")
636
  else:
637
+ hf_token = await _prompt_and_save_hf_token(prompt_session)
638
 
639
  # Create queues for communication
640
  submission_queue = asyncio.Queue()
 
653
  print(f"Loading MCP servers: {', '.join(config.mcpServers.keys())}")
654
  tool_router = ToolRouter(config.mcpServers, hf_token=hf_token, local_mode=True)
655
 
 
 
 
656
  # Session holder for interrupt/model/status access
657
  session_holder = [None]
658