feat: prepare embedder for Production

docker/dev.yml

@@ -72,7 +72,7 @@ services:
       parabol-network:
   text-embeddings-inference:
     container_name: text-embeddings-inference
-    image: ghcr.io/huggingface/text-embeddings-inference:cpu-0.6
+    image: ghcr.io/huggingface/text-embeddings-inference:cpu-1.1
     command:
       - "--model-id=llmrails/ember-v1"
     platform: linux/x86_64

packages/client/shared/gqlIds/EmbedderChannelId.ts

@@ -0,0 +1,9 @@
+export const EmbedderChannelId = {
+  join: (serverId: string) => `embedder:${serverId}`,
+  split: (id: string) => {
+    const [, serverId] = id.split(':')
+    return serverId
+  }
+}
+
+export default EmbedderChannelId

packages/embedder/addEmbeddingsMetadata.ts

@@ -0,0 +1,22 @@
+import RedisInstance from 'parabol-server/utils/RedisInstance'
+import {addEmbeddingsMetadataForRetrospectiveDiscussionTopic} from './addEmbeddingsMetadataForRet'
+import {EmbeddingObjectType, PubSubEmbedderMessage} from './embedder'
+
+export const addEmbeddingsMetadata = async (
+  redis: RedisInstance,
+  {objectType, startAt, endAt}: PubSubEmbedderMessage
+) => {
+  const ALL_OBJECT_TYPES: EmbeddingObjectType[] = ['retrospectiveDiscussionTopic']
+  const objectTypes = objectType ? [objectType] : ALL_OBJECT_TYPES
+
+  return Promise.all(
+    objectTypes.map((type) => {
+      switch (type) {
+        case 'retrospectiveDiscussionTopic':
+          return addEmbeddingsMetadataForRetrospectiveDiscussionTopic(redis, startAt, endAt)
+        default:
+          throw new Error(`Invalid object type: ${type}`)
+      }
+    })
+  )
+}

packages/embedder/addEmbeddingsMetadataForRet.ts

@@ -0,0 +1,83 @@
+import ms from 'ms'
+import getRethink from 'parabol-server/database/rethinkDriver'
+import getKysely from 'parabol-server/postgres/getKysely'
+import RedisInstance from 'parabol-server/utils/RedisInstance'
+import Redlock from 'redlock'
+
+const insertDiscussionsIntoMetadata = async (
+  discussions: {id: string; teamId: string; createdAt: Date}[]
+) => {
+  const pg = getKysely()
+  if (discussions.length === 0) return
+  const metadataRows = discussions.map(({id, teamId, createdAt}) => ({
+    refId: id,
+    objectType: 'retrospectiveDiscussionTopic' as const,
+    teamId,
+    // this is technically when the discussion was created. Discussions are mutable.
+    // The best solution would be a date range of min(commentUpdatedAt) to max(commentUpdatedAt)
+    refUpdatedAt: createdAt
+  }))
+
+  const PG_MAX_PARAMS = 65535
+  const metadataColParams = Object.keys(metadataRows[0]).length
+  const metadataBatchSize = Math.trunc(PG_MAX_PARAMS / metadataColParams)
+  const insertBatches = Array.from(
+    {length: Math.ceil(metadataRows.length / metadataBatchSize)},
+    (v, i) => metadataRows.slice(i * metadataBatchSize, i * metadataBatchSize + metadataBatchSize)
+  )
+  return Promise.all(
+    insertBatches.map((batch) => {
+      return pg
+        .insertInto('EmbeddingsMetadata')
+        .values(batch)
+        .onConflict((oc) => oc.doNothing())
+        .execute()
+    })
+  )
+}
+
+export const addEmbeddingsMetadataForRetrospectiveDiscussionTopic = async (
+  redis: RedisInstance,
+  startAt: Date | undefined,
+  endAt: Date | undefined
+) => {
+  const redlock = new Redlock([redis], {retryCount: 0})
+  try {
+    await redlock.acquire([`embedder_metadata_retrospectiveDiscussionTopic`], ms('10m'))
+  } catch {
+    // lock not acquired, another worker must be doing the job. abort
+    return
+  }
+  // load up the metadata table will all discussion topics that are a part of meetings ended within the given date range
+
+  const r = await getRethink()
+  const pg = getKysely()
+  const BATCH_SIZE = 1000
+  const rStartAt = startAt || r.minval
+  const rEndAt = endAt || r.maxval
+
+  let curStartAt = rStartAt
+  for (let i = 0; i < 1e6; i++) {
+    const endedMeetings = await r
+      .table('NewMeeting')
+      .between(curStartAt, rEndAt, {index: 'endedAt'})
+      .orderBy({index: 'endedAt'})
+      .filter({meetingType: 'retrospective'})
+      .limit(BATCH_SIZE)
+      .pluck('id', 'endedAt')
+      .run()
+    if (endedMeetings.length === 0) break
+    const endedMeetingIds = endedMeetings.map(({id}) => id!)
+    const endedMeetingDiscussions = await pg
+      .selectFrom('Discussion')
+      .select(['id', 'teamId', 'createdAt'])
+      .where('meetingId', 'in', endedMeetingIds)
+      .execute()
+    await insertDiscussionsIntoMetadata(endedMeetingDiscussions)
+
+    // assumes that fewer than BATCH_SIZE meetings share the same endedAt value.
+    // If this is not safe, we need to index on `endedAt + id`
+    const lastMeeting = endedMeetings[endedMeetings.length - 1]
+    curStartAt = lastMeeting.endedAt
+  }
+}

