feat(ai): Support restricting AI tags to a subset of existing tags (#2444)

* feat(ai): Support restricting AI tags to a subset of existing tags

Co-authored-by: Claude <noreply@anthropic.com>

4 files changed, 10 insertions, 0 deletions

diff --git a/packages/trpc/models/tags.ts b/packages/trpc/models/tags.ts
index 55532077..1d8f90b9 100644
--- a/packages/trpc/models/tags.ts
+++ b/packages/trpc/models/tags.ts
@@ -85,6 +85,7 @@ export class Tag {
     ctx: AuthedContext,
     opts: {
       nameContains?: string;
+      ids?: string[];
       attachedBy?: "ai" | "human" | "none";
       sortBy?: "name" | "usage" | "relevance";
       pagination?: {
@@ -119,6 +120,9 @@ export class Tag {
           opts.nameContains
             ? like(bookmarkTags.name, `%${opts.nameContains}%`)
             : undefined,
+          opts.ids && opts.ids.length > 0
+            ? inArray(bookmarkTags.id, opts.ids)
+            : undefined,
         ),
       )
       .groupBy(bookmarkTags.id, bookmarkTags.name)
diff --git a/packages/trpc/models/users.ts b/packages/trpc/models/users.ts
index 671f7d74..3340956a 100644
--- a/packages/trpc/models/users.ts
+++ b/packages/trpc/models/users.ts
@@ -447,6 +447,7 @@ export class User {
         autoTaggingEnabled: true,
         autoSummarizationEnabled: true,
         tagStyle: true,
+        curatedTagIds: true,
         inferredTagLang: true,
       },
     });
@@ -471,6 +472,7 @@ export class User {
       autoTaggingEnabled: settings.autoTaggingEnabled,
       autoSummarizationEnabled: settings.autoSummarizationEnabled,
       tagStyle: settings.tagStyle ?? "as-generated",
+      curatedTagIds: settings.curatedTagIds ?? null,
       inferredTagLang: settings.inferredTagLang,
     };
   }
@@ -500,6 +502,7 @@ export class User {
         autoTaggingEnabled: input.autoTaggingEnabled,
         autoSummarizationEnabled: input.autoSummarizationEnabled,
         tagStyle: input.tagStyle,
+        curatedTagIds: input.curatedTagIds,
         inferredTagLang: input.inferredTagLang,
       })
       .where(eq(users.id, this.user.id));
diff --git a/packages/trpc/routers/tags.ts b/packages/trpc/routers/tags.ts
index d4cfbe8c..5713c192 100644
--- a/packages/trpc/routers/tags.ts
+++ b/packages/trpc/routers/tags.ts
@@ -102,6 +102,7 @@ export const tagsAppRouter = router({
     .query(async ({ ctx, input }) => {
       return await Tag.getAll(ctx, {
         nameContains: input.nameContains,
+        ids: input.ids,
         attachedBy: input.attachedBy,
         sortBy: input.sortBy,
         pagination: input.limit
diff --git a/packages/trpc/routers/users.test.ts b/packages/trpc/routers/users.test.ts
index ccde4c86..d8ec90f9 100644
--- a/packages/trpc/routers/users.test.ts
+++ b/packages/trpc/routers/users.test.ts
@@ -167,6 +167,7 @@ describe("User Routes", () => {
       // AI Settings
       autoSummarizationEnabled: null,
       autoTaggingEnabled: null,
+      curatedTagIds: null,
       inferredTagLang: null,
       tagStyle: "titlecase-spaces",
     });
@@ -208,6 +209,7 @@ describe("User Routes", () => {
       // AI Settings
       autoSummarizationEnabled: true,
       autoTaggingEnabled: true,
+      curatedTagIds: null,
       inferredTagLang: "en",
       tagStyle: "lowercase-underscores",
     });