burtenshaw commited on
Commit
b750fb8
Β·
1 Parent(s): d0a5416

fix integration and improve logging

Browse files
Files changed (2) hide show
  1. app.py +237 -99
  2. mcp_server.py +46 -7
app.py CHANGED
@@ -82,25 +82,43 @@ app.add_middleware(CORSMiddleware, allow_origins=["*"])
82
 
83
  async def get_agent():
84
  """Get or create Agent instance"""
 
85
  global agent_instance
86
  if agent_instance is None and HF_TOKEN:
87
- agent_instance = Agent(
88
- model=HF_MODEL,
89
- provider=DEFAULT_PROVIDER,
90
- api_key=HF_TOKEN,
91
- servers=[
92
- {
93
- "type": "stdio",
94
- "config": {
95
- "command": "python",
96
- "args": ["mcp_server.py"],
97
- "cwd": ".", # Ensure correct working directory
98
- "env": {"HF_TOKEN": HF_TOKEN} if HF_TOKEN else {},
99
- },
100
- }
101
- ],
102
- )
103
- await agent_instance.load_tools()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
104
  return agent_instance
105
 
106
 
@@ -144,105 +162,225 @@ def extract_tags_from_text(text: str) -> List[str]:
144
 
145
  async def process_webhook_comment(webhook_data: Dict[str, Any]):
146
  """Process webhook to detect and add tags"""
147
- comment_content = webhook_data["comment"]["content"]
148
- discussion_title = webhook_data["discussion"]["title"]
149
- repo_name = webhook_data["repo"]["name"]
150
- discussion_num = webhook_data["discussion"]["num"]
151
- # Author is an object with "id" field
152
- comment_author = webhook_data["comment"]["author"].get("id", "unknown")
153
-
154
- # Extract potential tags from the comment and discussion title
155
- comment_tags = extract_tags_from_text(comment_content)
156
- title_tags = extract_tags_from_text(discussion_title)
157
- all_tags = list(set(comment_tags + title_tags))
158
-
159
- result_messages = []
160
-
161
- if not all_tags:
162
- result_messages.append("No recognizable tags found in the discussion.")
163
- else:
164
- agent = await get_agent()
165
- if not agent:
166
- msg = "Error: Agent not configured (missing HF_TOKEN)"
 
 
 
 
 
 
 
 
167
  result_messages.append(msg)
168
  else:
169
- # Process each tag
170
- for tag in all_tags:
 
 
 
 
 
 
 
 
171
  try:
172
- # Get response from agent
173
- responses = []
174
- prompt = (
175
- f"Add the tag '{tag}' to repository {repo_name} "
176
- "using add_new_tag"
177
- )
178
 
179
- async for item in agent.run(prompt):
180
- # Just collect the response content
181
- responses.append(str(item))
 
182
 
183
- response_text = " ".join(responses) if responses else "Completed"
 
184
 
185
- # Try to parse JSON from response if possible
186
- try:
187
- # Look for JSON in the response
188
- json_found = False
189
- for response_part in responses:
190
- response_str = str(response_part)
191
- if "{" in response_str and "}" in response_str:
192
- # Try to extract JSON from the response
193
- start_idx = response_str.find("{")
194
- end_idx = response_str.rfind("}") + 1
195
- json_str = response_str[start_idx:end_idx]
196
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
197
  try:
198
- json_response = json.loads(json_str)
199
- status = json_response.get("status")
200
- if status == "success":
201
- pr_url = json_response.get("pr_url", "")
202
- msg = f"Tag '{tag}': PR created - {pr_url}"
203
- elif status == "already_exists":
204
- msg = f"Tag '{tag}': Already exists"
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
205
  else:
206
- tag_msg = json_response.get(
207
- "message", "Processed"
208
  )
209
- msg = f"Tag '{tag}': {tag_msg}"
210
- json_found = True
211
- break
212
- except json.JSONDecodeError:
213
- continue
214
 
215
- if not json_found:
216
- # If no JSON found, use the response as is
217
- msg = f"Tag '{tag}': {response_text}"
218
 
