146 lines
5.0 KiB
TypeScript
146 lines
5.0 KiB
TypeScript
#!/usr/bin/env tsx
|
|
/**
|
|
* Re-embed all KnowledgeChunk rows with OpenAI (text-embedding-3-small @ 1024 dims)
|
|
* and tag metadata with provider openai (required for KB vector search filters).
|
|
*
|
|
* Usage: npm run re-embed-knowledge-base
|
|
*/
|
|
|
|
import { prisma } from "../src/lib/prisma"
|
|
import { embedBatch } from "../src/lib/kb/openaiEmbeddings"
|
|
import {
|
|
kbEmbeddingMetadata,
|
|
KB_EMBEDDING_DIMENSIONS,
|
|
} from "../src/lib/kb/embeddingMeta"
|
|
import { logger } from "../src/lib/logger"
|
|
|
|
const BATCH_SIZE = 32
|
|
|
|
interface KnowledgeChunkRow {
|
|
id: string
|
|
content: string
|
|
storeId: string
|
|
metadata: unknown
|
|
}
|
|
|
|
async function reEmbedKnowledgeBase() {
|
|
try {
|
|
logger.info("Starting knowledge base re-embedding (OpenAI)...")
|
|
|
|
const sql = `SELECT "id", "content", "storeId", "metadata"
|
|
FROM "KnowledgeChunk"
|
|
WHERE "content" IS NOT NULL AND "content" != ''
|
|
ORDER BY "createdAt" ASC`
|
|
const chunks = await prisma.$queryRawUnsafe<KnowledgeChunkRow[]>(sql)
|
|
|
|
if (chunks.length === 0) {
|
|
logger.info("No knowledge chunks found to re-embed.")
|
|
return
|
|
}
|
|
|
|
logger.info(`Found ${chunks.length} knowledge chunks to re-embed.`)
|
|
|
|
let processed = 0
|
|
let failed = 0
|
|
const totalBatches = Math.ceil(chunks.length / BATCH_SIZE)
|
|
|
|
for (let i = 0; i < chunks.length; i += BATCH_SIZE) {
|
|
const batch = chunks.slice(i, i + BATCH_SIZE)
|
|
const batchNumber = Math.floor(i / BATCH_SIZE) + 1
|
|
|
|
logger.info(
|
|
`Processing batch ${batchNumber}/${totalBatches} (${batch.length} chunks)...`,
|
|
)
|
|
|
|
try {
|
|
const texts = batch.map((c) => c.content)
|
|
const embeddings = await embedBatch(texts)
|
|
|
|
if (embeddings.length !== batch.length) {
|
|
logger.error(
|
|
`Embedding count mismatch: expected ${batch.length}, got ${embeddings.length}`,
|
|
)
|
|
failed += batch.length
|
|
continue
|
|
}
|
|
|
|
await Promise.all(
|
|
batch.map(async (chunk, index) => {
|
|
const embedding = embeddings[index]
|
|
|
|
if (!embedding || embedding.length === 0) {
|
|
logger.warn(`Skipping chunk ${chunk.id}: empty embedding`)
|
|
failed++
|
|
return
|
|
}
|
|
|
|
if (embedding.length !== KB_EMBEDDING_DIMENSIONS) {
|
|
logger.error(
|
|
`Chunk ${chunk.id}: expected ${KB_EMBEDDING_DIMENSIONS} dims, got ${embedding.length}`,
|
|
)
|
|
failed++
|
|
return
|
|
}
|
|
|
|
const prevMeta =
|
|
chunk.metadata &&
|
|
typeof chunk.metadata === "object" &&
|
|
!Array.isArray(chunk.metadata)
|
|
? (chunk.metadata as Record<string, unknown>)
|
|
: {}
|
|
const nextMeta = {
|
|
...prevMeta,
|
|
...kbEmbeddingMetadata(),
|
|
}
|
|
|
|
try {
|
|
await prisma.$executeRawUnsafe(
|
|
`
|
|
UPDATE "KnowledgeChunk"
|
|
SET "embedding" = $1::vector,
|
|
"metadata" = $2::jsonb,
|
|
"updatedAt" = NOW()
|
|
WHERE "id" = $3
|
|
`,
|
|
`[${embedding.join(",")}]`,
|
|
JSON.stringify(nextMeta),
|
|
chunk.id,
|
|
)
|
|
processed++
|
|
} catch (error) {
|
|
logger.error(`Failed to update chunk ${chunk.id}:`, error)
|
|
failed++
|
|
}
|
|
}),
|
|
)
|
|
} catch (error) {
|
|
logger.error(`Error processing batch ${batchNumber}:`, error)
|
|
failed += batch.length
|
|
}
|
|
}
|
|
|
|
logger.info(
|
|
`Re-embedding completed. Processed: ${processed}, Failed: ${failed}, Total: ${chunks.length}`,
|
|
)
|
|
} catch (error) {
|
|
logger.error("Fatal error during re-embedding:", error)
|
|
process.exit(1)
|
|
} finally {
|
|
await prisma.$disconnect()
|
|
}
|
|
}
|
|
|
|
if (require.main === module) {
|
|
reEmbedKnowledgeBase()
|
|
.then(() => {
|
|
logger.info("Re-embedding script finished successfully.")
|
|
process.exit(0)
|
|
})
|
|
.catch((error) => {
|
|
logger.error("Re-embedding script failed:", error)
|
|
process.exit(1)
|
|
})
|
|
}
|
|
|
|
export { reEmbedKnowledgeBase }
|