| |
| |
| |
|
|
| import os |
| import re |
| import json |
| from datetime import datetime |
| from typing import List, Dict, Any, Optional, Literal |
|
|
| from fastapi import FastAPI, Request, BackgroundTasks |
| from fastapi.middleware.cors import CORSMiddleware |
| import gradio as gr |
| import uvicorn |
| from pydantic import BaseModel |
| from fastapi.responses import RedirectResponse |
| from huggingface_hub.inference._mcp.agent import Agent |
| from dotenv import load_dotenv |
| load_dotenv() |
|
|
| HF_TOKEN=os.getenv("HF_TOKEN") |
| WEBHOOK_SECRET=os.getenv("WEBHOOK_SECRET") |
| HF_MODEL=os.getenv("HF_MODEL","HuggingFaceH4/zephyr-7b-beta") |
| DEFAULT_PROVIDER:Literal['hf-inference']="hf-inference" |
| HF_PROVIDER=os.getenv("HF_PROVIDER") |
| agent_instance: Optional[Agent]=None |
| tag_operations_store:List[Dict[str,Any]]=[] |
|
|
| RECOGNIZED_TAGS = { |
| "pytorch", |
| "tensorflow", |
| "jax", |
| "transformers", |
| "diffusers", |
| "text-generation", |
| "text-classification", |
| "question-answering", |
| "text-to-image", |
| "image-classification", |
| "object-detection", |
| " ", |
| "fill-mask", |
| "token-classification", |
| "translation", |
| "summarization", |
| "feature-extraction", |
| "sentence-similarity", |
| "zero-shot-classification", |
| "image-to-text", |
| "automatic-speech-recognition", |
| "audio-classification", |
| "voice-activity-detection", |
| "depth-estimation", |
| "image-segmentation", |
| "video-classification", |
| "reinforcement-learning", |
| "tabular-classification", |
| "tabular-regression", |
| "time-series-forecasting", |
| "graph-ml", |
| "robotics", |
| "computer-vision", |
| "nlp", |
| "cv", |
| "multimodal", |
| } |
|
|
|
|
| class WebhookEvent(BaseModel): |
| event: Dict[str, str] |
| comment: Dict[str, Any] |
| discussion: Dict[str, Any] |
| repo: Dict[str, str] |
|
|
|
|
| app = FastAPI(title="HF Tagging Bot") |
| app.add_middleware(CORSMiddleware, allow_origins=["*"]) |
|
|
|
|
| async def get_agent(): |
| """Get or create Agent instance""" |
| print("🤖 get_agent() called...") |
| global agent_instance |
| if agent_instance is None and HF_TOKEN: |
| print("🔧 Creating new Agent instance...") |
| print(f"🔑 HF_TOKEN present: {bool(HF_TOKEN)}") |
| print(f"🤖 Model: {HF_MODEL}") |
| print(f"🔗 Provider: {DEFAULT_PROVIDER}") |
|
|
| try: |
| agent_instance = Agent( |
| model=HF_MODEL, |
| provider=DEFAULT_PROVIDER, |
| api_key=HF_TOKEN, |
| servers=[ |
| { |
| "type": "stdio", |
| "command": "python", |
| "args": ["mcp_server.py"], |
| "cwd": ".", |
| "env": {"HF_TOKEN": HF_TOKEN} if HF_TOKEN else {}, |
| } |
| ], |
| ) |
| print("✅ Agent instance created successfully") |
| print("🔧 Loading tools...") |
| await agent_instance.load_tools() |
| print("✅ Tools loaded successfully") |
| except Exception as e: |
| print(f"❌ Error creating/loading agent: {str(e)}") |
| agent_instance = None |
| elif agent_instance is None: |
| print("❌ No HF_TOKEN available, cannot create agent") |
| else: |
| print("✅ Using existing agent instance") |
|
|
| return agent_instance |
|
|
|
|
| def extract_tags_from_text(text: str) -> List[str]: |
| """Extract potential tags from discussion text""" |
| text_lower = text.lower() |
| explicit_tags = [] |
| tag_pattern = r"tags?:\s*([a-zA-Z0-9-_,\s]+)" |
| matches = re.findall(tag_pattern, text_lower) |
| for match in matches: |
| tags = [tag.strip() for tag in match.split(",")] |
| explicit_tags.extend(tags) |
| hashtag_pattern = r"#([a-zA-Z0-9-_]+)" |
| hashtag_matches = re.findall(hashtag_pattern, text_lower) |
| explicit_tags.extend(hashtag_matches) |
|
|
| mentioned_tags = [] |
| for tag in RECOGNIZED_TAGS: |
| if tag in text_lower: |
| mentioned_tags.append(tag) |
|
|
| all_tags = list(set(explicit_tags + mentioned_tags)) |
|
|
| valid_tags = [] |
| for tag in all_tags: |
| if tag in RECOGNIZED_TAGS or tag in explicit_tags: |
| valid_tags.append(tag) |
|
|
| return valid_tags |
|
|
|
|
| async def process_webhook_comment(webhook_data: Dict[str, Any]): |
| """Process webhook to detect and add tags""" |
| print("🏷️ Starting process_webhook_comment...") |
|
|
| try: |
| comment_content = webhook_data["comment"]["content"] |
| discussion_title = webhook_data["discussion"]["title"] |
| repo_name = webhook_data["repo"]["name"] |
| discussion_num = webhook_data["discussion"]["num"] |
| comment_author = webhook_data["comment"]["author"].get("id", "unknown") |
|
|
| print(f"📝 Comment content: {comment_content}") |
| print(f"📰 Discussion title: {discussion_title}") |
| print(f"📦 Repository: {repo_name}") |
|
|
| comment_tags = extract_tags_from_text(comment_content) |
| title_tags = extract_tags_from_text(discussion_title) |
| all_tags = list(set(comment_tags + title_tags)) |
|
|
| print(f"🔍 Comment tags found: {comment_tags}") |
| print(f"🔍 Title tags found: {title_tags}") |
| print(f"🏷️ All unique tags: {all_tags}") |
|
|
| result_messages = [] |
|
|
| if not all_tags: |
| msg = "No recognizable tags found in the discussion." |
| print(f"❌ {msg}") |
| result_messages.append(msg) |
| else: |
| print("🤖 Getting agent instance...") |
| agent = await get_agent() |
| if not agent: |
| msg = "Error: Agent not configured (missing HF_TOKEN please check)" |
| print(f"❌ {msg}") |
| result_messages.append(msg) |
| else: |
| print("✅ Agent instance obtained successfully") |
|
|
| try: |
| user_prompt = f""" |
| I need to add the following tags to the repository '{repo_name}': {", ".join(all_tags)} |
| For each tag, please: |
| 1. Check if the tag already exists on the repository using get_current_tags |
| 2. If the tag doesn't exist, add it using add_new_tag |
| 3. Provide a summary of what was done for each tag |
| Please process all {len(all_tags)} tags: {", ".join(all_tags)} |
| """ |
|
|
| print("💬 Sending comprehensive prompt to agent...") |
| print(f"📝 Prompt: {user_prompt}") |
|
|
| conversation_result = [] |
|
|
| try: |
| async for item in agent.run(user_prompt): |
| item_str = str(item) |
| conversation_result.append(item_str) |
|
|
| if ( |
| "tool_call" in item_str.lower() |
| or "function" in item_str.lower() |
| ): |
| print(f"🔧 Agent using tools: {item_str[:200]}...") |
| elif "content" in item_str and len(item_str) < 500: |
| print(f"💭 Agent response: {item_str}") |
|
|
| full_response = " ".join(conversation_result) |
| print(f"📋 Agent conversation completed successfully") |
|
|
| for tag in all_tags: |
| tag_mentioned = tag.lower() in full_response.lower() |
|
|
| if ( |
| "already exists" in full_response.lower() |
| and tag_mentioned |
| ): |
| msg = f"Tag '{tag}': Already exists" |
| elif ( |
| "pr" in full_response.lower() |
| or "pull request" in full_response.lower() |
| ): |
| if tag_mentioned: |
| msg = f"Tag '{tag}': PR created successfully" |
| else: |
| msg = ( |
| f"Tag '{tag}': Processed " |
| "(PR may have been created)" |
| ) |
| elif "success" in full_response.lower() and tag_mentioned: |
| msg = f"Tag '{tag}': Successfully processed" |
| elif "error" in full_response.lower() and tag_mentioned: |
| msg = f"Tag '{tag}': Error during processing" |
| else: |
| msg = f"Tag '{tag}': Processed by agent" |
|
|
| print(f"✅ Result for tag '{tag}': {msg}") |
| result_messages.append(msg) |
|
|
| except Exception as agent_error: |
| print(f"⚠️ Agent streaming failed: {str(agent_error)}") |
| print("🔄 Falling back to direct MCP tool calls...") |
|
|
| try: |
| import sys |
| import importlib.util |
|
|
| spec = importlib.util.spec_from_file_location( |
| "mcp_server", "./mcp_server.py" |
| ) |
| mcp_module = importlib.util.module_from_spec(spec) |
| spec.loader.exec_module(mcp_module) |
| for tag in all_tags: |
| try: |
| print( |
| f"🔧 Directly calling get_current_tags for '{tag}'" |
| ) |
| current_tags_result = mcp_module.get_current_tags( |
| repo_name |
| ) |
| print( |
| f"📄 Current tags result: {current_tags_result}" |
| ) |
|
|
| import json |
|
|
| tags_data = json.loads(current_tags_result) |
|
|
| if tags_data.get("status") == "success": |
| current_tags = tags_data.get("current_tags", []) |
| if tag in current_tags: |
| msg = f"Tag '{tag}': Already exists" |
| print(f"✅ {msg}") |
| else: |
| print( |
| f"🔧 Directly calling add_new_tag for '{tag}'" |
| ) |
| add_result = mcp_module.add_new_tag( |
| repo_name, tag |
| ) |
| print(f"📄 Add tag result: {add_result}") |
|
|
| add_data = json.loads(add_result) |
| if add_data.get("status") == "success": |
| pr_url = add_data.get("pr_url", "") |
| msg = f"Tag '{tag}': PR created - {pr_url}" |
| elif ( |
| add_data.get("status") |
| == "already_exists" |
| ): |
| msg = f"Tag '{tag}': Already exists" |
| else: |
| msg = f"Tag '{tag}': {add_data.get('message', 'Processed')}" |
| print(f"✅ {msg}") |
| else: |
| error_msg = tags_data.get( |
| "error", "Unknown error" |
| ) |
| msg = f"Tag '{tag}': Error - {error_msg}" |
| print(f"❌ {msg}") |
|
|
| result_messages.append(msg) |
|
|
| except Exception as direct_error: |
| error_msg = f"Tag '{tag}': Direct call error - {str(direct_error)}" |
| print(f"❌ {error_msg}") |
| result_messages.append(error_msg) |
| except Exception as fallback_error: |
| error_msg = f"Fallback approach failed: {str(fallback_error)}" |
| print(f"❌ {error_msg}") |
| result_messages.append(error_msg) |
| except Exception as e: |
| error_msg = f"Error during agent processing {str(e)}" |
| print(f"❌ {error_msg}") |
| result_messages.append(error_msg) |
| base_url="https://huggingface.co" |
| discussion_url=f"{base_url}/{repo_name}/discussion/{discussion_num}" |
| interaction = { |
| "timestamp": datetime.now().isoformat(), |
| "repo": repo_name, |
| "discussion_title": discussion_title, |
| "discussion_num": discussion_num, |
| "discussion_url": discussion_url, |
| "original_comment": comment_content, |
| "comment_author": comment_author, |
| "detected_tags": all_tags, |
| "results": result_messages, |
| } |
| tag_operations_store.append(interaction) |
| final_result="|".join(result_messages) |
| print(f"💾 Stored interaction and returning result: {final_result}") |
| return final_result |
|
|
| except Exception as e: |
| error_msg = f"❌ Fatal error in process_webhook_comment: {str(e)}" |
| print(error_msg) |
| return error_msg |
|
|
|
|
| @app.get("/") |
| async def root(): |
| """Root endpoint with basic information""" |
| return { |
| "name":"HF Tagging Bot", |
| "status":"running", |
| "description":"Webhook listener for automatic model tagging", |
| "endpoints":{ |
| "webhook":"/webhook", |
| "health":"/health", |
| "operations":"/operations" |
| } |
| } |
|
|
|
|
| @app.get("/health") |
| async def health_check(): |
| """Health check endpoint for monitoring""" |
| agent=await get_agent() |
| return { |
| "status":"healthy", |
| "timestamp":datetime.now().isoformat(), |
| "components":{ |
| "webhook_secret":"configured" if WEBHOOK_SECRET else "missing", |
| "hf_token":"configured" if HF_TOKEN else "missing", |
| "mcp_agent":"ready" if agent else "not ready" |
| } |
| } |
|
|
| @app.get("/operations") |
| async def get_operations(): |
| """Get recent tag operations for monitoring""" |
| recent_ops=tag_operations_store[-50:] if tag_operations_store else [] |
| return { |
| "total_operations":len(tag_operations_store), |
| "recent_operations":recent_ops |
| } |
|
|
|
|
|
|
| @app.post("/webhook") |
| async def webhook_handler(request:Request, background_tasks:BackgroundTasks): |
| """ |
| Handle incoming webhooks from Hugging Face Hub |
| Following the pattern from: https://raw.githubusercontent.com/huggingface/hub-docs/refs/heads/main/docs/hub/webhooks-guide-discussion-bot.md |
| """ |
| print("🔔 Webhook received!") |
| webhook_secret=request.headers.get("X-webhook-Secret") |
| if webhook_secret!=WEBHOOK_SECRET: |
| print("❌ Invalid webhook secret") |
| return {"error":"incorrect secret"} |
| payload=await request.json() |
| print(f"📥 Received webhook payload: {json.dumps(payload, indent=2)}") |
| event=payload.get("event",{}) |
| scope=event.get("score") |
| action=event.get("action") |
|
|
| print(f"🔍 Event details - scope: {scope}, action: {action}") |
| scope_check = scope == "discussion" |
| action_check = action == "create" |
| not_pr = not payload["discussion"]["isPullRequest"] |
| scope_check = scope_check and not_pr |
| print(f"✅ not_pr: {not_pr}") |
| print(f"✅ scope_check: {scope_check}") |
| print(f"✅ action_check: {action_check}") |
| |
| if scope_check and action_check: |
| required_fields=['comment','discussion','repo'] |
| missing_fields=[field for field in required_fields if field not in payload] |
|
|
| if missing_fields: |
| error_msg = f"Missing required fields: {missing_fields}" |
| print(f"❌ {error_msg}") |
| return {"error": error_msg} |
| print(f"🚀 Processing webhook for repo: {payload['repo']['name']}") |
| background_tasks.add_task(process_webhook_comment,payload) |
| return {"status":"processing"} |
| print(f"⏭️ Ignoring webhook - scope: {scope}, action: {action}") |
| return {"status": "ignored"} |
|
|
| @app.post("/simulate_webhook") |
| async def simulate_webhook(repo_name:str,discussion_title:str,comment_content:str)->str: |
| """Simulate webhook for testing purposes""" |
| if not all([repo_name,discussion_title,comment_content]): |
| return "please fill in all fields" |
| mock_payload={ |
| "event":{"action":"create","scope":"discussion.comment"}, |
| "comment":{"content":comment_content,"author":{"id":"test-user"},"id":"mock-comment-id","hidden":False}, |
| "discussion":{"title":discussion_title,"num":len(tag_operations_store)+1,"id":"mock-comment-id","status":"open","isPullRequest":False}, |
| "repo":{"name":repo_name,"type":"model","private":False} |
| } |
| response=await process_webhook_comment(mock_payload) |
| return f"✅ Processed! Results: {response}" |
|
|
| def create_gradio_app(): |
| """Create Gradio interface""" |
| with gr.Blocks(title="HF Tagging Bot", theme=gr.themes.Soft()) as demo: |
| gr.Markdown("# 🏷️ HF Tagging Bot Dashboard") |
| gr.Markdown("*Automatically adds tags to models when mentioned in discussions*") |
|
|
| gr.Markdown(""" |
| ## How it works: |
| - Monitors HuggingFace Hub discussions |
| - Detects tag mentions in comments (e.g., "tag: pytorch", |
| "#transformers") |
| - Automatically adds recognized tags to the model repository |
| - Supports common ML tags like: pytorch, tensorflow, |
| text-generation, etc. |
| """) |
| with gr.Column(): |
| sim_repo=gr.Textbox(label="Repository",value="burtenshaw/play-mcp-repo-bot",placeholder="username/model-name") |
| sim_title= gr.Textbox(label="Discussion Title",value="Add pytorch tag",placeholder="Discussion title") |
| sim_comment=gr.Textbox(label="comment",lines=3,value="This model should have tags: pytorch, text-generation",placeholder="Comment mentioning tags ...") |
| sim_btn=gr.Button("🏷️ Test Tag Detection") |
| with gr.Column(): |
| sim_result=gr.Textbox(label="Result",lines=8) |
| sim_btn.click(fn=simulate_webhook,inputs=[sim_repo,sim_title,sim_comment],outputs=sim_result) |
| gr.Markdown(f"""## Recognized Tags: {",".join(sorted(RECOGNIZED_TAGS))}""") |
| return demo |
| gradio_app=create_gradio_app() |
| app=gr.mount_gradio_app(app,gradio_app,path="/gradio") |
|
|
| @app.get("/") |
| async def root_direct(): |
| return RedirectResponse(url="/gradio") |
|
|
| if __name__=="__main__": |
| print("🚀 Starting HF Tagging Bot...") |
| print("📊 Dashboard: http://localhost:7860/gradio") |
| print("🔗 Webhook: http://localhost:7860/webhook") |
| uvicorn.run("app:app",host="0.0.0.0",port=7860,reload=True) |
