kurihada/ai-desktop
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

完善小黑盒搜索功能，将小黑盒操作作为工具给大模型

Browse Source
This commit is contained in:
kurihada
2025-11-14 15:50:34 +08:00
parent 7b955de2f0
commit d86c4b21ae
46 changed files with 8477 additions and 5800 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/main/index.ts
+370 -21
View File
@@ -1,7 +1,7 @@
import { app, BrowserWindow, ipcMain, screen, globalShortcut, clipboard } from 'electron'
import { join } from 'path'
import { chromium, BrowserContext } from 'playwright'
import { existsSync, rmSync } from 'fs'
import { existsSync, rmSync, readFileSync, writeFileSync, mkdirSync } from 'fs'
import { ScraperFactory } from './scrapers'
import { XiaoheiheScrap } from './scrapers/xiaoheihe'
import { GenericScraper } from './scrapers/generic'
@@ -11,6 +11,7 @@ import { XiaoheiheService } from './platforms/xiaoheihe'
let floatingWindow: BrowserWindow | null = null
let settingsWindow: BrowserWindow | null = null
let chatWindow: BrowserWindow | null = null
let toolsPanelWindow: BrowserWindow | null = null
// Initialize scraper factory
const scraperFactory = new ScraperFactory()
@@ -25,14 +26,139 @@ platformServiceFactory.register(new XiaoheiheService())
let persistentContext: BrowserContext | null = null
const userDataDir = join(app.getPath('userData'), 'browser-data')
// Settings file path
const settingsDir = join(app.getPath('userData'), 'settings')
const settingsFilePath = join(settingsDir, 'config.json')
const loginInfoFilePath = join(settingsDir, 'login-info.json')
// Rate limiter for search operations
class SearchRateLimiter {
private lastSearchTime: number = 0
private searchCount: number = 0
private readonly minInterval: number = 3000 // 最小间隔 3 秒
private readonly maxSearchPerMinute: number = 10 // 每分钟最多 10 次
private readonly resetInterval: number = 60000 // 1 分钟重置计数
canSearch(): { allowed: boolean; waitTime?: number; reason?: string } {
const now = Date.now()
const timeSinceLastSearch = now - this.lastSearchTime
// 检查是否需要重置计数器
if (timeSinceLastSearch > this.resetInterval) {
this.searchCount = 0
}
// 检查是否超过频率限制
if (this.searchCount >= this.maxSearchPerMinute) {
const waitTime = this.resetInterval - timeSinceLastSearch
return {
allowed: false,
waitTime: Math.ceil(waitTime / 1000),
reason: '搜索过于频繁，请稍后再试'
}
}
// 检查最小间隔
if (timeSinceLastSearch < this.minInterval) {
const waitTime = this.minInterval - timeSinceLastSearch
return {
allowed: false,
waitTime: Math.ceil(waitTime / 1000),
reason: '请求过快，请稍后再试'
}
}
return { allowed: true }
}
recordSearch(): void {
this.lastSearchTime = Date.now()
this.searchCount++
}
reset(): void {
this.searchCount = 0
this.lastSearchTime = 0
}
}
const searchRateLimiter = new SearchRateLimiter()
// Ensure settings directory exists
function ensureSettingsDir(): void {
if (!existsSync(settingsDir)) {
mkdirSync(settingsDir, { recursive: true })
}
}
// Read settings from file
function readSettings(): any {
try {
ensureSettingsDir()
if (existsSync(settingsFilePath)) {
const data = readFileSync(settingsFilePath, 'utf-8')
return JSON.parse(data)
}
return {}
} catch (error) {
console.error('Failed to read settings:', error)
return {}
}
}
// Write settings to file
function writeSettings(settings: any): { success: boolean; error?: string } {
try {
ensureSettingsDir()
writeFileSync(settingsFilePath, JSON.stringify(settings, null, 2), 'utf-8')
return { success: true }
} catch (error) {
console.error('Failed to write settings:', error)
return {
success: false,
error: error instanceof Error ? error.message : '保存设置失败'
}
}
}
// Read login info from file
function readLoginInfo(): any {
try {
ensureSettingsDir()
if (existsSync(loginInfoFilePath)) {
const data = readFileSync(loginInfoFilePath, 'utf-8')
return JSON.parse(data)
}
return {}
} catch (error) {
console.error('Failed to read login info:', error)
return {}
}
}
// Write login info to file
function writeLoginInfo(loginInfo: any): { success: boolean; error?: string } {
try {
ensureSettingsDir()
writeFileSync(loginInfoFilePath, JSON.stringify(loginInfo, null, 2), 'utf-8')
return { success: true }
} catch (error) {
console.error('Failed to write login info:', error)
return {
success: false,
error: error instanceof Error ? error.message : '保存登录信息失败'
}
}
}
function createFloatingWindow(): void {
const { width } = screen.getPrimaryDisplay().workAreaSize
const { width, height } = screen.getPrimaryDisplay().workAreaSize
floatingWindow = new BrowserWindow({
width: 260,
height: 160,
x: width - 100,
y: 20,
width: 160,
height: 200,
x: width - 160,
y: Math.floor((height - 200) / 2),
frame: false,
transparent: true,
alwaysOnTop: true,
@@ -92,7 +218,8 @@ function createSettingsWindow(): void {
}
function createChatWindow(initialText?: string): void {
console.log('createChatWindow called with initialText:', initialText)
const startTime = Date.now()
console.log('[PERF] createChatWindow called at:', startTime)
// If chat window already exists, focus it and send new text if provided
if (chatWindow && !chatWindow.isDestroyed()) {
@@ -108,16 +235,18 @@ function createChatWindow(initialText?: string): void {
return
}
console.log('Creating new chat window')
console.log('[PERF] Creating new chat window')
chatWindow = new BrowserWindow({
width: 800,
height: 600,
title: 'AI 对话',
show: false, // Hide window during load for better perceived performance
webPreferences: {
preload: join(__dirname, '../preload/index.js'),
sandbox: false,
nodeIntegration: false,
contextIsolation: true
contextIsolation: true,
spellcheck: false // Disable spell check to avoid conflicts with IME
}
})
@@ -128,12 +257,20 @@ function createChatWindow(initialText?: string): void {
chatWindow.loadFile(join(__dirname, '../renderer/chat.html'))
}
// Show window when ready to avoid showing loading state
chatWindow.once('ready-to-show', () => {
const readyTime = Date.now()
console.log('[PERF] Chat window ready-to-show at:', readyTime)
console.log('[PERF] Time from create to ready:', readyTime - startTime, 'ms')
chatWindow?.show()
})
// Send initial text after page loads
if (initialText) {
console.log('Setting up did-finish-load listener for initial text')
chatWindow.webContents.once('did-finish-load', () => {
console.log('Chat window did-finish-load event fired')
// Add a small delay to ensure React components are mounted
// Add a small delay to ensure Vue components are mounted
setTimeout(() => {
if (chatWindow && !chatWindow.isDestroyed()) {
console.log('Sending initial text to new window:', initialText)
@@ -148,6 +285,49 @@ function createChatWindow(initialText?: string): void {
})
}
// Create tools panel window
function createToolsPanelWindow(): void {
console.log('Creating tools panel window')
// If tools panel already exists, focus it
if (toolsPanelWindow && !toolsPanelWindow.isDestroyed()) {
console.log('Tools panel window already exists, focusing')
toolsPanelWindow.focus()
return
}
toolsPanelWindow = new BrowserWindow({
width: 900,
height: 700,
title: '工具箱',
show: false,
webPreferences: {
preload: join(__dirname, '../preload/index.js'),
sandbox: false,
nodeIntegration: false,
contextIsolation: true,
spellcheck: false
}
})
// Load tools panel page
if (process.env['ELECTRON_RENDERER_URL']) {
toolsPanelWindow.loadURL(`${process.env['ELECTRON_RENDERER_URL']}/tools.html`)
} else {
toolsPanelWindow.loadFile(join(__dirname, '../renderer/tools.html'))
}
// Show window when ready
toolsPanelWindow.once('ready-to-show', () => {
console.log('Tools panel window ready-to-show')
toolsPanelWindow?.show()
})
toolsPanelWindow.on('closed', () => {
toolsPanelWindow = null
})
}
// Fetch article content using Playwright with factory pattern
async function fetchArticleContent(url: string): Promise<{
title: string
@@ -169,8 +349,10 @@ async function fetchArticleContent(url: string): Promise<{
hotScore: number
}
}> {
console.log('fetchArticleContent: Starting to fetch article from URL:', url)
let browser
try {
console.log('fetchArticleContent: Launching browser...')
browser = await chromium.launch({ headless: true })
const context = await browser.newContext({
userAgent:
@@ -178,25 +360,35 @@ async function fetchArticleContent(url: string): Promise<{
})
const page = await context.newPage()
console.log('fetchArticleContent: Navigating to URL...')
// Navigate to the URL
await page.goto(url, { waitUntil: 'networkidle', timeout: 30000 })
// Get appropriate scraper for this URL
console.log('fetchArticleContent: Getting scraper for URL...')
const scraper = scraperFactory.getScraper(url)
if (!scraper) {
throw new Error('No suitable scraper found for this URL')
}
console.log('fetchArticleContent: Using scraper:', scraper.constructor.name)
// Use scraper to extract content
const result = await scraper.scrape(page)
await browser.close()
console.log('fetchArticleContent: Scraping completed, article title:', result.title)
return result
} catch (error) {
if (browser) {
await browser.close()
}
console.error('fetchArticleContent: Error occurred:', error)
throw error
} finally {
// Always close browser in finally block to ensure cleanup
if (browser) {
try {
await browser.close()
console.log('fetchArticleContent: Browser closed')
} catch (closeError) {
console.error('Failed to close browser:', closeError)
}
}
}
}
@@ -244,7 +436,20 @@ async function waitForQrCodeLogin(): Promise<{
return { success: false, error: '浏览器上下文未初始化' }
}
const service = new XiaoheiheService()
return await service.waitForQrCodeLogin(persistentContext)
const result = await service.waitForQrCodeLogin(persistentContext)
// Save login info if successful
if (result.success && result.username) {
const savedLoginInfo = readLoginInfo()
savedLoginInfo['www.xiaoheihe.cn'] = {
username: result.username,
lastUpdate: new Date().toISOString()
}
writeLoginInfo(savedLoginInfo)
console.log('Login info saved for user:', result.username)
}
return result
} catch (error) {
console.error('Wait for QR code login error:', error)
return {
@@ -270,6 +475,26 @@ async function checkPlatformLoginFast(url: string): Promise<{
const context = await getPersistentContext()
const loginStatus = await service.checkLoginStatusFast(context)
// 如果已登录，尝试从本地文件读取用户名（如果 service 没有返回用户名）
if (loginStatus.isLoggedIn && !loginStatus.username) {
const savedLoginInfo = readLoginInfo()
const platformKey = url.replace(/https?:\/\//, '').split('/')[0]
if (savedLoginInfo[platformKey]?.username) {
loginStatus.username = savedLoginInfo[platformKey].username
}
}
// 如果已登录且有用户名，保存到本地文件
if (loginStatus.isLoggedIn && loginStatus.username) {
const savedLoginInfo = readLoginInfo()
const platformKey = url.replace(/https?:\/\//, '').split('/')[0]
savedLoginInfo[platformKey] = {
username: loginStatus.username,
lastUpdate: new Date().toISOString()
}
writeLoginInfo(savedLoginInfo)
}
return { success: true, ...loginStatus }
} catch (error) {
console.error('Check platform login fast error:', error)
@@ -288,6 +513,7 @@ async function checkPlatformLogin(url: string): Promise<{
username?: string
error?: string
}> {
let page: Awaited<ReturnType<BrowserContext['newPage']>> | undefined
try {
const service = platformServiceFactory.getService(url)
if (!service) {
@@ -295,11 +521,10 @@ async function checkPlatformLogin(url: string): Promise<{
}
const context = await getPersistentContext()
const page = await context.newPage()
page = await context.newPage()
await page.goto(url, { waitUntil: 'networkidle', timeout: 30000 })
const loginStatus = await service.checkLoginStatus(page)
await page.close()
return { success: true, ...loginStatus }
} catch (error) {
@@ -309,6 +534,15 @@ async function checkPlatformLogin(url: string): Promise<{
isLoggedIn: false,
error: error instanceof Error ? error.message : '检查登录状态失败'
}
} finally {
// Always close page in finally block to ensure cleanup
if (page) {
try {
await page.close()
} catch (closeError) {
console.error('Failed to close page:', closeError)
}
}
}
}
@@ -348,6 +582,72 @@ async function postCommentToPlatform(
}
}
// Search platform content
async function searchPlatform(
platform: string,
query: string
): Promise<{
success: boolean
results?: Array<{
title: string
url: string
author?: string
publishTime?: string
summary?: string
commentCount?: number
likeCount?: number
}>
error?: string
}> {
try {
console.log('searchPlatform: Starting search on platform:', platform)
console.log('searchPlatform: Query:', query)
// 检查搜索频率限制
const rateLimitCheck = searchRateLimiter.canSearch()
if (!rateLimitCheck.allowed) {
console.log('searchPlatform: Rate limit exceeded')
return {
success: false,
error: 'RATE_LIMIT_EXCEEDED',
results: []
}
}
// 根据平台名称构造一个URL来获取对应的服务
const platformUrls: Record<string, string> = {
xiaoheihe: 'https://www.xiaoheihe.cn'
}
const platformUrl = platformUrls[platform]
if (!platformUrl) {
return { success: false, error: '不支持的平台' }
}
const service = platformServiceFactory.getService(platformUrl)
if (!service) {
return { success: false, error: '未找到平台服务' }
}
console.log('searchPlatform: Getting persistent context...')
const context = await getPersistentContext()
console.log('searchPlatform: Context obtained, calling service.search...')
// 记录本次搜索
searchRateLimiter.recordSearch()
const result = await service.search(context, query)
console.log('searchPlatform: Search completed, found', result.results?.length || 0, 'results')
return result
} catch (error) {
console.error('searchPlatform: Exception occurred:', error)
return {
success: false,
error: error instanceof Error ? error.message : '搜索失败'
}
}
}
function registerGlobalShortcuts(): void {
// Register Command+K (Mac) or Ctrl+K (Windows/Linux)
const shortcut = process.platform === 'darwin' ? 'Command+K' : 'Control+K'
@@ -391,6 +691,18 @@ app.whenReady().then(() => {
createSettingsWindow()
})
// Handle open tools panel from renderer
ipcMain.on('open-tools-panel', () => {
createToolsPanelWindow()
})
// Handle close tools panel from renderer
ipcMain.on('close-tools-panel', () => {
if (toolsPanelWindow && !toolsPanelWindow.isDestroyed()) {
toolsPanelWindow.close()
}
})
// Handle open chat window from renderer
ipcMain.on('open-chat', (_, selectedText?: string) => {
console.log('open-chat event received, selectedText:', selectedText)
@@ -437,9 +749,22 @@ app.whenReady().then(() => {
return await checkPlatformLoginFast(url)
})
// Handle check platform login status
ipcMain.handle('check-platform-login', async (_, url: string) => {
return await checkPlatformLogin(url)
// Handle check platform login status (accepts platform name or URL)
ipcMain.handle('check-platform-login', async (_, arg: string | { platform: string }) => {
// Support both old format (URL string) and new format (object with platform)
let url: string
if (typeof arg === 'string') {
url = arg
} else {
const platformUrls: Record<string, string> = {
xiaoheihe: 'https://www.xiaoheihe.cn'
}
url = platformUrls[arg.platform]
if (!url) {
return { success: false, isLoggedIn: false, error: '不支持的平台' }
}
}
return await checkPlatformLoginFast(url)
})
// Handle post comment to platform
@@ -474,6 +799,12 @@ app.whenReady().then(() => {
console.log('User data directory deleted successfully')
}
// Clear saved login info from file
const savedLoginInfo = readLoginInfo()
delete savedLoginInfo['www.xiaoheihe.cn']
writeLoginInfo(savedLoginInfo)
console.log('Login info cleared for xiaoheihe')
return { success: true }
}
return { success: false, error: '不支持的平台' }
@@ -486,6 +817,24 @@ app.whenReady().then(() => {
}
})
// Handle search platform
ipcMain.handle(
'search-platform',
async (_, { platform, query }: { platform: string; query: string }) => {
return await searchPlatform(platform, query)
}
)
// Handle read settings
ipcMain.handle('read-settings', () => {
return readSettings()
})
// Handle write settings
ipcMain.handle('write-settings', (_, settings: any) => {
return writeSettings(settings)
})
createFloatingWindow()
registerGlobalShortcuts()