From 0647790cc21fefe2304158c8b371c480045c27bb Mon Sep 17 00:00:00 2001 From: Rzy <465386466@qq.com> Date: Sat, 25 Apr 2026 23:55:30 +0800 Subject: [PATCH] 22 --- doubao_test.go | 169 ++++++++++++ internal/collect/base.go | 2 + internal/collect/doubao.go | 424 +++++++++++++++++++++++++----- internal/collect/interface.go | 9 +- internal/collect/utils.go | 12 +- internal/collect/utils_test.go | 103 ++++++++ internal/collect/wenxin.go | 10 +- example_test.go => wenxin_test.go | 0 8 files changed, 657 insertions(+), 72 deletions(-) create mode 100644 doubao_test.go create mode 100644 internal/collect/utils_test.go rename example_test.go => wenxin_test.go (100%) diff --git a/doubao_test.go b/doubao_test.go new file mode 100644 index 0000000..aae3953 --- /dev/null +++ b/doubao_test.go @@ -0,0 +1,169 @@ +package collect + +import ( + "context" + "geo/internal/collect" + "geo/internal/config" + "testing" + "time" + + "github.com/gofiber/fiber/v2/log" +) + +var ( + doubaoCfg, _ = config.LoadConfig() + + doubaoManager = collect.NewCollectManager(context.Background(), doubaoCfg, log.DefaultLogger()) +) + +// TestDoubaoCollector_WaitLogin 测试豆包登录功能 +func TestDoubaoCollector_WaitLogin(t *testing.T) { + if testing.Short() { + t.Skip("跳过需要浏览器交互的测试") + } + + params := &collect.CollectParams{ + Headless: false, // 显示浏览器窗口以便扫码登录 + RequestID: "test_doubao_login_001", + Platform: "doubao", + } + + t.Log("开始测试豆包登录...") + t.Log("请在打开的浏览器窗口中完成豆包账号登录（扫码或输入账号密码）") + + success, msg := doubaoManager.WaitLogin("doubao", params) + + if !success { + t.Errorf("豆包登录失败: %s", msg) + return + } + + t.Logf("豆包登录成功: %s", msg) + t.Log("Cookie已保存，后续测试可以使用已登录状态") +} + +// TestDoubaoCollector_AskQuestion 测试豆包提问功能 +// 注意：此测试需要有效的登录状态 +func TestDoubaoCollector_AskQuestion(t *testing.T) { + if testing.Short() { + t.Skip("跳过需要浏览器交互的测试") + } + + // 设置收集参数 + params := &collect.CollectParams{ + Headless: false, // 显示浏览器以便调试 + RequestID: "test_doubao_001", + Platform: "doubao", + } + + // 定义提问内容 + question := "四川房地产软件排名" + t.Logf("向豆包提问: %s", question) + + // 调用管理器提问并获取答案 + result, err := doubaoManager.AskQuestion("doubao", params, question) + if err != nil { + t.Errorf("提问失败: %v", err) + return + } + + t.Logf("获取到答案:\n%s", result.Answer) + + // 验证答案非空 + if len(result.Answer) == 0 { + t.Error("答案为空") + } +} + +// TestDoubaoCollector_MultipleQuestions 测试豆包多次提问 +func TestDoubaoCollector_MultipleQuestions(t *testing.T) { + if testing.Short() { + t.Skip("跳过需要浏览器交互的测试") + } + + params := &collect.CollectParams{ + Headless: false, + RequestID: "test_doubao_multi_001", + Platform: "doubao", + } + + questions := []string{ + "什么是人工智能？", + "Python和Go的区别是什么？", + "如何学习编程？", + "中国有哪些著名的旅游景点？", + "健康饮食的基本原则是什么？", + "区块链技术的主要应用场景有哪些？", + "如何提高英语口语能力？", + "云计算的优势和劣势分别是什么？", + "环境保护的重要性体现在哪些方面？", + "未来十年最有前景的行业有哪些？", + } + + for i, question := range questions { + t.Logf("========== 第 %d/%d 个问题 ==========", i+1, len(questions)) + t.Logf("问题: %s", question) + + result, err := doubaoManager.AskQuestion("doubao", params, question) + if err != nil { + t.Errorf("第 %d 个问题提问失败: %v", i+1, err) + continue + } + + t.Logf("✓ 第 %d 个问题的答案长度: %d 字符", i+1, len(result.Answer)) + t.Logf("答案预览: %s...", result.Answer[:min(100, len(result.Answer))]) + + // 每个问题之间等待一下 + if i < len(questions)-1 { + t.Log("等待3秒后继续下一个问题...") + time.Sleep(3 * time.Second) + } + } + + t.Logf("========== 测试完成 ==========") + t.Logf("成功回答了 %d 个问题", len(questions)) +} + +// TestDoubaoCollector_SpeedTest 测试优化后的速度 +func TestDoubaoCollector_SpeedTest(t *testing.T) { + if testing.Short() { + t.Skip("跳过需要浏览器交互的测试") + } + + params := &collect.CollectParams{ + Headless: false, + RequestID: "test_doubao_speed", + Platform: "doubao", + } + + question := "1+1等于几？" + + startTime := time.Now() + t.Logf("开始提问: %s", question) + + result, err := doubaoManager.AskQuestion("doubao", params, question) + if err != nil { + t.Fatalf("提问失败: %v", err) + } + + elapsed := time.Since(startTime) + t.Logf("✓ 完成时间: %v", elapsed) + t.Logf("答案长度: %d 字符", len(result.Answer)) + t.Logf("分享链接: %s", result.ShareLink) + t.Logf("答案预览: %s...", result.Answer[:min(100, len(result.Answer))]) + + // 如果超过60秒，说明还有优化空间 + if elapsed > 60*time.Second { + t.Logf("⚠️ 警告: 耗时较长 (%v)，可能需要进一步优化", elapsed) + } else { + t.Logf("✅ 速度正常") + } +} + +// min 辅助函数 +func min(a, b int) int { + if a < b { + return a + } + return b +} diff --git a/internal/collect/base.go b/internal/collect/base.go index 91c160f..d22d4c5 100644 --- a/internal/collect/base.go +++ b/internal/collect/base.go @@ -21,6 +21,7 @@ type BaseCollector struct { Headless bool RequestID string Platform string + KeyWords []string Browser *rod.Browser Page *rod.Page @@ -52,6 +53,7 @@ func NewBaseCollector(ctx context.Context, params *CollectParams, config *config Headless: params.Headless, RequestID: params.RequestID, Platform: params.Platform, + KeyWords: params.KeyWords, Logger: baseLogger, config: config, MaxRetries: 3, diff --git a/internal/collect/doubao.go b/internal/collect/doubao.go index 98c437c..30349d4 100644 --- a/internal/collect/doubao.go +++ b/internal/collect/doubao.go @@ -7,6 +7,7 @@ import ( "strings" "time" + "github.com/atotto/clipboard" "github.com/go-rod/rod" "github.com/go-rod/rod/lib/proto" "github.com/gofiber/fiber/v2/log" @@ -32,24 +33,23 @@ func NewDoubaoCollector(ctx context.Context, params *CollectParams, cfg *config. // CheckLoginStatus 检查登录状态 func (c *DoubaoCollector) CheckLoginStatus() bool { - currentURL := c.GetCurrentURL() + c.LogInfo("检查豆包登录状态...") - // 检查是否在聊天页面 - 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 + // 方法3: 检查是否有登录按钮（如果存在说明未登录） + loginButtons, err := c.Page.Elements("button") + if err == nil { + for _, btn := range loginButtons { + text, _ := btn.Text() + trimmedText := strings.TrimSpace(text) + if trimmedText == "登录" || trimmedText == "Login" || strings.Contains(trimmedText, "登录") { + c.LogInfo(fmt.Sprintf("检测到登录按钮'%s'，说明未登录", trimmedText)) + return false + } } } - return false + c.LogInfo("未检测到登录状态相关元素") + return true } // WaitLogin 等待登录 @@ -59,20 +59,37 @@ func (c *DoubaoCollector) WaitLogin() (bool, string) { } defer c.Close() - c.Page.MustNavigate(c.LoginURL) + c.LogInfo("导航到豆包页面...") + c.Page.MustNavigate(c.ChatURL) c.Sleep(3) + // 截图查看初始状态 + c.Screenshot("doubao_initial") + if c.CheckLoginStatus() { + c.LogInfo("已登录，保存cookies") c.SaveCookies() return true, "already_logged_in" } + c.LogInfo("未登录，等待手动登录...") + c.Screenshot("doubao_need_login") + + // 最多等待300秒 for i := 0; i < 300; i++ { if c.CheckLoginStatus() { c.Sleep(2) c.SaveCookies() + c.Screenshot("doubao_login_success") + c.LogInfo("登录成功！") return true, "login_success" } + + // 每10秒输出一次提示 + if i%10 == 0 && i > 0 { + c.LogInfo(fmt.Sprintf("等待登录中... 已等待 %d 秒", i)) + } + time.Sleep(1 * time.Second) } @@ -90,7 +107,12 @@ func (c *DoubaoCollector) AskQuestion(question string) (*CollectResult, error) { return nil, fmt.Errorf("页面初始化失败: %v", err) } - c.Sleep(3) + // 检查是否登录 + if !c.CheckLoginStatus() { + return nil, fmt.Errorf("未登录，请先调用WaitLogin进行登录") + } + + c.LogInfo(fmt.Sprintf("开始提问: %s", question)) if err := c.inputQuestion(question); err != nil { return nil, fmt.Errorf("输入问题失败: %v", err) @@ -104,23 +126,30 @@ func (c *DoubaoCollector) AskQuestion(question string) (*CollectResult, error) { if err != nil { return nil, fmt.Errorf("获取答案失败: %v", err) } + answerStr, isExposure := HighlightKeywordsInHTML(answer, c.KeyWords) + // 获取分享链接 + shareLink := c.getShareLink() + + c.LogInfo(fmt.Sprintf("✓ 获取答案成功，长度: %d 字符", len(answer))) return &CollectResult{ - Answer: answer, - ShareLink: "", + Answer: answerStr, + ShareLink: shareLink, + IsExposure: isExposure, }, nil } // inputQuestion 输入问题 func (c *DoubaoCollector) inputQuestion(question string) error { + c.LogInfo("输入问题...") + + // 豆包的输入框选择器 - 使用精确的class匹配 inputSelectors := []string{ - "textarea[placeholder*='输入']", - "textarea[placeholder*='问']", - "textarea", - "[contenteditable='true']", - ".chat-input textarea", - "#input-box", - ".input-area textarea", + "textarea[placeholder*='发消息...']", + "[class*='input'] textarea", + "textarea.semi-input-textarea", + "textarea[placeholder='发消息...']", + "textarea[class*='semi-input-textarea']", } var inputBox *rod.Element @@ -129,6 +158,7 @@ func (c *DoubaoCollector) inputQuestion(question string) error { for _, selector := range inputSelectors { inputBox, err = c.WaitForElementVisible(selector, 10) if err == nil && inputBox != nil { + c.LogInfo(fmt.Sprintf("找到输入框: %s", selector)) break } } @@ -137,110 +167,382 @@ func (c *DoubaoCollector) inputQuestion(question string) error { 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.LogInfo(fmt.Sprintf("清空输入框失败: %v", err)) } - c.SleepMs(1000) + // 使用原生Input方法输入（更稳定） + inputBox.Input(question) + c.LogInfo(fmt.Sprintf("问题已输入: %s", question)) 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*='发送']", + "[aria-label*='send']", ".send-icon", + "button svg[path*='send']", + "button svg[path*='arrow']", } 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 } } + // 如果没找到，尝试遍历所有button元素 if sendBtn == nil { - sendBtn, err = c.Page.Element("button svg") - if err != nil { - return fmt.Errorf("未找到发送按钮") + c.LogInfo("通过选择器未找到发送按钮，尝试遍历所有button元素...") + allButtons, _ := c.Page.Elements("button") + for _, btn := range allButtons { + // 检查按钮是否可点击且可见 + visible, _ := btn.Visible() + if visible { + classAttr, _ := btn.Attribute("class") + text, _ := btn.Text() + + // 检查是否包含send、submit等关键词 + if classAttr != nil && (strings.Contains(strings.ToLower(*classAttr), "send") || + strings.Contains(strings.ToLower(*classAttr), "submit")) { + sendBtn = btn + c.LogInfo(fmt.Sprintf("通过class找到发送按钮: class=%s", *classAttr)) + break + } + + // 检查文本内容 + trimmedText := strings.TrimSpace(text) + if trimmedText == "发送" || trimmedText == "Send" { + sendBtn = btn + c.LogInfo(fmt.Sprintf("通过文本找到发送按钮: text=%s", trimmedText)) + break + } + } } } - c.SleepMs(500) + // 最后的fallback：查找最后一个button + if sendBtn == nil { + buttons, _ := c.Page.Elements("button") + if len(buttons) > 0 { + sendBtn = buttons[len(buttons)-1] + c.LogInfo("使用最后一个button作为发送按钮") + } + } - if err := c.JSClick(sendBtn); err != nil { + if sendBtn == nil { + return fmt.Errorf("未找到发送按钮") + } + + // 滚动到可见区域 + if err := sendBtn.ScrollIntoView(); err != nil { + c.LogInfo(fmt.Sprintf("滚动失败: %v", err)) + } + + // 点击发送按钮 + c.LogInfo("执行点击...") + if err := sendBtn.Click(proto.InputMouseButtonLeft, 1); err != nil { return fmt.Errorf("点击发送按钮失败: %v", err) } - c.SleepMs(2000) + c.LogInfo("已点击发送按钮") return nil } -// waitForAnswer 等待并获取答案 +// waitForAnswer 等待并获取答案（处理流式输出） func (c *DoubaoCollector) waitForAnswer() (string, error) { - timeout := 120 + c.LogInfo("等待AI回答...") + + timeout := 180 // 最大等待时间（秒） startTime := time.Now() - lastAnswerLength := 0 + + var lastAnswer string + var stableCount int // 稳定计数器 + const requiredStableCount = 5 // 需要连续5次内容不变才认为完成 + isAnswering := false // 标记是否正在回答中 for time.Since(startTime).Seconds() < float64(timeout) { + // 尝试多种方式查找答案容器 answerSelectors := []string{ - ".message-content", - ".response-text", + "[class*='message-content']", + "[class*='response-text']", "[class*='assistant'] [class*='content']", "[class*='bot'] [class*='message']", ".chat-message.bot", ".answer-box", + "[class*='answer']", + "[class*='reply']", + "[data-testid*='message']", + // 豆包特定的选择器 + "[class*='bubble']", + "[class*='chat-bubble']", + "[class*='msg-content']", + "[class*='text-content']", + ".markdown-body", + "[class*='markdown']", + // 更通用的选择器 + "div[class*='content']", + "div[class*='text']", + "article", + "section", } + var answerText string + for _, selector := range answerSelectors { answerElements, err := c.Page.Elements(selector) if err == nil && len(answerElements) > 0 { - lastAnswer := answerElements[len(answerElements)-1] + // 取最后一个元素（最新的回答） + lastAnswerElem := answerElements[len(answerElements)-1] - visible, _ := lastAnswer.Visible() + visible, _ := lastAnswerElem.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") + // 尝试获取HTML内容 + htmlContent, err := lastAnswerElem.HTML() + if err == nil && len(strings.TrimSpace(htmlContent)) > 30 { + // 清理HTML标签，只保留纯文本 + answerText = CleanHTMLTags(htmlContent) + c.LogInfo(fmt.Sprintf("找到答案容器: %s, 清理后文本长度: %d", selector, len(answerText))) + break + } - if !isGenerating { - currentLength := len(text) - if currentLength == lastAnswerLength && currentLength > 10 { - return strings.TrimSpace(text), nil - } - lastAnswerLength = currentLength - } + // 如果HTML获取失败，尝试获取文本 + text, err := lastAnswerElem.Text() + if err == nil && len(strings.TrimSpace(text)) > 30 { + answerText = strings.TrimSpace(text) + c.LogInfo(fmt.Sprintf("找到答案容器: %s, 文本长度: %d", selector, len(answerText))) + break } } } } - c.SleepMs(1500) + // 如果常规方法没找到，尝试查找所有包含较多文本的div + if answerText == "" { + allDivs, _ := c.Page.Elements("div") + for _, div := range allDivs { + visible, _ := div.Visible() + if !visible { + continue + } + + text, err := div.Text() + if err == nil { + trimmedText := strings.TrimSpace(text) + // 查找包含较多文本且不是输入框的div + if len(trimmedText) > 50 && len(trimmedText) < 5000 { + // 排除输入框相关的div + classAttr, _ := div.Attribute("class") + if classAttr != nil { + classLower := strings.ToLower(*classAttr) + if strings.Contains(classLower, "input") || + strings.Contains(classLower, "textarea") || + strings.Contains(classLower, "send") { + continue + } + } + + answerText = CleanHTMLTags(trimmedText) + c.LogInfo(fmt.Sprintf("通过遍历div找到答案，文本长度: %d", len(answerText))) + break + } + } + } + } + + // 检查是否获取到答案 + if answerText != "" && len(answerText) > 30 { + if !isAnswering { + c.LogInfo("检测到AI开始回答...") + isAnswering = true + } + + // 检查内容是否稳定（流式输出完成） + if answerText == lastAnswer { + stableCount++ + c.LogInfo(fmt.Sprintf("答案稳定中... (%d/%d), 长度: %d", stableCount, requiredStableCount, len(answerText))) + + // 如果内容稳定足够次数，说明回答完成 + if stableCount >= requiredStableCount { + c.LogInfo(fmt.Sprintf("✓ AI回答完成，最终长度: %d 字符", len(answerText))) + return answerText, nil + } + } else { + // 内容还在变化，重置计数器 + stableCount = 0 + lastAnswer = answerText + 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)) + // 截图帮助调试 + if elapsed%30 == 0 { + c.Screenshot(fmt.Sprintf("doubao_wait_answer_%d", elapsed)) + } + } } - return "", fmt.Errorf("等待答案超时") + return "", fmt.Errorf("等待答案超时（%d秒）", timeout) +} + +// getShareLink 尝试获取当前对话的分享链接 +func (c *DoubaoCollector) getShareLink() string { + c.LogInfo("尝试获取分享链接...") + + // 步骤1: 找到class包含message-action-button-main的div + actionDiv, err := c.Page.Element("div[class*='message-action-button-main']") + if err != nil || actionDiv == nil { + c.LogInfo("未找到message-action-button-main元素") + return "" + } + + c.LogInfo("找到message-action-button-main元素") + + // 步骤2: 在该div中找到所有button，取倒数第二个作为分享按钮 + buttons, err := actionDiv.Elements("button") + if err != nil || len(buttons) == 0 { + c.LogInfo("未找到button元素") + return "" + } + + if len(buttons) < 2 { + c.LogInfo(fmt.Sprintf("button数量不足(%d)，无法获取倒数第二个", len(buttons))) + return "" + } + + // 取倒数第二个button + shareBtn := buttons[len(buttons)-3] + c.LogInfo(fmt.Sprintf("找到分享按钮（倒数第2个，共%d个button）", len(buttons))) + + // 检查是否可点击，如果pointer-events为none，使用JavaScript点击 + visible, _ := shareBtn.Visible() + if !visible { + c.LogInfo("分享按钮不可见，尝试使用JavaScript点击") + // 使用立即执行函数，但返回一个空函数避免.apply错误 + _, err := c.Page.Eval(`(function(){Array.from(document.querySelectorAll('div[class*="message-action-button-main"] button')).slice(-2)[0].click();return function(){};})`) + if err != nil { + c.LogInfo(fmt.Sprintf("JavaScript点击失败: %v", err)) + return "" + } + } else { + // 正常点击 + if err := shareBtn.Click(proto.InputMouseButtonLeft, 1); err != nil { + c.LogInfo(fmt.Sprintf("点击分享按钮失败: %v，尝试JavaScript点击", err)) + // Fallback: 使用JavaScript点击 + _, err := c.Page.Eval(`(function(){Array.from(document.querySelectorAll('div[class*="message-action-button-main"] button')).slice(-2)[0].click();return function(){};})`) + if err != nil { + c.LogInfo(fmt.Sprintf("JavaScript点击也失败: %v", err)) + return "" + } + } + } + + c.SleepMs(500) + + // 步骤3: 找到内容为"复制链接"的span并点击 + copySpan, err := c.Page.ElementX("//span[contains(text(), '复制链接')]") + if err != nil || copySpan == nil { + c.LogInfo("未找到'复制链接'span元素") + return "" + } + + c.LogInfo("找到'复制链接'span元素，点击复制...") + if err := copySpan.Click(proto.InputMouseButtonLeft, 1); err != nil { + c.LogInfo(fmt.Sprintf("点击复制按钮失败: %v，尝试JavaScript点击", err)) + // Fallback: 使用JavaScript点击 + script := ` + (function() { + var spans = document.querySelectorAll('span'); + for (var i = 0; i < spans.length; i++) { + if (spans[i].textContent.includes('复制链接')) { + spans[i].click(); + return true; + } + } + return false; + })() + ` + result, err := c.Page.Eval(script) + if err != nil || result == nil { + c.LogInfo(fmt.Sprintf("JavaScript点击复制按钮失败: %v", err)) + return "" + } + } + c.SleepMs(500) + + // 步骤4: 从剪贴板获取内容 + clipboardContent := c.getClipboardContent() + if clipboardContent != "" { + c.LogInfo(fmt.Sprintf("从剪贴板获取到分享链接: %s", clipboardContent)) + return clipboardContent + } + + c.LogInfo("未能从剪贴板获取链接") + return "" +} + +// getClipboardContent 从剪贴板获取内容 +func (c *DoubaoCollector) getClipboardContent() string { + // 使用atotto/clipboard库读取系统剪贴板 + text, err := clipboard.ReadAll() + if err != nil { + c.LogInfo(fmt.Sprintf("读取剪贴板失败: %v", err)) + return "" + } + + if text == "" { + c.LogInfo("剪贴板内容为空") + return "" + } + + c.LogInfo(fmt.Sprintf("剪贴板原始内容: %s", text)) + return text +} + +// extractURL 从文本中提取 URL +func extractURL(text string) string { + // 简单的 URL 提取逻辑 + start := strings.Index(text, "https://") + if start == -1 { + start = strings.Index(text, "http://") + } + if start != -1 { + end := strings.Index(text[start:], " ") + if end == -1 { + return text[start:] + } + return text[start : start+end] + } + return "" } // SafeElement 安全地获取元素 diff --git a/internal/collect/interface.go b/internal/collect/interface.go index 3221d14..4a38b6a 100644 --- a/internal/collect/interface.go +++ b/internal/collect/interface.go @@ -17,8 +17,9 @@ type CollectorInterface interface { // CollectResult 收集结果 type CollectResult struct { - Answer string `json:"answer"` // AI回答内容 - ShareLink string `json:"share_link"` // 分享链接 + Answer string `json:"answer"` // AI回答内容 + ShareLink string `json:"share_link"` // 分享链接\ + IsExposure bool `json:"is_exposure"` // 是否曝光 } // NewCollector 创建收集器的工厂函数类型 @@ -40,8 +41,8 @@ type CollectorValue struct { type CollectParams struct { Headless bool // 是否无头模式 RequestID string // 请求ID - Platform string // 平台类型: wenxin, deepseek, doubao, qianwen - + Platform string + KeyWords []string } // CollectorMap 收集器注册表 diff --git a/internal/collect/utils.go b/internal/collect/utils.go index 58fbd02..2a53ee7 100644 --- a/internal/collect/utils.go +++ b/internal/collect/utils.go @@ -68,9 +68,10 @@ func CleanDivTags(html string) string { // htmlContent: 原始HTML内容 // pointKeys: 需要高亮的关键词列表 // 返回处理后的HTML内容，每个关键词会被不同颜色的span标签包裹 -func HighlightKeywordsInHTML(htmlContent string, pointKeys []string) string { +func HighlightKeywordsInHTML(htmlContent string, pointKeys []string) (string, bool) { + var isExposure bool if htmlContent == "" || len(pointKeys) == 0 { - return htmlContent + return htmlContent, isExposure } // 预定义的颜色列表（使用CSS颜色值） @@ -111,12 +112,17 @@ func HighlightKeywordsInHTML(htmlContent string, pointKeys []string) string { pattern := fmt.Sprintf(`(?i)(%s)`, escapedKeyword) re := regexp.MustCompile(pattern) + // 检查是否匹配到关键词 + if re.MatchString(result) { + isExposure = true + } + // 替换匹配的关键词为带颜色的span标签 replacement := fmt.Sprintf(`$1`, color) result = re.ReplaceAllString(result, replacement) } - return result + return result, isExposure } // HighlightKeywordsInText 在纯文本中高亮显示指定的关键词（先转换为HTML） diff --git a/internal/collect/utils_test.go b/internal/collect/utils_test.go new file mode 100644 index 0000000..e54e91c --- /dev/null +++ b/internal/collect/utils_test.go @@ -0,0 +1,103 @@ +package collect + +import ( + "strings" + "testing" +) + +// TestHighlightKeywordsInHTML 测试HTML内容关键词高亮功能 +func TestHighlightKeywordsInHTML(t *testing.T) { + html := `

