fix: remove sync_status, fix export command

Author: luoling8192

cspell.config.yaml

@@ -6,6 +6,7 @@ words:
   - guiiai
   - Logg
   - onnotice
+  - pgvector
   - tiktoken
   - unocss
   - xsai

packages/cli/src/commands/embed.ts

@@ -1,12 +1,7 @@
 import * as input from '@inquirer/prompts'
 import { useLogger } from '@tg-search/common'
 import { EmbeddingService } from '@tg-search/core'
-import {
-  getMessageCount,
-  getMessagesWithoutEmbedding,
-  getPartitionTables,
-  updateMessageEmbedding,
-} from '@tg-search/db'
+import { findMessagesByChatId, updateMessageEmbeddings } from '@tg-search/db'
 
 import { TelegramCommand } from '../command'
 
@@ -29,20 +24,24 @@ export class EmbedCommand extends TelegramCommand {
   }
 
   async execute(_args: string[], options: EmbedOptions): Promise<void> {
+    // Get chat ID (required)
+    const chatId = options.chatId || await input.input({
+      message: '请输入会话 ID：',
+      default: '',
+    })
+
+    if (!chatId) {
+      throw new Error('会话 ID 是必需的')
+    }
+
     // Get batch size
     const batchSize = options.batchSize || await input.input({
       message: '请输入每批处理的消息数量：',
       default: '1000',
     })
 
-    // Get chat ID
-    const chatId = options.chatId || await input.input({
-      message: '请输入会话 ID（留空处理所有会话）：',
-      default: '',
-    })
-
     // Get concurrency
-    const _concurrency = options.concurrency || await input.input({
+    const concurrency = options.concurrency || await input.input({
       message: '请输入并发数：',
       default: '4',
     })
@@ -51,48 +50,44 @@ export class EmbedCommand extends TelegramCommand {
     const embedding = new EmbeddingService()
 
     try {
-      // Get message count
-      const tables = await getPartitionTables()
-      let totalMessages = 0
-      let totalProcessed = 0
-      let failedEmbeddings = 0
-
-      for (const table of tables) {
-        const count = await getMessageCount(Number(table.tableName))
-        totalMessages += count
-      }
+      // Get all messages for the chat
+      const messages = await findMessagesByChatId(Number(chatId))
+      const messagesToEmbed = messages.filter(m => !m.embedding && m.content)
+      const totalMessages = messagesToEmbed.length
 
       logger.log(`共有 ${totalMessages} 条消息需要处理`)
 
       // Process messages in batches
-      while (true) {
-        const messages = await getMessagesWithoutEmbedding(
-          chatId ? Number(chatId) : undefined,
-          Number(batchSize),
-        )
-
-        if (messages.length === 0) {
-          break
-        }
+      let totalProcessed = 0
+      let failedEmbeddings = 0
 
-        logger.debug(`获取到 ${messages.length} 条消息`)
+      // Split messages into batches
+      for (let i = 0; i < messagesToEmbed.length; i += Number(batchSize)) {
+        const batch = messagesToEmbed.slice(i, i + Number(batchSize))
+        logger.debug(`处理第 ${i + 1} 到 ${i + batch.length} 条消息`)
 
         // Generate embeddings in parallel
-        const contents = messages.map(m => m.content || '')
+        const contents = batch.map(m => m.content!)
         const embeddings = await embedding.generateEmbeddings(contents)
 
-        // Update embeddings
-        for (let i = 0; i < messages.length; i++) {
-          const message = messages[i]
-          try {
-            await updateMessageEmbedding(message.id, message.chatId, embeddings[i])
-            totalProcessed++
-          }
-          catch (error) {
-            logger.withError(error).warn(`更新消息 ${message.id} 的向量嵌入失败`)
-            failedEmbeddings++
+        // Prepare updates
+        const updates = batch.map((message, index) => ({
+          id: message.id,
+          embedding: embeddings[index],
+        }))
+
+        try {
+          // Update embeddings in batches with concurrency control
+          for (let j = 0; j < updates.length; j += Number(concurrency)) {
+            const concurrentBatch = updates.slice(j, j + Number(concurrency))
+            await updateMessageEmbeddings(Number(chatId), concurrentBatch)
+            totalProcessed += concurrentBatch.length
           }
         }
+        catch (error) {
+          logger.withError(error).warn(`更新消息向量嵌入失败`)
+          failedEmbeddings += batch.length
+        }
 
         logger.log(`已处理 ${totalProcessed}/${totalMessages} 条消息，${failedEmbeddings} 条失败`)
       }

packages/cli/src/commands/export.ts

@@ -1,17 +1,19 @@
 import type { Dialog, TelegramMessageType } from '@tg-search/core'
+import type { MessageCreateInput, NewChat } from '@tg-search/db'
 
 import * as fs from 'node:fs/promises'
 import * as path from 'node:path'
 import * as input from '@inquirer/prompts'
-import { useLogger } from '@tg-search/common'
+import { getConfig, useLogger } from '@tg-search/common'
+import { createMessageBatch, updateChat } from '@tg-search/db'
 
 import { TelegramCommand } from '../command'
 
 const logger = useLogger()
 
 interface ExportOptions {
   chatId?: number
-  format?: 'json' | 'html'
+  format?: 'database' | 'html' | 'json'
   path?: string
   messageTypes?: TelegramMessageType[]
   startTime?: Date
@@ -20,6 +22,69 @@ interface ExportOptions {
   batchSize?: number
 }
 
+/**
+ * Process a batch of messages for database export
+ */
+async function processDatabaseBatch(
+  messages: MessageCreateInput[],
+  startIndex: number,
+): Promise<{ shouldStop: boolean, failedCount: number }> {
+  try {
+    // Create messages in batch
+    const result = await createMessageBatch(messages)
+    const firstMessage = messages[0]
+    const lastMessage = messages[messages.length - 1]
+
+    logger.debug(
+      `已保存 ${startIndex + 1} - ${startIndex + messages.length} 条消息 `
+      + `(${new Date(firstMessage.createdAt).toLocaleString()} - ${new Date(lastMessage.createdAt).toLocaleString()})`,
+    )
+
+    // If any message already exists, stop the export
+    if (!result.success || result.duplicateCount > 0) {
+      logger.debug('检测到已存在的消息，导出完成')
+      return { shouldStop: true, failedCount: 0 }
+    }
+
+    return { shouldStop: false, failedCount: 0 }
+  }
+  catch (error) {
+    logger.withError(error).error(`保存批次消息失败 (${startIndex + 1} - ${startIndex + messages.length})`)
+    return { shouldStop: false, failedCount: messages.length }
+  }
+}
+
+/**
+ * Save messages to JSON file
+ */
+async function saveToJsonFile(messages: any[], chatId: number, exportPath: string): Promise<boolean> {
+  try {
+    const fileName = `${chatId}_${new Date().toISOString().split('T')[0]}`
+    const filePath = path.join(exportPath, `${fileName}.json`)
+    await fs.writeFile(filePath, JSON.stringify(messages, null, 2))
+    logger.debug(`已保存 JSON 文件: ${filePath}`)
+    logger.log(`已导出到文件: ${filePath}`)
+    return true
+  }
+  catch (error) {
+    logger.withError(error).error('保存 JSON 文件失败')
+    return false
+  }
+}
+
+/**
+ * Update chat metadata
+ */
+async function updateChatMetadata(chat: Dialog): Promise<void> {
+  const chatInput: NewChat = {
+    id: chat.id,
+    type: chat.type,
+    title: chat.name,
+    lastSyncTime: new Date(),
+  }
+  await updateChat(chatInput)
+}
+
 /**
  * Export command to export messages from Telegram
  */
@@ -55,16 +120,22 @@ export class ExportCommand extends TelegramCommand {
     const format = options.format || await input.select({
       message: '请选择导出格式：',
       choices: [
-        { name: 'HTML', value: 'html' },
-        { name: 'JSON', value: 'json' },
+        { name: 'Database', value: 'database' },
       ],
     })
 
-    // Get export path
-    const exportPath = options.path || await input.input({
-      message: '请输入导出路径：',
-      default: './export',
-    })
+    // Get export path if not exporting to database
+    let exportPath = options.path
+    if (format !== 'database') {
+      exportPath = exportPath || await input.input({
+        message: '请输入导出路径：',
+        default: './export',
+      })
+
+      // Create export directory
+      await fs.mkdir(exportPath, { recursive: true })
+      logger.debug(`已创建导出目录: ${exportPath}`)
+    }
 
     // Get message types
     const messageTypes = options.messageTypes || await input.checkbox({
@@ -109,21 +180,14 @@ export class ExportCommand extends TelegramCommand {
       default: '0',
     })
 
-    // Get batch size
-    const batchSize = options.batchSize || await input.input({
-      message: '每多少条消息提醒一次继续？',
-      default: '3000',
-    })
-
-    // Create export directory
-    await fs.mkdir(exportPath, { recursive: true })
-    logger.debug(`已创建导出目录: ${exportPath}`)
+    // Get batch size from config
+    const batchSize = options.batchSize || getConfig().messageBatchSize
 
     // Export messages
     logger.log('正在导出消息...')
     let count = 0
     let failedCount = 0
-    const skippedCount = 0
+    let shouldStop = false
     const messages = []
 
     // Export messages
@@ -138,17 +202,72 @@ export class ExportCommand extends TelegramCommand {
       count++
       if (count % Number(batchSize) === 0) {
         logger.debug(`已处理 ${count} 条消息`)
+
+        // Save current batch to database if format is database
+        if (format === 'database') {
+          const batch = messages.slice(count - Number(batchSize), count)
+          const messageInputs: MessageCreateInput[] = batch.map(message => ({
+            id: message.id,
+            chatId: message.chatId,
+            type: message.type,
+            content: message.content,
+            fromId: message.fromId,
+            replyToId: message.replyToId,
+            forwardFromChatId: message.forwardFromChatId,
+            forwardFromMessageId: message.forwardFromMessageId,
+            views: message.views,
+            forwards: message.forwards,
+            createdAt: message.createdAt,
+          }))
+
+          const result = await processDatabaseBatch(messageInputs, count - Number(batchSize))
+          shouldStop = result.shouldStop
+          failedCount += result.failedCount
+
+          if (shouldStop) {
+            break
+          }
+        }
       }
     }
 
-    // Save to file
-    const fileName = `${chatId}_${new Date().toISOString().split('T')[0]}`
-    const filePath = path.join(exportPath, `${fileName}.${format}`)
-
+    // Save to file or database
     try {
       if (format === 'json') {
-        await fs.writeFile(filePath, JSON.stringify(messages, null, 2))
-        logger.debug(`已保存 JSON 文件: ${filePath}`)
+        const success = await saveToJsonFile(messages, chatId, exportPath!)
+        if (!success) {
+          failedCount = count
+        }
+      }
+      else if (format === 'database' && !shouldStop) {
+        // Update chat metadata
+        await updateChatMetadata(selectedChat)
+
+        // Save remaining messages
+        const remainingCount = count % Number(batchSize)
+        if (remainingCount > 0) {
+          const batch = messages.slice(-remainingCount)
+          const messageInputs: MessageCreateInput[] = batch.map(message => ({
+            id: message.id,
+            chatId: message.chatId,
+            type: message.type,
+            content: message.content,
+            fromId: message.fromId,
+            replyToId: message.replyToId,
+            forwardFromChatId: message.forwardFromChatId,
+            forwardFromMessageId: message.forwardFromMessageId,
+            views: message.views,
+            forwards: message.forwards,
+            createdAt: message.createdAt,
+          }))
+
+          const result = await processDatabaseBatch(messageInputs, count - remainingCount)
+          failedCount += result.failedCount
+        }
+
+        if (failedCount === 0) {
+          logger.log('已导出到数据库')
+        }
       }
       else {
         // TODO: 实现 HTML 导出
@@ -158,27 +277,13 @@ export class ExportCommand extends TelegramCommand {
     }
     catch (error) {
       failedCount = count
-      logger.withError(error).error(`保存文件失败: ${filePath}`)
-    }
-
-    if (failedCount === 0) {
-      logger.log(`已导出到文件: ${filePath}`)
+      logger.withError(error).error(`保存失败`)
     }
 
     const summary = failedCount > 0
-      ? `导出完成，共导出 ${count} 条消息，${failedCount} 条消息失败，${skippedCount} 条消息已存在`
-      : `导出完成，共导出 ${count} 条消息，${skippedCount} 条消息已存在`
+      ? `导出完成，共导出 ${count} 条消息，${failedCount} 条消息失败`
+      : `导出完成，共导出 ${count} 条消息`
     logger.log(summary)
-
-    // Ask if continue
-    const shouldContinue = await input.confirm({
-      message: '是否继续导出？',
-      default: false,
-    })
-
-    if (shouldContinue) {
-      await this.execute(_args, options)
-    }
   }
 }