@lorrylurui/code-intelligence-mcp 1.2.0 → 2.0.1

package/dist/workers/embeddingWorker.js

@@ -14,87 +14,132 @@
  * 大仓分片：
  * - 直接启动多个 worker 进程（同一 Redis）即可水平扩展，BullMQ 原生分布式协调
  */
-import { Worker } from 'bullmq';
-import Redis from 'ioredis';
+import { Worker, QueueEvents } from 'bullmq';
+import { Redis } from 'ioredis';
 import { env } from '../config/env.js';
-import { getMySqlPool } from '../db/mysql.js';
-import { createEmbeddingClient } from './embeddingClient.js';
+import { getPool } from '../db/postgres.js';
+import { createEmbeddingClient } from '../services/embeddingClient.js';
 import { indexedRowToEmbedText } from '../indexer/embedText.js';
+import { initCategoryEmbeddings, resolveCategory, } from '../indexer/categoryClassifier.js';
 import { SYMBOL_STATUS } from '../config/symbolStatus.js';
 async function processEmbedJob(job, pool) {
     const { semanticHash } = job.data;
-    const table = env.mysqlSymbolsTable;
+    // 优先使用 job payload 里的表名（跨项目场景），降级到 env（单项目场景）
+    const table = job.data.symbolsTable ?? env.symbolsTable;
+    const shortHash = semanticHash.slice(0, 10);
     const embedClient = createEmbeddingClient(env.embeddingServiceUrl);
+    const ts = () => new Date().toISOString();
     // Step 1: 缓存命中检查 —— 相同 semantic_hash 已有 online 向量
-    const [cached] = await pool.query(`SELECT embedding FROM ${table}
-         WHERE semantic_hash = ? AND status = ? AND embedding IS NOT NULL
+    const { rows: cached } = await pool.query(`SELECT embedding FROM ${table}
+         WHERE semantic_hash = $1 AND status = $2 AND embedding IS NOT NULL
          LIMIT 1`, [semanticHash, SYMBOL_STATUS.ONLINE]);
     let vector;
     if (cached.length > 0) {
         // Cache hit: 直接复用已有向量，0 次 API 调用
+        // pgvector 返回字符串 "[x1,x2,...]", JSON.parse 可直接解析
         vector =
             typeof cached[0].embedding === 'string'
                 ? JSON.parse(cached[0].embedding)
                 : cached[0].embedding;
-        console.error(`[worker] cache hit  hash=${semanticHash.slice(0, 10)}…`);
+        // cache hit 时只需把 pending 行的向量补齐（有可能是新增的同语义符号）
+        const cacheResult = await pool.query(`UPDATE ${table}
+             SET embedding = $1::vector, status = $2
+             WHERE semantic_hash = $3 AND status = $4`, [
+            `[${vector.join(',')}]`,
+            SYMBOL_STATUS.ONLINE,
+            semanticHash,
+            SYMBOL_STATUS.PENDING,
+        ]);
+        console.error(`[worker] ✅ cache hit   [${ts()}]  table=${table}  hash=${shortHash}…  updated ${cacheResult.rowCount ?? 0} row(s)  (0 API calls)`);
+        return { updatedRows: cacheResult.rowCount ?? 0 };
     }
-    else {
-        // Cache miss: 取一条 pending 行做 embedding
-        const [pending] = await pool.query(`SELECT name, type, category, path, description, content, meta
-             FROM ${table}
-             WHERE semantic_hash = ? AND status = ?
-             LIMIT 1`, [semanticHash, SYMBOL_STATUS.PENDING]);
-        if (pending.length === 0) {
-            // 所有行已被并发 worker 处理，幂等退出
-            return;
-        }
-        const row = pending[0];
-        const meta = typeof row.meta === 'string'
-            ? JSON.parse(row.meta)
-            : (row.meta ?? {});
-        const doc = indexedRowToEmbedText({ ...row, meta });
-        const vectors = await embedClient.embed([doc]);
-        vector = vectors[0];
-        console.error(`[worker] embedded   hash=${semanticHash.slice(0, 10)}… path=${row.path}:${row.name}`);
+    // Cache miss: 取一条 pending 行做 embedding
+    const { rows: pending } = await pool.query(`SELECT name, type, category, path, description, content, meta
+         FROM ${table}
+         WHERE semantic_hash = $1 AND status = $2
+         LIMIT 1`, [semanticHash, SYMBOL_STATUS.PENDING]);
+    if (pending.length === 0) {
+        // 所有行已被并发 worker 处理，幂等退出
+        console.error(`[worker] ⚠️  skip       [${ts()}]  table=${table}  hash=${shortHash}…  (no pending rows)`);
+        return { updatedRows: 0 };
     }
-    // Step 2: 批量写入 —— 批量更新所有拥有相同 semantic_hash 的 pending 行
-    // 一次 API 调用覆盖 N 个同义符号（大仓重复代码/多文件同函数场景收益明显）
-    await pool.query(`UPDATE ${table}
-         SET embedding = CAST(? AS JSON), status = ?
-         WHERE semantic_hash = ? AND status = ?`, [
-        JSON.stringify(vector),
+    const row = pending[0];
+    const meta = typeof row.meta === 'string' ? JSON.parse(row.meta) : (row.meta ?? {});
+    const rowObj = { ...row, meta };
+    console.error(`[worker] 🔄 embedding  [${ts()}]  table=${table}  hash=${shortHash}…  ${row.path}:${row.name}`);
+    // 与 reindex 保持一致：优先用 content（语义模板），降级用 indexedRowToEmbedText
+    const doc = row.content ?? indexedRowToEmbedText(rowObj);
+    const vectors = await embedClient.embed([doc]);
+    vector = vectors[0];
+    // 生成 category（规则 → embedding → LLM 三层融合）
+    const [resolvedRow] = await resolveCategory([rowObj], [vector]);
+    const resolvedCategory = resolvedRow.category ?? null;
+    // Step 2: 批量写入 —— 覆盖所有相同 semantic_hash 的 pending 行
+    const result = await pool.query(`UPDATE ${table}
+         SET embedding = $1::vector, status = $2, category = COALESCE($3, category)
+         WHERE semantic_hash = $4 AND status = $5`, [
+        `[${vector.join(',')}]`,
         SYMBOL_STATUS.ONLINE,
+        resolvedCategory,
         semanticHash,
         SYMBOL_STATUS.PENDING,
     ]);
+    console.error(`[worker] ✓  done        [${ts()}]  table=${table}  hash=${shortHash}…  category=${resolvedCategory ?? 'null'}  updated ${result.rowCount ?? 0} row(s)`);
+    return { updatedRows: result.rowCount ?? 0 };
 }
 /**
- * 启动 embedding worker，返回 Worker 实例（可用于优雅关闭）。
+ * 启动 embedding worker，返回包含 stop() 的句柄。
  */
-export function startEmbeddingWorker(opts = {}) {
+export async function startEmbeddingWorker(opts = {}) {
     const { concurrency = 5, rpmLimit = 100 } = opts;
     const connection = new Redis(env.redisUrl, {
         maxRetriesPerRequest: null,
         enableReadyCheck: false,
     });
-    const pool = getMySqlPool();
-    if (!pool) {
-        throw new Error('[embeddingWorker] MySQL pool unavailable — check env vars');
+    // 独立连接监听队列事件（BullMQ 要求不共用 Worker 连接）
+    const eventsConnection = new Redis(env.redisUrl, {
+        maxRetriesPerRequest: null,
+        enableReadyCheck: false,
+    });
+    const queueEvents = new QueueEvents('embedding', {
+        connection: eventsConnection,
+    });
+    const pool = getPool();
+    // 预热 category embeddings（仅在服务启动时调用一次）
+    if (env.embeddingServiceUrl) {
+        await initCategoryEmbeddings();
+        console.error('[embedding-worker] category embeddings initialized');
     }
     const worker = new Worker('embedding', (job) => processEmbedJob(job, pool), {
         connection,
         concurrency,
         // 全局限流：所有 worker 进程共享，防止触发 OpenAI rate limit
         limiter: { max: rpmLimit, duration: 60_000 },
+        // 完成后立即从 Redis 清除，避免 jobId 残留导致下次同 hash 无法入队
+        removeOnComplete: { count: 0 },
+        removeOnFail: { count: 100 },
     });
-    worker.on('completed', (job) => {
-        console.error(`[worker] ✓ job done  hash=${job.data.semanticHash.slice(0, 10)}…`);
+    // 累计统计：每次 drained 后重置
+    const stats = { completed: 0, updatedRows: 0 };
+    worker.on('completed', (_job, result) => {
+        stats.completed++;
+        stats.updatedRows += result?.updatedRows ?? 0;
     });
     worker.on('failed', (job, err) => {
-        console.error(`[worker] ✗ job fail  hash=${job?.data?.semanticHash?.slice(0, 10)}… err=${err.message}`);
+        console.error(`[worker] ✗  failed     [${new Date().toISOString()}]  table=${job?.data?.symbolsTable ?? env.symbolsTable}  hash=${job?.data?.semanticHash?.slice(0, 10)}…  err=${err.message}`);
     });
     worker.on('error', (err) => {
         console.error(`[worker] error: ${err.message}`);
     });
-    return worker;
+    // 队列清空时打汇总（全量 reindex 入队后监听，确认所有 embedding 已处理）
+    queueEvents.on('drained', () => {
+        console.error(`[worker] ✅ queue drained  [${new Date().toISOString()}]  completed=${stats.completed} jobs  rows_updated=${stats.updatedRows}`);
+        stats.completed = 0;
+        stats.updatedRows = 0;
+    });
+    const stop = async () => {
+        await worker.close();
+        await queueEvents.close();
+    };
+    return { worker, stop };
 }

package/package.json

@@ -1,6 +1,6 @@
 {
     "name": "@lorrylurui/code-intelligence-mcp",
-    "version": "1.2.0",
+    "version": "2.0.1",
     "private": false,
     "description": "MCP server 提供仓库内可复用代码块（ts/tsx/js/jsx/css/less）的索引和查询能力，支持基于代码上下文的智能推荐。",
     "type": "module",
@@ -23,7 +23,7 @@
         "embedding:dev": "cd embedding-service && python3 -m uvicorn app:app --host 127.0.0.1 --port 8765",
         "docker:up": "docker compose up -d",
         "docker:down": "docker compose down",
-        "docker:logs": "docker compose logs -f mysql"
+        "docker:logs": "docker compose logs -f postgres"
     },
     "dependencies": {
         "@babel/parser": "^7.29.2",
@@ -35,12 +35,14 @@
         "fast-glob": "^3.3.2",
         "ioredis": "^5.10.1",
         "mysql2": "^3.11.3",
+        "pg": "^8.20.0",
         "react": "^19.2.4",
         "ts-morph": "^25.0.0",
         "zod": "^3.23.8"
     },
     "devDependencies": {
         "@types/node": "^22.10.1",
+        "@types/pg": "^8.20.0",
         "tsx": "^4.19.2",
         "typescript": "^5.6.3"
     },