fix: dataloader mem leak

Signed-off-by: Matt Krick <matt.krick@gmail.com>

packages/embedder/JobQueueStream.ts

@@ -1,7 +1,7 @@
 import {Selectable, sql} from 'kysely'
 import ms from 'ms'
 import sleep from 'parabol-client/utils/sleep'
-import {DataLoaderWorker} from 'parabol-server/graphql/graphql'
+import RootDataLoader from 'parabol-server/dataloader/RootDataLoader'
 import 'parabol-server/initSentry'
 import getKysely from 'parabol-server/postgres/getKysely'
 import {DB} from 'parabol-server/postgres/pg'
@@ -13,10 +13,10 @@ import {updateJobState} from './indexing/updateJobState'
 
 type Job = Selectable<DB['EmbeddingsJobQueue']>
 export default class JobQueueStream implements AsyncIterator<Job> {
-  private dataLoader: DataLoaderWorker
+  private dataLoader: RootDataLoader
   private modelManager: ModelManager
 
-  constructor(modelManager: ModelManager, dataLoader: DataLoaderWorker) {
+  constructor(modelManager: ModelManager, dataLoader: RootDataLoader) {
     this.modelManager = modelManager
     this.dataLoader = dataLoader
   }
@@ -92,7 +92,7 @@ export default class JobQueueStream implements AsyncIterator<Job> {
         await pg
           .updateTable('EmbeddingsMetadata')
           .set({fullText})
-          .where('id', '=', metadata.id)
+          .where('id', '=', embeddingsMetadataId)
           .execute()
       }
     } catch (e) {
@@ -128,7 +128,9 @@ export default class JobQueueStream implements AsyncIterator<Job> {
     // Cannot use summarization strategy if generation model has same context length as embedding model
     // We must split the text & not tokens because the endpoint doesn't support decoding input tokens
     const chunks = isFullTextTooBig ? embeddingModel.splitText(fullText, 0) : [fullText]
-
+    if (isFullTextTooBig) {
+      console.log('we split!', chunks)
+    }
     await Promise.all(
       chunks.map(async (chunk, chunkNumber) => {
         const embeddingVector = await embeddingModel.getEmbedding(chunk)
@@ -150,7 +152,7 @@ export default class JobQueueStream implements AsyncIterator<Job> {
             chunkNumber: isFullTextTooBig ? chunkNumber : null
           })
           .onConflict((oc) =>
-            oc.doUpdateSet((eb) => ({
+            oc.column('embeddingsMetadataId').doUpdateSet((eb) => ({
               embedText: eb.ref('excluded.embedText'),
               embedding: eb.ref('excluded.embedding')
             }))

packages/embedder/embedder.ts

@@ -11,6 +11,7 @@ import {establishPrimaryEmbedder} from './establishPrimaryEmbedder'
 import {importHistoricalMetadata} from './importHistoricalMetadata'
 import {getRootDataLoader} from './indexing/getRootDataLoader'
 import {mergeAsyncIterators} from './mergeAsyncIterators'
+import RootDataLoader from 'parabol-server/dataloader/RootDataLoader'
 
 tracer.init({
   service: `embedder`,
@@ -81,7 +82,7 @@ const run = async () => {
   }
 
   const incomingStream = new RedisStream('embedderStream', 'embedderConsumerGroup', embedderChannel)
-  const dataLoader = getRootDataLoader()
+  const dataLoader = new RootDataLoader({maxBatchSize: 1000})
   const jobQueueStream = new JobQueueStream(modelManager, dataLoader)
 
   console.log(`\n⚡⚡⚡️️ Server ID: ${SERVER_ID}. Embedder is ready ⚡⚡⚡️️️`)

packages/embedder/indexing/createEmbeddingTextFrom.ts

@@ -1,12 +1,12 @@
 import {Selectable} from 'kysely'
-import {DataLoaderWorker} from 'parabol-server/graphql/graphql'
 import {DB} from 'parabol-server/postgres/pg'
 
+import RootDataLoader from 'parabol-server/dataloader/RootDataLoader'
 import {createTextFromRetrospectiveDiscussionTopic} from './retrospectiveDiscussionTopic'
 
 export const createEmbeddingTextFrom = async (
   embeddingsMetadata: Selectable<DB['EmbeddingsMetadata']>,
-  dataLoader: DataLoaderWorker
+  dataLoader: RootDataLoader
 ) => {
   switch (embeddingsMetadata.objectType) {
     case 'retrospectiveDiscussionTopic':

packages/embedder/indexing/retrospectiveDiscussionTopic.ts

@@ -1,7 +1,7 @@
-import getRethink, {RethinkSchema} from 'parabol-server/database/rethinkDriver'
+import {RethinkSchema} from 'parabol-server/database/rethinkDriver'
 import Comment from 'parabol-server/database/types/Comment'
 import {isMeetingRetrospective} from 'parabol-server/database/types/MeetingRetrospective'
-import {DataLoaderWorker} from 'parabol-server/graphql/graphql'
+import RootDataLoader from 'parabol-server/dataloader/RootDataLoader'
 import prettier from 'prettier'
 
 // Here's a generic reprentation of the text generated here:
@@ -21,14 +21,14 @@ import prettier from 'prettier'
 
 const IGNORE_COMMENT_USER_IDS = ['parabolAIUser']
 
-async function getPreferredNameByUserId(userId: string, dataLoader: DataLoaderWorker) {
+async function getPreferredNameByUserId(userId: string, dataLoader: RootDataLoader) {
   if (!userId) return 'Unknown'
   const user = await dataLoader.get('users').load(userId)
   return !user ? 'Unknown' : user.preferredName
 }
 
 async function formatThread(
-  dataLoader: DataLoaderWorker,
+  dataLoader: RootDataLoader,
   comments: Comment[],
   parentId: string | null = null,
   depth = 0
@@ -60,13 +60,12 @@ async function formatThread(
 
 export const createTextFromRetrospectiveDiscussionTopic = async (
   discussionId: string,
-  dataLoader: DataLoaderWorker,
+  dataLoader: RootDataLoader,
   textForReranking: boolean = false
 ) => {
   const discussion = await dataLoader.get('discussions').load(discussionId)
   if (!discussion) throw new Error(`Discussion not found: ${discussionId}`)
   const {discussionTopicId: reflectionGroupId, meetingId, summary: discussionSummary} = discussion
-  const r = await getRethink()
   const [newMeeting, reflectionGroup, reflections] = await Promise.all([
     dataLoader.get('newMeetings').load(meetingId),
     dataLoader.get('retroReflectionGroups').load(reflectionGroupId),
@@ -166,7 +165,7 @@ export const createTextFromRetrospectiveDiscussionTopic = async (
 
 export const newRetroDiscussionTopicsFromNewMeeting = async (
   newMeeting: RethinkSchema['NewMeeting']['type'],
-  dataLoader: DataLoaderWorker
+  dataLoader: RootDataLoader
 ) => {
   const discussPhase = newMeeting.phases.find((phase) => phase.phaseType === 'discuss')
   const orgId = (await dataLoader.get('teams').load(newMeeting.teamId))?.orgId