Playwright_Langgraph

import asyncio
import json
import logging
import os
from typing import TypedDict, List, Dict, Optional, Annotated, Any
from pydantic import BaseModel, Field
from playwright.async_api import async_playwright
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, START, END
from langgraph.prebuilt import ToolNode
import operator
from dotenv import load_dotenv
import os

# 必须在 os.getenv() 调用前执行
load_dotenv()

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger("UI_Framework")


# ================= 1. 数据模型与状态 =================
class StepPlan(BaseModel):
    tool_name: str
    args: Dict[str, Any]
    description: str = ""
    expected_result: str = ""


class TestTask(BaseModel):
    id: str
    prompt: str
    base_url: Optional[str] = None


class FrameworkState(TypedDict):
    task: TestTask
    plan: List[StepPlan]
    current_step_idx: int
    execution_log: Annotated[List[Dict], operator.add]  # 自动追加
    status: str  # pending, planning, executing, success, failed
    error: Optional[str]
    final_report: Optional[Dict]


# ================= 2. 浏览器上下文管理器 =================
class PlaywrightContext:
    def __init__(self):
        self.pw = None
        self.browser = None
        self.context = None
        self.page = None

    async def init(self, headless: bool = True):
        self.pw = await async_playwright().start()
        self.browser = await self.pw.chromium.launch(headless=headless)
        self.context = await self.browser.new_context(
            viewport={"width": 1440, "height": 900},
            ignore_https_errors=True,
            user_agent="LangGraph-UITestBot/1.0"
        )
        self.page = await self.context.new_page()
        return self.page

    async def close(self):
        if self.page: await self.page.close()
        if self.context: await self.context.close()
        if self.browser: await self.browser.close()
        if self.pw: await self.pw.stop()


browser_ctx = PlaywrightContext()


# ================= 3. Playwright 工具集（显式命名对齐 Prompt） =================
@tool("navigate")
async def pw_navigate(url: str) -> str:
    """导航到指定URL"""
    await browser_ctx.page.goto(url, wait_until="domcontentloaded", timeout=15000)
    return f"✅ 导航成功: {url}"

@tool("click")
async def pw_click(selector: str, timeout: int = 5000) -> str:
    """点击元素"""
    await browser_ctx.page.click(selector, timeout=timeout)
    return f"✅ 点击成功: {selector}"

@tool("fill")
async def pw_fill(selector: str, value: str) -> str:
    """填写表单"""
    await browser_ctx.page.fill(selector, value)
    return f"✅ 填写成功: {selector} = {value}"

@tool("wait_for")
async def pw_wait_for(selector: str, state: str = "visible", timeout: int = 10000) -> str:
    """等待元素出现/可点击"""
    await browser_ctx.page.wait_for_selector(selector, state=state, timeout=timeout)
    return f"✅ 等待成功: {selector} ({state})"

@tool("assert_text")
async def pw_assert_text(selector: str, expected: str) -> str:
    """断言元素包含指定文本"""
    actual = await browser_ctx.page.inner_text(selector)
    if expected.strip().lower() in actual.strip().lower():
        return f"✅ 断言通过: 期望'{expected}', 实际'{actual}'"
    raise AssertionError(f"❌ 断言失败: 期望'{expected}', 实际'{actual}'")

@tool("screenshot")
async def pw_screenshot(path: str = "tmp/screenshot.png") -> str:
    """全屏截图"""
    import os
    # ✅ 安全处理路径（防止传入纯文件名时 dirname 为空导致 makedirs 报错）
    dir_name = os.path.dirname(path)
    if dir_name:
        os.makedirs(dir_name, exist_ok=True)
    await browser_ctx.page.screenshot(path=path, full_page=True)
    return f"✅ 截图已保存: {path}"

@tool("select")
async def pw_select(selector: str, value: str) -> str:
    """下拉框选择"""
    await browser_ctx.page.select_option(selector, value=value)
    return f"✅ 选择成功: {selector} = {value}"