packages/embedder/ai_models/AbstractModel.ts

@@ -1,3 +1,7 @@
+import {sql} from 'kysely'
+import getKysely from 'parabol-server/postgres/getKysely'
+import {DB} from 'parabol-server/postgres/pg'
+
 export interface ModelConfig {
   model: string
   url: string
@@ -30,6 +34,8 @@ export interface EmbeddingModelParams {
   tableSuffix: string
 }
 
+export type EmbeddingsTable = Extract<keyof DB, `Embeddings_${string}`>
+
 export abstract class AbstractEmbeddingsModel extends AbstractModel {
   readonly embeddingDimensions: number
   readonly maxInputTokens: number
@@ -42,7 +48,68 @@ export abstract class AbstractEmbeddingsModel extends AbstractModel {
     this.tableName = `Embeddings_${modelParams.tableSuffix}`
   }
   protected abstract constructModelParams(config: EmbeddingModelConfig): EmbeddingModelParams
-  abstract getEmbedding(content: string): Promise<number[]>
+  abstract getEmbedding(content: string): Promise<number[] | Error>
+
+  abstract getTokens(content: string): Promise<number[] | Error>
+
+  async createTable() {
+    const pg = getKysely()
+    const hasTable =
+      (
+        await sql<number[]>`SELECT 1 FROM ${sql.id('pg_catalog', 'pg_tables')} WHERE ${sql.id(
+          'tablename'
+        )} = ${this.tableName}`.execute(pg)
+      ).rows.length > 0
+    if (hasTable) return undefined
+    const vectorDimensions = this.embeddingDimensions
+    console.log(`ModelManager: creating ${this.tableName} with ${vectorDimensions} dimensions`)
+    await sql`
+      DO $$
+        BEGIN
+        CREATE TABLE IF NOT EXISTS ${sql.id(this.tableName)} (
+          "id" INT GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
+          "embedText" TEXT,
+          "embedding" vector(${sql.raw(vectorDimensions.toString())}),
+          "embeddingsMetadataId" INTEGER UNIQUE NOT NULL,
+          FOREIGN KEY ("embeddingsMetadataId")
+            REFERENCES "EmbeddingsMetadata"("id")
+            ON DELETE CASCADE
+        );
+        CREATE INDEX IF NOT EXISTS "idx_${sql.raw(this.tableName)}_embedding_vector_cosign_ops"
+          ON ${sql.id(this.tableName)}
+          USING hnsw ("embedding" vector_cosine_ops);
+        END
+      $$;
+      CREATE OR REPLACE FUNCTION insert_metadata_in_queue_${sql.raw(this.tableName)} ()
+      RETURNS TRIGGER AS $$
+      BEGIN
+          INSERT INTO public."EmbeddingsJobQueue" ("model", "embeddingsMetadataId")
+              VALUES ('${sql.raw(this.tableName)}', NEW."embeddingsMetadataId");
+          RETURN NEW;
+      END;
+      $$ LANGUAGE plpgsql;
+
+      DROP TRIGGER IF EXISTS "embeddings_metadata_to_queue_${sql.raw(
+        this.tableName
+      )}" on "EmbeddingsMetadata";
+
+      CREATE TRIGGER "embeddings_metadata_to_queue_${sql.raw(this.tableName)}"
+
+      AFTER INSERT ON "EmbeddingsMetadata"
+      FOR EACH ROW
+      EXECUTE PROCEDURE insert_metadata_in_queue_${sql.raw(this.tableName)}();
+      `.execute(pg)
+
+    console.log(
+      `ModelManager: Queueing EmbeddingsMetadata into EmbeddingsJobQue for ${this.tableName}`
+    )
+    await sql`
+    INSERT INTO "EmbeddingsJobQueue" ("model", "embeddingsMetadataId")
+    SELECT '${sql.raw(this.tableName)}', "embeddingsMetadataId"
+    FROM "EmbeddingsMetadata"
+    ON CONFLICT DO NOTHING;
+    `.execute(pg)
+  }
 }
 
 export interface GenerationModelParams {

packages/embedder/ai_models/ModelManager.ts

@@ -1,5 +1,6 @@
-import {Kysely, sql} from 'kysely'
+import {sql} from 'kysely'
 
+import getKysely from 'parabol-server/postgres/getKysely'
 import {
   AbstractEmbeddingsModel,
   AbstractGenerationModel,
@@ -93,39 +94,37 @@ export class ModelManager {
     })
   }
 
-  async maybeCreateTables(pg: Kysely<any>) {
-    const maybePromises = this.embeddingModels.map(async (embeddingsModel) => {
-      const tableName = embeddingsModel.tableName
-      const hasTable =
-        (
-          await sql<number[]>`SELECT 1 FROM ${sql.id('pg_catalog', 'pg_tables')} WHERE ${sql.id(
-            'tablename'
-          )} = ${tableName}`.execute(pg)
-        ).rows.length > 0
-      if (hasTable) return undefined
-      const vectorDimensions = embeddingsModel.embeddingDimensions
-      console.log(`ModelManager: creating ${tableName} with ${vectorDimensions} dimensions`)
-      const query = sql`
-      DO $$
-  BEGIN
-  CREATE TABLE IF NOT EXISTS ${sql.id(tableName)} (
-    "id" INT GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
-    "embedText" TEXT,
-    "embedding" vector(${sql.raw(vectorDimensions.toString())}),
-    "embeddingsMetadataId" INTEGER NOT NULL,
-    FOREIGN KEY ("embeddingsMetadataId")
-      REFERENCES "EmbeddingsMetadata"("id")
-      ON DELETE CASCADE
-  );
-  CREATE INDEX IF NOT EXISTS "idx_${sql.raw(tableName)}_embedding_vector_cosign_ops"
-    ON ${sql.id(tableName)}
-    USING hnsw ("embedding" vector_cosine_ops);
-  END $$;
-
-      `
-      return query.execute(pg)
-    })
-    Promise.all(maybePromises)
+  async maybeCreateTables() {
+    return Promise.all(this.embeddingModels.map((model) => model.createTable()))
+  }
+  /*
+    Once a model is no longer used, don't schedule work for it in the job queue
+  */
+  async removeOldTriggers() {
+    const pg = getKysely()
+    const prefix = 'embeddings_metadata_to_queue_'
+    const triggers = await pg
+      .selectFrom('information_schema.triggers' as any)
+      .select('trigger_name')
+      .where('event_object_table', '=', 'EmbeddingsMetadata')
+      .where('trigger_name', 'like', `${prefix}%`)
+      .execute()
+    return Promise.all(
+      triggers.map(async ({trigger_name}) => {
+        // pgadmin lowercases triggers but PG doesn't. Lowercase it all just to be safe
+        const lowercaseTableName = trigger_name.slice(prefix.length).toLowerCase()
+        const isModelUsed = this.embeddingModels.some(
+          (model) => model.tableName.toLowerCase() === lowercaseTableName
+        )
+        if (isModelUsed) return
+        await sql`
+        DROP TRIGGER IF EXISTS ${sql.id(trigger_name)} on "EmbeddingsMetadata";
+        DROP FUNCTION IF EXISTS insert_metadata_in_queue_${sql.raw(lowercaseTableName)};`.execute(
+          pg
+        )
+        console.log(`Removed old trigger: ${trigger_name}`)
+      })
+    )
   }
 }
 

packages/embedder/ai_models/TextEmbeddingsInference.ts

@@ -27,6 +27,28 @@ export class TextEmbeddingsInference extends AbstractEmbeddingsModel {
     super(config)
   }
 
+  async getTokens(content: string) {
+    const fetchOptions = {
+      body: JSON.stringify({inputs: content}),
+      deadline: new Date(new Date().getTime() + MAX_REQUEST_TIME_S * 1000),
+      headers: {
+        Accept: 'application/json',
+        'Content-Type': 'application/json; charset=utf-8'
+      },
+      method: 'POST'
+    }
+
+    try {
+      const res = await fetchWithRetry(`${this.url}/tokenize`, fetchOptions)
+      const listOfTokens = (await res.json()) as number[][]
+      if (!listOfTokens) return new Error('listOfTokens is undefined')
+      if (listOfTokens.length !== 1 || !listOfTokens[0])
+        return new Error(`listOfTokens list length !== 1 (length: ${listOfTokens.length})`)
+      return listOfTokens[0]
+    } catch (e) {
+      return e instanceof Error ? e : new Error(e)
+    }
+  }
   public async getEmbedding(content: string) {
     const fetchOptions = {
       body: JSON.stringify({inputs: content}),
@@ -40,17 +62,14 @@ export class TextEmbeddingsInference extends AbstractEmbeddingsModel {
 
     try {
       const res = await fetchWithRetry(`${this.url}/embed`, fetchOptions)
-      const listOfVectors = (await res.json()) as Array<number[]>
-      if (!listOfVectors)
-        throw new Error('TextEmbeddingsInference.getEmbeddings(): listOfVectors is undefined')
+      const listOfVectors = (await res.json()) as number[][]
+      if (!listOfVectors) return new Error('listOfVectors is undefined')
       if (listOfVectors.length !== 1 || !listOfVectors[0])
-        throw new Error(
-          `TextEmbeddingsInference.getEmbeddings(): listOfVectors list length !== 1 (length: ${listOfVectors.length})`
-        )
+        return new Error(`listOfVectors list length !== 1 (length: ${listOfVectors.length})`)
       return listOfVectors[0]
     } catch (e) {
       console.log(`TextEmbeddingsInference.getEmbeddings() timeout: `, e)
-      throw e
+      return e instanceof Error ? e : new Error(e || 'Unknown Error')
     }
   }