219
- except Exception:
220
- msg = f"Tag '{tag}': Response parse error - {response_text}"
 
 
221
 
222
- result_messages.append(msg)
 
 
 
 
 
223
 
224
  except Exception as e:
225
- error_msg = f"Error processing tag '{tag}': {str(e)}"
 
226
  result_messages.append(error_msg)
227
 
228
- # Store the interaction
229
- base_url = "https://huggingface.co"
230
- discussion_url = f"{base_url}/{repo_name}/discussions/{discussion_num}"
231
-
232
- interaction = {
233
- "timestamp": datetime.now().isoformat(),
234
- "repo": repo_name,
235
- "discussion_title": discussion_title,
236
- "discussion_num": discussion_num,
237
- "discussion_url": discussion_url,
238
- "original_comment": comment_content,
239
- "comment_author": comment_author,
240
- "detected_tags": all_tags,
241
- "results": result_messages,
242
- }
243
-
244
- tag_operations_store.append(interaction)
245
- return " | ".join(result_messages)
 
 
 
 
 
 
 
246
 
247
 
248
  @app.post("/webhook")
 
82
 
83
  async def get_agent():
84
  """Get or create Agent instance"""
85
+ print("πŸ€– get_agent() called...")
86
  global agent_instance
87
  if agent_instance is None and HF_TOKEN:
88
+ print("πŸ”§ Creating new Agent instance...")
89
+ print(f"πŸ”‘ HF_TOKEN present: {bool(HF_TOKEN)}")
90
+ print(f"πŸ€– Model: {HF_MODEL}")
91
+ print(f"πŸ”— Provider: {DEFAULT_PROVIDER}")
92
+
93
+ try:
94
+ agent_instance = Agent(
95
+ model=HF_MODEL,
96
+ provider=DEFAULT_PROVIDER,
97
+ api_key=HF_TOKEN,
98
+ servers=[
99
+ {
100
+ "type": "stdio",
101
+ "config": {
102
+ "command": "python",
103
+ "args": ["mcp_server.py"],
104
+ "cwd": ".", # Ensure correct working directory
105
+ "env": {"HF_TOKEN": HF_TOKEN} if HF_TOKEN else {},
106
+ },
107
+ }
108
+ ],
109
+ )
110
+ print("βœ… Agent instance created successfully")
111
+ print("πŸ”§ Loading tools...")
112
+ await agent_instance.load_tools()
113
+ print("βœ… Tools loaded successfully")
114
+ except Exception as e:
115
+ print(f"❌ Error creating/loading agent: {str(e)}")
116
+ agent_instance = None
117
+ elif agent_instance is None:
118
+ print("❌ No HF_TOKEN available, cannot create agent")
119
+ else:
120
+ print("βœ… Using existing agent instance")
121
+
122
  return agent_instance
123
 
124
 
 
162
 
163
  async def process_webhook_comment(webhook_data: Dict[str, Any]):
164
  """Process webhook to detect and add tags"""
165
+ print("🏷️ Starting process_webhook_comment...")
166
+
167
+ try:
168
+ comment_content = webhook_data["comment"]["content"]
169
+ discussion_title = webhook_data["discussion"]["title"]
170
+ repo_name = webhook_data["repo"]["name"]
171
+ discussion_num = webhook_data["discussion"]["num"]
172
+ # Author is an object with "id" field
173
+ comment_author = webhook_data["comment"]["author"].get("id", "unknown")
174
+
175
+ print(f"πŸ“ Comment content: {comment_content}")
176
+ print(f"πŸ“° Discussion title: {discussion_title}")
177
+ print(f"πŸ“¦ Repository: {repo_name}")
178
+
179
+ # Extract potential tags from the comment and discussion title
180
+ comment_tags = extract_tags_from_text(comment_content)
181
+ title_tags = extract_tags_from_text(discussion_title)
182
+ all_tags = list(set(comment_tags + title_tags))
183
+
184
+ print(f"πŸ” Comment tags found: {comment_tags}")
185
+ print(f"πŸ” Title tags found: {title_tags}")
186
+ print(f"🏷️ All unique tags: {all_tags}")
187
+
188
+ result_messages = []
189
+
190
+ if not all_tags:
191
+ msg = "No recognizable tags found in the discussion."
192
+ print(f"❌ {msg}")
193
  result_messages.append(msg)
