From 0b801a7799065fcdfe9104e7d58a8a87e7de7511 Mon Sep 17 00:00:00 2001 From: Xiyun Hu Date: Tue, 17 Mar 2026 00:04:44 -0400 Subject: [PATCH] Improve visualizer agent structure --- .../requirements.txt | 2 + .../visualizer_agent.py | 1169 ++++++++++++++--- scripts/local_visualizer_smoke.sh | 265 ++++ 3 files changed, 1240 insertions(+), 196 deletions(-) create mode 100644 cyberneuro_multi-agent-neuroimaging-analysis/requirements.txt create mode 100755 scripts/local_visualizer_smoke.sh diff --git a/cyberneuro_multi-agent-neuroimaging-analysis/requirements.txt b/cyberneuro_multi-agent-neuroimaging-analysis/requirements.txt new file mode 100644 index 0000000..02ef72a --- /dev/null +++ b/cyberneuro_multi-agent-neuroimaging-analysis/requirements.txt @@ -0,0 +1,2 @@ +beautifulsoup4>=4.12.3 +ollama>=0.3.3 diff --git a/cyberneuro_multi-agent-neuroimaging-analysis/visualizer_agent.py b/cyberneuro_multi-agent-neuroimaging-analysis/visualizer_agent.py index e9dc5e9..68fe304 100644 --- a/cyberneuro_multi-agent-neuroimaging-analysis/visualizer_agent.py +++ b/cyberneuro_multi-agent-neuroimaging-analysis/visualizer_agent.py @@ -1,7 +1,13 @@ -import re import asyncio -from typing import Optional, List -from dataclasses import dataclass +import json +import logging +import re +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Callable, Dict, List, Optional, Tuple + +from bs4 import BeautifulSoup, Tag + # ============================================================ # CONFIGURATION @@ -14,261 +20,1033 @@ # Default model - change this to your preferred model DEFAULT_MODEL = "llama3.1:8b" +# Token budgets +MAX_PATCH_TOKENS = 800 +MAX_FULL_TOKENS = 4096 + +logger = logging.getLogger(__name__) + + # ============================================================ # DATA CLASSES # ============================================================ + @dataclass class ValidationResult: """Result of HTML validation.""" + valid: bool error: Optional[str] = None warnings: Optional[List[str]] = None + @dataclass -class EditResult: - """Result of visualization edit operation.""" +class VisualizerResponse: + """External response shape kept for MCP and FastAPI compatibility.""" + status: str # "success" or "error" html: Optional[str] = None message: Optional[str] = None warnings: Optional[List[str]] = None + +class EditIntent(Enum): + TITLE_CHANGE = "title_change" + SUBTITLE_CHANGE = "subtitle_change" + COLOR_CHANGE = "color_change" + AXIS_LABEL = "axis_label" + LEGEND_TOGGLE = "legend_toggle" + DATA_UPDATE = "data_update" + FONT_SIZE = "font_size" + ADD_ELEMENT = "add_element" + REMOVE_ELEMENT = "remove_element" + LAYOUT_CHANGE = "layout_change" + STYLE_CHANGE = "style_change" + COMPLEX = "complex" + + +@dataclass +class ClassifiedEdit: + intent: EditIntent + target_selector: Optional[str] = None + new_value: Optional[str] = None + parameters: dict = field(default_factory=dict) + confidence: float = 1.0 + + +@dataclass +class EditResult: + success: bool + html: str + method: str # "programmatic", "scoped_llm", "full_rewrite", "none" + changes_description: str + tokens_used: int = 0 + + +# ============================================================ +# PROMPTS +# ============================================================ + +INTENT_CLASSIFIER_SYSTEM_PROMPT = """You are an edit intent classifier. Given a user's request to modify an HTML visualization, return a JSON object with: +- \"intent\": one of [title_change, subtitle_change, color_change, axis_label, legend_toggle, data_update, font_size, add_element, remove_element, layout_change, style_change, complex] +- \"target_css\": CSS selector for the element to edit, or null +- \"new_value\": the new value if it's a simple replacement, or null +- \"confidence\": 0.0-1.0 + +Return ONLY the JSON object, nothing else.""" + + +SCOPED_PATCH_SYSTEM_PROMPT = """You are an HTML visualization editor. You will receive: +1. A user's edit request +2. A FRAGMENT of HTML - the specific section that needs editing + +Return ONLY the modified HTML fragment. Do not add explanations, markdown fences, or any text outside the HTML. The fragment must be valid HTML and must preserve the same root element tag and class structure. + +Rules: +- Make ONLY the change the user requested +- Preserve all existing classes, ids, and data attributes +- Preserve the dark theme (use CSS variables, not hardcoded light colors) +- Do not add external scripts or stylesheets +- If the requested change is unclear, make your best interpretation""" + + +SYSTEM_PROMPT = """You are an HTML visualization editor. Return ONLY the complete modified HTML card. Rules: +- Start with a div with class \"visualizationCard\" +- Do not use external scripts or stylesheets +- Preserve the dark theme using CSS variables +- Use the vc-header, vc-title, vc-subtitle, vc-body, vc-footer, vc-insight structure +- If data is missing, show a warning block instead of inventing values +- Make ONLY the change the user requested - do not alter anything else""" + + +# ============================================================ +# CLASSIFICATION CONSTANTS +# ============================================================ + +COLOR_NORMALIZATION_MAP: Dict[str, str] = { + "red": "#E24B4A", + "blue": "#378ADD", + "green": "#639922", + "yellow": "#EF9F27", + "orange": "#D85A30", + "purple": "#7F77DD", + "pink": "#D4537E", + "gray": "#888780", + "grey": "#888780", + "teal": "#1D9E75", + "coral": "#D85A30", + "amber": "#BA7517", + "black": "#2C2C2A", + "white": "#F1EFE8", +} + +COLOR_WORD_PATTERN = "|".join(sorted(COLOR_NORMALIZATION_MAP.keys(), key=len, reverse=True)) + +ELEMENT_SELECTOR_MAP: Dict[str, str] = { + "title": ".vc-title, h1, h2, .title, [class*='title']", + "subtitle": ".vc-subtitle, .subtitle, [class*='subtitle'], .description", + "description": ".vc-subtitle, .subtitle, [class*='subtitle'], .description", + "legend": ".vc-legend, [class*='legend']", + "grid": "[class*='grid'], .grid", + "tooltip": "[class*='tooltip'], .tooltip", + "label": "[class*='label'], .label", + "annotation": ".vc-insight, [class*='annotation']", + "footer": ".vc-footer, footer", + "header": ".vc-header, header", + "insight": ".vc-insight", +} + +INTENT_PATTERNS: Dict[EditIntent, List[re.Pattern[str]]] = { + EditIntent.TITLE_CHANGE: [ + re.compile(r"\b(?:change|update|set|make|rename|edit)\b[^\n]{0,120}\btitle\b", re.IGNORECASE), + re.compile(r"\btitle\b[^\n]{0,120}\b(?:to|should\s+be|as|say|be)\b", re.IGNORECASE), + ], + EditIntent.SUBTITLE_CHANGE: [ + re.compile(r"\b(?:change|update|set|make|rename|edit)\b[^\n]{0,120}\b(?:subtitle|description)\b", re.IGNORECASE), + re.compile(r"\b(?:subtitle|description)\b[^\n]{0,120}\b(?:to|should\s+be|as|say|be)\b", re.IGNORECASE), + ], + EditIntent.COLOR_CHANGE: [ + re.compile( + rf"\b(?:change|make|set|turn|color)\b[^\n]{{0,120}}\b(?:bars?|lines?|area|pie|slices?|background|fill|stroke|chart)\b[^\n]{{0,80}}(?:{COLOR_WORD_PATTERN}|#[0-9a-fA-F]{{3,8}})\b", + re.IGNORECASE, + ), + re.compile( + rf"(?:{COLOR_WORD_PATTERN}|#[0-9a-fA-F]{{3,8}})\b[^\n]{{0,80}}\b(?:bars?|lines?|area|pie|slices?|background|fill|stroke)\b", + re.IGNORECASE, + ), + ], + EditIntent.AXIS_LABEL: [ + re.compile(r"\b(?:change|update|set|rename|label)\b[^\n]{0,120}\b(?:x|y|horizontal|vertical)\s*-?\s*axis\b", re.IGNORECASE), + re.compile(r"\b(?:x|y|horizontal|vertical)\s*-?\s*axis\b[^\n]{0,120}\b(?:label|title|name)\b", re.IGNORECASE), + ], + EditIntent.LEGEND_TOGGLE: [ + re.compile(r"\b(?:show|hide|toggle|add|remove)\b[^\n]{0,60}\blegend\b", re.IGNORECASE), + ], + EditIntent.FONT_SIZE: [ + re.compile(r"\b(?:font|text)\s*size\b", re.IGNORECASE), + re.compile(r"\b(?:bigger|smaller|larger)\b[^\n]{0,40}\b(?:text|font)\b", re.IGNORECASE), + re.compile(r"\b\d+\s*px\b", re.IGNORECASE), + ], + EditIntent.ADD_ELEMENT: [ + re.compile(r"\b(?:add|include|insert)\b[^\n]{0,120}\b(?:title|subtitle|legend|grid|tooltip|label|annotation|footer|header|insight)\b", re.IGNORECASE), + ], + EditIntent.REMOVE_ELEMENT: [ + re.compile(r"\b(?:remove|delete|hide|drop)\b[^\n]{0,120}\b(?:title|subtitle|legend|grid|tooltip|label|annotation|footer|header|insight)\b", re.IGNORECASE), + ], + EditIntent.DATA_UPDATE: [ + re.compile(r"\b(?:update|change|replace)\b[^\n]{0,120}\b(?:data|values|numbers|dataset|series)\b", re.IGNORECASE), + ], + EditIntent.LAYOUT_CHANGE: [ + re.compile(r"\b(?:layout|arrange|alignment|position|move|resize|spacing)\b", re.IGNORECASE), + ], + EditIntent.STYLE_CHANGE: [ + re.compile(r"\b(?:style|theme|professional|cleaner|modern)\b", re.IGNORECASE), + ], +} + +INTENT_PRIORITY: List[EditIntent] = [ + EditIntent.TITLE_CHANGE, + EditIntent.SUBTITLE_CHANGE, + EditIntent.COLOR_CHANGE, + EditIntent.AXIS_LABEL, + EditIntent.REMOVE_ELEMENT, + EditIntent.ADD_ELEMENT, + EditIntent.LEGEND_TOGGLE, + EditIntent.FONT_SIZE, + EditIntent.DATA_UPDATE, + EditIntent.LAYOUT_CHANGE, + EditIntent.STYLE_CHANGE, +] + + # ============================================================ # VALIDATION FUNCTIONS # ============================================================ + def validate_visualization_html(html: str, allow_external_images: bool = False) -> ValidationResult: """ Validate HTML content meets visualization card requirements. """ warnings: List[str] = [] - + if not html or not isinstance(html, str): return ValidationResult(valid=False, error="HTML content is required") - + if len(html) > MAX_HTML_SIZE: return ValidationResult( - valid=False, - error=f"HTML content exceeds maximum size limit ({MAX_HTML_SIZE} characters)" + valid=False, + error=f"HTML content exceeds maximum size limit ({MAX_HTML_SIZE} characters)", ) - + if REQUIRED_ROOT_CLASS not in html: return ValidationResult( valid=False, - error=f"Missing required .{REQUIRED_ROOT_CLASS} root element" + error=f"Missing required .{REQUIRED_ROOT_CLASS} root element", ) - + # Block external script sources - external_script_pattern = re.compile(r']+src\s*=', re.IGNORECASE) + external_script_pattern = re.compile(r"]+src\s*=", re.IGNORECASE) if external_script_pattern.search(html): return ValidationResult( valid=False, - error="External script sources (' ) assert not result.valid, "Test 3 failed: Should reject external scripts" - print("✓ Test 3: Rejects external script sources") - + print("OK Test 3: Rejects external script sources") + # Test 4: Inline script allowed result = validate_visualization_html( '
' ) assert result.valid, f"Test 4 failed: {result.error}" - print("✓ Test 4: Allows inline scripts") - + print("OK Test 4: Allows inline scripts") + # Test 5: External stylesheet result = validate_visualization_html( '
' ) assert not result.valid, "Test 5 failed: Should reject external stylesheets" - print("✓ Test 5: Rejects external stylesheets") - + print("OK Test 5: Rejects external stylesheets") + # Test 6: Clean LLM output messy_output = "```html\n
test
\n```" cleaned = clean_llm_output(messy_output) assert cleaned == '
test
', f"Test 6 failed: {cleaned}" - print("✓ Test 6: Cleans markdown fences from output") - + print("OK Test 6: Cleans markdown fences from output") + print("\nAll validation tests passed!") async def quick_test(provider: str = "ollama", model: str = DEFAULT_MODEL): """Quick test of the edit function.""" print(f"\nQuick test using {provider} with model '{model}'...") - + result = await visualizer_edit_html( user_query="Add a footer that says 'Source: UK Biobank'", html=STARTER_HTML, provider=provider, - model=model + model=model, ) - + if result.status == "success": print("Success!") print(f"\nOutput HTML:\n{result.html}") else: print(f"Error: {result.message}") - + return result if __name__ == "__main__": import sys - + if len(sys.argv) > 1 and sys.argv[1] == "--test": # Run validation tests run_validation_tests() @@ -471,4 +1248,4 @@ async def quick_test(provider: str = "ollama", model: str = DEFAULT_MODEL): else: print("Usage:") print(" python visualizer_agent.py --test # Run validation tests") - print(" python visualizer_agent.py --quick # Quick LLM test (Ollama)") + print(" python visualizer_agent.py --quick # Quick LLM test (Ollama)") \ No newline at end of file diff --git a/scripts/local_visualizer_smoke.sh b/scripts/local_visualizer_smoke.sh new file mode 100755 index 0000000..ab0e21c --- /dev/null +++ b/scripts/local_visualizer_smoke.sh @@ -0,0 +1,265 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ROOT_DIR="$(cd "${SCRIPT_DIR}/.." && pwd)" +APP_DIR="${ROOT_DIR}/cyberneuro_multi-agent-neuroimaging-analysis" +MCP_DIR="${ROOT_DIR}/mcp_server" + +OLLAMA_URL="${OLLAMA_URL:-http://127.0.0.1:11434}" +FRONTEND_HOST="${FRONTEND_HOST:-127.0.0.1}" +FRONTEND_PORT="${FRONTEND_PORT:-5176}" +VISUALIZER_MODEL="${VISUALIZER_MODEL:-qwen2.5-coder:32b}" + +RUN_FRONTEND=1 +RUN_PY_QUICK=1 +WITH_MCP=0 +CHECK_ONLY=0 + +usage() { + cat <<'EOF' +Local visualizer smoke test. + +Usage: + bash scripts/local_visualizer_smoke.sh [options] + +Options: + --with-mcp Start MCP server (8010) and Node backend (8789) + --no-python-quick Skip python visualizer quick test + --no-frontend Skip frontend startup + --check-only Run health checks only, then exit + -h, --help Show help + +Environment overrides: + OLLAMA_URL (default: http://127.0.0.1:11434) + FRONTEND_HOST (default: 127.0.0.1) + FRONTEND_PORT (default: 5176) + VISUALIZER_MODEL (default: qwen2.5-coder:32b) +EOF +} + +while [[ $# -gt 0 ]]; do + case "$1" in + --with-mcp) + WITH_MCP=1 + ;; + --no-python-quick) + RUN_PY_QUICK=0 + ;; + --no-frontend) + RUN_FRONTEND=0 + ;; + --check-only) + CHECK_ONLY=1 + RUN_FRONTEND=0 + ;; + -h|--help) + usage + exit 0 + ;; + *) + echo "Unknown option: $1" >&2 + usage + exit 1 + ;; + esac + shift +done + +require_cmd() { + if ! command -v "$1" >/dev/null 2>&1; then + echo "ERROR: Missing required command: $1" >&2 + exit 1 + fi +} + +for cmd in curl python3 npm; do + require_cmd "$cmd" +done +if [[ "$WITH_MCP" -eq 1 ]]; then + require_cmd uv +fi + +if [[ ! -d "$APP_DIR" ]]; then + echo "ERROR: App directory not found: $APP_DIR" >&2 + exit 1 +fi + +TAGS_FILE="$(mktemp)" + +cleanup() { + local exit_code=$? + trap - EXIT INT TERM + + rm -f "$TAGS_FILE" >/dev/null 2>&1 || true + + for pid in "${FRONTEND_PID:-}" "${BACKEND_PID:-}" "${MCP_PID:-}"; do + if [[ -n "$pid" ]] && kill -0 "$pid" >/dev/null 2>&1; then + kill "$pid" >/dev/null 2>&1 || true + fi + done + + exit "$exit_code" +} +trap cleanup EXIT INT TERM + +pick_model() { + python3 - "$1" "$2" <<'PY' +import json +import sys + +requested = sys.argv[1] +path = sys.argv[2] + +with open(path, "r", encoding="utf-8") as f: + data = json.load(f) + +models = [m.get("name", "") for m in data.get("models", []) if m.get("name")] + +if requested in models: + print(requested) + raise SystemExit(0) + +preferred = [ + "qwen2.5-coder:32b", + "qwen2.5-coder:32b-instruct", + "llama3.1:8b", + "llama3", +] + +for name in preferred: + if name in models: + print(name) + raise SystemExit(0) + +for name in models: + if "coder" in name.lower(): + print(name) + raise SystemExit(0) + +print(models[0] if models else "") +PY +} + +wait_for_http() { + local url="$1" + local timeout="$2" + local i + + for ((i=1; i<=timeout; i++)); do + if curl -fsS "$url" >/dev/null 2>&1; then + return 0 + fi + sleep 1 + done + return 1 +} + +echo "[1/5] Checking Ollama at ${OLLAMA_URL} ..." +if ! curl -fsS "${OLLAMA_URL}/api/tags" > "$TAGS_FILE"; then + echo "ERROR: Ollama is not reachable at ${OLLAMA_URL}" >&2 + echo "Start Ollama, then verify with: curl ${OLLAMA_URL}/api/tags" >&2 + exit 1 +fi +echo "OK Ollama is reachable." + +SELECTED_MODEL="$(pick_model "$VISUALIZER_MODEL" "$TAGS_FILE")" +if [[ -z "$SELECTED_MODEL" ]]; then + echo "ERROR: No Ollama models are available. Pull a model first." >&2 + exit 1 +fi + +if [[ "$SELECTED_MODEL" != "$VISUALIZER_MODEL" ]]; then + echo "Requested model '${VISUALIZER_MODEL}' not found. Using '${SELECTED_MODEL}' instead." +else + echo "Using model '${SELECTED_MODEL}'." +fi + +if [[ "$RUN_PY_QUICK" -eq 1 ]]; then + echo "[2/5] Running python visualizer quick test ..." + if (cd "$APP_DIR" && python3 -c "import bs4, ollama" >/dev/null 2>&1); then + (cd "$APP_DIR" && python3 visualizer_agent.py --quick ollama "$SELECTED_MODEL") + else + echo "WARNING: Python deps missing for quick test." + echo "Install with: cd $APP_DIR && python3 -m pip install -r requirements.txt" + fi +else + echo "[2/5] Skipping python visualizer quick test (--no-python-quick)." +fi + +if [[ "$CHECK_ONLY" -eq 1 ]]; then + echo "Health checks passed. Exiting (--check-only)." + exit 0 +fi + +if [[ "$WITH_MCP" -eq 1 ]]; then + echo "[3/5] Starting MCP server on http://127.0.0.1:8010 ..." + ( + cd "$MCP_DIR" + uv sync >/dev/null + uvicorn mcp_server:http_app --host 127.0.0.1 --port 8010 + ) >/tmp/visualizer_mcp.log 2>&1 & + MCP_PID=$! + + if ! wait_for_http "http://127.0.0.1:8010/health" 30; then + echo "ERROR: MCP server failed to start. See /tmp/visualizer_mcp.log" >&2 + exit 1 + fi + echo "OK MCP server is healthy." + + echo "[4/5] Starting Node backend on http://127.0.0.1:8789 ..." + ( + cd "$APP_DIR" + FRONTEND_ORIGIN="http://${FRONTEND_HOST}:${FRONTEND_PORT}" \ + PORT=8789 \ + MCP_SERVER_URL="http://127.0.0.1:8010/mcp" \ + npm run backend + ) >/tmp/visualizer_backend.log 2>&1 & + BACKEND_PID=$! + + if ! wait_for_http "http://127.0.0.1:8789/health" 30; then + echo "ERROR: Backend failed to start. See /tmp/visualizer_backend.log" >&2 + exit 1 + fi + echo "OK backend is healthy." +else + echo "[3/5] Skipping MCP/backend startup (--with-mcp to enable)." +fi + +if [[ "$RUN_FRONTEND" -eq 1 ]]; then + if [[ ! -d "$APP_DIR/node_modules" ]]; then + echo "[4/5] Installing frontend dependencies ..." + (cd "$APP_DIR" && npm install) + else + echo "[4/5] Frontend dependencies already installed." + fi + + echo "[5/5] Starting frontend on http://${FRONTEND_HOST}:${FRONTEND_PORT} ..." + ( + cd "$APP_DIR" + npm run dev -- --host "$FRONTEND_HOST" --port "$FRONTEND_PORT" + ) & + FRONTEND_PID=$! + + if ! wait_for_http "http://${FRONTEND_HOST}:${FRONTEND_PORT}" 60; then + echo "ERROR: Frontend failed to start on port ${FRONTEND_PORT}." >&2 + exit 1 + fi + + cat <