package collect import ( "context" "fmt" "geo/internal/config" "log" "strings" "time" "github.com/go-rod/rod" "github.com/go-rod/rod/lib/proto" ) // DoubaoCollector 豆包收集器 type DoubaoCollector struct { *BaseCollector } // NewDoubaoCollector 创建豆包收集器 func NewDoubaoCollector(ctx context.Context, params *CollectParams, cfg *config.Config, logger *log.Logger) CollectorInterface { collector := &DoubaoCollector{ BaseCollector: NewBaseCollector(ctx, params, cfg, logger), } // 设置豆包的URL collector.LoginURL = "https://www.doubao.com/" collector.ChatURL = "https://www.doubao.com/chat/" return collector } // CheckLoginStatus 检查登录状态 func (c *DoubaoCollector) CheckLoginStatus() bool { currentURL := c.GetCurrentURL() // 检查是否在聊天页面 if strings.Contains(currentURL, "doubao.com") { // 查找用户信息元素 userInfo, err := c.SafeElement(".user-info, .avatar, [class*='user-profile']") if err == nil && userInfo != nil { return true } // 检查是否有输入框 inputBox, err := c.SafeElement("textarea, [contenteditable='true']") if err == nil && inputBox != nil { return true } } return false } // WaitLogin 等待登录 func (c *DoubaoCollector) WaitLogin() (bool, string) { c.LogInfo("开始等待豆包登录...") if err := c.SetupDriver(); err != nil { return false, fmt.Sprintf("浏览器启动失败: %v", err) } defer c.Close() // 访问豆包首页 c.Page.MustNavigate(c.LoginURL) c.Sleep(3) // 检查是否已登录 if c.CheckLoginStatus() { c.SaveCookies() c.LogInfo("已有登录状态") return true, "already_logged_in" } c.LogInfo("请登录豆包账号...") // 等待用户手动登录,最多300秒 for i := 0; i < 300; i++ { if c.CheckLoginStatus() { c.SaveCookies() c.LogInfo("登录成功") return true, "login_success" } time.Sleep(1 * time.Second) } return false, "登录超时,请检查网络或账号状态" } // AskQuestion 提问并获取答案 func (c *DoubaoCollector) 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() // 初始化页面 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 *DoubaoCollector) inputQuestion(question string) error { c.LogInfo("输入问题到豆包...") // 豆包的输入框选择器 inputSelectors := []string{ "textarea[placeholder*='输入']", "textarea[placeholder*='问']", "textarea", "[contenteditable='true']", ".chat-input textarea", "#input-box", ".input-area textarea", } 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) // 清空输入框 if err := c.ClearInput(inputBox); err != nil { c.LogInfo(fmt.Sprintf("清空输入框失败: %v", err)) } c.SleepMs(300) // 输入问题 if err := c.SetInputValue(inputBox, question); err != nil { inputBox.Input(question) } c.LogInfo(fmt.Sprintf("问题已输入")) c.SleepMs(1000) return nil } // clickSendButton 点击发送按钮 func (c *DoubaoCollector) clickSendButton() error { c.LogInfo("点击发送按钮...") // 发送按钮选择器 sendSelectors := []string{ "button[class*='send']", "button[class*='submit']", ".send-btn", ".submit-btn", "button svg[path*='send']", "[aria-label*='发送']", ".send-icon", } var sendBtn *rod.Element var err error for _, selector := range sendSelectors { sendBtn, err = c.WaitForElementClickable(selector, 5) if err == nil && sendBtn != nil { c.LogInfo(fmt.Sprintf("找到发送按钮: %s", selector)) break } } if sendBtn == nil { // 尝试查找发送图标 sendBtn, err = c.Page.Element("button svg") if err != nil { return fmt.Errorf("未找到发送按钮") } } c.SleepMs(500) // 点击发送按钮 if err := c.JSClick(sendBtn); err != nil { return fmt.Errorf("点击发送按钮失败: %v", err) } c.LogInfo("已点击发送按钮") c.SleepMs(2000) return nil } // waitForAnswer 等待并获取答案 func (c *DoubaoCollector) waitForAnswer() (string, error) { c.LogInfo("等待豆包回答...") timeout := 120 // 最大等待时间(秒) startTime := time.Now() lastAnswerLength := 0 for time.Since(startTime).Seconds() < float64(timeout) { // 查找答案区域 answerSelectors := []string{ ".message-content", ".response-text", "[class*='assistant'] [class*='content']", "[class*='bot'] [class*='message']", ".chat-message.bot", ".answer-box", } for _, selector := range answerSelectors { answerElements, err := c.Page.Elements(selector) if err == nil && len(answerElements) > 0 { // 获取最后一个答案元素 lastAnswer := answerElements[len(answerElements)-1] visible, _ := lastAnswer.Visible() if visible { text, err := lastAnswer.Text() if err == nil && len(strings.TrimSpace(text)) > 0 { // 检查是否正在生成 isGenerating := strings.Contains(text, "正在") || strings.Contains(text, "思考中") || strings.Contains(text, "typing") if !isGenerating { // 检查答案是否还在增长 currentLength := len(text) if currentLength == lastAnswerLength && currentLength > 10 { // 答案不再增长,认为已完成 c.LogInfo("获取到完整答案") return strings.TrimSpace(text), nil } lastAnswerLength = currentLength } } } } } c.SleepMs(1500) } return "", fmt.Errorf("等待答案超时") } // SafeElement 安全地获取元素 func (c *DoubaoCollector) 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) }