194
  else:
195
+ print("πŸ€– Getting agent instance...")
196
+ agent = await get_agent()
197
+ if not agent:
198
+ msg = "Error: Agent not configured (missing HF_TOKEN)"
199
+ print(f"❌ {msg}")
200
+ result_messages.append(msg)
201
+ else:
202
+ print("βœ… Agent instance obtained successfully")
203
+
204
+ # Process all tags in a single conversation with the agent
205
  try:
206
+ # Create a comprehensive prompt for the agent
207
+ user_prompt = f"""
208
+ I need to add the following tags to the repository '{repo_name}': {", ".join(all_tags)}
 
 
 
209
 
210
+ For each tag, please:
211
+ 1. Check if the tag already exists on the repository using get_current_tags
212
+ 2. If the tag doesn't exist, add it using add_new_tag
213
+ 3. Provide a summary of what was done for each tag
214
 
215
+ Please process all {len(all_tags)} tags: {", ".join(all_tags)}
216
+ """
217
 
218
+ print("πŸ’¬ Sending comprehensive prompt to agent...")
219
+ print(f"πŸ“ Prompt: {user_prompt}")
220
+
221
+ # Let the agent handle the entire conversation
222
+ conversation_result = []
 
 
 
 
 
 
223
 
224
+ try:
225
+ async for item in agent.run(user_prompt):
226
+ # The agent yields different types of items
227
+ item_str = str(item)
228
+ conversation_result.append(item_str)
229
+
230
+ # Log important events
231
+ if (
232
+ "tool_call" in item_str.lower()
233
+ or "function" in item_str.lower()
234
+ ):
235
+ print(f"πŸ”§ Agent using tools: {item_str[:200]}...")
236
+ elif "content" in item_str and len(item_str) < 500:
237
+ print(f"πŸ’­ Agent response: {item_str}")
238
+
239
+ # Extract the final response from the conversation
240
+ full_response = " ".join(conversation_result)
241
+ print(f"πŸ“‹ Agent conversation completed successfully")
242
+
243
+ # Try to extract meaningful results for each tag
244
+ for tag in all_tags:
245
+ tag_mentioned = tag.lower() in full_response.lower()
246
+
247
+ if (
248
+ "already exists" in full_response.lower()
249
+ and tag_mentioned
250
+ ):
251
+ msg = f"Tag '{tag}': Already exists"
252
+ elif (
253
+ "pr" in full_response.lower()
254
+ or "pull request" in full_response.lower()
255
+ ):
256
+ if tag_mentioned:
257
+ msg = f"Tag '{tag}': PR created successfully"
258
+ else:
259
+ msg = (
260
+ f"Tag '{tag}': Processed "
261
+ "(PR may have been created)"
262
+ )
263
+ elif "success" in full_response.lower() and tag_mentioned:
264
+ msg = f"Tag '{tag}': Successfully processed"
265
+ elif "error" in full_response.lower() and tag_mentioned:
266
+ msg = f"Tag '{tag}': Error during processing"
267
+ else:
268
+ msg = f"Tag '{tag}': Processed by agent"
269
+
270
+ print(f"βœ… Result for tag '{tag}': {msg}")
271
+ result_messages.append(msg)
272
+
273
+ except Exception as agent_error:
274
+ print(f"⚠️ Agent streaming failed: {str(agent_error)}")
275
+ print("πŸ”„ Falling back to direct MCP tool calls...")
276
+
277
+ # Import the MCP server functions directly as fallback
278
+ try:
279
+ import sys
280
+ import importlib.util
281
+
282
+ # Load the MCP server module
283
+ spec = importlib.util.spec_from_file_location(
284
+ "mcp_server", "./mcp_server.py"
285
+ )
286
+ mcp_module = importlib.util.module_from_spec(spec)
287
+ spec.loader.exec_module(mcp_module)
288
+
289
+ # Use the MCP tools directly for each tag
290
+ for tag in all_tags:
291
  try:
