Fu-Jie_openwebui-extensions/plugins/actions/infographic/verify_generation.py

import os
import sys
import json
import requests
from datetime import datetime
from dotenv import load_dotenv
import pathlib

# Load .env from the same directory as this script
env_path = pathlib.Path(__file__).parent / ".env"
load_dotenv(dotenv_path=env_path)
# =================================================================
# Configuration
# =================================================================

API_KEY = os.getenv("OPENAI_API_KEY")
BASE_URL = os.getenv("OPENAI_BASE_URL")
MODEL = os.getenv("OPENAI_MODEL", "gpt-4o")  # Default to gpt-4o if not set

if not API_KEY or not BASE_URL:
    print(
        "Error: OPENAI_API_KEY and OPENAI_BASE_URL environment variables must be set."
    )
    sys.exit(1)

# =================================================================
# Prompts (Extracted from infographic_cn.py)
# =================================================================

SYSTEM_PROMPT_INFOGRAPHIC_ASSISTANT = """
You are a professional infographic design expert who can analyze user-provided text content and convert it into AntV Infographic syntax format.

## Infographic Syntax Specification

Infographic syntax is a Mermaid-like declarative syntax for describing infographic templates, data, and themes.

### Syntax Rules
- Entry uses `infographic <template-name>`
- Key-value pairs are separated by spaces, **absolutely NO colons allowed**
- Use two spaces for indentation
- Object arrays use `-` with line breaks

⚠️ **IMPORTANT WARNING: This is NOT YAML format!**
- ❌ Wrong: `children:` `items:` `data:` (with colons)
- ✅ Correct: `children` `items` `data` (without colons)

### Template Library & Selection Guide

#### 1. List & Hierarchy (Text-heavy)
-   **Linear & Short (Steps/Phases)** -> `list-row-horizontal-icon-arrow`
-   **Linear & Long (Rankings/Details)** -> `list-vertical`
-   **Grouped / Parallel (Features/Catalog)** -> `list-grid`
-   **Hierarchical (Org Chart/Taxonomy)** -> `tree-vertical` or `tree-horizontal`
-   **Central Idea (Brainstorming)** -> `mindmap`

#### 2. Sequence & Relationship (Flow-based)
-   **Time-based (History/Plan)** -> `sequence-roadmap-vertical-simple`
-   **Process Flow (Complex)** -> `sequence-zigzag` or `sequence-horizontal`
-   **Resource Flow / Distribution** -> `relation-sankey`
-   **Circular Relationship** -> `relation-circle`

#### 3. Comparison & Analysis
-   **Binary Comparison (A vs B)** -> `compare-binary`
-   **SWOT Analysis** -> `compare-swot`
-   **Multi-item Comparison Table** -> `compare-table`
-   **Quadrant Analysis (Importance vs Urgency)** -> `quadrant-quarter`

#### 4. Charts & Data (Metric-heavy)
-   **Key Metrics / Data Cards** -> `statistic-card`
-   **Distribution / Comparison** -> `chart-bar` or `chart-column`
-   **Trend over Time** -> `chart-line` or `chart-area`
-   **Proportion / Part-to-Whole** -> `chart-pie` or `chart-doughnut`

### Infographic Syntax Guide

#### 1. Structure
-   **Entry**: `infographic <template-name>`
-   **Blocks**: `data`, `theme`, `design` (optional)
-   **Format**: Key-value pairs separated by spaces, 2-space indentation.
-   **Arrays**: Object arrays use `-` (newline), simple arrays use inline values.

#### 2. Data Block (`data`)
-   `title`: Main title
-   `desc`: Subtitle or description
-   `items`: List of data items
-     - `label`: Item title
-     - `value`: Numerical value (required for Charts/Stats)
-     - `desc`: Item description (optional)
-     - `icon`: Icon name (e.g., `mdi/rocket-launch`)
-     - `time`: Time label (Optional, for Roadmap/Sequence)
-     - `children`: Nested items (ONLY for Tree/Mindmap/Sankey/SWOT)
-     - `illus`: Illustration name (ONLY for Quadrant)

#### 3. Theme Block (`theme`)
-   `colorPrimary`: Main color (Hex)
-   `colorBg`: Background color (Hex)
-   `palette`: Color list (Space separated)
-   `textColor`: Text color (Hex)
-   `stylize`: Style effect (e.g., `rough`, `flat`)

### Examples

#### Chart (Bar Chart)
infographic chart-bar
data
  title Revenue Growth
  desc Monthly revenue in 2024
  items
    - label Jan
      value 1200
    - label Feb
      value 1500
    - label Mar
      value 1800

#### Comparison (SWOT)
infographic compare-swot
data
  title Project SWOT
  items
    - label Strengths
      children
        - label Strong team
        - label Innovative tech
    - label Weaknesses
      children
        - label Limited budget
    - label Opportunities
      children
        - label Emerging market
    - label Threats
      children
        - label High competition

#### Relationship (Sankey)
infographic relation-sankey
data
  title Energy Flow
  items
    - label Solar
      value 100
      children
        - label Grid
          value 60
        - label Battery
          value 40
    - label Wind
      value 80
      children
        - label Grid
          value 80

#### Quadrant (Importance vs Urgency)
infographic quadrant-quarter
data
  title Task Management
  items
    - label Critical Bug
      desc Fix immediately
      illus mdi/bug
    - label Feature Request
      desc Plan for next sprint
      illus mdi/star

### Output Rules
1.  **Strict Syntax**: Follow the indentation and formatting rules exactly.
2.  **No Explanations**: Output ONLY the syntax code block.
3.  **Language**: Use the user's requested language for content.
"""

