File size: 6,005 Bytes
955ef08 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 |
# --------------------------------------------------------------------------------------
# 1. Imports and Configuration
# ---------------------------------------------------------------------------------------
import os
import json
from fastmcp import FastMCP
from huggingface_hub import HfApi, model_info, ModelCard, ModelCardData
from huggingface_hub.utils import HfHubHTTPError
from huggingface_hub import CommitOperationAdd
import traceback
from dotenv import load_dotenv
load_dotenv()
HF_TOKEN=os.getenv("HF_TOKEN")
hf_api=HfApi(token=HF_TOKEN) if HF_TOKEN else None
mcp=FastMCP("hf-tagging-bot")
# -------------------------------------------------------------------------------
# 2. Get Current Tags Tool
# -------------------------------------------------------------------------------
#@mcp.tool(name="get_current_tags")
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":"success",
"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
# --------------------------------------------------------------------------------------------------------
# 3. 3. Add New Tag Tool
# --------------------------------------------------------------------------------------------------------
#@mcp.tool(name="add_new_tag")
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:
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}")
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
updated_tags=current_tags+[new_tag]
print(f"🆕 will update tags from {current_tags} to {updated_tags}")
try:
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:
print(f"📄 creating new model card (none exists)")
card=ModelCard("")
card.data=ModelCardData()
card_dict=card.data.to_dict()
card_dict['tags']=updated_tags
card.data=ModelCardData(**card_dict)
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)}
🤖This is a pull request created by the HuggingFace Hub Tagging Bot.
"""
print(f"🚀 creating PR with title: {pr_title}")
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
)
pr_url_attr=commit_info.pr_url
pr_url=pr_url_attr if hasattr(commit_info,"prl_url") else str(commit_info)
print(f"✅ PR created successfully! URL: {pr_url}")
result={
"status":"success",
"repo_id":repo_id,
"tag":new_tag,
"prl_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)}")
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()
|