Spaces:
Paused
Paused
| #!/usr/bin/env python | |
| """ | |
| Weave MCP Server - Report creation and editing functionality for W&B reports. | |
| """ | |
| from pathlib import Path | |
| from typing import Dict, List, Optional, Union | |
| import re | |
| import wandb_workspaces.reports.v2 as wr | |
| from dotenv import load_dotenv | |
| import wandb | |
| from wandb_mcp_server.utils import get_rich_logger | |
| # Load environment variables | |
| load_dotenv(dotenv_path=Path(__file__).parent.parent.parent / ".env") | |
| # Configure logging | |
| logger = get_rich_logger(__name__) | |
| CREATE_WANDB_REPORT_TOOL_DESCRIPTION = """Create a new Weights & Biases Report and add text and HTML-rendered charts. Useful to save/document analysis and other findings. | |
| Only call this tool if the user explicitly asks to create a report or save to wandb/weights & biases. | |
| Always provide the returned report link to the user. | |
| <plots_html_usage_guide> | |
| - If the analsis has generated plots then they can be logged to a Weights & Biases report via converting them to html. | |
| - All charts should be properly rendered in raw HTML, do not use any placeholders for any chart, render everything. | |
| - All charts should be beautiful, tasteful and well proportioned. | |
| - Plot html code should use SVG chart elements that should render properly in any modern browser. | |
| - Include interactive hover effects where it makes sense. | |
| - If the analysis contains multiple charts, break up the html into one section of html per chart. | |
| - Ensure that the axis labels are properly set and aligned for each chart. | |
| - Always use valid markdown for the report text. | |
| </plots_html_usage_guide> | |
| <plots_html_format_guide> | |
| **IMPORTANT: plots_html Parameter Format** | |
| - The plots_html parameter accepts either: | |
| 1. A dictionary where keys are chart names and values are HTML strings: {"chart1": "<html>...</html>", "chart2": "<html>...</html>"} | |
| 2. A single HTML string (will be automatically wrapped with key "chart") | |
| - Do NOT pass raw HTML as a JSON string - pass it directly as an HTML string | |
| - If you have multiple charts, use the dictionary format for better organization | |
| - The tool will provide feedback about how your input was processed | |
| </plots_html_format_guide> | |
| Args: | |
| entity_name: str, The W&B entity (team or username) - required | |
| project_name: str, The W&B project name - required | |
| title: str, Title of the W&B Report - required | |
| description: str, Optional description of the W&B Report | |
| markdown_report_text: str, beuatifully formatted markdown text for the report body | |
| plots_html: str, Optional dict of plot name and html string of any charts created as part of an analysis | |
| Returns: | |
| str, The url to the report | |
| Example: | |
| ```python | |
| # Create a simple report | |
| report = create_report( | |
| entity_name="my-team", | |
| project_name="my-project", | |
| title="Model Analysis Report", | |
| description="Analysis of our latest model performance", | |
| markdown_report_text=''' | |
| # Model Analysis Report | |
| [TOC] | |
| ## Performance Summary | |
| Our model achieved 95% accuracy on the test set. | |
| ### Key Metrics | |
| Precision: 0.92 | |
| Recall: 0.89 | |
| ''' | |
| ) | |
| ``` | |
| """ | |
| def create_report( | |
| entity_name: str, | |
| project_name: str, | |
| title: str, | |
| description: Optional[str] = None, | |
| markdown_report_text: Optional[str] = None, | |
| plots_html: Optional[Union[Dict[str, str], str]] = None, | |
| ) -> Dict[str, str]: | |
| """ | |
| Create a new Weights & Biases Report and add text and charts. Useful to save/document analysis and other findings. | |
| Args: | |
| entity_name: The W&B entity (team or username) | |
| project_name: The W&B project name | |
| title: Title of the W&B Report | |
| description: Optional description of the W&B Report | |
| markdown_report_text: Optional markdown text for the report body | |
| plots_html: Optional dict of plot name and html string, or single HTML string | |
| Returns: | |
| Dict with 'url' and 'processing_details' keys | |
| """ | |
| import json | |
| # Process plots_html and collect warnings | |
| processed_plots_html = None | |
| processing_warnings = [] | |
| if isinstance(plots_html, str): | |
| try: | |
| # First try to parse as JSON (dictionary) | |
| processed_plots_html = json.loads(plots_html) | |
| processing_warnings.append("Successfully parsed plots_html as JSON dictionary") | |
| except json.JSONDecodeError: | |
| # If it's not valid JSON, treat as raw HTML and wrap in dictionary | |
| if plots_html.strip(): # Only if not empty | |
| processed_plots_html = {"chart": plots_html} | |
| processing_warnings.append("plots_html was not valid JSON, treated as raw HTML and wrapped with key 'chart'") | |
| else: | |
| processed_plots_html = None | |
| processing_warnings.append("plots_html was empty string, no charts will be included") | |
| elif isinstance(plots_html, dict): | |
| processed_plots_html = plots_html | |
| processing_warnings.append(f"Successfully processed plots_html dictionary with {len(plots_html)} chart(s)") | |
| elif plots_html is None: | |
| processing_warnings.append("No plots_html provided, report will contain only text content") | |
| else: | |
| processing_warnings.append(f"Unexpected plots_html type: {type(plots_html)}, no charts will be included") | |
| processed_plots_html = None | |
| # Get the current API key from context | |
| from wandb_mcp_server.api_client import WandBApiManager | |
| api_key = WandBApiManager.get_api_key() | |
| # Store original environment key | |
| import os | |
| old_key = os.environ.get("WANDB_API_KEY") | |
| try: | |
| # Set API key temporarily for wandb.init() | |
| if api_key: | |
| os.environ["WANDB_API_KEY"] = api_key | |
| wandb.init( | |
| entity=entity_name, project=project_name, job_type="mcp_report_creation" | |
| ) | |
| # Initialize the report | |
| report = wr.Report( | |
| entity=entity_name, | |
| project=project_name, | |
| title=title, | |
| description=description or "", | |
| width="fluid", | |
| ) | |
| # Log plots | |
| plots_dict = {} | |
| if processed_plots_html: | |
| for plot_name, html in processed_plots_html.items(): | |
| wandb.log({plot_name: wandb.Html(html)}) | |
| plots_dict[plot_name] = html | |
| pg = [] | |
| for k, v in plots_dict.items(): | |
| pg.append( | |
| wr.PanelGrid( | |
| panels=[ | |
| wr.MediaBrowser( | |
| media_keys=[k], | |
| num_columns=1, | |
| layout=wr.Layout(w=20, h=20), # , x=5, y=5) | |
| ), | |
| ] | |
| ) | |
| ) | |
| else: | |
| pg = None | |
| blocks = parse_report_content_enhanced(markdown_report_text or "") | |
| # Add blocks if provided | |
| if pg: | |
| report.blocks = blocks + pg | |
| else: | |
| report.blocks = blocks | |
| logger.info(f"Report blocks: {report.blocks}") | |
| # Save the report | |
| report.save() | |
| wandb.finish() | |
| logger.info(f"Created report: {title}") | |
| return { | |
| "url": report.url, | |
| "processing_details": processing_warnings | |
| } | |
| except Exception as e: | |
| logger.error(f"Error creating report: {e}") | |
| # Include processing details in the error for better debugging | |
| error_msg = f"Error creating report: {e}" | |
| if processing_warnings: | |
| error_msg += f"\n\nProcessing details: {'; '.join(processing_warnings)}" | |
| raise Exception(error_msg) | |
| finally: | |
| # Restore original environment variable | |
| if old_key: | |
| os.environ["WANDB_API_KEY"] = old_key | |
| elif "WANDB_API_KEY" in os.environ: | |
| del os.environ["WANDB_API_KEY"] | |
| def edit_report( | |
| report_url: str, | |
| title: Optional[str] = None, | |
| description: Optional[str] = None, | |
| blocks: Optional[List[Union[wr.H1, wr.H2, wr.H3, wr.PanelGrid]]] = None, | |
| append_blocks: bool = False, | |
| ) -> wr.Report: | |
| """ | |
| Edit an existing W&B report. | |
| Args: | |
| report_url: The URL of the report to edit | |
| title: Optional new title for the report | |
| description: Optional new description for the report | |
| blocks: Optional list of blocks to update or append | |
| append_blocks: If True, append new blocks to existing ones. If False, replace existing blocks. | |
| """ | |
| try: | |
| # Load the existing report | |
| report = wr.Report.from_url(report_url) | |
| # Update title if provided | |
| if title: | |
| report.title = title | |
| # Update description if provided | |
| if description: | |
| report.description = description | |
| # Update blocks if provided | |
| if blocks: | |
| if append_blocks: | |
| # Append new blocks to existing ones | |
| report.blocks = (report.blocks or []) + blocks | |
| else: | |
| # Replace existing blocks | |
| report.blocks = blocks | |
| # Save the changes | |
| report.save() | |
| logger.info(f"Updated report: {report.title}") | |
| return report | |
| except Exception as e: | |
| logger.error(f"Error editing report: {e}") | |
| raise | |
| def parse_report_content_enhanced( | |
| text: str, | |
| ) -> List[Union[wr.H1, wr.H2, wr.H3, wr.P, wr.TableOfContents]]: | |
| """ | |
| Parse markdown-like text into W&B report blocks with paragraph grouping. | |
| """ | |
| blocks = [] | |
| lines = text.strip().split("\n") | |
| current_paragraph = [] | |
| for line in lines: | |
| # Check if this is a special line (header or TOC) | |
| h1_match = re.match(r"^# (.+)$", line) | |
| h2_match = re.match(r"^## (.+)$", line) | |
| h3_match = re.match(r"^### (.+)$", line) | |
| is_toc = line.strip().lower() == "[toc]" | |
| # If we hit a special line and have paragraph content, finalize the paragraph | |
| if (h1_match or h2_match or h3_match or is_toc) and current_paragraph: | |
| blocks.append(wr.P("\n".join(current_paragraph))) | |
| current_paragraph = [] | |
| # Handle the current line | |
| if h1_match: | |
| blocks.append(wr.H1(h1_match.group(1))) | |
| elif h2_match: | |
| blocks.append(wr.H2(h2_match.group(1))) | |
| elif h3_match: | |
| blocks.append(wr.H3(h3_match.group(1))) | |
| elif is_toc: | |
| blocks.append(wr.TableOfContents()) | |
| else: | |
| if line.strip(): # Only add non-empty lines | |
| current_paragraph.append(line) | |
| # Don't forget any remaining paragraph content | |
| if current_paragraph: | |
| blocks.append(wr.P("\n".join(current_paragraph))) | |
| return blocks | |