USER_PROMPT_GENERATE_INFOGRAPHIC = """
请分析以下文本内容，将其核心信息转换为 AntV Infographic 语法格式。

---
**用户上下文信息:**
用户姓名: {user_name}
当前日期时间: {current_date_time_str}
用户语言: {user_language}
---

**文本内容:**
{long_text_content}

请根据文本特点选择最合适的信息图模板，并输出规范的 infographic 语法。注意保持正确的缩进格式（两个空格）。
"""

# =================================================================
# Test Cases
# =================================================================

TEST_CASES = [
    {
        "name": "List Grid (Features)",
        "content": """
        MiniMax 2025 模型矩阵全解析：
        1. 极致 MoE 架构优化：国内首批 MoE 路线坚定者。
        2. Video-01 视频生成：杀手锏级多模态能力。
        3. 情感陪伴与角色扮演：源自星野基因。
        4. 超长上下文与精准召回：支持 128k+ 窗口。
        """,
    },
    {
        "name": "Tree Vertical (Hierarchy)",
        "content": """
        公司组织架构：
        - 研发部
            - 后端组
            - 前端组
        - 市场部
            - 销售组
            - 推广组
        """,
    },
    {
        "name": "Statistic Card (Metrics)",
        "content": """
        2024年Q4 核心指标：
        - 总用户数：1,234,567
        - 日活跃用户：89%
        - 营收增长：+45%
        """,
    },
    {
        "name": "Mindmap (Brainstorming)",
        "content": """
        人工智能的应用领域：
        - 生成式 AI：文本生成、图像生成、视频生成
        - 预测性 AI：股市预测、天气预报
        - 决策 AI：自动驾驶、游戏博弈
        """,
    },
    {
        "name": "SWOT Analysis",
        "content": """
        某初创咖啡品牌的 SWOT 分析：
        优势：产品口味独特，选址精准。
        劣势：品牌知名度低，资金压力大。
        机会：线上外卖市场增长，年轻人对精品咖啡需求增加。
        威胁：行业巨头价格战，原材料成本上涨。
        """,
    },
    {
        "name": "Sankey (Relationship)",
        "content": """
        家庭月度开支流向：
        总收入 10000 元。
        其中 4000 元用于房贷。
        3000 元用于日常消费（包括 2000 元餐饮，1000 元交通）。
        2000 元用于教育培训。
        1000 元存入银行。
        """,
    },
    {
        "name": "Quadrant (Analysis)",
        "content": """
        个人任务象限分析：
        - 紧急且重要：修复线上紧急 Bug，准备下午的客户会议。
        - 重要不紧急：制定年度学习计划，健身锻炼。
        - 紧急不重要：回复琐碎的邮件，接听推销电话。
        - 不紧急不重要：刷社交媒体，看无聊的综艺。
        """,
    },
    {
        "name": "Chart Bar (Data)",
        "content": """
        2023年各季度营收情况（亿元）：
        第一季度：12.5
        第二季度：15.8
        第三季度：14.2
        第四季度：18.9
        """,
    },
]

