Spaces:
Running on CPU Upgrade
Running on CPU Upgrade
Commit ·
7edb225
1
Parent(s): 085cd02
feat: prompt for HF token interactively when not found
Browse filesInstead 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>
- 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 |
-
#
|
|
|
|
|
|
|
|
|
|
| 588 |
hf_token = _get_hf_token()
|
| 589 |
if hf_token:
|
| 590 |
print("HF token loaded")
|
| 591 |
else:
|
| 592 |
-
|
| 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 |
|