1 files changed, 28 insertions, 8 deletions
diff --git a/packages/trpc/routers/admin.ts b/packages/trpc/routers/admin.ts
index e4985b5c..91f4a34f 100644
--- a/packages/trpc/routers/admin.ts
+++ b/packages/trpc/routers/admin.ts
@@ -1,5 +1,5 @@
 import { TRPCError } from "@trpc/server";
-import { count, eq, sum } from "drizzle-orm";
+import { count, eq, or, sum } from "drizzle-orm";
 import { z } from "zod";
 
 import { assets, bookmarkLinks, bookmarks, users } from "@karakeep/db/schema";
@@ -129,11 +129,21 @@ export const adminAppRouter = router({
         ctx.db
           .select({ value: count() })
           .from(bookmarks)
-          .where(eq(bookmarks.taggingStatus, "pending")),
+          .where(
+            or(
+              eq(bookmarks.taggingStatus, "pending"),
+              eq(bookmarks.summarizationStatus, "pending"),
+            ),
+          ),
         ctx.db
           .select({ value: count() })
           .from(bookmarks)
-          .where(eq(bookmarks.taggingStatus, "failure")),
+          .where(
+            or(
+              eq(bookmarks.taggingStatus, "failure"),
+              eq(bookmarks.summarizationStatus, "failure"),
+            ),
+          ),
 
         // Tidy Assets
         TidyAssetsQueue.stats(),
@@ -233,7 +243,8 @@ export const adminAppRouter = router({
   reRunInferenceOnAllBookmarks: adminProcedure
     .input(
       z.object({
-        taggingStatus: z.enum(["success", "failure", "all"]),
+        type: z.enum(["tag", "summarize"]),
+        status: z.enum(["success", "failure", "all"]),
       }),
     )
     .mutation(async ({ input, ctx }) => {
@@ -241,13 +252,22 @@ export const adminAppRouter = router({
         columns: {
           id: true,
         },
-        ...(input.taggingStatus === "all"
-          ? {}
-          : { where: eq(bookmarks.taggingStatus, input.taggingStatus) }),
+        ...{
+          tag:
+            input.status === "all"
+              ? {}
+              : { where: eq(bookmarks.taggingStatus, input.status) },
+          summarize:
+            input.status === "all"
+              ? {}
+              : { where: eq(bookmarks.summarizationStatus, input.status) },
+        }[input.type],
       });
 
       await Promise.all(
-        bookmarkIds.map((b) => OpenAIQueue.enqueue({ bookmarkId: b.id })),
+        bookmarkIds.map((b) =>
+          OpenAIQueue.enqueue({ bookmarkId: b.id, type: input.type }),
+        ),
       );
     }),
   tidyAssets: adminProcedure.mutation(async () => {