Spaces:
Running
Running
#!/usr/bin/env python3 | |
""" | |
Simplified MCP Server for HuggingFace Hub Tagging Operations using FastMCP | |
""" | |
import os | |
import json | |
from fastmcp import FastMCP | |
from huggingface_hub import HfApi, model_info, ModelCard, ModelCardData | |
from huggingface_hub.utils import HfHubHTTPError | |
from dotenv import load_dotenv | |
load_dotenv() | |
# Configuration | |
HF_TOKEN = os.getenv("HF_TOKEN") | |
# Initialize HF API client | |
hf_api = HfApi(token=HF_TOKEN) if HF_TOKEN else None | |
# Create the FastMCP server | |
mcp = FastMCP("hf-tagging-bot") | |
def get_current_tags(repo_id: str) -> str: | |
"""Get current tags from a HuggingFace model repository""" | |
print(f"π§ get_current_tags called with repo_id: {repo_id}") | |
if not hf_api: | |
error_result = {"error": "HF token not configured"} | |
json_str = json.dumps(error_result) | |
print(f"β No HF API token - returning: {json_str}") | |
return json_str | |
try: | |
print(f"π‘ Fetching model info for: {repo_id}") | |
info = model_info(repo_id=repo_id, token=HF_TOKEN) | |
current_tags = info.tags if info.tags else [] | |
print(f"π·οΈ Found {len(current_tags)} tags: {current_tags}") | |
result = { | |
"status": "success", | |
"repo_id": repo_id, | |
"current_tags": current_tags, | |
"count": len(current_tags), | |
} | |
json_str = json.dumps(result) | |
print(f"β get_current_tags returning: {json_str}") | |
return json_str | |
except Exception as e: | |
print(f"β Error in get_current_tags: {str(e)}") | |
error_result = {"status": "error", "repo_id": repo_id, "error": str(e)} | |
json_str = json.dumps(error_result) | |
print(f"β get_current_tags error returning: {json_str}") | |
return json_str | |
def add_new_tag(repo_id: str, new_tag: str) -> str: | |
"""Add a new tag to a HuggingFace model repository via PR""" | |
print(f"π§ add_new_tag called with repo_id: {repo_id}, new_tag: {new_tag}") | |
if not hf_api: | |
error_result = {"error": "HF token not configured"} | |
json_str = json.dumps(error_result) | |
print(f"β No HF API token - returning: {json_str}") | |
return json_str | |
try: | |
# Get current model info and tags | |
print(f"π‘ Fetching current model info for: {repo_id}") | |
info = model_info(repo_id=repo_id, token=HF_TOKEN) | |
current_tags = info.tags if info.tags else [] | |
print(f"π·οΈ Current tags: {current_tags}") | |
# Check if tag already exists | |
if new_tag in current_tags: | |
print(f"β οΈ Tag '{new_tag}' already exists in {current_tags}") | |
result = { | |
"status": "already_exists", | |
"repo_id": repo_id, | |
"tag": new_tag, | |
"message": f"Tag '{new_tag}' already exists", | |
} | |
json_str = json.dumps(result) | |
print(f"π·οΈ add_new_tag (already exists) returning: {json_str}") | |
return json_str | |
# Add the new tag to existing tags | |
updated_tags = current_tags + [new_tag] | |
print(f"π Will update tags from {current_tags} to {updated_tags}") | |
# Create model card content with updated tags | |
try: | |
# Load existing model card | |
print(f"π Loading existing model card...") | |
card = ModelCard.load(repo_id, token=HF_TOKEN) | |
if not hasattr(card, "data") or card.data is None: | |
card.data = ModelCardData() | |
except HfHubHTTPError: | |
# Create new model card if none exists | |
print(f"π Creating new model card (none exists)") | |
card = ModelCard("") | |
card.data = ModelCardData() | |
# Update tags - create new ModelCardData with updated tags | |
card_dict = card.data.to_dict() | |
card_dict["tags"] = updated_tags | |
card.data = ModelCardData(**card_dict) | |
# Create a pull request with the updated model card | |
pr_title = f"Add '{new_tag}' tag" | |
pr_description = f""" | |
## Add tag: {new_tag} | |
This PR adds the `{new_tag}` tag to the model repository. | |
**Changes:** | |
- Added `{new_tag}` to model tags | |
- Updated from {len(current_tags)} to {len(updated_tags)} tags | |
**Current tags:** {", ".join(current_tags) if current_tags else "None"} | |
**New tags:** {", ".join(updated_tags)} | |
""" | |
print(f"π Creating PR with title: {pr_title}") | |
# Create commit with updated model card using CommitOperationAdd | |
from huggingface_hub import CommitOperationAdd | |
commit_info = hf_api.create_commit( | |
repo_id=repo_id, | |
operations=[ | |
CommitOperationAdd( | |
path_in_repo="README.md", path_or_fileobj=str(card).encode("utf-8") | |
) | |
], | |
commit_message=pr_title, | |
commit_description=pr_description, | |
token=HF_TOKEN, | |
create_pr=True, | |
) | |
# Extract PR URL from commit info | |
pr_url_attr = commit_info.pr_url | |
pr_url = pr_url_attr if hasattr(commit_info, "pr_url") else str(commit_info) | |
print(f"β PR created successfully! URL: {pr_url}") | |
result = { | |
"status": "success", | |
"repo_id": repo_id, | |
"tag": new_tag, | |
"pr_url": pr_url, | |
"previous_tags": current_tags, | |
"new_tags": updated_tags, | |
"message": f"Created PR to add tag '{new_tag}'", | |
} | |
json_str = json.dumps(result) | |
print(f"β add_new_tag success returning: {json_str}") | |
return json_str | |
except Exception as e: | |
print(f"β Error in add_new_tag: {str(e)}") | |
print(f"β Error type: {type(e)}") | |
import traceback | |
print(f"β Traceback: {traceback.format_exc()}") | |
error_result = { | |
"status": "error", | |
"repo_id": repo_id, | |
"tag": new_tag, | |
"error": str(e), | |
} | |
json_str = json.dumps(error_result) | |
print(f"β add_new_tag error returning: {json_str}") | |
return json_str | |
if __name__ == "__main__": | |
mcp.run() | |