feat: use pgvecto.rs (#3605)
This commit is contained in:
Jason Rasmussen
@@ -1,6 +1,6 @@
|
||||
import { AssetType } from '@app/infra/entities';
|
||||
import { Transform } from 'class-transformer';
|
||||
import { IsArray, IsBoolean, IsEnum, IsNotEmpty, IsString } from 'class-validator';
|
||||
import { IsBoolean, IsEnum, IsNotEmpty, IsString } from 'class-validator';
|
||||
import { Optional, toBoolean } from '../../domain.util';
|
||||
|
||||
export class SearchDto {
|
||||
@@ -23,58 +23,6 @@ export class SearchDto {
|
||||
@Optional()
|
||||
type?: AssetType;
|
||||
|
||||
@IsBoolean()
|
||||
@Optional()
|
||||
@Transform(toBoolean)
|
||||
isFavorite?: boolean;
|
||||
|
||||
@IsBoolean()
|
||||
@Optional()
|
||||
@Transform(toBoolean)
|
||||
isArchived?: boolean;
|
||||
|
||||
@IsString()
|
||||
@IsNotEmpty()
|
||||
@Optional()
|
||||
'exifInfo.city'?: string;
|
||||
|
||||
@IsString()
|
||||
@IsNotEmpty()
|
||||
@Optional()
|
||||
'exifInfo.state'?: string;
|
||||
|
||||
@IsString()
|
||||
@IsNotEmpty()
|
||||
@Optional()
|
||||
'exifInfo.country'?: string;
|
||||
|
||||
@IsString()
|
||||
@IsNotEmpty()
|
||||
@Optional()
|
||||
'exifInfo.make'?: string;
|
||||
|
||||
@IsString()
|
||||
@IsNotEmpty()
|
||||
@Optional()
|
||||
'exifInfo.model'?: string;
|
||||
|
||||
@IsString()
|
||||
@IsNotEmpty()
|
||||
@Optional()
|
||||
'exifInfo.projectionType'?: string;
|
||||
|
||||
@IsString({ each: true })
|
||||
@IsArray()
|
||||
@Optional()
|
||||
@Transform(({ value }) => value.split(','))
|
||||
'smartInfo.objects'?: string[];
|
||||
|
||||
@IsString({ each: true })
|
||||
@IsArray()
|
||||
@Optional()
|
||||
@Transform(({ value }) => value.split(','))
|
||||
'smartInfo.tags'?: string[];
|
||||
|
||||
@IsBoolean()
|
||||
@Optional()
|
||||
@Transform(toBoolean)
|
||||
|
||||
@@ -1,29 +1,20 @@
|
||||
import { BadRequestException } from '@nestjs/common';
|
||||
import { SystemConfigKey } from '@app/infra/entities';
|
||||
import {
|
||||
albumStub,
|
||||
assetStub,
|
||||
asyncTick,
|
||||
authStub,
|
||||
faceStub,
|
||||
newAlbumRepositoryMock,
|
||||
newAssetRepositoryMock,
|
||||
newJobRepositoryMock,
|
||||
newMachineLearningRepositoryMock,
|
||||
newPersonRepositoryMock,
|
||||
newSearchRepositoryMock,
|
||||
newSmartInfoRepositoryMock,
|
||||
newSystemConfigRepositoryMock,
|
||||
searchStub,
|
||||
personStub,
|
||||
} from '@test';
|
||||
import { plainToInstance } from 'class-transformer';
|
||||
import { mapAsset } from '../asset';
|
||||
import { JobName } from '../job';
|
||||
import {
|
||||
IAlbumRepository,
|
||||
IAssetRepository,
|
||||
IJobRepository,
|
||||
IMachineLearningRepository,
|
||||
IPersonRepository,
|
||||
ISearchRepository,
|
||||
ISmartInfoRepository,
|
||||
ISystemConfigRepository,
|
||||
} from '../repositories';
|
||||
import { SearchDto } from './dto';
|
||||
@@ -33,401 +24,126 @@ jest.useFakeTimers();
|
||||
|
||||
describe(SearchService.name, () => {
|
||||
let sut: SearchService;
|
||||
let albumMock: jest.Mocked<IAlbumRepository>;
|
||||
let assetMock: jest.Mocked<IAssetRepository>;
|
||||
let configMock: jest.Mocked<ISystemConfigRepository>;
|
||||
let jobMock: jest.Mocked<IJobRepository>;
|
||||
let personMock: jest.Mocked<IPersonRepository>;
|
||||
let machineMock: jest.Mocked<IMachineLearningRepository>;
|
||||
let searchMock: jest.Mocked<ISearchRepository>;
|
||||
let personMock: jest.Mocked<IPersonRepository>;
|
||||
let smartInfoMock: jest.Mocked<ISmartInfoRepository>;
|
||||
|
||||
beforeEach(async () => {
|
||||
albumMock = newAlbumRepositoryMock();
|
||||
beforeEach(() => {
|
||||
assetMock = newAssetRepositoryMock();
|
||||
configMock = newSystemConfigRepositoryMock();
|
||||
jobMock = newJobRepositoryMock();
|
||||
personMock = newPersonRepositoryMock();
|
||||
machineMock = newMachineLearningRepositoryMock();
|
||||
searchMock = newSearchRepositoryMock();
|
||||
|
||||
sut = new SearchService(albumMock, assetMock, jobMock, machineMock, personMock, searchMock, configMock);
|
||||
|
||||
searchMock.checkMigrationStatus.mockResolvedValue({ assets: false, albums: false, faces: false });
|
||||
|
||||
delete process.env.TYPESENSE_ENABLED;
|
||||
await sut.init();
|
||||
});
|
||||
|
||||
const disableSearch = () => {
|
||||
searchMock.setup.mockClear();
|
||||
searchMock.checkMigrationStatus.mockClear();
|
||||
jobMock.queue.mockClear();
|
||||
process.env.TYPESENSE_ENABLED = 'false';
|
||||
};
|
||||
|
||||
afterEach(() => {
|
||||
sut.teardown();
|
||||
personMock = newPersonRepositoryMock();
|
||||
smartInfoMock = newSmartInfoRepositoryMock();
|
||||
sut = new SearchService(configMock, machineMock, personMock, smartInfoMock, assetMock);
|
||||
});
|
||||
|
||||
it('should work', () => {
|
||||
expect(sut).toBeDefined();
|
||||
});
|
||||
|
||||
describe('request dto', () => {
|
||||
it('should convert smartInfo.tags to a string list', () => {
|
||||
const instance = plainToInstance(SearchDto, { 'smartInfo.tags': 'a,b,c' });
|
||||
expect(instance['smartInfo.tags']).toEqual(['a', 'b', 'c']);
|
||||
});
|
||||
describe('searchPerson', () => {
|
||||
it('should pass options to search', async () => {
|
||||
const { name } = personStub.withName;
|
||||
|
||||
it('should handle empty smartInfo.tags', () => {
|
||||
const instance = plainToInstance(SearchDto, {});
|
||||
expect(instance['smartInfo.tags']).toBeUndefined();
|
||||
});
|
||||
await sut.searchPerson(authStub.user1, { name, withHidden: false });
|
||||
|
||||
it('should convert smartInfo.objects to a string list', () => {
|
||||
const instance = plainToInstance(SearchDto, { 'smartInfo.objects': 'a,b,c' });
|
||||
expect(instance['smartInfo.objects']).toEqual(['a', 'b', 'c']);
|
||||
});
|
||||
expect(personMock.getByName).toHaveBeenCalledWith(authStub.user1.id, name, { withHidden: false });
|
||||
|
||||
it('should handle empty smartInfo.objects', () => {
|
||||
const instance = plainToInstance(SearchDto, {});
|
||||
expect(instance['smartInfo.objects']).toBeUndefined();
|
||||
});
|
||||
});
|
||||
await sut.searchPerson(authStub.user1, { name, withHidden: true });
|
||||
|
||||
describe(`init`, () => {
|
||||
it('should skip when search is disabled', async () => {
|
||||
disableSearch();
|
||||
await sut.init();
|
||||
|
||||
expect(searchMock.setup).not.toHaveBeenCalled();
|
||||
expect(searchMock.checkMigrationStatus).not.toHaveBeenCalled();
|
||||
expect(jobMock.queue).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should skip schema migration if not needed', async () => {
|
||||
await sut.init();
|
||||
|
||||
expect(searchMock.setup).toHaveBeenCalled();
|
||||
expect(jobMock.queue).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should do schema migration if needed', async () => {
|
||||
searchMock.checkMigrationStatus.mockResolvedValue({ assets: true, albums: true, faces: true });
|
||||
await sut.init();
|
||||
|
||||
expect(searchMock.setup).toHaveBeenCalled();
|
||||
expect(jobMock.queue.mock.calls).toEqual([
|
||||
[{ name: JobName.SEARCH_INDEX_ASSETS }],
|
||||
[{ name: JobName.SEARCH_INDEX_ALBUMS }],
|
||||
[{ name: JobName.SEARCH_INDEX_FACES }],
|
||||
]);
|
||||
expect(personMock.getByName).toHaveBeenCalledWith(authStub.user1.id, name, { withHidden: true });
|
||||
});
|
||||
});
|
||||
|
||||
describe('getExploreData', () => {
|
||||
it('should throw bad request exception if search is disabled', async () => {
|
||||
disableSearch();
|
||||
await expect(sut.getExploreData(authStub.admin)).rejects.toBeInstanceOf(BadRequestException);
|
||||
expect(searchMock.explore).not.toHaveBeenCalled();
|
||||
});
|
||||
it('should get assets by city and tag', async () => {
|
||||
assetMock.getAssetIdByCity.mockResolvedValueOnce({
|
||||
fieldName: 'exifInfo.city',
|
||||
items: [{ value: 'Paris', data: assetStub.image.id }],
|
||||
});
|
||||
assetMock.getAssetIdByTag.mockResolvedValueOnce({
|
||||
fieldName: 'smartInfo.tags',
|
||||
items: [{ value: 'train', data: assetStub.imageFrom2015.id }],
|
||||
});
|
||||
assetMock.getByIds.mockResolvedValueOnce([assetStub.image, assetStub.imageFrom2015]);
|
||||
const expectedResponse = [
|
||||
{ fieldName: 'exifInfo.city', items: [{ value: 'Paris', data: mapAsset(assetStub.image) }] },
|
||||
{ fieldName: 'smartInfo.tags', items: [{ value: 'train', data: mapAsset(assetStub.imageFrom2015) }] },
|
||||
];
|
||||
|
||||
it('should return explore data if feature flag SEARCH is set', async () => {
|
||||
searchMock.explore.mockResolvedValue([{ fieldName: 'name', items: [{ value: 'image', data: assetStub.image }] }]);
|
||||
assetMock.getByIds.mockResolvedValue([assetStub.image]);
|
||||
const result = await sut.getExploreData(authStub.user1);
|
||||
|
||||
await expect(sut.getExploreData(authStub.admin)).resolves.toEqual([
|
||||
{
|
||||
fieldName: 'name',
|
||||
items: [{ value: 'image', data: mapAsset(assetStub.image) }],
|
||||
},
|
||||
]);
|
||||
|
||||
expect(searchMock.explore).toHaveBeenCalledWith(authStub.admin.id);
|
||||
expect(assetMock.getByIds).toHaveBeenCalledWith([assetStub.image.id]);
|
||||
expect(result).toEqual(expectedResponse);
|
||||
});
|
||||
});
|
||||
|
||||
describe('search', () => {
|
||||
// it('should throw an error is search is disabled', async () => {
|
||||
// sut['enabled'] = false;
|
||||
it('should throw an error if query is missing', async () => {
|
||||
await expect(sut.search(authStub.user1, { q: '' })).rejects.toThrow('Missing query');
|
||||
});
|
||||
|
||||
// await expect(sut.search(authStub.admin, {})).rejects.toBeInstanceOf(BadRequestException);
|
||||
|
||||
// expect(searchMock.searchAlbums).not.toHaveBeenCalled();
|
||||
// expect(searchMock.searchAssets).not.toHaveBeenCalled();
|
||||
// });
|
||||
|
||||
it('should search assets and albums using text search', async () => {
|
||||
searchMock.searchAssets.mockResolvedValue(searchStub.withImage);
|
||||
searchMock.searchAlbums.mockResolvedValue(searchStub.emptyResults);
|
||||
assetMock.getByIds.mockResolvedValue([assetStub.image]);
|
||||
|
||||
await expect(sut.search(authStub.admin, {})).resolves.toEqual({
|
||||
it('should search by metadata if `clip` option is false', async () => {
|
||||
const dto: SearchDto = { q: 'test query', clip: false };
|
||||
assetMock.searchMetadata.mockResolvedValueOnce([assetStub.image]);
|
||||
const expectedResponse = {
|
||||
albums: {
|
||||
total: 0,
|
||||
count: 0,
|
||||
page: 1,
|
||||
items: [],
|
||||
facets: [],
|
||||
distances: [],
|
||||
},
|
||||
assets: {
|
||||
total: 1,
|
||||
count: 1,
|
||||
page: 1,
|
||||
items: [mapAsset(assetStub.image)],
|
||||
facets: [],
|
||||
distances: [],
|
||||
},
|
||||
});
|
||||
};
|
||||
|
||||
// expect(searchMock.searchAssets).toHaveBeenCalledWith('*', { userId: authStub.admin.id });
|
||||
expect(searchMock.searchAlbums).toHaveBeenCalledWith('*', { userId: authStub.admin.id });
|
||||
const result = await sut.search(authStub.user1, dto);
|
||||
|
||||
expect(result).toEqual(expectedResponse);
|
||||
expect(assetMock.searchMetadata).toHaveBeenCalledWith(dto.q, authStub.user1.id, { numResults: 250 });
|
||||
expect(smartInfoMock.searchCLIP).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should search assets and albums using vector search', async () => {
|
||||
searchMock.vectorSearch.mockResolvedValue(searchStub.emptyResults);
|
||||
searchMock.searchAlbums.mockResolvedValue(searchStub.emptyResults);
|
||||
machineMock.encodeText.mockResolvedValue([123]);
|
||||
|
||||
await expect(sut.search(authStub.admin, { clip: true, query: 'foo' })).resolves.toEqual({
|
||||
it('should search by CLIP if `clip` option is true', async () => {
|
||||
const dto: SearchDto = { q: 'test query', clip: true };
|
||||
const embedding = [1, 2, 3];
|
||||
smartInfoMock.searchCLIP.mockResolvedValueOnce([assetStub.image]);
|
||||
machineMock.encodeText.mockResolvedValueOnce(embedding);
|
||||
const expectedResponse = {
|
||||
albums: {
|
||||
total: 0,
|
||||
count: 0,
|
||||
page: 1,
|
||||
items: [],
|
||||
facets: [],
|
||||
distances: [],
|
||||
},
|
||||
assets: {
|
||||
total: 0,
|
||||
count: 0,
|
||||
page: 1,
|
||||
items: [],
|
||||
total: 1,
|
||||
count: 1,
|
||||
items: [mapAsset(assetStub.image)],
|
||||
facets: [],
|
||||
distances: [],
|
||||
},
|
||||
});
|
||||
};
|
||||
|
||||
expect(machineMock.encodeText).toHaveBeenCalledWith(expect.any(String), { text: 'foo' }, expect.any(Object));
|
||||
expect(searchMock.vectorSearch).toHaveBeenCalledWith([123], {
|
||||
userId: authStub.admin.id,
|
||||
clip: true,
|
||||
query: 'foo',
|
||||
});
|
||||
expect(searchMock.searchAlbums).toHaveBeenCalledWith('foo', {
|
||||
userId: authStub.admin.id,
|
||||
clip: true,
|
||||
query: 'foo',
|
||||
});
|
||||
});
|
||||
});
|
||||
const result = await sut.search(authStub.user1, dto);
|
||||
|
||||
describe('handleIndexAssets', () => {
|
||||
it('should call done, even when there are no assets', async () => {
|
||||
await sut.handleIndexAssets();
|
||||
|
||||
expect(searchMock.importAssets).toHaveBeenCalledWith([], true);
|
||||
expect(result).toEqual(expectedResponse);
|
||||
expect(smartInfoMock.searchCLIP).toHaveBeenCalledWith({ ownerId: authStub.user1.id, embedding, numResults: 100 });
|
||||
expect(assetMock.searchMetadata).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should index all the assets', async () => {
|
||||
assetMock.getAll.mockResolvedValue({
|
||||
items: [assetStub.image],
|
||||
hasNextPage: false,
|
||||
});
|
||||
it('should throw an error if clip is requested but disabled', async () => {
|
||||
const dto: SearchDto = { q: 'test query', clip: true };
|
||||
configMock.load
|
||||
.mockResolvedValueOnce([{ key: SystemConfigKey.MACHINE_LEARNING_ENABLED, value: false }])
|
||||
.mockResolvedValueOnce([{ key: SystemConfigKey.MACHINE_LEARNING_CLIP_ENABLED, value: false }]);
|
||||
|
||||
await sut.handleIndexAssets();
|
||||
|
||||
expect(searchMock.importAssets.mock.calls).toEqual([
|
||||
[[assetStub.image], false],
|
||||
[[], true],
|
||||
]);
|
||||
});
|
||||
|
||||
it('should skip if search is disabled', async () => {
|
||||
sut['enabled'] = false;
|
||||
|
||||
await sut.handleIndexAssets();
|
||||
|
||||
expect(searchMock.importAssets).not.toHaveBeenCalled();
|
||||
expect(searchMock.importAlbums).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('handleIndexAsset', () => {
|
||||
it('should skip if search is disabled', () => {
|
||||
sut['enabled'] = false;
|
||||
sut.handleIndexAsset({ ids: [assetStub.image.id] });
|
||||
});
|
||||
|
||||
it('should index the asset', () => {
|
||||
sut.handleIndexAsset({ ids: [assetStub.image.id] });
|
||||
});
|
||||
});
|
||||
|
||||
describe('handleIndexAlbums', () => {
|
||||
it('should skip if search is disabled', async () => {
|
||||
sut['enabled'] = false;
|
||||
await sut.handleIndexAlbums();
|
||||
});
|
||||
|
||||
it('should index all the albums', async () => {
|
||||
albumMock.getAll.mockResolvedValue([albumStub.empty]);
|
||||
|
||||
await sut.handleIndexAlbums();
|
||||
|
||||
expect(searchMock.importAlbums).toHaveBeenCalledWith([albumStub.empty], true);
|
||||
});
|
||||
});
|
||||
|
||||
describe('handleIndexAlbum', () => {
|
||||
it('should skip if search is disabled', () => {
|
||||
sut['enabled'] = false;
|
||||
sut.handleIndexAlbum({ ids: [albumStub.empty.id] });
|
||||
});
|
||||
|
||||
it('should index the album', () => {
|
||||
sut.handleIndexAlbum({ ids: [albumStub.empty.id] });
|
||||
});
|
||||
});
|
||||
|
||||
describe('handleRemoveAlbum', () => {
|
||||
it('should skip if search is disabled', () => {
|
||||
sut['enabled'] = false;
|
||||
sut.handleRemoveAlbum({ ids: ['album1'] });
|
||||
});
|
||||
|
||||
it('should remove the album', () => {
|
||||
sut.handleRemoveAlbum({ ids: ['album1'] });
|
||||
});
|
||||
});
|
||||
|
||||
describe('handleRemoveAsset', () => {
|
||||
it('should skip if search is disabled', () => {
|
||||
sut['enabled'] = false;
|
||||
sut.handleRemoveAsset({ ids: ['asset1'] });
|
||||
});
|
||||
|
||||
it('should remove the asset', () => {
|
||||
sut.handleRemoveAsset({ ids: ['asset1'] });
|
||||
});
|
||||
});
|
||||
|
||||
describe('handleIndexFaces', () => {
|
||||
it('should call done, even when there are no faces', async () => {
|
||||
personMock.getAllFaces.mockResolvedValue([]);
|
||||
|
||||
await sut.handleIndexFaces();
|
||||
|
||||
expect(searchMock.importFaces).toHaveBeenCalledWith([], true);
|
||||
});
|
||||
|
||||
it('should index all the faces', async () => {
|
||||
personMock.getAllFaces.mockResolvedValue([faceStub.face1]);
|
||||
|
||||
await sut.handleIndexFaces();
|
||||
|
||||
expect(searchMock.importFaces.mock.calls).toEqual([
|
||||
[
|
||||
[
|
||||
{
|
||||
id: 'asset-id|person-1',
|
||||
ownerId: 'user-id',
|
||||
assetId: 'asset-id',
|
||||
personId: 'person-1',
|
||||
embedding: [1, 2, 3, 4],
|
||||
},
|
||||
],
|
||||
false,
|
||||
],
|
||||
[[], true],
|
||||
]);
|
||||
});
|
||||
|
||||
it('should skip if search is disabled', async () => {
|
||||
sut['enabled'] = false;
|
||||
|
||||
await sut.handleIndexFaces();
|
||||
|
||||
expect(searchMock.importFaces).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('handleIndexAsset', () => {
|
||||
it('should skip if search is disabled', async () => {
|
||||
sut['enabled'] = false;
|
||||
await sut.handleIndexFace({ assetId: 'asset-1', personId: 'person-1' });
|
||||
|
||||
expect(searchMock.importFaces).not.toHaveBeenCalled();
|
||||
expect(personMock.getFacesByIds).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should index the face', async () => {
|
||||
personMock.getFacesByIds.mockResolvedValue([faceStub.face1]);
|
||||
|
||||
await sut.handleIndexFace({ assetId: 'asset-1', personId: 'person-1' });
|
||||
|
||||
expect(personMock.getFacesByIds).toHaveBeenCalledWith([{ assetId: 'asset-1', personId: 'person-1' }]);
|
||||
});
|
||||
});
|
||||
|
||||
describe('handleRemoveFace', () => {
|
||||
it('should skip if search is disabled', () => {
|
||||
sut['enabled'] = false;
|
||||
sut.handleRemoveFace({ assetId: 'asset-1', personId: 'person-1' });
|
||||
});
|
||||
|
||||
it('should remove the face', () => {
|
||||
sut.handleRemoveFace({ assetId: 'asset-1', personId: 'person-1' });
|
||||
});
|
||||
});
|
||||
|
||||
describe('flush', () => {
|
||||
it('should flush queued album updates', async () => {
|
||||
albumMock.getByIds.mockResolvedValue([albumStub.empty]);
|
||||
|
||||
sut.handleIndexAlbum({ ids: ['album1'] });
|
||||
|
||||
jest.runOnlyPendingTimers();
|
||||
|
||||
await asyncTick(4);
|
||||
|
||||
expect(albumMock.getByIds).toHaveBeenCalledWith(['album1']);
|
||||
expect(searchMock.importAlbums).toHaveBeenCalledWith([albumStub.empty], false);
|
||||
});
|
||||
|
||||
it('should flush queued album deletes', async () => {
|
||||
sut.handleRemoveAlbum({ ids: ['album1'] });
|
||||
|
||||
jest.runOnlyPendingTimers();
|
||||
|
||||
await asyncTick(4);
|
||||
|
||||
expect(searchMock.deleteAlbums).toHaveBeenCalledWith(['album1']);
|
||||
});
|
||||
|
||||
it('should flush queued asset updates', async () => {
|
||||
assetMock.getByIds.mockResolvedValue([assetStub.image]);
|
||||
|
||||
sut.handleIndexAsset({ ids: ['asset1'] });
|
||||
|
||||
jest.runOnlyPendingTimers();
|
||||
|
||||
await asyncTick(4);
|
||||
|
||||
expect(assetMock.getByIds).toHaveBeenCalledWith(['asset1']);
|
||||
expect(searchMock.importAssets).toHaveBeenCalledWith([assetStub.image], false);
|
||||
});
|
||||
|
||||
it('should flush queued asset deletes', async () => {
|
||||
sut.handleRemoveAsset({ ids: ['asset1'] });
|
||||
|
||||
jest.runOnlyPendingTimers();
|
||||
|
||||
await asyncTick(4);
|
||||
|
||||
expect(searchMock.deleteAssets).toHaveBeenCalledWith(['asset1']);
|
||||
await expect(sut.search(authStub.user1, dto)).rejects.toThrow('CLIP is not enabled');
|
||||
await expect(sut.search(authStub.user1, dto)).rejects.toThrow('CLIP is not enabled');
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -1,396 +1,99 @@
|
||||
import { AlbumEntity, AssetEntity, AssetFaceEntity } from '@app/infra/entities';
|
||||
import { AssetEntity } from '@app/infra/entities';
|
||||
import { Inject, Injectable, Logger } from '@nestjs/common';
|
||||
import { mapAlbumWithAssets } from '../album';
|
||||
import { AssetResponseDto, mapAsset } from '../asset';
|
||||
import { AuthUserDto } from '../auth';
|
||||
import { usePagination } from '../domain.util';
|
||||
import { IAssetFaceJob, IBulkEntityJob, JOBS_ASSET_PAGINATION_SIZE, JobName } from '../job';
|
||||
import { PersonResponseDto } from '../person/person.dto';
|
||||
import { PersonResponseDto } from '../person';
|
||||
import {
|
||||
AssetFaceId,
|
||||
IAlbumRepository,
|
||||
IAssetRepository,
|
||||
IJobRepository,
|
||||
IMachineLearningRepository,
|
||||
IPersonRepository,
|
||||
ISearchRepository,
|
||||
ISmartInfoRepository,
|
||||
ISystemConfigRepository,
|
||||
OwnedFaceEntity,
|
||||
SearchCollection,
|
||||
SearchExploreItem,
|
||||
SearchResult,
|
||||
SearchStrategy,
|
||||
} from '../repositories';
|
||||
import { FeatureFlag, SystemConfigCore } from '../system-config';
|
||||
import { SearchDto, SearchPeopleDto } from './dto';
|
||||
import { SearchResponseDto } from './response-dto';
|
||||
|
||||
interface SyncQueue {
|
||||
upsert: Set<string>;
|
||||
delete: Set<string>;
|
||||
}
|
||||
|
||||
@Injectable()
|
||||
export class SearchService {
|
||||
private logger = new Logger(SearchService.name);
|
||||
private enabled = false;
|
||||
private timer: NodeJS.Timeout | null = null;
|
||||
private configCore: SystemConfigCore;
|
||||
|
||||
private albumQueue: SyncQueue = {
|
||||
upsert: new Set(),
|
||||
delete: new Set(),
|
||||
};
|
||||
|
||||
private assetQueue: SyncQueue = {
|
||||
upsert: new Set(),
|
||||
delete: new Set(),
|
||||
};
|
||||
|
||||
private faceQueue: SyncQueue = {
|
||||
upsert: new Set(),
|
||||
delete: new Set(),
|
||||
};
|
||||
|
||||
constructor(
|
||||
@Inject(IAlbumRepository) private albumRepository: IAlbumRepository,
|
||||
@Inject(IAssetRepository) private assetRepository: IAssetRepository,
|
||||
@Inject(IJobRepository) private jobRepository: IJobRepository,
|
||||
@Inject(ISystemConfigRepository) configRepository: ISystemConfigRepository,
|
||||
@Inject(IMachineLearningRepository) private machineLearning: IMachineLearningRepository,
|
||||
@Inject(IPersonRepository) private personRepository: IPersonRepository,
|
||||
@Inject(ISearchRepository) private searchRepository: ISearchRepository,
|
||||
@Inject(ISystemConfigRepository) configRepository: ISystemConfigRepository,
|
||||
@Inject(ISmartInfoRepository) private smartInfoRepository: ISmartInfoRepository,
|
||||
@Inject(IAssetRepository) private assetRepository: IAssetRepository,
|
||||
) {
|
||||
this.configCore = SystemConfigCore.create(configRepository);
|
||||
}
|
||||
|
||||
teardown() {
|
||||
if (this.timer) {
|
||||
clearInterval(this.timer);
|
||||
this.timer = null;
|
||||
}
|
||||
}
|
||||
|
||||
async init() {
|
||||
this.enabled = await this.configCore.hasFeature(FeatureFlag.SEARCH);
|
||||
if (!this.enabled) {
|
||||
return;
|
||||
}
|
||||
|
||||
this.logger.log('Running bootstrap');
|
||||
await this.searchRepository.setup();
|
||||
|
||||
const migrationStatus = await this.searchRepository.checkMigrationStatus();
|
||||
if (migrationStatus[SearchCollection.ASSETS]) {
|
||||
this.logger.debug('Queueing job to re-index all assets');
|
||||
await this.jobRepository.queue({ name: JobName.SEARCH_INDEX_ASSETS });
|
||||
}
|
||||
if (migrationStatus[SearchCollection.ALBUMS]) {
|
||||
this.logger.debug('Queueing job to re-index all albums');
|
||||
await this.jobRepository.queue({ name: JobName.SEARCH_INDEX_ALBUMS });
|
||||
}
|
||||
if (migrationStatus[SearchCollection.FACES]) {
|
||||
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 searchPerson(authUser: AuthUserDto, dto: SearchPeopleDto): Promise<PersonResponseDto[]> {
|
||||
return this.personRepository.getByName(authUser.id, dto.name, { withHidden: dto.withHidden });
|
||||
}
|
||||
|
||||
async getExploreData(authUser: AuthUserDto): Promise<SearchExploreItem<AssetResponseDto>[]> {
|
||||
await this.configCore.requireFeature(FeatureFlag.SEARCH);
|
||||
|
||||
const results = await this.searchRepository.explore(authUser.id);
|
||||
const lookup = await this.getLookupMap(
|
||||
results.reduce(
|
||||
(ids: string[], result: SearchExploreItem<AssetEntity>) => [
|
||||
...ids,
|
||||
...result.items.map((item) => item.data.id),
|
||||
],
|
||||
[],
|
||||
),
|
||||
);
|
||||
const options = { maxFields: 12, minAssetsPerField: 5 };
|
||||
const results = await Promise.all([
|
||||
this.assetRepository.getAssetIdByCity(authUser.id, options),
|
||||
this.assetRepository.getAssetIdByTag(authUser.id, options),
|
||||
]);
|
||||
const assetIds = new Set<string>(results.flatMap((field) => field.items.map((item) => item.data)));
|
||||
const assets = await this.assetRepository.getByIds(Array.from(assetIds));
|
||||
const assetMap = new Map<string, AssetResponseDto>(assets.map((asset) => [asset.id, mapAsset(asset)]));
|
||||
|
||||
return results.map(({ fieldName, items }) => ({
|
||||
fieldName,
|
||||
items: items
|
||||
.map(({ value, data }) => ({ value, data: lookup[data.id] }))
|
||||
.filter(({ data }) => !!data)
|
||||
.map(({ value, data }) => ({ value, data: mapAsset(data) })),
|
||||
items: items.map(({ value, data }) => ({ value, data: assetMap.get(data) as AssetResponseDto })),
|
||||
}));
|
||||
}
|
||||
|
||||
async search(authUser: AuthUserDto, dto: SearchDto): Promise<SearchResponseDto> {
|
||||
const { machineLearning } = await this.configCore.getConfig();
|
||||
await this.configCore.requireFeature(FeatureFlag.SEARCH);
|
||||
|
||||
const query = dto.q || dto.query || '*';
|
||||
const query = dto.q || dto.query;
|
||||
if (!query) {
|
||||
throw new Error('Missing query');
|
||||
}
|
||||
const hasClip = machineLearning.enabled && machineLearning.clip.enabled;
|
||||
const strategy = dto.clip && hasClip ? SearchStrategy.CLIP : SearchStrategy.TEXT;
|
||||
const filters = { userId: authUser.id, ...dto };
|
||||
if (dto.clip && !hasClip) {
|
||||
throw new Error('CLIP is not enabled');
|
||||
}
|
||||
const strategy = dto.clip ? SearchStrategy.CLIP : SearchStrategy.TEXT;
|
||||
|
||||
let assets: AssetEntity[] = [];
|
||||
|
||||
let assets: SearchResult<AssetEntity>;
|
||||
switch (strategy) {
|
||||
case SearchStrategy.CLIP:
|
||||
const {
|
||||
machineLearning: { clip },
|
||||
} = await this.configCore.getConfig();
|
||||
const embedding = await this.machineLearning.encodeText(machineLearning.url, { text: query }, clip);
|
||||
assets = await this.searchRepository.vectorSearch(embedding, filters);
|
||||
const embedding = await this.machineLearning.encodeText(
|
||||
machineLearning.url,
|
||||
{ text: query },
|
||||
machineLearning.clip,
|
||||
);
|
||||
assets = await this.smartInfoRepository.searchCLIP({ ownerId: authUser.id, embedding, numResults: 100 });
|
||||
break;
|
||||
case SearchStrategy.TEXT:
|
||||
assets = await this.assetRepository.searchMetadata(query, authUser.id, { numResults: 250 });
|
||||
default:
|
||||
assets = await this.searchRepository.searchAssets(query, filters);
|
||||
break;
|
||||
}
|
||||
|
||||
const albums = await this.searchRepository.searchAlbums(query, filters);
|
||||
const lookup = await this.getLookupMap(assets.items.map((asset) => asset.id));
|
||||
|
||||
return {
|
||||
albums: { ...albums, items: albums.items.map(mapAlbumWithAssets) },
|
||||
albums: {
|
||||
total: 0,
|
||||
count: 0,
|
||||
items: [],
|
||||
facets: [],
|
||||
},
|
||||
assets: {
|
||||
...assets,
|
||||
items: assets.items
|
||||
.map((item) => lookup[item.id])
|
||||
.filter((item) => !!item)
|
||||
.map((asset) => mapAsset(asset)),
|
||||
total: assets.length,
|
||||
count: assets.length,
|
||||
items: assets.map((asset) => mapAsset(asset)),
|
||||
facets: [],
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
searchPerson(authUser: AuthUserDto, dto: SearchPeopleDto): Promise<PersonResponseDto[]> {
|
||||
return this.personRepository.getByName(authUser.id, dto.name, { withHidden: dto.withHidden });
|
||||
}
|
||||
|
||||
async handleIndexAlbums() {
|
||||
if (!this.enabled) {
|
||||
return false;
|
||||
}
|
||||
|
||||
const albums = this.patchAlbums(await this.albumRepository.getAll());
|
||||
this.logger.log(`Indexing ${albums.length} albums`);
|
||||
await this.searchRepository.importAlbums(albums, true);
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
async handleIndexAssets() {
|
||||
if (!this.enabled) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// TODO: do this in batches based on searchIndexVersion
|
||||
const assetPagination = usePagination(JOBS_ASSET_PAGINATION_SIZE, (pagination) =>
|
||||
this.assetRepository.getAll(pagination, { isVisible: true }),
|
||||
);
|
||||
|
||||
for await (const assets of assetPagination) {
|
||||
this.logger.debug(`Indexing ${assets.length} assets`);
|
||||
|
||||
const patchedAssets = this.patchAssets(assets);
|
||||
await this.searchRepository.importAssets(patchedAssets, false);
|
||||
}
|
||||
|
||||
await this.searchRepository.importAssets([], true);
|
||||
|
||||
this.logger.debug('Finished re-indexing all assets');
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
async handleIndexFaces() {
|
||||
if (!this.enabled) {
|
||||
return false;
|
||||
}
|
||||
await this.searchRepository.deleteAllFaces();
|
||||
|
||||
// TODO: do this in batches based on searchIndexVersion
|
||||
const faces = this.patchFaces(await this.personRepository.getAllFaces());
|
||||
this.logger.log(`Indexing ${faces.length} faces`);
|
||||
|
||||
const chunkSize = 1000;
|
||||
for (let i = 0; i < faces.length; i += chunkSize) {
|
||||
await this.searchRepository.importFaces(faces.slice(i, i + chunkSize), false);
|
||||
}
|
||||
|
||||
await this.searchRepository.importFaces([], true);
|
||||
|
||||
this.logger.debug('Finished re-indexing all faces');
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
handleIndexAlbum({ ids }: IBulkEntityJob) {
|
||||
if (!this.enabled) {
|
||||
return false;
|
||||
}
|
||||
|
||||
for (const id of ids) {
|
||||
this.albumQueue.upsert.add(id);
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
handleIndexAsset({ ids }: IBulkEntityJob) {
|
||||
if (!this.enabled) {
|
||||
return false;
|
||||
}
|
||||
|
||||
for (const id of ids) {
|
||||
this.assetQueue.upsert.add(id);
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
async handleIndexFace({ assetId, personId }: IAssetFaceJob) {
|
||||
if (!this.enabled) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// immediately push to typesense
|
||||
await this.searchRepository.importFaces(await this.idsToFaces([{ assetId, personId }]), false);
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
handleRemoveAlbum({ ids }: IBulkEntityJob) {
|
||||
if (!this.enabled) {
|
||||
return false;
|
||||
}
|
||||
|
||||
for (const id of ids) {
|
||||
this.albumQueue.delete.add(id);
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
handleRemoveAsset({ ids }: IBulkEntityJob) {
|
||||
if (!this.enabled) {
|
||||
return false;
|
||||
}
|
||||
|
||||
for (const id of ids) {
|
||||
this.assetQueue.delete.add(id);
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
handleRemoveFace({ assetId, personId }: IAssetFaceJob) {
|
||||
if (!this.enabled) {
|
||||
return false;
|
||||
}
|
||||
|
||||
this.faceQueue.delete.add(this.asKey({ assetId, personId }));
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
private async flush() {
|
||||
if (this.albumQueue.upsert.size > 0) {
|
||||
const ids = [...this.albumQueue.upsert.keys()];
|
||||
const items = await this.idsToAlbums(ids);
|
||||
this.logger.debug(`Flushing ${items.length} album upserts`);
|
||||
await this.searchRepository.importAlbums(items, false);
|
||||
this.albumQueue.upsert.clear();
|
||||
}
|
||||
|
||||
if (this.albumQueue.delete.size > 0) {
|
||||
const ids = [...this.albumQueue.delete.keys()];
|
||||
this.logger.debug(`Flushing ${ids.length} album deletes`);
|
||||
await this.searchRepository.deleteAlbums(ids);
|
||||
this.albumQueue.delete.clear();
|
||||
}
|
||||
|
||||
if (this.assetQueue.upsert.size > 0) {
|
||||
const ids = [...this.assetQueue.upsert.keys()];
|
||||
const items = await this.idsToAssets(ids);
|
||||
this.logger.debug(`Flushing ${items.length} asset upserts`);
|
||||
await this.searchRepository.importAssets(items, false);
|
||||
this.assetQueue.upsert.clear();
|
||||
}
|
||||
|
||||
if (this.assetQueue.delete.size > 0) {
|
||||
const ids = [...this.assetQueue.delete.keys()];
|
||||
this.logger.debug(`Flushing ${ids.length} asset deletes`);
|
||||
await this.searchRepository.deleteAssets(ids);
|
||||
this.assetQueue.delete.clear();
|
||||
}
|
||||
|
||||
if (this.faceQueue.upsert.size > 0) {
|
||||
const ids = [...this.faceQueue.upsert.keys()].map((key) => this.asParts(key));
|
||||
const items = await this.idsToFaces(ids);
|
||||
this.logger.debug(`Flushing ${items.length} face upserts`);
|
||||
await this.searchRepository.importFaces(items, false);
|
||||
this.faceQueue.upsert.clear();
|
||||
}
|
||||
|
||||
if (this.faceQueue.delete.size > 0) {
|
||||
const ids = [...this.faceQueue.delete.keys()];
|
||||
this.logger.debug(`Flushing ${ids.length} face deletes`);
|
||||
await this.searchRepository.deleteFaces(ids);
|
||||
this.faceQueue.delete.clear();
|
||||
}
|
||||
}
|
||||
|
||||
private async idsToAlbums(ids: string[]): Promise<AlbumEntity[]> {
|
||||
const entities = await this.albumRepository.getByIds(ids);
|
||||
return this.patchAlbums(entities);
|
||||
}
|
||||
|
||||
private async idsToAssets(ids: string[]): Promise<AssetEntity[]> {
|
||||
const entities = await this.assetRepository.getByIds(ids);
|
||||
return this.patchAssets(entities.filter((entity) => entity.isVisible));
|
||||
}
|
||||
|
||||
private async idsToFaces(ids: AssetFaceId[]): Promise<OwnedFaceEntity[]> {
|
||||
return this.patchFaces(await this.personRepository.getFacesByIds(ids));
|
||||
}
|
||||
|
||||
private patchAssets(assets: AssetEntity[]): AssetEntity[] {
|
||||
return assets;
|
||||
}
|
||||
|
||||
private patchAlbums(albums: AlbumEntity[]): AlbumEntity[] {
|
||||
return albums.map((entity) => ({ ...entity, assets: [] }));
|
||||
}
|
||||
|
||||
private patchFaces(faces: AssetFaceEntity[]): OwnedFaceEntity[] {
|
||||
const results: OwnedFaceEntity[] = [];
|
||||
for (const face of faces) {
|
||||
if (face.personId) {
|
||||
results.push({
|
||||
id: this.asKey(face as AssetFaceId),
|
||||
ownerId: face.asset.ownerId,
|
||||
assetId: face.assetId,
|
||||
personId: face.personId,
|
||||
embedding: face.embedding,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
return results;
|
||||
}
|
||||
|
||||
private asKey(face: AssetFaceId): string {
|
||||
return `${face.assetId}|${face.personId}`;
|
||||
}
|
||||
|
||||
private asParts(key: string): AssetFaceId {
|
||||
const [assetId, personId] = key.split('|');
|
||||
return { assetId, personId };
|
||||
}
|
||||
|
||||
private async getLookupMap(assetIds: string[]) {
|
||||
const assets = await this.assetRepository.getByIds(assetIds);
|
||||
const lookup: Record<string, AssetEntity> = {};
|
||||
for (const asset of assets) {
|
||||
lookup[asset.id] = asset;
|
||||
}
|
||||
return lookup;
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user