Added embedding of nodes with better handling of paths in Codebase

Author: preeesha

rocket-chatter-ingestion-server/src/core/llm.ts

@@ -3,7 +3,6 @@ import { OPENAI_KEY } from "../constants"
 
 const openai = new OpenAI({ apiKey: OPENAI_KEY })
 
-
 export namespace LLM {
 	export async function generateEmbeddings(data: string): Promise<number[]> {
 		try {

rocket-chatter-ingestion-server/src/process/prepare/codebase.ts

@@ -1,13 +1,26 @@
-import { existsSync, mkdirSync, rmSync, writeFileSync } from "fs"
+import {
+	existsSync,
+	mkdirSync,
+	readFileSync,
+	readdirSync,
+	rmSync,
+	writeFileSync,
+} from "fs"
 import { glob } from "glob"
+import path from "path"
+
 import { DBNode } from "../../core/dbNode"
 import { IFileProcessor } from "./processor/file.types"
 import { SourceFile } from "./sourceFile"
 import { ISourceFile } from "./sourceFile.types"
 
 export class Codebase {
 	private _path: string
-	private _dataPath: string
+	private _dataDirPath: string
+	private _dataDirName: string
+	private _embeddingsDirPath: string
+	private _embeddingsDirName: string
+
 	private _batchSize: number
 	private _fileProcessor: IFileProcessor
 
@@ -19,7 +32,11 @@ export class Codebase {
 		if (path.endsWith("/")) path = path.replace(/\/+$/, "")
 
 		this._path = path
-		this._dataPath = ""
+		this._dataDirName = ""
+		this._dataDirPath = ""
+		this._embeddingsDirName = ""
+		this._embeddingsDirPath = ""
+
 		this._batchSize = batchSize
 		this._fileProcessor = fileProcessor
 
@@ -28,32 +45,36 @@ export class Codebase {
 		this.makeFilesBatches()
 	}
 
-	private makePath(path: string): string {
-		return `${this._path}/${path}`
-	}
-
 	private initializeDataDirectory(removeExisting = true): void {
-		this._dataPath = "data"
+		this._dataDirName = `data-${Date.now()}`
+		this._dataDirPath = path.resolve(this._path, this._dataDirName)
+
+		this._embeddingsDirName = `${this._dataDirName}/embeddings`
+		this._embeddingsDirPath = path.resolve(this._path, this._embeddingsDirName)
+
+		/* Handle data directory */
+		if (removeExisting && existsSync(this._dataDirPath))
+			rmSync(this._dataDirPath, { recursive: true })
+		mkdirSync(this._dataDirPath)
 
-		const path = this.makePath(this._dataPath)
-		if (removeExisting && existsSync(path)) rmSync(path, { recursive: true })
-		mkdirSync(path)
+		/* Handle embeddings directory */
+		mkdirSync(this._embeddingsDirPath)
 	}
 
 	private prepareFilesMetadata() {
 		const extensions = ["ts", "tsx", "js", "jsx"]
 
-		// console.log(`🕒 Preparing metadata for files: *.${extensions.join(", *.")}`)
+		console.log(`🕒 Preparing metadata for files: *.${extensions.join(", *.")}`)
 		{
 			const globPatterns = extensions.map((x) => `**/*.${x}`)
 			for (const pattern of globPatterns) {
 				const files = glob
-					.sync(this.makePath(pattern))
+					.sync(`${this._path}/${pattern}`)
 					.map((x) => new SourceFile(x))
 				this._files.push(...files)
 			}
 		}
-		// console.log(`✅ Prepared metadata for ${this._files.length} files\n`)
+		console.log(`✅ Prepared metadata for ${this._files.length} files\n`)
 	}
 
 	private makeFilesBatches() {
@@ -74,7 +95,7 @@ export class Codebase {
 		if (entries.length === 0) return 0
 		const batch = Object.fromEntries(entries)
 		writeFileSync(
-			this.makePath(`${this._dataPath}/${fileName}`),
+			path.resolve(this._dataDirPath, fileName),
 			JSON.stringify(batch, null, 2)
 		)
 
@@ -88,7 +109,7 @@ export class Codebase {
 	): Promise<number> {
 		let nNodesProcessed = 0
 
-		// console.log(`🕒 Processing ${start}-${end} files`)
+		console.log(`🕒 Processing ${start}-${end} files`)
 		{
 			let nodes: Record<string, DBNode> = {}
 
@@ -114,9 +135,9 @@ export class Codebase {
 
 			nNodesProcessed = Object.keys(nodes).length
 		}
-		// console.log(
-		// 	`✅ Processed ${start}-${end} files (${nNodesProcessed} nodes)\n`
-		// )
+		console.log(
+			`✅ Processed ${start}-${end} files (${nNodesProcessed} nodes)\n`
+		)
 
 		return nNodesProcessed
 	}
@@ -133,14 +154,66 @@ export class Codebase {
 	 * @returns Promise<void>
 	 */
 	async process(): Promise<void> {
-		// console.log("🕒 Preparing Nodes\n")
+		console.log("🕒 Preparing Nodes\n")
 
 		let nodesProcessed = 0
 		for (const [index, batch] of this._batches.entries()) {
 			const [start, end] = batch
 			nodesProcessed += await this.processFilesBatch(index, start, end)
 		}
 
-		// console.log(`✅ Prepared ${nodesProcessed} nodes`)
+		console.log(`✅ Prepared ${nodesProcessed} nodes`)
+	}
+
+	async embed(): Promise<void> {
+		console.log("🕒 Preparing Embeddings")
+
+		if (existsSync(this._embeddingsDirPath))
+			rmSync(this._embeddingsDirPath, { recursive: true })
+		mkdirSync(this._embeddingsDirPath)
+
+		const files = readdirSync(this._dataDirPath)
+			.filter((x) => x.endsWith(".json"))
+			.map((x) => path.resolve(this._dataDirPath, x)) // convert path like "batch-1.json" to "./data/batch-1.json"
+
+		const embeddingsPerNode = 2
+		const maxAllowedEmbeddingsPerMinute = 2800 // openai limitation for embeddings
+		const nFilesPerBatch = Math.floor(
+			maxAllowedEmbeddingsPerMinute / this._batchSize / embeddingsPerNode
+		)
+
+		let batch = 0
+		for (let i = 0; i < files.length; i += nFilesPerBatch) {
+			const start = i
+			const end = Math.min(i + nFilesPerBatch, files.length)
+
+			console.log(`\n🕒 Embedding ${start}-${end} files`)
+
+			let nodes: Record<string, DBNode> = {}
+			for (const file of files.slice(start, end)) {
+				// to convert file content from a plain string to js object
+				const data = JSON.parse(readFileSync(file, "utf-8"))
+				Object.assign(nodes, data)
+			}
+
+			const jobs = Object.values(nodes).map(async (x) => {
+				nodes[x.id] = await DBNode.fillEmbeddings(new DBNode(x))
+			})
+			await Promise.all(jobs)
+
+			writeFileSync(
+				`${this._embeddingsDirPath}/batch-${batch++}.json`,
+				JSON.stringify(nodes, null, 2)
+			)
+
+			console.log(`✅ Embedded ${start}-${end} files\n`)
+
+			if (i + nFilesPerBatch < files.length) {
+				console.log(`🕒 Waiting for 60 seconds`)
+				await new Promise((resolve) => setTimeout(resolve, 60 * 1000))
+			}
+		}
+
+		console.log(`✅ Prepared embeddings for nodes`)
 	}
 }

rocket-chatter-ingestion-server/src/use.ts

@@ -3,7 +3,8 @@ import { FileProcessor } from "./process/prepare/processor/file"
 
 async function main() {
 	const codebase = new Codebase("./project", new FileProcessor())
-	codebase.process()
+	await codebase.process()
+	await codebase.embed()
 }
 
 main()