feat(ui): add regeneration option for embeddings generation

2026-02-12 15:52:47 +00:00 · 2025-12-14 03:56:40 +03:00
parent 7e55decf64
commit 51526b239f
9 changed files with 327 additions and 64 deletions
--- a/messages/en.json
+++ b/messages/en.json
@@ -346,6 +346,7 @@
      "confirm": "Import",
      "delete": "Delete",
      "generateEmbeddings": "Generate Embeddings",
+      "regenerateEmbeddings": "Regenerate all embeddings",
      "pending": "pending",
      "embeddingsSuccess": "{count} embeddings generated",
      "embeddingsResult": "Generated: {success}, Failed: {failed}"
--- a/package-lock.json
+++ b/package-lock.json
@@ -21,6 +21,7 @@
        "@radix-ui/react-dropdown-menu": "^2.1.16",
        "@radix-ui/react-label": "^2.1.8",
        "@radix-ui/react-popover": "^1.1.15",
+        "@radix-ui/react-progress": "^1.1.8",
        "@radix-ui/react-scroll-area": "^1.2.10",
        "@radix-ui/react-select": "^2.2.6",
        "@radix-ui/react-separator": "^1.1.8",
@@ -4016,6 +4017,30 @@
        }
      }
    },
+    "node_modules/@radix-ui/react-progress": {
+      "version": "1.1.8",
+      "resolved": "https://registry.npmjs.org/@radix-ui/react-progress/-/react-progress-1.1.8.tgz",
+      "integrity": "sha512-+gISHcSPUJ7ktBy9RnTqbdKW78bcGke3t6taawyZ71pio1JewwGSJizycs7rLhGTvMJYCQB1DBK4KQsxs7U8dA==",
+      "license": "MIT",
+      "dependencies": {
+        "@radix-ui/react-context": "1.1.3",
+        "@radix-ui/react-primitive": "2.1.4"
+      },
+      "peerDependencies": {
+        "@types/react": "*",
+        "@types/react-dom": "*",
+        "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc",
+        "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc"
+      },
+      "peerDependenciesMeta": {
+        "@types/react": {
+          "optional": true
+        },
+        "@types/react-dom": {
+          "optional": true
+        }
+      }
+    },
    "node_modules/@radix-ui/react-roving-focus": {
      "version": "1.1.11",
      "resolved": "https://registry.npmjs.org/@radix-ui/react-roving-focus/-/react-roving-focus-1.1.11.tgz",
--- a/package.json
+++ b/package.json
@@ -29,6 +29,7 @@
    "@radix-ui/react-dropdown-menu": "^2.1.16",
    "@radix-ui/react-label": "^2.1.8",
    "@radix-ui/react-popover": "^1.1.15",
+    "@radix-ui/react-progress": "^1.1.8",
    "@radix-ui/react-scroll-area": "^1.2.10",
    "@radix-ui/react-select": "^2.2.6",
    "@radix-ui/react-separator": "^1.1.8",
--- a/src/app/admin/page.tsx
+++ b/src/app/admin/page.tsx
@@ -38,15 +38,25 @@ export default async function AdminPage() {
    isAISearchEnabled(),
  ]);
  
-  // Count prompts without embeddings (JSON null check requires separate query)
+  // Count prompts without embeddings and total public prompts
  let promptsWithoutEmbeddings = 0;
+  let totalPublicPrompts = 0;
  if (aiSearchEnabled) {
-    promptsWithoutEmbeddings = await db.prompt.count({
-      where: {
-        isPrivate: false,
-        embedding: { equals: Prisma.DbNull },
-      },
-    });
+    [promptsWithoutEmbeddings, totalPublicPrompts] = await Promise.all([
+      db.prompt.count({
+        where: {
+          isPrivate: false,
+          deletedAt: null,
+          embedding: { equals: Prisma.DbNull },
+        },
+      }),
+      db.prompt.count({
+        where: {
+          isPrivate: false,
+          deletedAt: null,
+        },
+      }),
+    ]);
  }

  // Fetch data for tables
