KevinMidboe/immich
1
0
Fork 0
mirror of https://github.com/KevinMidboe/immich.git synced 2025-10-29 17:40:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(web,server)!: configure machine learning via the UI (#3768)

Browse Source
This commit is contained in:
Jason Rasmussen
2023-08-25 00:15:03 -04:00
committed by GitHub
parent 2cccef174a
commit 8211afb726
52 changed files with 831 additions and 649 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
server/src/domain/search/response-dto/index.ts
View File

@@ -1,3 +1,2 @@
export * from './search-config-response.dto';
export * from './search-explore.response.dto';
export * from './search-response.dto';

3
server/src/domain/search/response-dto/search-config-response.dto.ts
View File

@@ -1,3 +0,0 @@
export class SearchConfigResponseDto {
enabled!: boolean;
}

94
server/src/domain/search/search.service.spec.ts
View File

@@ -1,5 +1,3 @@
import { BadRequestException } from '@nestjs/common';
import { ConfigService } from '@nestjs/config';
import {
albumStub,
assetStub,
@@ -12,12 +10,14 @@ import {
newJobRepositoryMock,
newMachineLearningRepositoryMock,
newSearchRepositoryMock,
newSystemConfigRepositoryMock,
searchStub,
} from '@test';
import { plainToInstance } from 'class-transformer';
import { IAlbumRepository } from '../album/album.repository';
import { IAssetRepository } from '../asset/asset.repository';
import { IFaceRepository } from '../facial-recognition';
import { ISystemConfigRepository } from '../index';
import { JobName } from '../job';
import { IJobRepository } from '../job/job.repository';
import { IMachineLearningRepository } from '../smart-info';
@@ -31,29 +31,26 @@ describe(SearchService.name, () => {
let sut: SearchService;
let albumMock: jest.Mocked<IAlbumRepository>;
let assetMock: jest.Mocked<IAssetRepository>;
let configMock: jest.Mocked<ISystemConfigRepository>;
let faceMock: jest.Mocked<IFaceRepository>;
let jobMock: jest.Mocked<IJobRepository>;
let machineMock: jest.Mocked<IMachineLearningRepository>;
let searchMock: jest.Mocked<ISearchRepository>;
let configMock: jest.Mocked<ConfigService>;
const makeSut = (value?: string) => {
if (value) {
configMock.get.mockReturnValue(value);
}
return new SearchService(albumMock, assetMock, faceMock, jobMock, machineMock, searchMock, configMock);
};
beforeEach(() => {
beforeEach(async () => {
albumMock = newAlbumRepositoryMock();
assetMock = newAssetRepositoryMock();
configMock = newSystemConfigRepositoryMock();
faceMock = newFaceRepositoryMock();
jobMock = newJobRepositoryMock();
machineMock = newMachineLearningRepositoryMock();
searchMock = newSearchRepositoryMock();
configMock = { get: jest.fn() } as unknown as jest.Mocked<ConfigService>;
sut = makeSut();
sut = new SearchService(albumMock, assetMock, configMock, faceMock, jobMock, machineMock, searchMock);
searchMock.checkMigrationStatus.mockResolvedValue({ assets: false, albums: false, faces: false });
await sut.init();
});
afterEach(() => {
@@ -86,45 +83,18 @@ describe(SearchService.name, () => {
});
});
describe('isEnabled', () => {
it('should be enabled by default', () => {
expect(sut.isEnabled()).toBe(true);
});
it('should be disabled via an env variable', () => {
const sut = makeSut('false');
expect(sut.isEnabled()).toBe(false);
});
});
describe('getConfig', () => {
it('should return the config', () => {
expect(sut.getConfig()).toEqual({ enabled: true });
});
it('should return the config when search is disabled', () => {
const sut = makeSut('false');
expect(sut.getConfig()).toEqual({ enabled: false });
});
});
describe(`init`, () => {
it('should skip when search is disabled', async () => {
const sut = makeSut('false');
// it('should skip when search is disabled', async () => {
// await sut.init();
await sut.init();
// expect(searchMock.setup).not.toHaveBeenCalled();
// expect(searchMock.checkMigrationStatus).not.toHaveBeenCalled();
// expect(jobMock.queue).not.toHaveBeenCalled();
expect(searchMock.setup).not.toHaveBeenCalled();
expect(searchMock.checkMigrationStatus).not.toHaveBeenCalled();
expect(jobMock.queue).not.toHaveBeenCalled();
sut.teardown();
});
// sut.teardown();
// });
it('should skip schema migration if not needed', async () => {
searchMock.checkMigrationStatus.mockResolvedValue({ assets: false, albums: false, faces: false });
await sut.init();
expect(searchMock.setup).toHaveBeenCalled();
@@ -145,14 +115,14 @@ describe(SearchService.name, () => {
});
describe('search', () => {
it('should throw an error is search is disabled', async () => {
const sut = makeSut('false');
// it('should throw an error is search is disabled', async () => {
// sut['enabled'] = false;
await expect(sut.search(authStub.admin, {})).rejects.toBeInstanceOf(BadRequestException);
// await expect(sut.search(authStub.admin, {})).rejects.toBeInstanceOf(BadRequestException);
expect(searchMock.searchAlbums).not.toHaveBeenCalled();
expect(searchMock.searchAssets).not.toHaveBeenCalled();
});
// expect(searchMock.searchAlbums).not.toHaveBeenCalled();
// expect(searchMock.searchAssets).not.toHaveBeenCalled();
// });
it('should search assets and albums', async () => {
searchMock.searchAssets.mockResolvedValue(searchStub.emptyResults);
@@ -205,7 +175,7 @@ describe(SearchService.name, () => {
});
it('should skip if search is disabled', async () => {
const sut = makeSut('false');
sut['enabled'] = false;
await sut.handleIndexAssets();
@@ -216,7 +186,7 @@ describe(SearchService.name, () => {
describe('handleIndexAsset', () => {
it('should skip if search is disabled', () => {
const sut = makeSut('false');
sut['enabled'] = false;
sut.handleIndexAsset({ ids: [assetStub.image.id] });
});
@@ -227,7 +197,7 @@ describe(SearchService.name, () => {
describe('handleIndexAlbums', () => {
it('should skip if search is disabled', () => {
const sut = makeSut('false');
sut['enabled'] = false;
sut.handleIndexAlbums();
});
@@ -242,7 +212,7 @@ describe(SearchService.name, () => {
describe('handleIndexAlbum', () => {
it('should skip if search is disabled', () => {
const sut = makeSut('false');
sut['enabled'] = false;
sut.handleIndexAlbum({ ids: [albumStub.empty.id] });
});
@@ -253,7 +223,7 @@ describe(SearchService.name, () => {
describe('handleRemoveAlbum', () => {
it('should skip if search is disabled', () => {
const sut = makeSut('false');
sut['enabled'] = false;
sut.handleRemoveAlbum({ ids: ['album1'] });
});
@@ -264,7 +234,7 @@ describe(SearchService.name, () => {
describe('handleRemoveAsset', () => {
it('should skip if search is disabled', () => {
const sut = makeSut('false');
sut['enabled'] = false;
sut.handleRemoveAsset({ ids: ['asset1'] });
});
@@ -305,7 +275,7 @@ describe(SearchService.name, () => {
});
it('should skip if search is disabled', async () => {
const sut = makeSut('false');
sut['enabled'] = false;
await sut.handleIndexFaces();
@@ -315,7 +285,7 @@ describe(SearchService.name, () => {
describe('handleIndexAsset', () => {
it('should skip if search is disabled', () => {
const sut = makeSut('false');
sut['enabled'] = false;
sut.handleIndexFace({ assetId: 'asset-1', personId: 'person-1' });
expect(searchMock.importFaces).not.toHaveBeenCalled();
@@ -333,7 +303,7 @@ describe(SearchService.name, () => {
describe('handleRemoveFace', () => {
it('should skip if search is disabled', () => {
const sut = makeSut('false');
sut['enabled'] = false;
sut.handleRemoveFace({ assetId: 'asset-1', personId: 'person-1' });
});

47
server/src/domain/search/search.service.ts
View File

@@ -1,18 +1,17 @@
import { AlbumEntity, AssetEntity, AssetFaceEntity } from '@app/infra/entities';
import { BadRequestException, Inject, Injectable, Logger } from '@nestjs/common';
import { ConfigService } from '@nestjs/config';
import { Inject, Injectable, Logger } from '@nestjs/common';
import { mapAlbumWithAssets } from '../album';
import { IAlbumRepository } from '../album/album.repository';
import { AssetResponseDto, mapAsset } from '../asset';
import { IAssetRepository } from '../asset/asset.repository';
import { AuthUserDto } from '../auth';
import { MACHINE_LEARNING_ENABLED } from '../domain.constant';
import { usePagination } from '../domain.util';
import { AssetFaceId, IFaceRepository } from '../facial-recognition';
import { IAssetFaceJob, IBulkEntityJob, IJobRepository, JobName, JOBS_ASSET_PAGINATION_SIZE } from '../job';
import { IMachineLearningRepository } from '../smart-info';
import { FeatureFlag, ISystemConfigRepository, SystemConfigCore } from '../system-config';
import { SearchDto } from './dto';
import { SearchConfigResponseDto, SearchResponseDto } from './response-dto';
import { SearchResponseDto } from './response-dto';
import {
ISearchRepository,
OwnedFaceEntity,
@@ -30,8 +29,9 @@ interface SyncQueue {
@Injectable()
export class SearchService {
private logger = new Logger(SearchService.name);
private enabled: boolean;
private enabled = false;
private timer: NodeJS.Timer | null = null;
private configCore: SystemConfigCore;
private albumQueue: SyncQueue = {
upsert: new Set(),
@@ -51,16 +51,13 @@ export class SearchService {
constructor(
@Inject(IAlbumRepository) private albumRepository: IAlbumRepository,
@Inject(IAssetRepository) private assetRepository: IAssetRepository,
@Inject(ISystemConfigRepository) configRepository: ISystemConfigRepository,
@Inject(IFaceRepository) private faceRepository: IFaceRepository,
@Inject(IJobRepository) private jobRepository: IJobRepository,
@Inject(IMachineLearningRepository) private machineLearning: IMachineLearningRepository,
@Inject(ISearchRepository) private searchRepository: ISearchRepository,
configService: ConfigService,
) {
this.enabled = configService.get('TYPESENSE_ENABLED') !== 'false';
if (this.enabled) {
this.timer = setInterval(() => this.flush(), 5_000);
}
this.configCore = new SystemConfigCore(configRepository);
}
teardown() {
@@ -70,17 +67,8 @@ export class SearchService {
}
}
isEnabled() {
return this.enabled;
}
getConfig(): SearchConfigResponseDto {
return {
enabled: this.enabled,
};
}
async init() {
this.enabled = await this.configCore.hasFeature(FeatureFlag.SEARCH);
if (!this.enabled) {
return;
}
@@ -101,10 +89,13 @@ export class SearchService {
this.logger.debug('Queueing job to re-index all faces');
await this.jobRepository.queue({ name: JobName.SEARCH_INDEX_FACES });
}
this.timer = setInterval(() => this.flush(), 5_000);
}
async getExploreData(authUser: AuthUserDto): Promise<SearchExploreItem<AssetResponseDto>[]> {
this.assertEnabled();
await this.configCore.requireFeature(FeatureFlag.SEARCH);
const results = await this.searchRepository.explore(authUser.id);
const lookup = await this.getLookupMap(
results.reduce(
@@ -126,16 +117,18 @@ export class SearchService {
}
async search(authUser: AuthUserDto, dto: SearchDto): Promise<SearchResponseDto> {
this.assertEnabled();
const { machineLearning } = await this.configCore.getConfig();
await this.configCore.requireFeature(FeatureFlag.SEARCH);
const query = dto.q || dto.query || '*';
const strategy = dto.clip && MACHINE_LEARNING_ENABLED ? SearchStrategy.CLIP : SearchStrategy.TEXT;
const hasClip = machineLearning.enabled && machineLearning.clipEncodeEnabled;
const strategy = dto.clip && hasClip ? SearchStrategy.CLIP : SearchStrategy.TEXT;
const filters = { userId: authUser.id, ...dto };
let assets: SearchResult<AssetEntity>;
switch (strategy) {
case SearchStrategy.CLIP:
const clip = await this.machineLearning.encodeText(query);
const clip = await this.machineLearning.encodeText(machineLearning.url, query);
assets = await this.searchRepository.vectorSearch(clip, filters);
break;
case SearchStrategy.TEXT:
@@ -333,12 +326,6 @@ export class SearchService {
}
}
private assertEnabled() {
if (!this.enabled) {
throw new BadRequestException('Search is disabled');
}
}
private async idsToAlbums(ids: string[]): Promise<AlbumEntity[]> {
const entities = await this.albumRepository.getByIds(ids);
return this.patchAlbums(entities);