tag-this-repo / server.py
burtenshaw
first commit
05c2ac8
raw
history blame
6.25 kB
import os
from datetime import datetime
from typing import List, Dict, Any, Optional
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()
# Configuration
WEBHOOK_SECRET = os.getenv("WEBHOOK_SECRET", "your-webhook-secret")
HF_TOKEN = os.getenv("HF_TOKEN")
HF_MODEL = os.getenv("HF_MODEL", "microsoft/DialoGPT-medium")
HF_PROVIDER = os.getenv("HF_PROVIDER", "huggingface")
# Simple storage for processed comments
comments_store: List[Dict[str, Any]] = []
# Agent instance
agent_instance: Optional[Agent] = None
class WebhookEvent(BaseModel):
event: Dict[str, str]
comment: Dict[str, Any]
discussion: Dict[str, Any]
repo: Dict[str, str]
app = FastAPI(title="HF Discussion Bot")
app.add_middleware(CORSMiddleware, allow_origins=["*"])
async def get_agent():
"""Get or create Agent instance"""
global agent_instance
if agent_instance is None and HF_TOKEN:
agent_instance = Agent(
model=HF_MODEL,
provider=HF_PROVIDER,
api_key=HF_TOKEN,
servers=[
{
"type": "stdio",
"config": {"command": "python", "args": ["mcp_server.py"]},
}
],
)
await agent_instance.load_tools()
return agent_instance
async def process_webhook_comment(webhook_data: Dict[str, Any]):
"""Process webhook using Agent with MCP tools"""
comment_content = webhook_data["comment"]["content"]
discussion_title = webhook_data["discussion"]["title"]
repo_name = webhook_data["repo"]["name"]
discussion_num = webhook_data["discussion"]["num"]
agent = await get_agent()
if not agent:
ai_response = "Error: Agent not configured (missing HF_TOKEN)"
else:
# Use Agent to respond to the discussion
prompt = f"""
Please respond to this HuggingFace discussion comment using the available tools.
Repository: {repo_name}
Discussion: {discussion_title} (#{discussion_num})
Comment: {comment_content}
First use generate_discussion_response to create a helpful response, then use post_discussion_comment to post it.
"""
try:
response_parts = []
async for item in agent.run(prompt):
# Collect the agent's response
if hasattr(item, "content") and item.content:
response_parts.append(item.content)
elif isinstance(item, str):
response_parts.append(item)
ai_response = (
" ".join(response_parts) if response_parts else "No response generated"
)
except Exception as e:
ai_response = f"Error using agent: {str(e)}"
# Store the interaction with reply link
discussion_url = f"https://huggingface.co/{repo_name}/discussions/{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,
"ai_response": ai_response,
"comment_author": webhook_data["comment"]["author"],
}
comments_store.append(interaction)
return ai_response
@app.post("/webhook")
async def webhook_handler(request: Request, background_tasks: BackgroundTasks):
"""Handle HF Hub webhooks"""
webhook_secret = request.headers.get("X-Webhook-Secret")
if webhook_secret != WEBHOOK_SECRET:
return {"error": "Invalid webhook secret"}
payload = await request.json()
event = payload.get("event", {})
if event.get("action") == "create" and event.get("scope") == "discussion.comment":
background_tasks.add_task(process_webhook_comment, payload)
return {"status": "processing"}
return {"status": "ignored"}
async def simulate_webhook(
repo_name: str, discussion_title: str, comment_content: str
) -> str:
"""Simulate webhook for testing"""
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": "test-user",
"created_at": datetime.now().isoformat(),
},
"discussion": {
"title": discussion_title,
"num": len(comments_store) + 1,
},
"repo": {"name": repo_name},
}
response = await process_webhook_comment(mock_payload)
return f"βœ… Processed! AI Response: {response}"
def create_gradio_app():
"""Create Gradio interface"""
with gr.Blocks(title="HF Discussion Bot", theme=gr.themes.Soft()) as demo:
gr.Markdown("# πŸ€– HF Discussion Bot Dashboard")
gr.Markdown("*Powered by HuggingFace Tiny Agents + FastMCP*")
with gr.Column():
sim_repo = gr.Textbox(label="Repository", value="microsoft/DialoGPT-medium")
sim_title = gr.Textbox(label="Discussion Title", value="Test Discussion")
sim_comment = gr.Textbox(
label="Comment",
lines=3,
value="How do I use this model?",
)
sim_btn = gr.Button("πŸ“€ Test Webhook")
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],
)
return demo
# Mount Gradio app
gradio_app = create_gradio_app()
app = gr.mount_gradio_app(app, gradio_app, path="/gradio")
if __name__ == "__main__":
print("πŸš€ Starting HF Discussion Bot with Tiny Agents...")
print("πŸ“Š Dashboard: http://localhost:8001/gradio")
print("πŸ”— Webhook: http://localhost:8001/webhook")
uvicorn.run("server:app", host="0.0.0.0", port=8001, reload=True)