generated from kgod/ai-review-template
228 lines
6.9 KiB
Python
228 lines
6.9 KiB
Python
"""节点实现"""
|
|
import asyncio
|
|
import hashlib
|
|
|
|
from langchain_core.output_parsers import PydanticOutputParser
|
|
from pydantic import BaseModel, Field
|
|
from langchain_core.messages import HumanMessage
|
|
from src.bash_model import GeneralLlm, AnalyseLlm
|
|
from .state import NextPageState
|
|
from .prompts import FIND_PAGINATION_AREA_PROMPT, FIND_NEXT_BUTTON_PROMPT
|
|
|
|
|
|
# 结构化输出模型
|
|
class SelectorResult(BaseModel):
|
|
"""选择器结果"""
|
|
selector: str | None = Field(description="CSS选择器，找不到则为None")
|
|
reason: str = Field(description="选择或找不到的原因")
|
|
|
|
|
|
async def fetch_page(state: NextPageState) -> dict:
|
|
"""节点1: 获取页面 HTML"""
|
|
url = state["url"]
|
|
page = state["page"]
|
|
|
|
print(f"[节点1] 访问页面: {url}")
|
|
|
|
await page.goto(url, wait_until="load", timeout=60000)
|
|
await asyncio.sleep(3)
|
|
|
|
# 获取清理后的 HTMLa
|
|
html = await page.evaluate("""
|
|
() => {
|
|
const clone = document.body.cloneNode(true);
|
|
clone.querySelectorAll('style, script, noscript, svg, link').forEach(el => el.remove());
|
|
return clone.innerHTML;
|
|
}
|
|
""")
|
|
|
|
print(f"[节点1] 获取 HTML 完成，长度: {len(html)}")
|
|
|
|
return {
|
|
"html": html,
|
|
"failed_selectors": [],
|
|
"retry_count": 0
|
|
}
|
|
|
|
|
|
async def find_pagination_area(state: NextPageState) -> dict:
|
|
"""节点2: 找分页区域"""
|
|
html = state["html"]
|
|
page = state["page"]
|
|
|
|
print("[节点2] 分析分页区域...")
|
|
|
|
# LLM 分析
|
|
prompt = FIND_PAGINATION_AREA_PROMPT.format(html=html)
|
|
|
|
chain = AnalyseLlm | PydanticOutputParser(pydantic_object=SelectorResult)
|
|
result = await chain.ainvoke([HumanMessage(content=prompt)])
|
|
|
|
#
|
|
# llm = AnalyseLlm.with_structured_output(SelectorResult)
|
|
# result = await llm.ainvoke([HumanMessage(content=prompt)])
|
|
|
|
if not result.selector:
|
|
# 无分页控件，直接成功
|
|
print(f"[节点2] 无分页控件: {result.reason}")
|
|
return {
|
|
"has_pagination": 0,
|
|
"is_valid": True
|
|
}
|
|
|
|
print(f"[节点2] 找到分页区域: {result.selector}")
|
|
|
|
# 获取区域 HTML
|
|
try:
|
|
pagination_html = await page.inner_html(result.selector)
|
|
except Exception as e:
|
|
# 选择器无效，当作无分页
|
|
print(f"[节点2] 选择器无效: {result.selector}, {e}")
|
|
return {
|
|
"has_pagination": 0,
|
|
"is_valid": True
|
|
}
|
|
|
|
return {
|
|
"has_pagination": 1,
|
|
"pagination_selector": result.selector,
|
|
"pagination_html": pagination_html
|
|
}
|
|
|
|
|
|
async def find_next_button(state: NextPageState) -> dict:
|
|
"""节点3: 找下一页按钮"""
|
|
pagination_html = state["pagination_html"]
|
|
pagination_selector = state["pagination_selector"]
|
|
failed_selectors = state.get("failed_selectors", [])
|
|
retry_count = state.get("retry_count", 0)
|
|
page = state["page"]
|
|
|
|
print(f"[节点3] 分析下一页按钮... (重试: {retry_count})")
|
|
|
|
# LLM 分析
|
|
prompt = FIND_NEXT_BUTTON_PROMPT.format(
|
|
pagination_html=pagination_html,
|
|
failed_selectors="\n".join(failed_selectors) if failed_selectors else "无"
|
|
)
|
|
|
|
chain = AnalyseLlm | PydanticOutputParser(pydantic_object=SelectorResult)
|
|
result = await chain.ainvoke([HumanMessage(content=prompt)])
|
|
|
|
# llm = AnalyseLlm.with_structured_output(SelectorResult)
|
|
# result = await llm.ainvoke([HumanMessage(content=prompt)])
|
|
|
|
next_selector = result.selector
|
|
if not next_selector:
|
|
# 有分页控件但无可点击的下一页（仅一页），统一设为无分页
|
|
print(f"[节点3] 无可点击的下一页: {result.reason}")
|
|
return {
|
|
"selector": None,
|
|
"has_pagination": 0,
|
|
"is_valid": True
|
|
}
|
|
|
|
# 组合完整选择器
|
|
selector = f"{pagination_selector} {next_selector}"
|
|
|
|
# 检查是否已失败过
|
|
if selector in failed_selectors:
|
|
print(f"[节点3] 选择器已失败过: {selector}")
|
|
return {
|
|
"selector": None,
|
|
"failed_selectors": failed_selectors,
|
|
"retry_count": retry_count + 1
|
|
}
|
|
|
|
print(f"[节点3] 选择器: {selector}")
|
|
|
|
# 验证元素存在
|
|
element = await page.query_selector(selector)
|
|
if not element:
|
|
print("[节点3] 选择器未匹配到元素")
|
|
return {
|
|
"selector": None,
|
|
"failed_selectors": failed_selectors + [selector],
|
|
"retry_count": retry_count + 1
|
|
}
|
|
|
|
return {"selector": selector}
|
|
|
|
|
|
async def validate_next(state: NextPageState) -> dict:
|
|
"""节点4: 验证下一页"""
|
|
selector = state["selector"]
|
|
page = state["page"]
|
|
url = state["url"]
|
|
failed_selectors = state.get("failed_selectors", [])
|
|
retry_count = state.get("retry_count", 0)
|
|
|
|
print(f"[节点4] 验证选择器: {selector}")
|
|
|
|
context = page.context
|
|
|
|
# 1. 记录状态
|
|
before_url = page.url
|
|
tabs_before = len(context.pages)
|
|
content_before = await page.content()
|
|
hash_before = hashlib.md5(content_before.encode()).hexdigest()
|
|
|
|
# 2. 点击（先 hover 触发可能的显示效果）
|
|
try:
|
|
element = page.locator(selector).first
|
|
await element.scroll_into_view_if_needed()
|
|
await asyncio.sleep(0.2)
|
|
await element.hover()
|
|
await asyncio.sleep(0.3)
|
|
await element.click()
|
|
except Exception as e:
|
|
print(f"[节点4] 点击失败: {e}")
|
|
return {
|
|
"is_valid": False,
|
|
"change_type": "no_change",
|
|
"failed_selectors": failed_selectors + [selector],
|
|
"retry_count": retry_count + 1
|
|
}
|
|
|
|
await asyncio.sleep(3)
|
|
|
|
# 3. 检测变化
|
|
tabs_after = len(context.pages)
|
|
after_url = page.url
|
|
content_after = await page.content()
|
|
hash_after = hashlib.md5(content_after.encode()).hexdigest()
|
|
|
|
if tabs_after > tabs_before:
|
|
change_type = "new_tab"
|
|
# 关闭新标签
|
|
await context.pages[-1].close()
|
|
elif after_url != before_url:
|
|
change_type = "url_change"
|
|
elif hash_after != hash_before:
|
|
change_type = "content_change"
|
|
else:
|
|
change_type = "no_change"
|
|
|
|
print(f"[节点4] 变化类型: {change_type}")
|
|
|
|
# 4. 恢复状态
|
|
if change_type in ("url_change", "content_change"):
|
|
await page.goto(url, wait_until="load", timeout=60000)
|
|
await asyncio.sleep(2)
|
|
|
|
# 5. 判断结果
|
|
if change_type == "no_change":
|
|
print("[节点4] 验证失败: 点击后无变化")
|
|
return {
|
|
"is_valid": False,
|
|
"change_type": change_type,
|
|
"failed_selectors": failed_selectors + [selector],
|
|
"retry_count": retry_count + 1
|
|
}
|
|
else:
|
|
print("[节点4] 验证成功!")
|
|
return {
|
|
"is_valid": True,
|
|
"change_type": change_type
|
|
}
|