feature: Support regenerating AI tags only for failed bookmarks

hoarder-app · Oct 13, 2024 · 2ccc15e · 2ccc15e
1 parent fb57f25
commit 2ccc15e
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 11 deletions.
diff --git a/apps/web/components/dashboard/admin/AdminActions.tsx b/apps/web/components/dashboard/admin/AdminActions.tsx
@@ -101,7 +101,16 @@ export default function AdminActions() {
         <ActionButton
           variant="destructive"
           loading={isInferencePending}
-          onClick={() => reRunInferenceOnAllBookmarks()}
+          onClick={() =>
+            reRunInferenceOnAllBookmarks({ taggingStatus: "failure" })
+          }
+        >
+          Regenerate AI Tags for Failed Bookmarks Only
+        </ActionButton>
+        <ActionButton
+          variant="destructive"
+          loading={isInferencePending}
+          onClick={() => reRunInferenceOnAllBookmarks({ taggingStatus: "all" })}
         >
           Regenerate AI Tags for All Bookmarks
         </ActionButton>

diff --git a/packages/trpc/routers/admin.ts b/packages/trpc/routers/admin.ts
@@ -145,17 +145,26 @@ export const adminAppRouter = router({
 
     await Promise.all(bookmarkIds.map((b) => triggerSearchReindex(b.id)));
   }),
-  reRunInferenceOnAllBookmarks: adminProcedure.mutation(async ({ ctx }) => {
-    const bookmarkIds = await ctx.db.query.bookmarks.findMany({
-      columns: {
-        id: true,
-      },
-    });
+  reRunInferenceOnAllBookmarks: adminProcedure
+    .input(
+      z.object({
+        taggingStatus: z.enum(["success", "failure", "all"]),
+      }),
+    )
+    .mutation(async ({ input, ctx }) => {
+      const bookmarkIds = await ctx.db.query.bookmarks.findMany({
+        columns: {
+          id: true,
+        },
+        ...(input.taggingStatus === "all"
+          ? {}
+          : { where: eq(bookmarks.taggingStatus, input.taggingStatus) }),
+      });
 
-    await Promise.all(
-      bookmarkIds.map((b) => OpenAIQueue.enqueue({ bookmarkId: b.id })),
-    );
-  }),
+      await Promise.all(
+        bookmarkIds.map((b) => OpenAIQueue.enqueue({ bookmarkId: b.id })),
+      );
+    }),
   tidyAssets: adminProcedure.mutation(async () => {
     await TidyAssetsQueue.enqueue({
       cleanDanglingAssets: true,