Merge pull request #1923 from giselles-ai/1827-persist-dvs-id-embeddings

Link document embeddings to their parent vector stores

Author: shige

apps/studio.giselles.ai/drizzle/schema.ts

@@ -442,6 +442,11 @@ export const documentEmbeddings = pgTable(
 	"document_embeddings",
 	{
 		dbId: serial("db_id").primaryKey(),
+		documentVectorStoreDbId: integer("document_vector_store_db_id")
+			.notNull()
+			.references(() => documentVectorStores.dbId, {
+				onDelete: "cascade",
+			}),
 		documentVectorStoreSourceDbId: integer("document_vector_store_source_db_id")
 			.notNull()
 			.references(() => documentVectorStoreSources.dbId, {
@@ -479,6 +484,7 @@ export const documentEmbeddings = pgTable(
 				sql`(${table.embedding}::halfvec(3072)) halfvec_cosine_ops`,
 			)
 			.where(sql`${table.embeddingDimensions} = 3072`),
+		index("doc_embs_store_idx").on(table.documentVectorStoreDbId),
 	],
 );
 
@@ -489,6 +495,10 @@ export const documentEmbeddingsRelations = relations(
 			fields: [documentEmbeddings.documentVectorStoreSourceDbId],
 			references: [documentVectorStoreSources.dbId],
 		}),
+		documentVectorStore: one(documentVectorStores, {
+			fields: [documentEmbeddings.documentVectorStoreDbId],
+			references: [documentVectorStores.dbId],
+		}),
 	}),
 );
 

apps/studio.giselles.ai/lib/vector-stores/document/database.ts

@@ -70,6 +70,7 @@ export async function updateDocumentVectorStoreSourceStatus(
 }
 
 interface InsertEmbeddingsParams {
+	storeDbId: number;
 	sourceDbId: number;
 	embeddingProfileId: EmbeddingProfileId;
 	dimensions: EmbeddingDimensions;
@@ -85,6 +86,7 @@ export async function insertDocumentEmbeddings(
 	params: InsertEmbeddingsParams,
 ): Promise<void> {
 	const {
+		storeDbId,
 		sourceDbId,
 		embeddingProfileId,
 		dimensions,
@@ -98,6 +100,7 @@ export async function insertDocumentEmbeddings(
 
 	// Convert embeddings to the format expected by the database
 	const values = embeddings.map((emb) => ({
+		documentVectorStoreDbId: storeDbId,
 		documentVectorStoreSourceDbId: sourceDbId,
 		embeddingProfileId,
 		embeddingDimensions: dimensions,

apps/studio.giselles.ai/lib/vector-stores/document/ingest/ingest-document.ts

@@ -211,6 +211,7 @@ export async function ingestDocument(
 
 				// Insert embeddings into database
 				await insertDocumentEmbeddings({
+					storeDbId: source.documentVectorStoreDbId,
 					sourceDbId: source.dbId,
 					embeddingProfileId,
 					dimensions: embeddingResult.dimensions as EmbeddingDimensions,

apps/studio.giselles.ai/migrations/0051_glossy_nocturne.sql

@@ -0,0 +1,3 @@
+ALTER TABLE "document_embeddings" ADD COLUMN "document_vector_store_db_id" integer NOT NULL;--> statement-breakpoint
+ALTER TABLE "document_embeddings" ADD CONSTRAINT "document_embeddings_document_vector_store_db_id_document_vector_stores_db_id_fk" FOREIGN KEY ("document_vector_store_db_id") REFERENCES "public"."document_vector_stores"("db_id") ON DELETE cascade ON UPDATE no action;--> statement-breakpoint
+CREATE INDEX "doc_embs_store_idx" ON "document_embeddings" USING btree ("document_vector_store_db_id");