package collect import ( "context" "fmt" "geo/internal/config" "log" "strings" "time" "github.com/go-rod/rod" "github.com/go-rod/rod/lib/proto" ) // WenxinCollector 文心一言收集器 type WenxinCollector struct { *BaseCollector } // NewWenxinCollector 创建文心一言收集器 func NewWenxinCollector(ctx context.Context, params *CollectParams, cfg *config.Config, logger *log.Logger) CollectorInterface { collector := &WenxinCollector{ BaseCollector: NewBaseCollector(ctx, params, cfg, logger), } // 设置文心一言的URL collector.LoginURL = "https://passport.baidu.com/v2/?login" collector.ChatURL = "https://yiyan.baidu.com/" return collector } // CheckLoginStatus 检查登录状态 func (c *WenxinCollector) CheckLoginStatus() bool { currentURL := c.GetCurrentURL() // 如果在登录页面,说明未登录 if strings.Contains(currentURL, "passport.baidu.com") { return false } // 检查页面上是否存在内容为"登录"或"Login"的button,如果存在说明未登录 loginButtons, err := c.Page.Elements("button") if err == nil { for _, btn := range loginButtons { text, _ := btn.Text() trimmedText := strings.TrimSpace(text) if trimmedText == "登录" || trimmedText == "Login" { c.LogInfo(fmt.Sprintf("检测到页面上有'%s'按钮,说明未登录", trimmedText)) return false } } } // 如果没有找到"登录"或"Login"按钮,说明已登录 return true } // WaitLogin 等待登录 func (c *WenxinCollector) WaitLogin() (bool, string) { c.LogInfo("开始等待文心一言登录...") if err := c.SetupDriver(); err != nil { return false, fmt.Sprintf("浏览器启动失败: %v", err) } defer c.Close() // 访问聊天页面 c.Page.MustNavigate(c.ChatURL) c.Sleep(3) // 检查是否已登录 if c.CheckLoginStatus() { c.SaveCookies() c.LogInfo("已有登录状态") return true, "already_logged_in" } c.LogInfo("检测到未登录,请在当前页面完成登录(扫码或输入账号密码)...") // 不跳转页面,在当前页面循环检查登录按钮是否存在 // 最多等待300秒 for i := 0; i < 5000; i++ { // 检查页面上是否还存在"登录"或"Login"按钮 loginButtonExists := false buttons, err := c.Page.Elements("button") if err == nil { for _, btn := range buttons { text, _ := btn.Text() trimmedText := strings.TrimSpace(text) if trimmedText == "登录" || trimmedText == "Login" { loginButtonExists = true break } } } // 如果登录按钮不存在,说明已登录 if !loginButtonExists { c.Sleep(2) // 等待页面稳定 c.SaveCookies() c.LogInfo("登录成功:登录按钮已消失") return true, "login_success" } // 每秒检查一次 time.Sleep(1 * time.Second) // 每30秒输出一次提示 if i > 0 && i%30 == 0 { c.LogInfo(fmt.Sprintf("等待登录中... 已等待 %d 秒", i)) } } return false, "登录超时,请检查网络或账号状态" } // AskQuestion 提问并获取答案 func (c *WenxinCollector) AskQuestion(question string) (string, error) { c.LogInfo(fmt.Sprintf("开始提问: %s", question)) // 初始化浏览器 if err := c.SetupDriver(); err != nil { return "", fmt.Errorf("浏览器启动失败: %v", err) } defer c.Close() //初始化页面(加载cookies和检查登录) if err := c.InitPage(); err != nil { return "", fmt.Errorf("页面初始化失败,请先调用WaitLogin登录: %v", err) } // 等待页面完全加载 c.Sleep(3) // 查找输入框并输入问题 if err := c.inputQuestion(question); err != nil { return "", fmt.Errorf("输入问题失败: %v", err) } // 点击发送按钮 if err := c.clickSendButton(); err != nil { return "", fmt.Errorf("点击发送按钮失败: %v", err) } // 等待并获取答案 answer, err := c.waitForAnswer() if err != nil { return "", fmt.Errorf("获取答案失败: %v", err) } c.LogInfo(fmt.Sprintf("成功获取答案,长度: %d 字符", len(answer))) return answer, nil } // inputQuestion 输入问题 func (c *WenxinCollector) inputQuestion(question string) error { c.LogInfo("输入问题...") // 文心一言的输入框选择器 - 根据实际页面结构调整 inputSelectors := []string{ "[contenteditable='true']", "div[contenteditable]", ".editable__T7WAW4uW", "[class*='editable']", } var inputBox *rod.Element var err error for _, selector := range inputSelectors { inputBox, err = c.WaitForElementVisible(selector, 10) if err == nil && inputBox != nil { c.LogInfo(fmt.Sprintf("找到输入框: %s", selector)) break } } if inputBox == nil { return fmt.Errorf("未找到输入框") } // 点击获取焦点 if err := inputBox.Click(proto.InputMouseButtonLeft, 1); err != nil { return fmt.Errorf("点击输入框失败: %v", err) } c.SleepMs(500) // fallback: 使用Focus + Input inputBox.Focus() c.SleepMs(200) inputBox.Input(question) c.LogInfo(fmt.Sprintf("问题已输入: %s", question)) c.SleepMs(1000) return nil } // clickSendButton 点击发送按钮 func (c *WenxinCollector) clickSendButton() error { c.LogInfo("点击发送按钮...") // 使用正则匹配包含"send"的class(防CSS混淆) allElements, err := c.Page.Elements("*") if err != nil { return fmt.Errorf("获取页面元素失败: %v", err) } var sendBtn *rod.Element for _, elem := range allElements { classAttr, _ := elem.Attribute("class") if classAttr != nil && strings.Contains(strings.ToLower(*classAttr), "send") { // 检查是否是可点击的元素(button、div等) tagName, _ := elem.Property("tagName") if tagName.Str() == "BUTTON" || tagName.Str() == "DIV" { sendBtn = elem c.LogInfo(fmt.Sprintf("通过正则找到发送按钮: class=%s, tag=%s", *classAttr, tagName.Str())) break } } } if sendBtn == nil { // fallback: 尝试查找最后一个button buttons, _ := c.Page.Elements("button") if len(buttons) > 0 { sendBtn = buttons[len(buttons)-1] c.LogInfo("使用最后一个button作为发送按钮") } } if sendBtn == nil { return fmt.Errorf("未找到发送按钮") } c.SleepMs(500) // 滚动到可见区域 if err := sendBtn.ScrollIntoView(); err != nil { c.LogInfo(fmt.Sprintf("滚动失败: %v", err)) } c.SleepMs(300) // 点击发送按钮 c.LogInfo("执行点击...") if err := sendBtn.Click(proto.InputMouseButtonLeft, 1); err != nil { return fmt.Errorf("点击发送按钮失败: %v", err) } c.LogInfo("已点击发送按钮") c.SleepMs(1000) // 检测是否发送成功:检查send按钮是否消失或变成pause按钮 maxWaitTime := 10 // 最多等待10秒 for i := 0; i < maxWaitTime*2; i++ { // 检查是否存在pause开头的按钮(表示正在生成) pauseExists, err := c.hasPauseButton() if err == nil && pauseExists { c.LogInfo("✓ 检测到pause按钮,消息发送成功,AI正在回答...") return nil } // 检查send按钮是否还存在 sendExists, _ := c.hasSendButton() if !sendExists { c.LogInfo("✓ send按钮已消失,消息发送成功") return nil } c.SleepMs(500) } c.LogInfo("⚠ 无法确认消息是否发送成功,但已尽力尝试") return nil } // hasSendButton 检查是否存在send开头的按钮 func (c *WenxinCollector) hasSendButton() (bool, error) { allElements, err := c.Page.Elements("*") if err != nil { return false, err } for _, elem := range allElements { classAttr, _ := elem.Attribute("class") if classAttr != nil && strings.Contains(strings.ToLower(*classAttr), "send") { tagName, _ := elem.Property("tagName") if tagName.Str() == "BUTTON" || tagName.Str() == "DIV" { return true, nil } } } return false, nil } // hasPauseButton 检查是否存在pause开头的按钮 func (c *WenxinCollector) hasPauseButton() (bool, error) { allElements, err := c.Page.Elements("*") if err != nil { return false, err } for _, elem := range allElements { classAttr, _ := elem.Attribute("class") if classAttr != nil && strings.Contains(strings.ToLower(*classAttr), "pause") { tagName, _ := elem.Property("tagName") if tagName.Str() == "BUTTON" || tagName.Str() == "DIV" { return true, nil } } } return false, nil } // waitForAnswer 等待并获取答案(处理流式输出) func (c *WenxinCollector) waitForAnswer() (string, error) { c.LogInfo("等待AI回答...") timeout := 180 // 最大等待时间(秒),流式输出可能需要更长时间 startTime := time.Now() var lastAnswer string var stableCount int // 稳定计数器,连续N次内容不变则认为完成 const requiredStableCount = 5 // 需要连续5次内容不变才认为完成 isAnswering := false // 标记是否正在回答中 for time.Since(startTime).Seconds() < float64(timeout) { // 检查是否存在pause按钮(表示正在生成答案) pauseExists, _ := c.hasPauseButton() if pauseExists { if !isAnswering { c.LogInfo("检测到pause按钮,AI正在生成回答...") isAnswering = true } } else if isAnswering { // pause按钮消失,可能回答完成了 c.LogInfo("pause按钮消失,检查回答是否完成...") // 再等待几次确认内容稳定 if stableCount >= requiredStableCount && lastAnswer != "" { c.LogInfo(fmt.Sprintf("✓ AI回答完成,最终长度: %d 字符", len(lastAnswer))) return lastAnswer, nil } } // 直接通过ID查找答案容器 answerElem, err := c.Page.Element("#answer_text_id") var answerText string if err == nil && answerElem != nil { // 获取整个HTML内容 htmlContent, err := answerElem.HTML() if err == nil && len(strings.TrimSpace(htmlContent)) > 30 { // 清理HTML标签,只保留纯文本 answerText = CleanHTMLTags(htmlContent) c.LogInfo(fmt.Sprintf("找到答案容器,清理后文本长度: %d", len(answerText))) } else { // 如果HTML获取失败,尝试获取文本 textContent, _ := answerElem.Text() answerText = strings.TrimSpace(textContent) c.LogInfo(fmt.Sprintf("找到答案容器,文本长度: %d", len(answerText))) } } else { c.LogInfo("未找到#answer_text_id元素") } // 检查是否获取到答案 if answerText != "" && len(answerText) > 30 { // 检查内容是否稳定(流式输出完成) if answerText == lastAnswer { stableCount++ c.LogInfo(fmt.Sprintf("答案稳定中... (%d/%d), 长度: %d", stableCount, requiredStableCount, len(answerText))) // 如果pause按钮不存在且内容稳定,说明回答完成 if !pauseExists && stableCount >= requiredStableCount { c.LogInfo(fmt.Sprintf("✓ AI回答完成,最终长度: %d 字符", len(answerText))) return answerText, nil } } else { // 内容还在变化,重置计数器 stableCount = 0 lastAnswer = answerText if pauseExists { c.LogInfo(fmt.Sprintf("检测到流式输出,当前长度: %d 字符", len(answerText))) } } } c.SleepMs(1500) // 每1.5秒检查一次 // 每10秒输出一次等待状态 elapsed := int(time.Since(startTime).Seconds()) if elapsed > 0 && elapsed%10 == 0 { c.LogInfo(fmt.Sprintf("等待AI回答中... 已等待 %d 秒", elapsed)) } } return "", fmt.Errorf("等待答案超时(%d秒)", timeout) } // SafeElement 安全地获取元素 func (c *WenxinCollector) SafeElement(selector string) (*rod.Element, error) { exists, _, err := c.Page.Has(selector) if err != nil { return nil, err } if !exists { return nil, nil } return c.Page.Element(selector) }