package collect import ( "context" "fmt" "geo/internal/config" "strings" "time" "github.com/go-rod/rod" "github.com/go-rod/rod/lib/proto" "github.com/gofiber/fiber/v2/log" ) // DoubaoCollector 豆包收集器 type DoubaoCollector struct { *BaseCollector } // NewDoubaoCollector 创建豆包收集器 func NewDoubaoCollector(ctx context.Context, params *CollectParams, cfg *config.Config, logger log.AllLogger) 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) { 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() return true, "already_logged_in" } for i := 0; i < 300; i++ { if c.CheckLoginStatus() { c.Sleep(2) c.SaveCookies() return true, "login_success" } time.Sleep(1 * time.Second) } return false, "登录超时" } // AskQuestion 提问并获取答案 func (c *DoubaoCollector) AskQuestion(question string) (*CollectResult, error) { if err := c.SetupDriver(); err != nil { return nil, fmt.Errorf("浏览器启动失败: %v", err) } defer c.Close() if err := c.InitPage(); err != nil { return nil, fmt.Errorf("页面初始化失败: %v", err) } c.Sleep(3) if err := c.inputQuestion(question); err != nil { return nil, fmt.Errorf("输入问题失败: %v", err) } if err := c.clickSendButton(); err != nil { return nil, fmt.Errorf("点击发送按钮失败: %v", err) } answer, err := c.waitForAnswer() if err != nil { return nil, fmt.Errorf("获取答案失败: %v", err) } return &CollectResult{ Answer: answer, ShareLink: "", }, nil } // inputQuestion 输入问题 func (c *DoubaoCollector) inputQuestion(question string) error { 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 { 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 { // Ignore clear error } c.SleepMs(300) if err := c.SetInputValue(inputBox, question); err != nil { inputBox.Input(question) } c.SleepMs(1000) return nil } // clickSendButton 点击发送按钮 func (c *DoubaoCollector) clickSendButton() error { 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 { 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.SleepMs(2000) return nil } // waitForAnswer 等待并获取答案 func (c *DoubaoCollector) waitForAnswer() (string, error) { 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 { 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) }