Browse Source

Add vectorStoreRelationId field

Shun Miyazawa 1 year ago
parent
commit
d2b97240ee

+ 12 - 2
apps/app/src/features/openai/server/models/vector-store-file-relation.ts

@@ -5,6 +5,7 @@ import { type Model, type Document, Schema } from 'mongoose';
 import { getOrCreateModel } from '~/server/util/mongoose-utils';
 import { getOrCreateModel } from '~/server/util/mongoose-utils';
 
 
 export interface VectorStoreFileRelation {
 export interface VectorStoreFileRelation {
+  vectorStoreRelationId: mongoose.Types.ObjectId;
   pageId: mongoose.Types.ObjectId;
   pageId: mongoose.Types.ObjectId;
   fileIds: string[];
   fileIds: string[];
   isAttachedToVectorStore: boolean;
   isAttachedToVectorStore: boolean;
@@ -18,7 +19,7 @@ interface VectorStoreFileRelationModel extends Model<VectorStoreFileRelation> {
 }
 }
 
 
 export const prepareVectorStoreFileRelations = (
 export const prepareVectorStoreFileRelations = (
-    pageId: Types.ObjectId, fileId: string, relationsMap: Map<string, VectorStoreFileRelation>,
+    vectorStoreRelationId: Types.ObjectId, pageId: Types.ObjectId, fileId: string, relationsMap: Map<string, VectorStoreFileRelation>,
 ): Map<string, VectorStoreFileRelation> => {
 ): Map<string, VectorStoreFileRelation> => {
   const pageIdStr = pageId.toHexString();
   const pageIdStr = pageId.toHexString();
   const existingData = relationsMap.get(pageIdStr);
   const existingData = relationsMap.get(pageIdStr);
@@ -30,6 +31,7 @@ export const prepareVectorStoreFileRelations = (
   // If the data doesn't exist, create a new one and add it to the map
   // If the data doesn't exist, create a new one and add it to the map
   else {
   else {
     relationsMap.set(pageIdStr, {
     relationsMap.set(pageIdStr, {
+      vectorStoreRelationId,
       pageId,
       pageId,
       fileIds: [fileId],
       fileIds: [fileId],
       isAttachedToVectorStore: false,
       isAttachedToVectorStore: false,
@@ -40,6 +42,11 @@ export const prepareVectorStoreFileRelations = (
 };
 };
 
 
 const schema = new Schema<VectorStoreFileRelationDocument, VectorStoreFileRelationModel>({
 const schema = new Schema<VectorStoreFileRelationDocument, VectorStoreFileRelationModel>({
+  vectorStoreRelationId: {
+    type: Schema.Types.ObjectId,
+    ref: 'VectorStore',
+    required: true,
+  },
   pageId: {
   pageId: {
     type: Schema.Types.ObjectId,
     type: Schema.Types.ObjectId,
     ref: 'Page',
     ref: 'Page',
@@ -63,7 +70,10 @@ schema.statics.upsertVectorStoreFileRelations = async function(vectorStoreFileRe
       return {
       return {
         updateOne: {
         updateOne: {
           filter: { pageId: data.pageId },
           filter: { pageId: data.pageId },
-          update: { $addToSet: { fileIds: { $each: data.fileIds } } },
+          update: {
+            $set: { vectorStoreId: data.vectorStoreRelationId },
+            $addToSet: { fileIds: { $each: data.fileIds } },
+          },
           upsert: true,
           upsert: true,
         },
         },
       };
       };