133 lines
3.7 KiB
Go
133 lines
3.7 KiB
Go
package collect
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
"geo/internal/config"
|
||
"os"
|
||
"path/filepath"
|
||
"sync"
|
||
|
||
"github.com/go-rod/rod"
|
||
"github.com/go-rod/rod/lib/launcher"
|
||
"github.com/gofiber/fiber/v2/log"
|
||
)
|
||
|
||
// CollectManager 收集管理器 - 单浏览器临时Page模式
|
||
type CollectManager struct {
|
||
ctx context.Context
|
||
config *config.Config
|
||
logger log.AllLogger
|
||
|
||
browser *rod.Browser // 全局唯一的浏览器实例
|
||
mu sync.RWMutex // 保护浏览器访问的互斥锁
|
||
|
||
collectors map[string]CollectorInterface // 平台名称 -> Collector 实例
|
||
}
|
||
|
||
// NewCollectManager 创建收集管理器（服务启动时调用）
|
||
func NewCollectManager(ctx context.Context, cfg *config.Config, logger log.AllLogger) *CollectManager {
|
||
manager := &CollectManager{
|
||
ctx: ctx,
|
||
config: cfg,
|
||
logger: logger,
|
||
collectors: make(map[string]CollectorInterface),
|
||
}
|
||
|
||
// 注册所有平台的 Collector
|
||
manager.registerCollectors()
|
||
|
||
// 服务启动时立即创建全局浏览器
|
||
manager.initBrowser()
|
||
|
||
return manager
|
||
}
|
||
|
||
// initBrowser 初始化全局浏览器实例
|
||
func (m *CollectManager) initBrowser() {
|
||
m.logger.Info("🚀 正在初始化浏览器...")
|
||
|
||
// 创建用户数据目录（所有平台共用）
|
||
userDataDir := filepath.Join(m.config.Sys.ChromeDataDir, "global", "main")
|
||
os.MkdirAll(userDataDir, 0755)
|
||
|
||
// 使用 launcher 启动 Chrome
|
||
l := launcher.New().
|
||
Bin(m.config.Sys.ChromePath).
|
||
UserDataDir(userDataDir).
|
||
Headless(false). // 有头模式，显示窗口
|
||
Leakless(true). // 让 Chrome 进程独立运行
|
||
Set("disable-blink-features", "AutomationControlled").
|
||
Set("window-size", "1920,1080").
|
||
Set("start-maximized", "").
|
||
Set("lang", "zh-CN")
|
||
|
||
url, err := l.Launch()
|
||
if err != nil {
|
||
m.logger.Errorf("❌ 启动浏览器失败: %v", err)
|
||
panic(fmt.Sprintf("启动浏览器失败: %v", err))
|
||
}
|
||
|
||
m.logger.Infof("✅ 浏览器启动成功: %s", url)
|
||
|
||
// 连接到浏览器
|
||
m.browser = rod.New().Context(m.ctx).ControlURL(url).MustConnect()
|
||
|
||
m.logger.Info("✅ 浏览器连接成功")
|
||
}
|
||
|
||
// Close 关闭管理器（注意：不关闭浏览器，让浏览器进程继续运行）
|
||
func (m *CollectManager) Close() {
|
||
m.mu.Lock()
|
||
defer m.mu.Unlock()
|
||
|
||
m.logger.Info("🔒 CollectManager 关闭（浏览器保持运行）...")
|
||
|
||
// 注意：这里不调用 m.browser.Close()
|
||
// 让浏览器进程继续运行，下次服务启动时可以复用
|
||
|
||
m.logger.Info("✅ Manager 已关闭，浏览器仍在后台运行")
|
||
}
|
||
|
||
// registerCollectors 注册所有平台的 Collector
|
||
func (m *CollectManager) registerCollectors() {
|
||
// 文心一言
|
||
params := &CollectParams{Platform: "wenxin"}
|
||
m.collectors["wenxin"] = NewWenxinCollector(m.ctx, params, m.config, m.logger, m.browser, nil)
|
||
|
||
// DeepSeek
|
||
params = &CollectParams{Platform: "deepseek"}
|
||
m.collectors["deepseek"] = NewDeepseekCollector(m.ctx, params, m.config, m.logger, m.browser, nil)
|
||
|
||
// 豆包
|
||
params = &CollectParams{Platform: "doubao"}
|
||
m.collectors["doubao"] = NewDoubaoCollector(m.ctx, params, m.config, m.logger, m.browser, nil)
|
||
|
||
// 通义千问
|
||
params = &CollectParams{Platform: "qianwen"}
|
||
m.collectors["qianwen"] = NewQianwenCollector(m.ctx, params, m.config, m.logger, m.browser, nil)
|
||
}
|
||
|
||
// GetCollector 获取指定平台的 Collector
|
||
func (m *CollectManager) GetCollector(platform string) (CollectorInterface, error) {
|
||
m.mu.RLock()
|
||
defer m.mu.RUnlock()
|
||
|
||
collector, ok := m.collectors[platform]
|
||
if !ok {
|
||
return nil, fmt.Errorf("不支持的平台: %s", platform)
|
||
}
|
||
|
||
return collector, nil
|
||
}
|
||
|
||
// AskQuestion 向指定平台提问（每次创建临时 Page）
|
||
func (m *CollectManager) AskQuestion(platform string, question string) (*CollectResult, error) {
|
||
collector, err := m.GetCollector(platform)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
return collector.AskQuestion(question)
|
||
}
|