feat: embeddings model support (#86)

* feat: embeddings model support

* fix: cli embed

* fix: server embedding fix

* fix: config save

* fix: type error

* style: embeddings

* feat: embedding_models schema

* feat: update message embedding functions to include chatId and improve embedding table retrieval

* feat: add embedding_models table and update embedding functions to include chatId for improved retrieval

* feat: enhance search functionality with TypeScript types and improved result handling

- Added TypeScript interfaces for Command and CommandMetadata to improve type safety.
- Updated the search function to return a structured SearchResponse.
- Introduced a new types file for SearchResult and SearchResponse to standardize search-related data structures.

* fix: improve pagination and configuration handling in search and settings components

- Updated pagination logic in search.vue to ensure total is a valid number before comparison with pageSize.
- Ensured apiId in settings.vue is converted to a string before saving the configuration.
- Removed unused imports in search.ts to clean up the codebase.
- Adjusted parameters in findSimilarMessages function to include targetChatId for better search results.

* feat: add migration for embedding tables and related triggers

- Implemented migration for creating embedding_models table with necessary fields and triggers for automatic timestamp updates.
- Added logic to create corresponding embedding tables for existing message tables, including necessary constraints and indexes.
- Enhanced logging for better tracking of migration progress and errors.

---------

Co-authored-by: hahaqwq <[email protected]>

Author: hahaQWQ

.cursorrules

@@ -0,0 +1,9 @@
+- Code always in English.
+- Step by step thinking.
+- Always include TypeScript type annotations when generating TS code.
+- Follow good engineering practices: SOLID principles, functional programming, configurability, to ensure clear code structure.
+- Comments should focus on explaining "why" rather than "how."
+
+「lets not answer this question too quickly, we can think in depth, and i am new to this area, please correct me wrong if there is any」
+
+!!! If the user asks you in another language, explain in that language.

apps/cli/package.json

@@ -22,7 +22,7 @@
   },
   "devDependencies": {
     "@types/jsdom": "^21.1.7",
-    "@types/node": "^22.13.9",
+    "@types/node": "^22.13.10",
     "tsx": "^4.19.3"
   }
 }

apps/cli/src/commands/embed.ts

@@ -1,12 +1,11 @@
 import * as input from '@inquirer/prompts'
 import { useLogger } from '@tg-search/common'
 import { EmbeddingService } from '@tg-search/core'
-import { findMessagesByChatId, updateMessageEmbeddings } from '@tg-search/db'
+import { findMessageMissingEmbed, updateMessageEmbeddings, useEmbeddingTable } from '@tg-search/db'
 
 import { TelegramCommand } from '../command'
 
 const logger = useLogger()