# 工具列表保持不变
TOOLS = [pw_navigate, pw_click, pw_fill, pw_wait_for, pw_assert_text, pw_screenshot, pw_select]
tool_node = ToolNode(TOOLS)


# ================= 4. LLM 规划器（替换原代码） =================
class StepPlanList(BaseModel):
    """强制根节点为对象，兼容 DashScope/JSON Mode"""
    steps: List[StepPlan]

planner_llm = ChatOpenAI(
    api_key=os.getenv("API_KEY"),
    base_url=os.getenv("BASE_URL", "https://dashscope.yuncs.com/compatible-mode/v1"),
    model="qwen-plus",
    temperature=0.0
).with_structured_output(
    StepPlanList,
    method="function_calling"  # ✅ 阿里云/通义千问对 function_calling 支持更稳定
)

async def plan_node(state: FrameworkState) -> Dict:
    # ✅ 修复：TypedDict 运行时是 dict，必须用 [] 访问
    task_dict = state["task"]
    prompt = task_dict.get("prompt", "")
    base_url = task_dict.get("base_url")
    if base_url:
        prompt = f"🌐 目标网站: {base_url}\n📝 任务: {prompt}"

    system_msg = """你是资深 UI 自动化测试架构师。请将自然语言测试任务分解为可执行步骤。
    🔑 规则:
    1. 严格使用工具: navigate, click, fill, wait_for, assert_text, screenshot, select
    2. 选择器优先使用: data-testid, id, role(name="xxx"), 或稳定 CSS 路径
    3. 第一步必须是 navigate (除非任务明确说明已在页面)
    4. 复杂交互拆分为多步，每步只做一个动作"""

    try:
        plan_obj = await planner_llm.ainvoke([
            {"role": "system", "content": system_msg},
            {"role": "user", "content": prompt}
        ])
        return {
            "plan": [p.model_dump() for p in plan_obj.steps],
            "status": "planning",
            "current_step_idx": 0,
            "execution_log": [{"step": "plan", "msg": f"🤖 已生成 {len(plan_obj.steps)} 个测试步骤"}]
        }
    except Exception as e:
        return {"status": "failed", "error": f"⚠️ 规划失败: {str(e)}", "execution_log": []}


# ================= 5. 执行器（单步推进 + 错误拦截） =================
async def execute_node(state: FrameworkState) -> Dict:
    plan = state["plan"]
    idx = state["current_step_idx"]
    logs = []

    if idx >= len(plan):
        return {"status": "success", "execution_log": [{"step": "complete", "msg": "🏁 所有步骤执行完毕"}]}

    step = plan[idx]
    tool_name = step.get("tool_name")
    args = step.get("args", {})

    target_tool = next((t for t in TOOLS if t.name == tool_name), None)
    if not target_tool:
        return {"status": "failed", "error": f"🔧 未知工具: {tool_name}", "current_step_idx": idx}

    try:
        result = await target_tool.ainvoke(args)
        logs.append({"step": idx, "tool": tool_name, "result": result, "status": "success"})
        return {"current_step_idx": idx + 1, "execution_log": logs}
    except Exception as e:
        logs.append({"step": idx, "tool": tool_name, "error": str(e), "status": "failed"})
        return {
            "status": "failed",
            "error": f"❌ 步骤 {idx} ({tool_name}) 执行失败: {str(e)}",
            "execution_log": logs,
            "current_step_idx": idx
        }


# ================= 6. 验证与清理 =================
async def validate_node(state: FrameworkState) -> Dict:
    return {"execution_log": [{"step": "validate", "msg": "🔍 执行最终状态校验"}]}


async def teardown_node(state: FrameworkState) -> Dict:
    await browser_ctx.close()
    return {"execution_log": [{"step": "teardown", "msg": "🛑 浏览器资源已释放"}]}


# ================= 7. 组装 LangGraph =================
workflow = StateGraph(FrameworkState)
workflow.add_node("plan", plan_node)
workflow.add_node("execute", execute_node)
workflow.add_node("validate", validate_node)
workflow.add_node("teardown", teardown_node)

