feat(server): lru cache for query embeddings (#19181)

This commit is contained in:
Mert
2025-06-16 11:03:49 -04:00
committed by GitHub
parent 176d53c1b3
commit 7347f64958
3 changed files with 85 additions and 19 deletions

View File

@@ -1,4 +1,5 @@
import { BadRequestException, Injectable } from '@nestjs/common';
import { LRUMap } from 'mnemonist';
import { AssetMapOptions, AssetResponseDto, MapAsset, mapAsset } from 'src/dtos/asset-response.dto';
import { AuthDto } from 'src/dtos/auth.dto';
import { mapPerson, PersonResponseDto } from 'src/dtos/person.dto';
@@ -24,6 +25,8 @@ import { isSmartSearchEnabled } from 'src/utils/misc';
@Injectable()
export class SearchService extends BaseService {
private embeddingCache = new LRUMap<string, string>(100);
async searchPerson(auth: AuthDto, dto: SearchPeopleDto): Promise<PersonResponseDto[]> {
const people = await this.personRepository.getByName(auth.user.id, dto.name, { withHidden: dto.withHidden });
return people.map((person) => mapPerson(person));
@@ -98,16 +101,21 @@ export class SearchService extends BaseService {
throw new BadRequestException('Smart search is not enabled');
}
const userIds = await this.getUserIdsToSearch(auth);
const embedding = await this.machineLearningRepository.encodeText(machineLearning.urls, dto.query, {
modelName: machineLearning.clip.modelName,
language: dto.language,
});
const userIds = this.getUserIdsToSearch(auth);
const key = machineLearning.clip.modelName + dto.query + dto.language;
let embedding = this.embeddingCache.get(key);
if (!embedding) {
embedding = await this.machineLearningRepository.encodeText(machineLearning.urls, dto.query, {
modelName: machineLearning.clip.modelName,
language: dto.language,
});
this.embeddingCache.set(key, embedding);
}
const page = dto.page ?? 1;
const size = dto.size || 100;
const { hasNextPage, items } = await this.searchRepository.searchSmart(
{ page, size },
{ ...dto, userIds, embedding },
{ ...dto, userIds: await userIds, embedding },
);
return this.mapResponse(items, hasNextPage ? (page + 1).toString() : null, { auth });