292
+ print(
293
+ f"πŸ”§ Directly calling get_current_tags for '{tag}'"
294
+ )
295
+ current_tags_result = mcp_module.get_current_tags(
296
+ repo_name
297
+ )
298
+ print(
299
+ f"πŸ“„ Current tags result: {current_tags_result}"
300
+ )
301
+
302
+ # Parse the JSON result
303
+ import json
304
+
305
+ tags_data = json.loads(current_tags_result)
306
+
307
+ if tags_data.get("status") == "success":
308
+ current_tags = tags_data.get("current_tags", [])
309
+ if tag in current_tags:
310
+ msg = f"Tag '{tag}': Already exists"
311
+ print(f"βœ… {msg}")
312
+ else:
313
+ print(
314
+ f"πŸ”§ Directly calling add_new_tag for '{tag}'"
315
+ )
316
+ add_result = mcp_module.add_new_tag(
317
+ repo_name, tag
318
+ )
319
+ print(f"πŸ“„ Add tag result: {add_result}")
320
+
321
+ add_data = json.loads(add_result)
322
+ if add_data.get("status") == "success":
323
+ pr_url = add_data.get("pr_url", "")
324
+ msg = f"Tag '{tag}': PR created - {pr_url}"
325
+ elif (
326
+ add_data.get("status")
327
+ == "already_exists"
328
+ ):
329
+ msg = f"Tag '{tag}': Already exists"
330
+ else:
331
+ msg = f"Tag '{tag}': {add_data.get('message', 'Processed')}"
332
+ print(f"βœ… {msg}")
333
  else:
334
+ error_msg = tags_data.get(
335
+ "error", "Unknown error"
336
  )
337
+ msg = f"Tag '{tag}': Error - {error_msg}"
338
+ print(f"❌ {msg}")
 
 
 
339
 
340
+ result_messages.append(msg)
 
 
341
 
342
+ except Exception as direct_error:
343
+ error_msg = f"Tag '{tag}': Direct call error - {str(direct_error)}"
344
+ print(f"❌ {error_msg}")
345
+ result_messages.append(error_msg)
346
 
347
+ except Exception as fallback_error:
348
+ error_msg = (
349
+ f"Fallback approach failed: {str(fallback_error)}"
350
+ )
351
+ print(f"❌ {error_msg}")
352
+ result_messages.append(error_msg)
353
 
354
  except Exception as e:
355
+ error_msg = f"Error during agent processing: {str(e)}"
356
+ print(f"❌ {error_msg}")
357
  result_messages.append(error_msg)
358
 
359
+ # Store the interaction
360
+ base_url = "https://huggingface.co"
361
+ discussion_url = f"{base_url}/{repo_name}/discussions/{discussion_num}"
362
+
363
+ interaction = {
364
+ "timestamp": datetime.now().isoformat(),
365
+ "repo": repo_name,
366
+ "discussion_title": discussion_title,
367
+ "discussion_num": discussion_num,
368
+ "discussion_url": discussion_url,
369
+ "original_comment": comment_content,
370
+ "comment_author": comment_author,
371
+ "detected_tags": all_tags,
372
+ "results": result_messages,
373
+ }
374
+
375
+ tag_operations_store.append(interaction)
376
+ final_result = " | ".join(result_messages)
377
+ print(f"πŸ’Ύ Stored interaction and returning result: {final_result}")
378
+ return final_result
379
+
380
+ except Exception as e:
381
+ error_msg = f"❌ Fatal error in process_webhook_comment: {str(e)}"
382
+ print(error_msg)
383
+ return error_msg
384
 
385
 
386
  @app.post("/webhook")
mcp_server.py CHANGED
@@ -25,12 +25,19 @@ mcp = FastMCP("hf-tagging-bot")
25
  @mcp.tool()
26
  def get_current_tags(repo_id: str) -> str:
27
  """Get current tags from a HuggingFace model repository"""
 
 
28
  if not hf_api:
29
- return json.dumps({"error": "HF token not configured"})
 
 
 
30
 
31
  try:
 
32
  info = model_info(repo_id=repo_id, token=HF_TOKEN)
