Spaces:
Runtime error
Runtime error
File size: 7,644 Bytes
878d9df d50d17c 878d9df d50d17c 878d9df d50d17c 878d9df d50d17c 878d9df cdab719 1011949 24f6bbf 1011949 ed9307c 1011949 0ca4abd |
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 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 |
#!/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
import gradio as gr
import uvicorn
load_dotenv()
#add logging
import logging
logger = logging.getLogger(__name__)
logging.basicConfig(filename='/home/user/tag_server.log', encoding='utf-8', level=logging.DEBUG)
logger.debug(' This message should go to the log file')
logger.info(' So should this')
# 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-server")
@mcp.tool()
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}")
global g_results
g_results = {
"status": "success",
"repo_id": repo_id,
"current_tags": current_tags,
"count": len(current_tags),
}
json_str = json.dumps(g_results)
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
@mcp.tool()
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
def create_gradio_app():
"""Create Gradio interface"""
with gr.Blocks(title="HF Tagging Bot", theme=gr.themes.Soft( font=[gr.themes.GoogleFont("Roboto"), "Open Sans", "sans-serif"]) ) as demo:
gr.Markdown("# π·οΈ HF Tagging Bot Dashboard")
with gr.Column():
sim_repo = gr.Textbox(
label="Repository",
value="burtenshaw/play-mcp-repo-bot",
placeholder="username/model-name",
)
sim_btn = gr.Button("π·οΈ Test Tag Detection")
with gr.Column():
sim_result = gr.Textbox(label="Result", lines=8)
sim_btn.click(
fn=get_current_tags,
inputs=[sim_repo,],
outputs=sim_result,
)
return demo
# Mount Gradio app
gradio_app = create_gradio_app()
#app_gradio = gr.mount_gradio_app(mcp, gradio_app, path="/gradio")
if __name__ == "__main__":
print("π Starting HF Tagging Bot server")
print("π Dashboard: http://localhost:7860/gradio")
print(" Starting http://localhost:443/mcp")
mcp.run(
transport="streamable-http",
host="0.0.0.0",
port=8000,
path="/mcp/",
log_level="debug",
)
|