在四川房地产软件领域，根据功能深度、本地化服务、技术实力及性价比等维度评测，以下软件表现突出且排名靠前：

1. 云案场

核心优势：

适用场景：

2. 明源云客

核心优势：

适用场景：

3. 用友地产CRM / 金蝶我家云售楼版

核心优势：

适用场景：

4. 元度云案场

核心优势：

适用场景：

5. 贝壳找房/链家网

核心优势：

适用场景：

排名依据与选型建议

  1. 功能深度：云案场与明源云客在全流程覆盖与风控领域表现突出，适合大型房企；用友/金蝶强于业财一体化。
  2. 本地化服务：云案场在四川本地响应速度与案例经验占优。
  3. 性价比：元度云案场实施成本低，适合中小型房企；云案场提供灵活模块组合，适配不同规模需求。
  4. 技术实力：云案场、明源云客等获等保认证，数据安全有保障。

建议：

` + keyWords := []string{"云案场", "关键词2"} + result := HighlightKeywordsInText(html, keyWords) + t.Log(result) +} + +// TestHighlightKeywordsInHTML_ColorAssignment 测试颜色分配逻辑 +func TestHighlightKeywordsInHTML_ColorAssignment(t *testing.T) { + // 创建一个包含所有关键词的HTML内容 + keywords := make([]string, 20) + htmlParts := make([]string, 20) + for i := 0; i < 20; i++ { + keyword := "关键词" + string(rune('A'+i)) + keywords[i] = keyword + htmlParts[i] = "

" + keyword + "

" + } + + htmlContent := strings.Join(htmlParts, "") + + result := HighlightKeywordsInHTML(htmlContent, keywords) + + // 验证所有关键词都被处理（应该都有span标签） + spanCount := strings.Count(result, `