33
  current_tags = info.tags if info.tags else []
 
34
 
35
  result = {
36
  "status": "success",
@@ -38,45 +45,63 @@ def get_current_tags(repo_id: str) -> str:
38
  "current_tags": current_tags,
39
  "count": len(current_tags),
40
  }
41
- return json.dumps(result)
 
 
42
 
43
  except Exception as e:
 
44
  error_result = {"status": "error", "repo_id": repo_id, "error": str(e)}
45
- return json.dumps(error_result)
 
 
46
 
47
 
48
  @mcp.tool()
49
  def add_new_tag(repo_id: str, new_tag: str) -> str:
50
  """Add a new tag to a HuggingFace model repository via PR"""
 
 
51
  if not hf_api:
52
- return json.dumps({"error": "HF token not configured"})
 
 
 
53
 
54
  try:
55
  # Get current model info and tags
 
56
  info = model_info(repo_id=repo_id, token=HF_TOKEN)
57
  current_tags = info.tags if info.tags else []
 
58
 
59
  # Check if tag already exists
60
  if new_tag in current_tags:
 
61
  result = {
62
  "status": "already_exists",
63
  "repo_id": repo_id,
64
  "tag": new_tag,
65
  "message": f"Tag '{new_tag}' already exists",
66
  }
67
- return json.dumps(result)
 
 
68
 
69
  # Add the new tag to existing tags
70
  updated_tags = current_tags + [new_tag]
 
71
 
72
  # Create model card content with updated tags
73
  try:
74
  # Load existing model card
 
75
  card = ModelCard.load(repo_id, token=HF_TOKEN)
76
  if not hasattr(card, "data") or card.data is None:
77
  card.data = ModelCardData()
78
  except HfHubHTTPError:
79
  # Create new model card if none exists
 
80
  card = ModelCard("")
81
  card.data = ModelCardData()
82
 
@@ -100,6 +125,8 @@ This PR adds the `{new_tag}` tag to the model repository.
100
  **New tags:** {", ".join(updated_tags)}