# =================================================================
# Helper Functions
# =================================================================


def generate_infographic(content):
    now = datetime.now()
    current_date_time_str = now.strftime("%Y年%m月%d日 %H:%M:%S")

    formatted_user_prompt = USER_PROMPT_GENERATE_INFOGRAPHIC.format(
        user_name="TestUser",
        current_date_time_str=current_date_time_str,
        user_language="zh-CN",
        long_text_content=content,
    )

    headers = {"Authorization": f"Bearer {API_KEY}", "Content-Type": "application/json"}

    payload = {
        "model": MODEL,
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT_INFOGRAPHIC_ASSISTANT},
            {"role": "user", "content": formatted_user_prompt},
        ],
        "stream": False,
    }

    try:
        response = requests.post(
            f"{BASE_URL}/chat/completions", headers=headers, json=payload
        )
        response.raise_for_status()
        return response.json()["choices"][0]["message"]["content"]
    except Exception as e:
        print(f"API Request Failed: {e}")
        return None


def validate_syntax(syntax):
    if not syntax:
        return False, "Empty output"

    # Basic checks
    if "infographic" not in syntax.lower():
        return False, "Missing 'infographic' keyword"

    if "data" not in syntax.lower() and "items" not in syntax.lower():
        return False, "Missing 'data' or 'items' block"

    # Check for colons in keys (simple heuristic)
    lines = syntax.split("\n")
    for line in lines:
        stripped = line.strip()
        if not stripped:
            continue
        # Ignore lines that are likely values or descriptions containing colons
        first_word = stripped.split()[0]
        if first_word in ["title", "desc", "label", "value", "time", "icon"]:
            continue  # These can have colons in value

        # Check for key: pattern at start of line
        if re.match(r"^\w+:", stripped):
            return False, f"Found colon in key: {stripped}"

    return True, "Syntax looks valid"


# =================================================================
# Main Execution
# =================================================================

import re


def main():
    print(f"Starting Infographic Generation Verification...")
    print(f"API: {BASE_URL}")
    print(f"Model: {MODEL}")
    print("-" * 50)

    results = []

    for case in TEST_CASES:
        print(f"\nTesting Case: {case['name']}")
        print("Generating...")

        output = generate_infographic(case["content"])

        if output:
            # Clean output (remove markdown code blocks if present)
            clean_output = output
            if "```" in output:
                match = re.search(
                    r"```(?:infographic|mermaid)?\s*(.*?)\s*```", output, re.DOTALL
                )
                if match:
                    clean_output = match.group(1).strip()

            print(f"Output Preview:\n{clean_output[:200]}...")

            is_valid, message = validate_syntax(clean_output)
            if is_valid:
                print(f"✅ Validation Passed: {message}")
                results.append({"name": case["name"], "status": "PASS"})
            else:
                print(f"❌ Validation Failed: {message}")
                print(f"Full Output:\n{output}")
                results.append({"name": case["name"], "status": "FAIL"})
        else:
            print("❌ Generation Failed")
            results.append({"name": case["name"], "status": "ERROR"})

    print("\n" + "=" * 50)
    print("Summary")
    print("=" * 50)
    for res in results:
        print(f"{res['name']}: {res['status']}")


if __name__ == "__main__":
    main()