Spaces:
Sleeping
Sleeping
| # ---------------------------------------------------------------------------------- | |
| # 1. Agent Configuration | |
| # ---------------------------------------------------------------------------------- | |
| 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 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","microsoft/DialoGPT-medium") | |
| DEFAULT_PROVIDER:Literal['hf-inference']="hf-inference" | |
| HF_PROVIDER=os.getenv("HF_PROVIDER",DEFAULT_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", | |
| "config": { | |
| "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)" | |
| 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 | |
| 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" | |
| } | |
| } | |
| 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" | |
| } | |
| } | |
| 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 | |
| } | |
| 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"} | |
| 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="Sandhya2002/tag-boat") | |
| 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") | |
| if __name__=="__main__": | |
| print("🚀 Starting HF Tagging Bot...") | |
| print("📊 Dashboard: http://localhost:7860/gradio") | |
| print("🔗 Webhook: http://localhost:7860/webhook") | |
| unicorn.run("app:app",host="0.0.0.0",port=7860,reload=True) | |