101
  """
102
 
 
 
103
  # Create commit with updated model card using CommitOperationAdd
104
  from huggingface_hub import CommitOperationAdd
105
 
@@ -120,6 +147,8 @@ This PR adds the `{new_tag}` tag to the model repository.
120
  pr_url_attr = commit_info.pr_url
121
  pr_url = pr_url_attr if hasattr(commit_info, "pr_url") else str(commit_info)
122
 
 
 
123
  result = {
124
  "status": "success",
125
  "repo_id": repo_id,
@@ -129,16 +158,26 @@ This PR adds the `{new_tag}` tag to the model repository.
129
  "new_tags": updated_tags,
130
  "message": f"Created PR to add tag '{new_tag}'",
131
  }
132
- return json.dumps(result)
 
 
133
 
134
  except Exception as e:
 
 
 
 
 
 
135
  error_result = {
136
  "status": "error",
137
  "repo_id": repo_id,
138
  "tag": new_tag,
139
  "error": str(e),
140
  }
141
- return json.dumps(error_result)
 
 
142
 
143
 
144
  if __name__ == "__main__":
 
25
  @mcp.tool()
26
  def get_current_tags(repo_id: str) -> str:
27
  """Get current tags from a HuggingFace model repository"""
28
+ print(f"πŸ”§ get_current_tags called with repo_id: {repo_id}")
29
+
30
  if not hf_api:
31
+ error_result = {"error": "HF token not configured"}
32
+ json_str = json.dumps(error_result)
33
+ print(f"❌ No HF API token - returning: {json_str}")
34
+ return json_str
35
 
36
  try:
37
+ print(f"πŸ“‘ Fetching model info for: {repo_id}")
38
  info = model_info(repo_id=repo_id, token=HF_TOKEN)
39
  current_tags = info.tags if info.tags else []
40
+ print(f"🏷️ Found {len(current_tags)} tags: {current_tags}")
41
 
42
  result = {
43
  "status": "success",
 
45
  "current_tags": current_tags,
46
  "count": len(current_tags),
47
  }
48
+ json_str = json.dumps(result)
49
+ print(f"βœ… get_current_tags returning: {json_str}")
50
+ return json_str
51
 
52
  except Exception as e:
53
+ print(f"❌ Error in get_current_tags: {str(e)}")
54
  error_result = {"status": "error", "repo_id": repo_id, "error": str(e)}
55
+ json_str = json.dumps(error_result)
56
+ print(f"❌ get_current_tags error returning: {json_str}")
57
+ return json_str
58
 
59
 
60
  @mcp.tool()
61
  def add_new_tag(repo_id: str, new_tag: str) -> str:
62
  """Add a new tag to a HuggingFace model repository via PR"""
63
+ print(f"πŸ”§ add_new_tag called with repo_id: {repo_id}, new_tag: {new_tag}")
64
+
65
  if not hf_api:
66
+ error_result = {"error": "HF token not configured"}
67
+ json_str = json.dumps(error_result)
68
+ print(f"❌ No HF API token - returning: {json_str}")
69
+ return json_str
70
 
71
  try:
72
  # Get current model info and tags
73
+ print(f"πŸ“‘ Fetching current model info for: {repo_id}")
74
  info = model_info(repo_id=repo_id, token=HF_TOKEN)
75
  current_tags = info.tags if info.tags else []
76
+ print(f"🏷️ Current tags: {current_tags}")
77
 
78
  # Check if tag already exists
79
  if new_tag in current_tags:
80
+ print(f"⚠️ Tag '{new_tag}' already exists in {current_tags}")
81
  result = {
82
  "status": "already_exists",
83
  "repo_id": repo_id,
84
  "tag": new_tag,
85
  "message": f"Tag '{new_tag}' already exists",
86
  }
87
+ json_str = json.dumps(result)
88
+ print(f"🏷️ add_new_tag (already exists) returning: {json_str}")
89
+ return json_str
90
 
91
  # Add the new tag to existing tags
92
  updated_tags = current_tags + [new_tag]
93
+ print(f"πŸ†• Will update tags from {current_tags} to {updated_tags}")
94
 
95
  # Create model card content with updated tags
96
  try:
97
  # Load existing model card
98
+ print(f"πŸ“„ Loading existing model card...")
99
  card = ModelCard.load(repo_id, token=HF_TOKEN)
100
  if not hasattr(card, "data") or card.data is None:
101
  card.data = ModelCardData()
102
  except HfHubHTTPError:
103
  # Create new model card if none exists
104
+ print(f"πŸ“„ Creating new model card (none exists)")
105
  card = ModelCard("")
106
  card.data = ModelCardData()
107
 
 
125
  **New tags:** {", ".join(updated_tags)}
126
  """
127
 
128
+ print(f"πŸš€ Creating PR with title: {pr_title}")
129
+
130
  # Create commit with updated model card using CommitOperationAdd
131
  from huggingface_hub import CommitOperationAdd
132
 
 
147
  pr_url_attr = commit_info.pr_url
148
  pr_url = pr_url_attr if hasattr(commit_info, "pr_url") else str(commit_info)
149
 
150
+ print(f"βœ… PR created successfully! URL: {pr_url}")
151
+
152
  result = {
153
  "status": "success",
154
  "repo_id": repo_id,
 
158
  "new_tags": updated_tags,
159
  "message": f"Created PR to add tag '{new_tag}'",
160
  }
161
+ json_str = json.dumps(result)
162
+ print(f"βœ… add_new_tag success returning: {json_str}")
163
+ return json_str
164
 
165
  except Exception as e:
166
+ print(f"❌ Error in add_new_tag: {str(e)}")
167
+ print(f"❌ Error type: {type(e)}")
168
+ import traceback
169
+
170
+ print(f"❌ Traceback: {traceback.format_exc()}")
171
+
172
  error_result = {
173
  "status": "error",
174
  "repo_id": repo_id,
175
  "tag": new_tag,
176
  "error": str(e),
177
  }
178
+ json_str = json.dumps(error_result)
179
+ print(f"❌ add_new_tag error returning: {json_str}")
180
+ return json_str
181
 
182
 
183
  if __name__ == "__main__":