Schedule document suggestions + bug fixes

apps/workers/src/workers/index.ts

@@ -1,5 +1,6 @@
-import { queues } from '@latitude-data/core/queues'
 import { captureException } from '$/utils/sentry'
+import { setupSchedules } from '@latitude-data/core/jobs'
+import { queues } from '@latitude-data/core/queues'
 import { Worker } from 'bullmq'
 
 import { defaultWorker } from './worker-definitions/defaultWorker'
@@ -17,6 +18,8 @@ const WORKERS = [defaultWorker]
 export default async function startWorkers() {
   const connection = await queues({ enableOfflineQueue: true })
 
+  await setupSchedules(connection)
+
   return WORKERS.flatMap((w) =>
     w.queues.map((q) => {
       const worker = new Worker(q, w.processor, {

apps/workers/src/workers/worker-definitions/defaultWorker.ts

@@ -4,6 +4,7 @@ import { buildProcessor } from '../_shared'
 
 const defaultWorkerQueues = [
   Queues.defaultQueue,
+  Queues.maintenanceQueue,
   Queues.eventHandlersQueue,
   Queues.eventsQueue,
   Queues.liveEvaluationsQueue,

packages/core/src/events/events.d.ts

@@ -8,7 +8,7 @@ import type {
   DocumentSuggestion,
   DocumentVersion,
   Evaluation,
-  EvaluationResult,
+  EvaluationDto,
   EvaluationResultDto,
   LogSources,
   MagicLinkToken,
@@ -319,8 +319,8 @@ export type DocumentCreatedEvent = LatitudeEventGeneric<
 export type EvaluationResultCreatedEvent = LatitudeEventGeneric<
   'evaluationResultCreated',
   {
-    evaluationResult: EvaluationResult
-    evaluation: Evaluation
+    evaluationResult: EvaluationResultDto
+    evaluation: EvaluationDto
     documentLog: DocumentLog
     workspaceId: number
   }

packages/core/src/events/handlers/requestDocumentSuggestionJob.test.ts

@@ -0,0 +1,163 @@
+import { beforeEach, describe, expect, it, MockInstance, vi } from 'vitest'
+import {
+  DocumentLog,
+  EvaluationDto,
+  EvaluationMetadataType,
+  EvaluationResultableType,
+  EvaluationResultDto,
+  LogSources,
+  Workspace,
+} from '../../browser'
+import * as jobs from '../../jobs'
+import * as factories from '../../tests/factories'
+import { requestDocumentSuggestionJob } from './requestDocumentSuggestionJob'
+
+describe('requestDocumentSuggestionJob', () => {
+  let mocks: {
+    enqueueGenerateDocumentSuggestionJob: MockInstance
+  }
+
+  let workspace: Workspace
+  let evaluation: EvaluationDto
+  let result: EvaluationResultDto
+  let documentLog: DocumentLog
+
+  beforeEach(async () => {
+    vi.clearAllMocks()
+    vi.resetAllMocks()
+    vi.restoreAllMocks()
+
+    const {
+      workspace: w,
+      providers,
+      project,
+      user,
+    } = await factories.createProject()
+    workspace = w
+    const provider = providers[0]!
+
+    const commit = await factories.createCommit({
+      projectId: project.id,
+      user: user,
+    })
+
+    const { documentVersion: d } = await factories.createDocumentVersion({
+      workspace: workspace,
+      user: user,
+      commit: commit,
+      path: 'prompt',
+      content: factories.helpers.createPrompt({ provider }),
+    })
+    const document = d
+
+    evaluation = await factories.createEvaluation({
+      workspace: workspace,
+      user: user,
+      metadataType: EvaluationMetadataType.LlmAsJudgeAdvanced,
+      resultType: EvaluationResultableType.Number,
+      resultConfiguration: { minValue: 0, maxValue: 100 },
+    })
+
+    await factories.createConnectedEvaluation({
+      workspace: workspace,
+      user: user,
+      evaluationUuid: evaluation.uuid,
+      documentUuid: document.documentUuid,
+      live: true,
+    })
+
+    const { documentLog: l } = await factories.createDocumentLog({
+      document,
+      commit,
+      source: LogSources.Playground,
+    })
+    documentLog = l
+
+    const providerLog = await factories.createProviderLog({
+      workspace: workspace,
+      documentLogUuid: documentLog.uuid,
+      providerId: provider.id,
+      providerType: provider.provider,
+    })
+
+    const { evaluationResult: r } = await factories.createEvaluationResult({
+      evaluation: evaluation,
+      documentLog: documentLog,
+      evaluatedProviderLog: providerLog,
+      result: '31',
+    })
+    result = r
+
+    mocks = { enqueueGenerateDocumentSuggestionJob: vi.fn() }
+    vi.spyOn(jobs, 'setupJobs').mockResolvedValue({
+      defaultQueue: {
+        jobs: {
+          enqueueGenerateDocumentSuggestionJob:
+            mocks.enqueueGenerateDocumentSuggestionJob,
+        },
+      },
+    } as any)
+  })
+
+  it('not enqueues generate suggestion job when result is not from playground', async () => {
+    documentLog.source = LogSources.API
+    result.source = LogSources.API
+
+    await requestDocumentSuggestionJob({
+      data: {
+        type: 'evaluationResultCreated',
+        data: {
+          workspaceId: workspace.id,
+          evaluationResult: result,
+          evaluation: evaluation,
+          documentLog: documentLog,
+        },
+      },
+    })
+
+    expect(mocks.enqueueGenerateDocumentSuggestionJob).not.toHaveBeenCalled()
+  })
+
+  it('not enqueues generate suggestion job when result has passed', async () => {
+    result.result = (evaluation.resultConfiguration as any).maxValue
+
+    await requestDocumentSuggestionJob({
+      data: {
+        type: 'evaluationResultCreated',
+        data: {
+          workspaceId: workspace.id,
+          evaluationResult: result,
+          evaluation: evaluation,
+          documentLog: documentLog,
+        },
+      },
+    })
+
+    expect(mocks.enqueueGenerateDocumentSuggestionJob).not.toHaveBeenCalled()
+  })
+
+  it('enqueues generate suggestion job', async () => {
+    await requestDocumentSuggestionJob({
+      data: {
+        type: 'evaluationResultCreated',
+        data: {
+          workspaceId: workspace.id,
+          evaluationResult: result,
+          evaluation: evaluation,
+          documentLog: documentLog,
+        },
+      },
+    })
+
+    expect(mocks.enqueueGenerateDocumentSuggestionJob).toHaveBeenCalledOnce()
+    expect(mocks.enqueueGenerateDocumentSuggestionJob).toHaveBeenCalledWith(
+      {
+        workspaceId: workspace.id,
+        commitId: documentLog.commitId,
+        documentUuid: documentLog.documentUuid,
+        evaluationId: result.evaluationId,
+      },
+      { attempts: 1, deduplication: { id: expect.any(String) } },
+    )
+  })
+})

packages/core/src/events/handlers/requestDocumentSuggestionJob.ts

@@ -1,8 +1,9 @@
 import { LogSources } from '../../browser'
 import { unsafelyFindWorkspace } from '../../data-access'
 import { setupJobs } from '../../jobs'
+import { generateDocumentSuggestionJobKey } from '../../jobs/job-definitions'
 import { NotFoundError } from '../../lib'
-import { EvaluationResultsRepository } from '../../repositories'
+import { hasEvaluationResultPassed } from '../../services/evaluationResults'
 import { EvaluationResultCreatedEvent } from '../events'
 
 export const requestDocumentSuggestionJob = async ({
@@ -12,25 +13,35 @@ export const requestDocumentSuggestionJob = async ({
 }) => {
   const {
     workspaceId,
-    evaluationResult: { id: resultId },
-    documentLog: { commitId, documentUuid },
+    evaluationResult: result,
+    evaluation,
+    documentLog,
   } = event.data
 
   const workspace = await unsafelyFindWorkspace(workspaceId)
   if (!workspace) throw new NotFoundError(`Workspace not found ${workspaceId}`)
 
-  const resultsRepository = new EvaluationResultsRepository(workspace.id)
-  const result = await resultsRepository.find(resultId).then((r) => r.unwrap())
-  if (result.source !== LogSources.Playground) return
+  if (documentLog.source !== LogSources.Playground) return
+  if (hasEvaluationResultPassed({ result, evaluation })) return
 
   const queues = await setupJobs()
   queues.defaultQueue.jobs.enqueueGenerateDocumentSuggestionJob(
     {
       workspaceId: workspace.id,
-      commitId: commitId,
-      documentUuid: documentUuid,
-      evaluationId: result.evaluationId,
+      commitId: documentLog.commitId,
+      documentUuid: documentLog.documentUuid,
+      evaluationId: evaluation.id,
+    },
+    {
+      attempts: 1,
+      deduplication: {
+        id: generateDocumentSuggestionJobKey({
+          workspaceId: workspace.id,
+          commitId: documentLog.commitId,
+          documentUuid: documentLog.documentUuid,
+          evaluationId: evaluation.id,
+        }),
+      },
     },
-    { attempts: 1 },
   )
 }

packages/core/src/jobs/constants.ts

@@ -1,6 +1,7 @@
 // TODO: Review if we can remove this declarations
 export enum Queues {
   defaultQueue = 'defaultQueue',
+  maintenanceQueue = 'maintenanceQueue',
   eventsQueue = 'eventsQueue',
   eventHandlersQueue = 'eventHandlersQueue',
   liveEvaluationsQueue = 'liveEvaluationsQueue',
@@ -24,6 +25,8 @@ export enum Jobs {
   createDocumentLogFromSpanJob = 'createDocumentLogFromSpanJob',
   createDocumentLogsFromSpansJob = 'createDocumentLogsFromSpansJob',
   generateDocumentSuggestionJob = 'generateDocumentSuggestionJob',
+  requestDocumentSuggestionsJob = 'requestDocumentSuggestionsJob',
+  cleanDocumentSuggestionsJob = 'cleanDocumentSuggestionsJob',
 }
 
 export const QUEUES = {
@@ -41,8 +44,13 @@ export const QUEUES = {
       'runEvaluationJob',
       'uploadDocumentLogsJob',
       'generateDocumentSuggestionJob',
+      'requestDocumentSuggestionsJob',
     ],
   },
+  [Queues.maintenanceQueue]: {
+    name: Queues.maintenanceQueue,
+    jobs: ['cleanDocumentSuggestionsJob'],
+  },
   [Queues.eventsQueue]: {
     name: Queues.eventsQueue,
     jobs: ['publishEventJob', 'createEventJob', 'publishToAnalyticsJob'],

packages/core/src/jobs/index.ts

@@ -1,12 +1,26 @@
+import Redis from 'ioredis'
 import { setupQueues } from './queues'
 
 export { Worker } from 'bullmq'
 
 let queues: Awaited<ReturnType<typeof setupQueues>>
 
-export async function setupJobs() {
-  if (queues) return queues
-  queues = await setupQueues()
-
+export async function setupJobs(connection?: Redis) {
+  if (!queues) queues = await setupQueues(connection)
   return queues
 }
+
+export async function setupSchedules(connection?: Redis) {
+  if (!queues) queues = await setupQueues(connection)
+
+  // Every day at 8 AM
+  await queues.defaultQueue.jobs.scheduleRequestDocumentSuggestionsJob(
+    '0 0 8 * * *',
+    { attempts: 1 },
+  )
+  // Every day at 2 AM
+  await queues.maintenanceQueue.jobs.scheduleCleanDocumentSuggestionsJob(
+    '0 0 2 * * *',
+    { attempts: 1 },
+  )
+}