refactor(server): jobs (#2023)

* refactor: job to domain

* chore: regenerate open api

* chore: tests

* fix: missing breaks

* fix: get asset with missing exif data

---------

Co-authored-by: Alex Tran <alex.tran1502@gmail.com>

server/libs/domain/src/smart-info/machine-learning.interface.ts

@@ -5,7 +5,7 @@ export interface MachineLearningInput {
 }
 
 export interface IMachineLearningRepository {
-  tagImage(input: MachineLearningInput): Promise<string[]>;
+  classifyImage(input: MachineLearningInput): Promise<string[]>;
   detectObjects(input: MachineLearningInput): Promise<string[]>;
   encodeImage(input: MachineLearningInput): Promise<number[]>;
   encodeText(input: string): Promise<number[]>;

server/libs/domain/src/smart-info/smart-info.service.spec.ts

@@ -1,6 +1,13 @@
 import { AssetEntity } from '@app/infra/db/entities';
-import { newJobRepositoryMock, newMachineLearningRepositoryMock, newSmartInfoRepositoryMock } from '../../test';
-import { IJobRepository } from '../job';
+import {
+  assetEntityStub,
+  newAssetRepositoryMock,
+  newJobRepositoryMock,
+  newMachineLearningRepositoryMock,
+  newSmartInfoRepositoryMock,
+} from '../../test';
+import { IAssetRepository, WithoutProperty } from '../asset';
+import { IJobRepository, JobName } from '../job';
 import { IMachineLearningRepository } from './machine-learning.interface';
 import { ISmartInfoRepository } from './smart-info.repository';
 import { SmartInfoService } from './smart-info.service';
@@ -12,35 +19,63 @@ const asset = {
 
 describe(SmartInfoService.name, () => {
   let sut: SmartInfoService;
+  let assetMock: jest.Mocked<IAssetRepository>;
   let jobMock: jest.Mocked<IJobRepository>;
   let smartMock: jest.Mocked<ISmartInfoRepository>;
   let machineMock: jest.Mocked<IMachineLearningRepository>;
 
   beforeEach(async () => {
+    assetMock = newAssetRepositoryMock();
     smartMock = newSmartInfoRepositoryMock();
     jobMock = newJobRepositoryMock();
     machineMock = newMachineLearningRepositoryMock();
-    sut = new SmartInfoService(jobMock, smartMock, machineMock);
+    sut = new SmartInfoService(assetMock, jobMock, smartMock, machineMock);
   });
 
   it('should work', () => {
     expect(sut).toBeDefined();
   });
 
+  describe('handleQueueObjectTagging', () => {
+    it('should queue the assets without tags', async () => {
+      assetMock.getWithout.mockResolvedValue([assetEntityStub.image]);
+
+      await sut.handleQueueObjectTagging({ force: false });
+
+      expect(jobMock.queue.mock.calls).toEqual([
+        [{ name: JobName.CLASSIFY_IMAGE, data: { asset: assetEntityStub.image } }],
+        [{ name: JobName.DETECT_OBJECTS, data: { asset: assetEntityStub.image } }],
+      ]);
+      expect(assetMock.getWithout).toHaveBeenCalledWith(WithoutProperty.OBJECT_TAGS);
+    });
+
+    it('should queue all the assets', async () => {
+      assetMock.getAll.mockResolvedValue([assetEntityStub.image]);
+
+      await sut.handleQueueObjectTagging({ force: true });
+
+      expect(jobMock.queue.mock.calls).toEqual([
+        [{ name: JobName.CLASSIFY_IMAGE, data: { asset: assetEntityStub.image } }],
+        [{ name: JobName.DETECT_OBJECTS, data: { asset: assetEntityStub.image } }],
+      ]);
+      expect(assetMock.getAll).toHaveBeenCalled();
+    });
+  });
+
   describe('handleTagImage', () => {
     it('should skip assets without a resize path', async () => {
-      await sut.handleTagImage({ asset: { resizePath: '' } as AssetEntity });
+      await sut.handleClassifyImage({ asset: { resizePath: '' } as AssetEntity });
 
       expect(smartMock.upsert).not.toHaveBeenCalled();
-      expect(machineMock.tagImage).not.toHaveBeenCalled();
+      expect(machineMock.classifyImage).not.toHaveBeenCalled();
     });
 
     it('should save the returned tags', async () => {
-      machineMock.tagImage.mockResolvedValue(['tag1', 'tag2', 'tag3']);
+      machineMock.classifyImage.mockResolvedValue(['tag1', 'tag2', 'tag3']);
 
-      await sut.handleTagImage({ asset });
+      await sut.handleClassifyImage({ asset });
 
-      expect(machineMock.tagImage).toHaveBeenCalledWith({ thumbnailPath: 'path/to/resize.ext' });
+      expect(machineMock.classifyImage).toHaveBeenCalledWith({ thumbnailPath: 'path/to/resize.ext' });
       expect(smartMock.upsert).toHaveBeenCalledWith({
         assetId: 'asset-1',
         tags: ['tag1', 'tag2', 'tag3'],
@@ -48,19 +83,19 @@ describe(SmartInfoService.name, () => {
     });
 
     it('should handle an error with the machine learning pipeline', async () => {
-      machineMock.tagImage.mockRejectedValue(new Error('Unable to read thumbnail'));
+      machineMock.classifyImage.mockRejectedValue(new Error('Unable to read thumbnail'));
 
-      await sut.handleTagImage({ asset });
+      await sut.handleClassifyImage({ asset });
 
       expect(smartMock.upsert).not.toHaveBeenCalled();
     });
 
     it('should no update the smart info if no tags were returned', async () => {
-      machineMock.tagImage.mockResolvedValue([]);
+      machineMock.classifyImage.mockResolvedValue([]);
 
-      await sut.handleTagImage({ asset });
+      await sut.handleClassifyImage({ asset });
 
-      expect(machineMock.tagImage).toHaveBeenCalled();
+      expect(machineMock.classifyImage).toHaveBeenCalled();
       expect(smartMock.upsert).not.toHaveBeenCalled();
     });
   });
@@ -102,4 +137,53 @@ describe(SmartInfoService.name, () => {
       expect(smartMock.upsert).not.toHaveBeenCalled();
     });
   });
+
+  describe('handleQueueEncodeClip', () => {
+    it('should queue the assets without clip embeddings', async () => {
+      assetMock.getWithout.mockResolvedValue([assetEntityStub.image]);
+
+      await sut.handleQueueEncodeClip({ force: false });
+
+      expect(jobMock.queue).toHaveBeenCalledWith({ name: JobName.ENCODE_CLIP, data: { asset: assetEntityStub.image } });
+      expect(assetMock.getWithout).toHaveBeenCalledWith(WithoutProperty.CLIP_ENCODING);
+    });
+
+    it('should queue all the assets', async () => {
+      assetMock.getAll.mockResolvedValue([assetEntityStub.image]);
+
+      await sut.handleQueueEncodeClip({ force: true });
+
+      expect(jobMock.queue).toHaveBeenCalledWith({ name: JobName.ENCODE_CLIP, data: { asset: assetEntityStub.image } });
+      expect(assetMock.getAll).toHaveBeenCalled();
+    });
+  });
+
+  describe('handleEncodeClip', () => {
+    it('should skip assets without a resize path', async () => {
+      await sut.handleEncodeClip({ asset: { resizePath: '' } as AssetEntity });
+
+      expect(smartMock.upsert).not.toHaveBeenCalled();
+      expect(machineMock.encodeImage).not.toHaveBeenCalled();
+    });
+
+    it('should save the returned objects', async () => {
+      machineMock.encodeImage.mockResolvedValue([0.01, 0.02, 0.03]);
+
+      await sut.handleEncodeClip({ asset });
+
+      expect(machineMock.encodeImage).toHaveBeenCalledWith({ thumbnailPath: 'path/to/resize.ext' });
+      expect(smartMock.upsert).toHaveBeenCalledWith({
+        assetId: 'asset-1',
+        clipEmbedding: [0.01, 0.02, 0.03],
+      });
+    });
+
+    it('should handle an error with the machine learning pipeline', async () => {
+      machineMock.encodeImage.mockRejectedValue(new Error('Unable to read thumbnail'));
+
+      await sut.handleEncodeClip({ asset });
+
+      expect(smartMock.upsert).not.toHaveBeenCalled();
+    });
+  });
 });

server/libs/domain/src/smart-info/smart-info.service.ts

@@ -1,6 +1,7 @@
 import { MACHINE_LEARNING_ENABLED } from '@app/common';
 import { Inject, Injectable, Logger } from '@nestjs/common';
-import { IAssetJob, IJobRepository, JobName } from '../job';
+import { IAssetRepository, WithoutProperty } from '../asset';
+import { IAssetJob, IBaseJob, IJobRepository, JobName } from '../job';
 import { IMachineLearningRepository } from './machine-learning.interface';
 import { ISmartInfoRepository } from './smart-info.repository';
 
@@ -9,26 +10,24 @@ export class SmartInfoService {
   private logger = new Logger(SmartInfoService.name);
 
   constructor(
+    @Inject(IAssetRepository) private assetRepository: IAssetRepository,
     @Inject(IJobRepository) private jobRepository: IJobRepository,
     @Inject(ISmartInfoRepository) private repository: ISmartInfoRepository,
     @Inject(IMachineLearningRepository) private machineLearning: IMachineLearningRepository,
   ) {}
 
-  async handleTagImage(data: IAssetJob) {
-    const { asset } = data;
-
-    if (!MACHINE_LEARNING_ENABLED || !asset.resizePath) {
-      return;
-    }
-
+  async handleQueueObjectTagging({ force }: IBaseJob) {
     try {
-      const tags = await this.machineLearning.tagImage({ thumbnailPath: asset.resizePath });
-      if (tags.length > 0) {
-        await this.repository.upsert({ assetId: asset.id, tags });
-        await this.jobRepository.queue({ name: JobName.SEARCH_INDEX_ASSET, data: { ids: [asset.id] } });
+      const assets = force
+        ? await this.assetRepository.getAll()
+        : await this.assetRepository.getWithout(WithoutProperty.OBJECT_TAGS);
+
+      for (const asset of assets) {
+        await this.jobRepository.queue({ name: JobName.CLASSIFY_IMAGE, data: { asset } });
+        await this.jobRepository.queue({ name: JobName.DETECT_OBJECTS, data: { asset } });
       }
     } catch (error: any) {
-      this.logger.error(`Unable to run image tagging pipeline: ${asset.id}`, error?.stack);
+      this.logger.error(`Unable to queue object tagging`, error?.stack);
     }
   }
 
@@ -50,6 +49,38 @@ export class SmartInfoService {
     }
   }
 
+  async handleClassifyImage(data: IAssetJob) {
+    const { asset } = data;
+
+    if (!MACHINE_LEARNING_ENABLED || !asset.resizePath) {
+      return;
+    }
+
+    try {
+      const tags = await this.machineLearning.classifyImage({ thumbnailPath: asset.resizePath });
+      if (tags.length > 0) {
+        await this.repository.upsert({ assetId: asset.id, tags });
+        await this.jobRepository.queue({ name: JobName.SEARCH_INDEX_ASSET, data: { ids: [asset.id] } });
+      }
+    } catch (error: any) {
+      this.logger.error(`Unable to run image tagging pipeline: ${asset.id}`, error?.stack);
+    }
+  }
+
+  async handleQueueEncodeClip({ force }: IBaseJob) {
+    try {
+      const assets = force
+        ? await this.assetRepository.getAll()
+        : await this.assetRepository.getWithout(WithoutProperty.CLIP_ENCODING);
+
+      for (const asset of assets) {
+        await this.jobRepository.queue({ name: JobName.ENCODE_CLIP, data: { asset } });
+      }
+    } catch (error: any) {
+      this.logger.error(`Unable to queue clip encoding`, error?.stack);
+    }
+  }
+
   async handleEncodeClip(data: IAssetJob) {
     const { asset } = data;