Update mcp_server.py

mcp_server.py

@@ -1,184 +1,185 @@
-#!/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")
-
-
-@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}")
-
-        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
-
-
-@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
-
-
-if __name__ == "__main__":
-    mcp.run()
+#!/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")
+if HF_TOKEN:
+    HF_TOKEN = HF_TOKEN.strip() 
+# Initialize HF API client
+hf_api = HfApi(token=HF_TOKEN) if HF_TOKEN else None
+
+# Create the FastMCP server
+mcp = FastMCP("hf-tagging-bot")
+
+
+@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}")
+
+        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
+
+
+@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
+
+
+if __name__ == "__main__":
+    mcp.run()