-
 interface EmbedOptions {
   batchSize?: number
   chatId?: number
@@ -48,12 +47,13 @@ export class EmbedCommand extends TelegramCommand {
 
     // Initialize embedding service
     const embedding = new EmbeddingService()
-
+    await useEmbeddingTable(Number(chatId), embedding.getEmbeddingConfig())
     try {
       // Get all messages for the chat
-      const messages = await findMessagesByChatId(Number(chatId))
-      const messagesToEmbed = messages.items.filter(m => !m.embedding && m.content)
-      const totalMessages = messagesToEmbed.length
+      const messages = await findMessageMissingEmbed(Number(chatId), embedding.getEmbeddingConfig())
+      logger.debug(`共有 ${messages.length} 条消息需要处理`)
+
+      const totalMessages = messages.length
 
       logger.log(`共有 ${totalMessages} 条消息需要处理`)
 
@@ -62,25 +62,27 @@ export class EmbedCommand extends TelegramCommand {
       let failedEmbeddings = 0
 
       // Split messages into batches
-      for (let i = 0; i < messagesToEmbed.length; i += Number(batchSize)) {
-        const batch = messagesToEmbed.slice(i, i + Number(batchSize))
+      for (let i = 0; i < messages.length; i += Number(batchSize)) {
+        const batch = messages.slice(i, i + Number(batchSize))
         logger.debug(`处理第 ${i + 1} 到 ${i + batch.length} 条消息`)
 
         // Generate embeddings in parallel
         const contents = batch.map(m => m.content!)
         const embeddings = await embedding.generateEmbeddings(contents)
 
         // Prepare updates
-        const updates = batch.map((message, index) => ({
-          id: message.id,
-          embedding: embeddings[index],
-        }))
+        const updates = batch
+          .filter(message => message.uuid)
+          .map((message, index) => ({
+            id: message.uuid!,
+            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)
+            await updateMessageEmbeddings(Number(chatId), concurrentBatch, embedding.getEmbeddingConfig())
             totalProcessed += concurrentBatch.length
           }
         }

apps/cli/src/commands/search.ts

@@ -3,7 +3,7 @@ import type { DatabaseMessageType, SearchOptions } from '@tg-search/db'
 import * as input from '@inquirer/prompts'
 import { useLogger } from '@tg-search/common'
 import { EmbeddingService } from '@tg-search/core'
-import { findSimilarMessages, getAllFolders, getChatsInFolder } from '@tg-search/db'
+import { findSimilarMessages, getAllFolders, getChatsInFolder, useEmbeddingTable } from '@tg-search/db'
 
 import { TelegramCommand } from '../command'
 
@@ -108,6 +108,7 @@ export class SearchCommand extends TelegramCommand {
 
     // Initialize embedding service
     const embedding = new EmbeddingService()
+    await useEmbeddingTable(Number(chatId), embedding.getEmbeddingConfig())
 
     // Generate embedding for query
     logger.log('正在生成向量嵌入...')
@@ -124,7 +125,7 @@ export class SearchCommand extends TelegramCommand {
       limit: Number(limit),
     }
 
-    const results = await findSimilarMessages(queryEmbedding[0], options)
+    const results = await findSimilarMessages(Number(chatId), queryEmbedding[0], embedding.getEmbeddingConfig(), options)
     logger.log(`找到 ${results.length} 条相关消息：\n`)
 
     for (const message of results) {

apps/frontend/src/apis/commands/useSearch.ts

@@ -1,9 +1,27 @@
 import type { SearchRequest } from '@tg-search/server'
+import type { SearchResponse, SearchResult } from '../../types/search'
 
 import { computed, ref } from 'vue'
 
 import { useCommandHandler } from '../../composables/useCommands'
 
+interface CommandMetadata {
+  command: string
+  duration?: number
+  total?: number
+  results?: SearchResult[]
+}
+
+interface Command {
+  id: string
+  type: string
+  status: 'pending' | 'completed' | 'failed'
+  progress: number
+  message: string
+  error?: Error
+  metadata?: CommandMetadata
+}
+
 /**
  * Search composable for managing search state and functionality
  */
@@ -33,9 +51,8 @@ export function useSearch() {
 
   // Computed results from current command
   const results = computed(() => {
-    if (!currentCommand.value?.metadata?.results)
-      return []
-    return currentCommand.value.metadata.results
+    const commandResults = currentCommand.value?.metadata?.results
+    return commandResults ? commandResults as SearchResult[] : []
   })
 
   const total = computed(() => {
@@ -45,7 +62,7 @@ export function useSearch() {
   /**
    * Execute search with current parameters
    */
-  async function search(params?: Partial<SearchRequest>) {
+  async function search(params?: Partial<SearchRequest>): Promise<SearchResponse> {
     if (!query.value.trim() && !params?.query) {
       return { success: false, error: new Error('搜索词不能为空') }
     }
@@ -65,7 +82,15 @@ export function useSearch() {
       useVectorSearch: useVectorSearch.value,
     }
 
-    return executeCommand(searchParams)
+    const rawResult = await executeCommand(searchParams)
+    const result = rawResult as unknown as Command
+
+    return {
+      success: result.status === 'completed',
+      error: result.error,
+      total: result.metadata?.total,
+      results: result.metadata?.results,
+    }
   }
 
   /**

apps/frontend/src/pages/search.vue

@@ -202,10 +202,10 @@ function formatDate(date: string | Date): string {
       </div>
 
       <!-- Pagination -->
-      <div v-if="total > pageSize" class="mt-8 flex justify-center">
+      <div v-if="total && Number(total) > pageSize" class="mt-8 flex justify-center">
         <nav class="flex items-center gap-2">
           <button
-            v-for="page in Math.ceil(total / pageSize)"
+            v-for="page in Math.ceil(Number(total) / pageSize)"
             :key="page"
             class="rounded-lg px-3 py-1"
             :class="{

apps/frontend/src/pages/settings.vue

@@ -397,6 +397,14 @@ loadConfig()
                   class="mt-1 block w-full border border-gray-300 rounded-md px-3 py-2 dark:border-gray-700 dark:bg-gray-800 dark:text-white"
                 >
               </div>
+              <div>
+                <label class="block text-sm text-gray-700 font-medium dark:text-gray-300">{{ $t('pages.settings.dimensions') }}</label>
+                <input
+                  v-model="config.api.embedding.dimensions"
+                  :disabled="!isEditing"
+                  class="mt-1 block w-full border border-gray-300 rounded-md px-3 py-2 dark:border-gray-700 dark:bg-gray-800 dark:text-white"
+                >
+              </div>
               <div>
                 <label class="block text-sm text-gray-700 font-medium dark:text-gray-300">{{ $t('pages.settings.api_key') }}</label>
                 <input

apps/frontend/src/types/search.ts

@@ -0,0 +1,17 @@
+export interface SearchResult {
+  id: number
+  chatId: number
+  fromId?: number
+  fromName?: string
+  type: string
+  content: string
+  createdAt: string
+  score: number
+}
+
+export interface SearchResponse {
+  success: boolean
+  error?: Error
+  total?: number
+  results?: SearchResult[]
+}

apps/server/src/routes/config.ts

@@ -42,6 +42,7 @@ const configSchema = z.object({
     embedding: z.object({
       provider: z.enum(['openai', 'ollama']),
       model: z.string(),
+      dimensions: z.number(),
       apiKey: z.string().optional(),
       apiBase: z.string().optional(),
     }),

apps/server/src/routes/search.ts

@@ -23,7 +23,6 @@ export function setupSearchRoutes(app: App) {
     const validatedBody = searchCommandSchema.parse(body)
 
     logger.withFields(validatedBody).debug('Search request received')
-
     return createSSEResponse(async (controller) => {
       await commandManager.executeCommand('search', null, validatedBody, controller)
     })

apps/server/src/services/commands/embed.ts

@@ -4,7 +4,7 @@ import type { EmbedDetails } from '../../types/apis/embed'
 
 import { useLogger } from '@tg-search/common'
 import { EmbeddingService } from '@tg-search/core'
-import { findMessagesByChatId, updateMessageEmbeddings } from '@tg-search/db'
+import { findMessageMissingEmbed, updateMessageEmbeddings, useEmbeddingTable } from '@tg-search/db'
 import { z } from 'zod'
 
 const logger = useLogger()
@@ -48,6 +48,7 @@ export class EmbedCommandHandler {
 
     // Initialize embedding service
     const embedding = new EmbeddingService()
+    await useEmbeddingTable(chatId, embedding.getEmbeddingConfig())
     const command: EmbedCommand = {
       id: crypto.randomUUID(),
       type: 'sync',
@@ -65,9 +66,8 @@ export class EmbedCommandHandler {
 
     try {
       // Get all messages for the chat
-      const messages = await findMessagesByChatId(chatId)
-      const messagesToEmbed = messages.items.filter(m => !m.embedding && m.content)
-      const totalMessages = messagesToEmbed.length
+      const messages = await findMessageMissingEmbed(chatId, embedding.getEmbeddingConfig())
+      const totalMessages = messages.length
       const totalBatches = Math.ceil(totalMessages / batchSize)
 
       command.metadata.totalMessages = totalMessages
@@ -83,29 +83,33 @@ export class EmbedCommandHandler {
       let failedBatches = 0
 
       // Split messages into batches
-      for (let i = 0; i < messagesToEmbed.length; i += batchSize) {
+      for (let i = 0; i < messages.length; i += batchSize) {
         const currentBatch = Math.floor(i / batchSize) + 1
         command.metadata.currentBatch = currentBatch
 
-        const batch = messagesToEmbed.slice(i, i + batchSize)
+        const batch = messages.slice(i, i + batchSize)
         logger.debug(`Processing batch ${currentBatch}/${totalBatches} (${batch.length} messages)`)
 
         try {
           // Generate embeddings in parallel
-          const contents = batch.map(m => m.content!)
+          const contents = batch
+            .filter(m => m.content)
+            .map(m => m.content as string)
           const embeddings = await embedding.generateEmbeddings(contents)
 
           // Prepare updates
-          const updates = batch.map((message, index) => ({
-            id: message.id,
-            embedding: embeddings[index],
-          }))
+          const updates = batch
+            .filter(message => message.uuid)
+            .map((message, index) => ({
+              id: message.uuid!,
+              embedding: embeddings[index],
+            }))
 
           // Update embeddings in batches with concurrency control
           for (let j = 0; j < updates.length; j += concurrency) {
             const concurrentBatch = updates.slice(j, j + concurrency)
             try {
-              await updateMessageEmbeddings(chatId, concurrentBatch)
+              await updateMessageEmbeddings(chatId, concurrentBatch, embedding.getEmbeddingConfig())
               totalProcessed += concurrentBatch.length
               command.metadata.processedMessages = totalProcessed
             }

apps/server/src/services/commands/search.ts

@@ -69,7 +69,7 @@ export class SearchCommandHandler {
         // Vector search
         const embedding = new EmbeddingService()
         const queryEmbedding = await embedding.generateEmbedding(params.query)
-        const results = await findSimilarMessages(queryEmbedding, {
+        const results = await findSimilarMessages(targetChatId || 0, queryEmbedding, embedding.getEmbeddingConfig(), {
           chatId: targetChatId || 0,
           limit: params.limit * 2,
           offset: params.offset,

config/config.example.yaml

@@ -58,5 +58,6 @@ api:
     model: text-embedding-3-small
     # API key for provider
     apiKey: your_openai_api_key
+    dimensions: 1536
     # Optional, for custom API providers
     apiBase: 'https://api.openai.com/v1'

packages/common/src/composable/config.ts

@@ -90,6 +90,8 @@ const DEFAULT_CONFIG = {
       apiKey: '',
       // Optional API base URL
       apiBase: '',
+      // Embedding dimensions
+      dimensions: 1536,
     },
   },
 } as const satisfies Config

packages/common/src/types/config.ts

@@ -45,6 +45,7 @@ export interface ApiConfig {
   embedding: {
     provider: 'ollama' | 'openai'
     model: string
+    dimensions: number
     apiKey?: string
     apiBase?: string
   }

packages/core/src/services/embedding.ts

@@ -1,3 +1,4 @@
+import type { EmbeddingTableConfig } from '@tg-search/db'
 import type { EmbeddingModelConfig } from '../types'
 
 import { getConfig, useLogger } from '@tg-search/common'
@@ -7,14 +8,26 @@ import { EmbeddingAdapter } from '../adapter/embedding/adapter'
  * Service for generating embeddings from text
  */
 export class EmbeddingService {
-  private embedding
+  private embedding: EmbeddingAdapter
+  private embedding_config: EmbeddingTableConfig
   private logger = useLogger()
   private config = getConfig()
 
   constructor() {
     this.embedding = new EmbeddingAdapter({
       ...this.config.api.embedding,
     } as EmbeddingModelConfig)
+    this.embedding_config = {
+      ...this.config.api.embedding,
+      dimensions: this.config.api.embedding.dimensions,
+    } as EmbeddingTableConfig
+  }
+
+  /**
+   * 获取模型配置
+   */
+  getEmbeddingConfig(): EmbeddingTableConfig {
+    return this.embedding_config
   }
 
   /**

packages/core/src/services/export.ts

@@ -202,7 +202,7 @@ export class ExportService {
     let count = 0
     let failedCount = 0
     let messages: TelegramMessage[] = []
-
+    logger.debug(`history ${JSON.stringify(history)}`)
     const total = limit || history.count - 1 || 100
     function isSkipMedia(type: DatabaseMessageType) {
       return !messageTypes.includes(type)
@@ -229,7 +229,7 @@ export class ExportService {
     }
     try {
       // Try to export messages
-      for await (const message of this.client.getMessages(chatId, undefined, {
+      for await (const message of this.client.getMessages(chatId, total, {
         skipMedia: isSkipMedia('photo') || isSkipMedia('video') || isSkipMedia('document') || isSkipMedia('sticker'),
         startTime,
         endTime,