diff options
Diffstat (limited to 'apps/workers')
| -rw-r--r-- | apps/workers/workers/inference/tagging.ts | 34 |
1 files changed, 33 insertions, 1 deletions
diff --git a/apps/workers/workers/inference/tagging.ts b/apps/workers/workers/inference/tagging.ts index b3006193..668c1d5e 100644 --- a/apps/workers/workers/inference/tagging.ts +++ b/apps/workers/workers/inference/tagging.ts @@ -85,6 +85,7 @@ async function buildPrompt( bookmark: NonNullable<Awaited<ReturnType<typeof fetchBookmark>>>, tagStyle: ZTagStyle, inferredTagLang: string, + curatedTags?: string[], ): Promise<string | null> { const prompts = await fetchCustomPrompts(bookmark.userId, "text"); if (bookmark.link) { @@ -110,6 +111,7 @@ Description: ${bookmark.link.description ?? ""} Content: ${content ?? ""}`, serverConfig.inference.contextLength, tagStyle, + curatedTags, ); } @@ -120,6 +122,7 @@ Content: ${content ?? ""}`, bookmark.text.text ?? "", serverConfig.inference.contextLength, tagStyle, + curatedTags, ); } @@ -133,6 +136,7 @@ async function inferTagsFromImage( abortSignal: AbortSignal, tagStyle: ZTagStyle, inferredTagLang: string, + curatedTags?: string[], ): Promise<InferenceResponse | null> { const { asset, metadata } = await readAsset({ userId: bookmark.userId, @@ -160,6 +164,7 @@ async function inferTagsFromImage( inferredTagLang, await fetchCustomPrompts(bookmark.userId, "images"), tagStyle, + curatedTags, ), metadata.contentType, base64, @@ -235,6 +240,7 @@ async function inferTagsFromPDF( abortSignal: AbortSignal, tagStyle: ZTagStyle, inferredTagLang: string, + curatedTags?: string[], ) { const prompt = await buildTextPrompt( inferredTagLang, @@ -242,6 +248,7 @@ async function inferTagsFromPDF( `Content: ${bookmark.asset.content}`, serverConfig.inference.contextLength, tagStyle, + curatedTags, ); setSpanAttributes({ "inference.model": serverConfig.inference.textModel, @@ -261,8 +268,14 @@ async function inferTagsFromText( abortSignal: AbortSignal, tagStyle: ZTagStyle, inferredTagLang: string, + curatedTags?: string[], ) { - const prompt = await buildPrompt(bookmark, tagStyle, inferredTagLang); + const prompt = await buildPrompt( + bookmark, + tagStyle, + inferredTagLang, + curatedTags, + ); if (!prompt) { return null; } @@ -285,6 +298,7 @@ async function inferTags( abortSignal: AbortSignal, tagStyle: ZTagStyle, inferredTagLang: string, + curatedTags?: string[], ) { setSpanAttributes({ "user.id": bookmark.userId, @@ -306,6 +320,7 @@ async function inferTags( abortSignal, tagStyle, inferredTagLang, + curatedTags, ); } else if (bookmark.asset) { switch (bookmark.asset.assetType) { @@ -317,6 +332,7 @@ async function inferTags( abortSignal, tagStyle, inferredTagLang, + curatedTags, ); break; case "pdf": @@ -327,6 +343,7 @@ async function inferTags( abortSignal, tagStyle, inferredTagLang, + curatedTags, ); break; default: @@ -507,6 +524,7 @@ export async function runTagging( columns: { autoTaggingEnabled: true, tagStyle: true, + curatedTagIds: true, inferredTagLang: true, }, }); @@ -518,6 +536,19 @@ export async function runTagging( return; } + // Resolve curated tag names if configured + let curatedTagNames: string[] | undefined; + if (userSettings?.curatedTagIds && userSettings.curatedTagIds.length > 0) { + const tags = await db.query.bookmarkTags.findMany({ + where: and( + eq(bookmarkTags.userId, bookmark.userId), + inArray(bookmarkTags.id, userSettings.curatedTagIds), + ), + columns: { name: true }, + }); + curatedTagNames = tags.map((t) => t.name); + } + logger.info( `[inference][${jobId}] Starting an inference job for bookmark with id "${bookmark.id}"`, ); @@ -529,6 +560,7 @@ export async function runTagging( job.abortSignal, userSettings?.tagStyle ?? "as-generated", userSettings?.inferredTagLang ?? serverConfig.inference.inferredTagLang, + curatedTagNames, ); if (tags === null) { |