@@ -220,7 +230,8 @@ export default async function AdminPage() {
        <TabsContent value="prompts">
          <PromptsManagement 
            aiSearchEnabled={aiSearchEnabled} 
-            promptsWithoutEmbeddings={promptsWithoutEmbeddings} 
+            promptsWithoutEmbeddings={promptsWithoutEmbeddings}
+            totalPublicPrompts={totalPublicPrompts}
          />
        </TabsContent>

--- a/src/app/api/admin/embeddings/route.ts
+++ b/src/app/api/admin/embeddings/route.ts
@@ -1,8 +1,8 @@
-import { NextResponse } from "next/server";
+import { NextRequest, NextResponse } from "next/server";
 import { auth } from "@/lib/auth";
 import { generateAllEmbeddings, isAISearchEnabled } from "@/lib/ai/embeddings";

-export async function POST() {
+export async function POST(request: NextRequest) {
  try {
    const session = await auth();
    if (!session?.user || session.user.role !== "ADMIN") {
@@ -17,11 +17,40 @@ export async function POST() {
      );
    }

-    const result = await generateAllEmbeddings();
+    // Check if regenerate mode
+    const { searchParams } = new URL(request.url);
+    const regenerate = searchParams.get("regenerate") === "true";

-    return NextResponse.json({
-      message: "Embeddings generated",
-      ...result,
+    // Create a streaming response for progress updates
+    const encoder = new TextEncoder();
+    const stream = new ReadableStream({
+      async start(controller) {
+        try {
+          const result = await generateAllEmbeddings(
+            (current, total, success, failed) => {
+              const progress = JSON.stringify({ current, total, success, failed, done: false });
+              controller.enqueue(encoder.encode(`data: ${progress}\n\n`));
+            },
+            regenerate
+          );
+          
+          const final = JSON.stringify({ ...result, done: true });
+          controller.enqueue(encoder.encode(`data: ${final}\n\n`));
+          controller.close();
+        } catch (error) {
+          const errorMsg = JSON.stringify({ error: "Failed to generate embeddings", done: true });
+          controller.enqueue(encoder.encode(`data: ${errorMsg}\n\n`));
+          controller.close();
+        }
+      },
+    });
+
+    return new Response(stream, {
+      headers: {
+        "Content-Type": "text/event-stream",
+        "Cache-Control": "no-cache",
+        "Connection": "keep-alive",
+      },
    });
  } catch (error) {
    console.error("Generate embeddings error:", error);
--- a/src/components/admin/ai-search-settings.tsx
+++ b/src/components/admin/ai-search-settings.tsx
@@ -4,25 +4,37 @@ import { useState } from "react";
 import { useTranslations } from "next-intl";
 import { Button } from "@/components/ui/button";
 import { Card, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card";
-import { Sparkles, Loader2, CheckCircle, AlertCircle } from "lucide-react";
+import { Progress } from "@/components/ui/progress";
+import { Sparkles, Loader2, CheckCircle, AlertCircle, RefreshCw } from "lucide-react";
 import { toast } from "sonner";

 interface AISearchSettingsProps {
  enabled: boolean;
  promptsWithoutEmbeddings: number;
+  totalPrompts: number;
 }

-export function AISearchSettings({ enabled, promptsWithoutEmbeddings }: AISearchSettingsProps) {
+interface ProgressState {
+  current: number;
+  total: number;
+  success: number;
+  failed: number;
+}
+
+export function AISearchSettings({ enabled, promptsWithoutEmbeddings, totalPrompts }: AISearchSettingsProps) {
  const t = useTranslations("admin");
  const [isGenerating, setIsGenerating] = useState(false);
+  const [progress, setProgress] = useState<ProgressState | null>(null);
  const [result, setResult] = useState<{ success: number; failed: number } | null>(null);

-  const handleGenerateEmbeddings = async () => {
+  const handleGenerateEmbeddings = async (regenerate: boolean = false) => {
    setIsGenerating(true);
    setResult(null);
+    setProgress(null);

    try {
-      const response = await fetch("/api/admin/embeddings", {
+      const url = regenerate ? "/api/admin/embeddings?regenerate=true" : "/api/admin/embeddings";
+      const response = await fetch(url, {
        method: "POST",
      });

@@ -31,13 +43,45 @@ export function AISearchSettings({ enabled, promptsWithoutEmbeddings }: AISearch
        throw new Error(data.error || "Failed to generate embeddings");
      }

-      const data = await response.json();
-      setResult({ success: data.success, failed: data.failed });
-      toast.success(t("aiSearch.generateSuccess", { count: data.success }));
+      // Read the stream
+      const reader = response.body?.getReader();
+      if (!reader) throw new Error("No response body");
+
+      const decoder = new TextDecoder();
+      
+      while (true) {
+        const { done, value } = await reader.read();
+        if (done) break;
+        
+        const text = decoder.decode(value);
+        const lines = text.split("\n\n").filter(line => line.startsWith("data: "));
+        
+        for (const line of lines) {
+          const jsonStr = line.replace("data: ", "");
+          try {
+            const data = JSON.parse(jsonStr);
+            
+            if (data.done) {
+              setResult({ success: data.success, failed: data.failed });
+              toast.success(t("aiSearch.generateSuccess", { count: data.success }));
+            } else {
+              setProgress({
+                current: data.current,
+                total: data.total,
+                success: data.success,
+                failed: data.failed,
+              });
+            }
+          } catch {
+            // Ignore parse errors
+          }
+        }
+      }
    } catch (error) {
      toast.error(error instanceof Error ? error.message : "Failed to generate embeddings");
    } finally {
      setIsGenerating(false);
+      setProgress(null);
    }
  };

@@ -45,6 +89,8 @@ export function AISearchSettings({ enabled, promptsWithoutEmbeddings }: AISearch
    return null;
  }

+  const progressPercent = progress ? Math.round((progress.current / progress.total) * 100) : 0;
+
  return (
    <Card>
      <CardHeader>
@@ -62,7 +108,22 @@ export function AISearchSettings({ enabled, promptsWithoutEmbeddings }: AISearch
          </div>
        </div>

-        {result && (
+        {/* Progress bar */}
+        {progress && (
+          <div className="space-y-2">
+            <Progress value={progressPercent} className="h-2" />
+            <div className="flex justify-between text-xs text-muted-foreground">
+              <span>{progress.current} / {progress.total}</span>
+              <span>{progressPercent}%</span>
+            </div>
+            <div className="flex gap-4 text-xs">
+              <span className="text-green-600">✓ {progress.success}</span>
+              {progress.failed > 0 && <span className="text-red-600">✗ {progress.failed}</span>}
+            </div>
+          </div>
+        )}
+
+        {result && !progress && (
          <div className="flex items-center gap-2 text-sm">
            {result.failed === 0 ? (
              <CheckCircle className="h-4 w-4 text-green-500" />
@@ -75,23 +136,39 @@ export function AISearchSettings({ enabled, promptsWithoutEmbeddings }: AISearch
          </div>
        )}

-        <Button
-          onClick={handleGenerateEmbeddings}
-          disabled={isGenerating || promptsWithoutEmbeddings === 0}
-          className="w-full"
-        >
-          {isGenerating ? (
-            <>
-              <Loader2 className="h-4 w-4 mr-2 animate-spin" />
-              {t("aiSearch.generating")}
-            </>
-          ) : (
-            <>
-              <Sparkles className="h-4 w-4 mr-2" />
-              {t("aiSearch.generateButton")}
-            </>
-          )}
-        </Button>
+        <div className="flex gap-2">
+          <Button
+            onClick={() => handleGenerateEmbeddings(false)}
+            disabled={isGenerating || promptsWithoutEmbeddings === 0}
+            className="flex-1"
+          >
+            {isGenerating && !progress ? (
+              <>
+                <Loader2 className="h-4 w-4 mr-2 animate-spin" />
+                {t("aiSearch.generating")}
+              </>
+            ) : isGenerating && progress ? (
+              <>
+                <Loader2 className="h-4 w-4 mr-2 animate-spin" />
+                {progress.current}/{progress.total}
+              </>
+            ) : (
+              <>
+                <Sparkles className="h-4 w-4 mr-2" />
+                {t("aiSearch.generateButton")}
+              </>
+            )}
+          </Button>
+          
+          <Button
+            variant="outline"
+            onClick={() => handleGenerateEmbeddings(true)}
+            disabled={isGenerating || totalPrompts === 0}
+            title={t("aiSearch.regenerateTooltip")}
+          >
+            <RefreshCw className="h-4 w-4" />
+          </Button>
+        </div>
      </CardContent>
    </Card>
  );
--- a/src/components/admin/prompts-management.tsx
+++ b/src/components/admin/prompts-management.tsx
@@ -3,9 +3,10 @@
 import { useState } from "react";
 import { useRouter } from "next/navigation";
 import { useTranslations } from "next-intl";
-import { Upload, Trash2, Loader2, CheckCircle, AlertCircle, Sparkles, Download } from "lucide-react";
+import { Upload, Trash2, Loader2, CheckCircle, AlertCircle, Sparkles, Download, RefreshCw } from "lucide-react";
 import { Button } from "@/components/ui/button";
 import { Card, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card";
+import { Progress } from "@/components/ui/progress";
 import {
  AlertDialog,
  AlertDialogAction,
@@ -26,12 +27,20 @@ interface ImportResult {
  errors: string[];
 }

+interface ProgressState {
+  current: number;
+  total: number;
+  success: number;
+  failed: number;
+}
+
 interface PromptsManagementProps {
  aiSearchEnabled: boolean;
  promptsWithoutEmbeddings: number;
+  totalPublicPrompts: number;
 }

-export function PromptsManagement({ aiSearchEnabled, promptsWithoutEmbeddings }: PromptsManagementProps) {
+export function PromptsManagement({ aiSearchEnabled, promptsWithoutEmbeddings, totalPublicPrompts }: PromptsManagementProps) {
  const router = useRouter();
  const t = useTranslations("admin");
  const [loading, setLoading] = useState(false);
@@ -41,6 +50,7 @@ export function PromptsManagement({ aiSearchEnabled, promptsWithoutEmbeddings }:
  const [showDeleteConfirm, setShowDeleteConfirm] = useState(false);
  const [importResult, setImportResult] = useState<ImportResult | null>(null);
  const [embeddingResult, setEmbeddingResult] = useState<{ success: number; failed: number } | null>(null);
+  const [embeddingProgress, setEmbeddingProgress] = useState<ProgressState | null>(null);

  const handleImport = async () => {
    setLoading(true);
@@ -92,25 +102,60 @@ export function PromptsManagement({ aiSearchEnabled, promptsWithoutEmbeddings }:
    }
  };

-  const handleGenerateEmbeddings = async () => {
+  const handleGenerateEmbeddings = async (regenerate: boolean = false) => {
    setGenerating(true);
    setEmbeddingResult(null);
+    setEmbeddingProgress(null);

    try {
-      const res = await fetch("/api/admin/embeddings", { method: "POST" });
-      const data = await res.json();
+      const url = regenerate ? "/api/admin/embeddings?regenerate=true" : "/api/admin/embeddings";
+      const res = await fetch(url, { method: "POST" });

      if (!res.ok) {
+        const data = await res.json();
        throw new Error(data.error || "Failed to generate embeddings");
      }

-      setEmbeddingResult({ success: data.success, failed: data.failed });
-      toast.success(t("prompts.embeddingsSuccess", { count: data.success }));
-      router.refresh();
+      // Read the stream
+      const reader = res.body?.getReader();
+      if (!reader) throw new Error("No response body");
+
+      const decoder = new TextDecoder();
+      
+      while (true) {
+        const { done, value } = await reader.read();
+        if (done) break;
+        
+        const text = decoder.decode(value);
+        const lines = text.split("\n\n").filter(line => line.startsWith("data: "));
+        
+        for (const line of lines) {
+          const jsonStr = line.replace("data: ", "");
+          try {
+            const data = JSON.parse(jsonStr);
+            
+            if (data.done) {
+              setEmbeddingResult({ success: data.success, failed: data.failed });
+              toast.success(t("prompts.embeddingsSuccess", { count: data.success }));
+              router.refresh();
+            } else {
+              setEmbeddingProgress({
+                current: data.current,
+                total: data.total,
+                success: data.success,
+                failed: data.failed,
+              });
+            }
+          } catch {
+            // Ignore parse errors
+          }
+        }
+      }
    } catch (error) {
      toast.error(error instanceof Error ? error.message : "Failed to generate embeddings");
    } finally {
      setGenerating(false);
+      setEmbeddingProgress(null);
    }
  };

@@ -182,14 +227,44 @@ export function PromptsManagement({ aiSearchEnabled, promptsWithoutEmbeddings }:
              <Button
                size="sm"
                variant="outline"
-                onClick={handleGenerateEmbeddings}
+                onClick={() => handleGenerateEmbeddings(false)}
                disabled={loading || deleting || generating || promptsWithoutEmbeddings === 0}
              >
-                {generating ? <Loader2 className="h-4 w-4 animate-spin" /> : <><Sparkles className="h-4 w-4 mr-2" />{t("prompts.generateEmbeddings")}</>}
+                {generating && !embeddingProgress ? (
+                  <Loader2 className="h-4 w-4 animate-spin" />
+                ) : generating && embeddingProgress ? (
+                  <><Loader2 className="h-4 w-4 animate-spin mr-2" />{embeddingProgress.current}/{embeddingProgress.total}</>
+                ) : (
+                  <><Sparkles className="h-4 w-4 mr-2" />{t("prompts.generateEmbeddings")}</>
+                )}
+              </Button>
+              <Button
+                size="sm"
+                variant="ghost"
+                onClick={() => handleGenerateEmbeddings(true)}
+                disabled={loading || deleting || generating || totalPublicPrompts === 0}
+                title={t("prompts.regenerateEmbeddings")}
+              >
+                <RefreshCw className="h-4 w-4" />
              </Button>
            </div>

-            {embeddingResult && (
+            {/* Progress bar */}
+            {embeddingProgress && (
+              <div className="space-y-2">
+                <Progress value={Math.round((embeddingProgress.current / embeddingProgress.total) * 100)} className="h-2" />
+                <div className="flex justify-between text-xs text-muted-foreground">
+                  <span>{embeddingProgress.current} / {embeddingProgress.total}</span>
+                  <span>{Math.round((embeddingProgress.current / embeddingProgress.total) * 100)}%</span>
+                </div>
+                <div className="flex gap-4 text-xs">
+                  <span className="text-green-600">✓ {embeddingProgress.success}</span>
+                  {embeddingProgress.failed > 0 && <span className="text-red-600">✗ {embeddingProgress.failed}</span>}
+                </div>
+              </div>
+            )}
+
+            {embeddingResult && !embeddingProgress && (
              <div className="flex items-center gap-2 text-sm text-muted-foreground">
                {embeddingResult.failed === 0 ? <CheckCircle className="h-4 w-4 text-green-500" /> : <AlertCircle className="h-4 w-4 text-amber-500" />}
                <span>{t("prompts.embeddingsResult", { success: embeddingResult.success, failed: embeddingResult.failed })}</span>
--- a/src/components/ui/progress.tsx
+++ b/src/components/ui/progress.tsx
@@ -0,0 +1,28 @@
+"use client";
+
+import * as React from "react";
+import * as ProgressPrimitive from "@radix-ui/react-progress";
+
+import { cn } from "@/lib/utils";
+
+const Progress = React.forwardRef<
+  React.ElementRef<typeof ProgressPrimitive.Root>,
+  React.ComponentPropsWithoutRef<typeof ProgressPrimitive.Root>
+>(({ className, value, ...props }, ref) => (
+  <ProgressPrimitive.Root
+    ref={ref}
+    className={cn(
+      "relative h-4 w-full overflow-hidden rounded-full bg-secondary",
+      className
+    )}
+    {...props}
+  >
+    <ProgressPrimitive.Indicator
+      className="h-full w-full flex-1 bg-primary transition-all"
+      style={{ transform: `translateX(-${100 - (value || 0)}%)` }}
+    />
+  </ProgressPrimitive.Root>
+));
+Progress.displayName = ProgressPrimitive.Root.displayName;
+
+export { Progress };
--- a/src/lib/ai/embeddings.ts
+++ b/src/lib/ai/embeddings.ts
@@ -50,19 +50,23 @@ export async function generatePromptEmbedding(promptId: string): Promise<void> {
    prompt.content,
  ].join("\n\n").trim();

-  try {
-    const embedding = await generateEmbedding(textToEmbed);
-    
-    await db.prompt.update({
-      where: { id: promptId },
-      data: { embedding },
-    });
-  } catch (error) {
-    console.error(`Failed to generate embedding for prompt ${promptId}:`, error);
-  }
+  const embedding = await generateEmbedding(textToEmbed);
+  
+  await db.prompt.update({
+    where: { id: promptId },
+    data: { embedding },
+  });
 }

-export async function generateAllEmbeddings(): Promise<{ success: number; failed: number }> {
+// Delay helper to avoid rate limits
+function delay(ms: number): Promise<void> {
+  return new Promise(resolve => setTimeout(resolve, ms));
+}
+
+export async function generateAllEmbeddings(
+  onProgress?: (current: number, total: number, success: number, failed: number) => void,
+  regenerate: boolean = false
+): Promise<{ success: number; failed: number; total: number }> {
  const config = await getConfig();
  if (!config.features.aiSearch) {
    throw new Error("AI Search is not enabled");
@@ -70,26 +74,38 @@ export async function generateAllEmbeddings(): Promise<{ success: number; failed

  const prompts = await db.prompt.findMany({
    where: { 
-      embedding: { equals: Prisma.DbNull },
+      ...(regenerate ? {} : { embedding: { equals: Prisma.DbNull } }),
      isPrivate: false,
      deletedAt: null,
    },
    select: { id: true },
  });

+  const total = prompts.length;
  let success = 0;
  let failed = 0;

-  for (const prompt of prompts) {
+  for (let i = 0; i < prompts.length; i++) {
+    const prompt = prompts[i];
    try {
      await generatePromptEmbedding(prompt.id);
      success++;
    } catch {
      failed++;
    }
+    
+    // Report progress
+    if (onProgress) {
+      onProgress(i + 1, total, success, failed);
+    }
+    
+    // Rate limit: wait 200ms between requests to avoid hitting API limits
+    if (i < prompts.length - 1) {
+      await delay(200);
+    }
  }

-  return { success, failed };
+  return { success, failed, total };
 }

 function cosineSimilarity(a: number[], b: number[]): number {