workflow.set_entry_point("plan")
workflow.add_edge("plan", "execute")
workflow.add_edge("validate", "teardown")
workflow.add_edge("teardown", END)


def route_after_execute(state: FrameworkState) -> str:
    if state.get("status") == "failed":
        return "validate"
    if state["current_step_idx"] < len(state["plan"]):
        return "execute"
    return "validate"


workflow.add_conditional_edges("execute", route_after_execute)
graph = workflow.compile()


# ================= 8. 框架运行器（支持批量/配置驱动） =================
class UITestFramework:
    def __init__(self):
        self.graph = graph

    async def run_task(self, prompt: str, base_url: str = None, channel: str=None,headless: bool = True, task_id: str = "auto",
                       **kwargs) -> Dict:
        # ✅ 兼容配置中的 "id" 字段，覆盖默认 task_id
        task_id = kwargs.pop("id", task_id)
        print("task_id >>",task_id)

        # ✅ 过滤其他未定义参数（防配置污染）
        unknown_keys = [k for k in kwargs if k not in ("task_id",)]
        if unknown_keys:
            import logging
            logging.warning(f"⚠️ 忽略未定义参数: {unknown_keys}")

        await browser_ctx.init(headless=headless)
        initial_state = {
            "task": {"id": task_id, "prompt": prompt, "base_url": base_url,"channel":channel},
            "plan": [],
            "current_step_idx": 0,
            "execution_log": [],
            "status": "pending",
            "error": None,
            "final_report": None
        }
        try:
            print("initial_state >>",initial_state)
            return await self.graph.ainvoke(initial_state)
        except Exception as e:
            await browser_ctx.close()
            return {"status": "critical_failed", "error": str(e), "execution_log": []}

    async def run_from_yaml(self, yaml_path: str, headless: bool = True) -> List[Dict]:
        import yaml
        with open(yaml_path, "r", encoding="utf-8") as f:
            config = yaml.safe_load(f)

        results = []
        for task in config.get("tasks", []):
            logger.info(f"🚀 执行: {task.get('prompt', 'Unknown')[:60]}...")
            res = await self.run_task(
                prompt=task["prompt"],
                base_url=task.get("base_url"),
                headless=headless,
                task_id=task.get("id", f"task_{len(results) + 1}")
            )
            results.append(res)
        return results


# ================= 9. 使用示例 =================
class UITestFramework:
    def __init__(self):
        self.graph = graph

    async def run_task(self, prompt: str, base_url: str = None, channel: str = None,
                       headless: bool = True, task_id: str = "auto", **kwargs) -> Dict:
        task_id = kwargs.pop("id", task_id)

        try:
            await browser_ctx.init(headless=headless)
        except Exception as e:
            return {"status": "critical_failed", "error": f"浏览器启动失败: {str(e)}", "execution_log": []}

        initial_state = {
            "task": {"id": task_id, "prompt": prompt, "base_url": base_url, "channel": channel},
            "plan": [],
            "current_step_idx": 0,
            "execution_log": [],
            "status": "pending",
            "error": None,
            "final_report": None
        }
        try:
            return await self.graph.ainvoke(initial_state)
        except Exception as e:
            await browser_ctx.close()
            return {"status": "critical_failed", "error": str(e), "execution_log": []}


# ================= 9. 使用示例（替换 main） =================
async def main():
    framework = UITestFramework()

    tasks_config = {
        "tasks": [
            {
                "id": "login_test",
                "prompt": "访问首页，点击 Log In 按钮，aaa@163.com,然后登录",
                "base_url": "https://www.baidu.com"
            }
        ]
    }

    for t in tasks_config["tasks"]:
        print(f"\n🚀 开始执行: {t['id']}")
        res = await framework.run_task(**t, headless=False)

        print(f"📊 状态: {res['status'].upper()}")
        if res.get("error"):
            print(f"🔴 错误详情: {res['error']}")
        print("📝 执行日志:")
        for log in res.get("execution_log", []):
            print(f"  {log.get('msg', log)}")


if __name__ == "__main__":
    asyncio.run(main())