[OPIK-2497] [FE] Add batch tag update UI for traces, spans and threads

YarivHashaiComet · YarivHashaiComet · commit 63319efbbd8b · 2025-11-03T17:55:46.000+02:00
diff --git a/apps/opik-frontend/src/api/traces/useSpanBatchUpdateMutation.ts b/apps/opik-frontend/src/api/traces/useSpanBatchUpdateMutation.ts
@@ -0,0 +1,56 @@
+import { useMutation, useQueryClient } from "@tanstack/react-query";
+import { AxiosError } from "axios";
+import get from "lodash/get";
+
+import api, { SPANS_KEY, SPANS_REST_ENDPOINT } from "@/api/api";
+import { Span } from "@/types/traces";
+import { useToast } from "@/components/ui/use-toast";
+
+type UseSpanBatchUpdateMutationParams = {
+  projectId: string;
+  spanIds: string[];
+  span: Partial<Span>;
+  mergeTags?: boolean;
+};
+
+const useSpanBatchUpdateMutation = () => {
+  const queryClient = useQueryClient();
+  const { toast } = useToast();
+
+  return useMutation({
+    mutationFn: async ({
+      spanIds,
+      span,
+      mergeTags,
+    }: UseSpanBatchUpdateMutationParams) => {
+      const { data } = await api.patch(SPANS_REST_ENDPOINT + "batch", {
+        ids: spanIds,
+        update: span,
+        merge_tags: mergeTags,
+      });
+
+      return data;
+    },
+    onError: (error: AxiosError) => {
+      const message = get(
+        error,
+        ["response", "data", "message"],
+        error.message,
+      );
+
+      toast({
+        title: "Error",
+        description: message,
+        variant: "destructive",
+      });
+    },
+    onSettled: (data, error, variables) => {
+      queryClient.invalidateQueries({
+        queryKey: [SPANS_KEY, { projectId: variables.projectId }],
+      });
+    },
+  });
+};
+
+export default useSpanBatchUpdateMutation;
+
diff --git a/apps/opik-frontend/src/api/traces/useThreadBatchUpdateMutation.ts b/apps/opik-frontend/src/api/traces/useThreadBatchUpdateMutation.ts
@@ -0,0 +1,60 @@
+import { useMutation, useQueryClient } from "@tanstack/react-query";
+import { AxiosError } from "axios";
+import get from "lodash/get";
+
+import api, { THREADS_KEY, TRACES_REST_ENDPOINT } from "@/api/api";
+import { useToast } from "@/components/ui/use-toast";
+
+type UseThreadBatchUpdateMutationParams = {
+  projectId: string;
+  threadIds: string[];
+  thread: {
+    tags?: string[];
+  };
+  mergeTags?: boolean;
+};
+
+const useThreadBatchUpdateMutation = () => {
+  const queryClient = useQueryClient();
+  const { toast } = useToast();
+
+  return useMutation({
+    mutationFn: async ({
+      threadIds,
+      thread,
+      mergeTags,
+    }: UseThreadBatchUpdateMutationParams) => {
+      const { data } = await api.patch(
+        TRACES_REST_ENDPOINT + "threads/batch",
+        {
+          ids: threadIds,
+          update: thread,
+          merge_tags: mergeTags,
+        },
+      );
+
+      return data;
+    },
+    onError: (error: AxiosError) => {
+      const message = get(
+        error,
+        ["response", "data", "message"],
+        error.message,
+      );
+
+      toast({
+        title: "Error",
+        description: message,
+        variant: "destructive",
+      });
+    },
+    onSettled: (data, error, variables) => {
+      queryClient.invalidateQueries({
+        queryKey: [THREADS_KEY, { projectId: variables.projectId }],
+      });
+    },
+  });
+};
+
+export default useThreadBatchUpdateMutation;
+
diff --git a/apps/opik-frontend/src/api/traces/useTraceBatchUpdateMutation.ts b/apps/opik-frontend/src/api/traces/useTraceBatchUpdateMutation.ts
@@ -0,0 +1,56 @@
+import { useMutation, useQueryClient } from "@tanstack/react-query";
+import { AxiosError } from "axios";
+import get from "lodash/get";
+
+import api, { TRACES_KEY, TRACES_REST_ENDPOINT } from "@/api/api";
+import { Trace } from "@/types/traces";
+import { useToast } from "@/components/ui/use-toast";
+
+type UseTraceBatchUpdateMutationParams = {
+  projectId: string;
+  traceIds: string[];
+  trace: Partial<Trace>;
+  mergeTags?: boolean;
+};
+
+const useTraceBatchUpdateMutation = () => {
+  const queryClient = useQueryClient();
+  const { toast } = useToast();
+
+  return useMutation({
+    mutationFn: async ({
+      traceIds,
+      trace,
+      mergeTags,
+    }: UseTraceBatchUpdateMutationParams) => {
+      const { data } = await api.patch(TRACES_REST_ENDPOINT + "batch", {
+        ids: traceIds,
+        update: trace,
+        merge_tags: mergeTags,
+      });
+
+      return data;
+    },
+    onError: (error: AxiosError) => {
+      const message = get(
+        error,
+        ["response", "data", "message"],
+        error.message,
+      );
+
+      toast({
+        title: "Error",
+        description: message,
+        variant: "destructive",
+      });
+    },
+    onSettled: (data, error, variables) => {
+      queryClient.invalidateQueries({
+        queryKey: [TRACES_KEY, { projectId: variables.projectId }],
+      });
+    },
+  });
+};
+
+export default useTraceBatchUpdateMutation;
+
diff --git a/apps/opik-frontend/src/components/pages-shared/traces/AddTagDialog/AddTagDialog.tsx b/apps/opik-frontend/src/components/pages-shared/traces/AddTagDialog/AddTagDialog.tsx
@@ -1,5 +1,5 @@
 import React, { useState } from "react";
-import { Trace, Span } from "@/types/traces";
+import { Trace, Span, Thread } from "@/types/traces";
 import { TRACE_DATA_TYPE } from "@/hooks/useTracesOrSpansList";
 import {
   Dialog,
@@ -11,16 +11,23 @@ import {
 import { Button } from "@/components/ui/button";
 import { Input } from "@/components/ui/input";
 import { useToast } from "@/components/ui/use-toast";
-import useTraceUpdateMutation from "@/api/traces/useTraceUpdateMutation";
-import useSpanUpdateMutation from "@/api/traces/useSpanUpdateMutation";
+import useTraceBatchUpdateMutation from "@/api/traces/useTraceBatchUpdateMutation";
+import useSpanBatchUpdateMutation from "@/api/traces/useSpanBatchUpdateMutation";
+import useThreadBatchUpdateMutation from "@/api/traces/useThreadBatchUpdateMutation";
 import useAppStore from "@/store/AppStore";
 
+export enum TAG_ENTITY_TYPE {
+  traces = "traces",
+  spans = "spans",
+  threads = "threads",
+}
+
 type AddTagDialogProps = {
-  rows: Array<Trace | Span>;
+  rows: Array<Trace | Span | Thread>;
   open: boolean | number;
   setOpen: (open: boolean | number) => void;
   projectId: string;
-  type: TRACE_DATA_TYPE;
+  type: TRACE_DATA_TYPE | TAG_ENTITY_TYPE.threads;
   onSuccess?: () => void;
 };
 
@@ -35,9 +42,9 @@ const AddTagDialog: React.FunctionComponent<AddTagDialogProps> = ({
   const { toast } = useToast();
   const workspaceName = useAppStore((state) => state.activeWorkspaceName);
   const [newTag, setNewTag] = useState<string>("");
-  const traceUpdateMutation = useTraceUpdateMutation();
-  const spanUpdateMutation = useSpanUpdateMutation();
-  const MAX_ENTITIES = 10;
+  const traceBatchUpdateMutation = useTraceBatchUpdateMutation();
+  const spanBatchUpdateMutation = useSpanBatchUpdateMutation();
+  const threadBatchUpdateMutation = useThreadBatchUpdateMutation();
 
   const handleClose = () => {
     setOpen(false);
@@ -47,54 +54,55 @@ const AddTagDialog: React.FunctionComponent<AddTagDialogProps> = ({
   const handleAddTag = () => {
     if (!newTag) return;
 
-    const promises: Promise<unknown>[] = [];
-
-    rows.forEach((row) => {
-      const currentTags = row.tags || [];
-
-      if (currentTags.includes(newTag)) return;
-
-      const newTags = [...currentTags, newTag];
-
-      if (type === TRACE_DATA_TYPE.traces) {
-        promises.push(
-          traceUpdateMutation.mutateAsync({
-            projectId,
-            traceId: row.id,
-            trace: {
-              workspace_name: workspaceName,
-              project_id: projectId,
-              tags: newTags,
-            },
-          }),
-        );
-      } else {
-        const span = row as Span;
-        const parentId = span.parent_span_id;
+    let mutationPromise;
+    let entityName;
 
-        promises.push(
-          spanUpdateMutation.mutateAsync({
-            projectId,
-            spanId: span.id,
-            span: {
-              workspace_name: workspaceName,
-              project_id: projectId,
-              ...(parentId && { parent_span_id: parentId }),
-              trace_id: span.trace_id,
-              tags: newTags,
-            },
-          }),
-        );
-      }
-    });
+    if (type === TRACE_DATA_TYPE.traces) {
+      const ids = rows.map((row) => row.id);
+      mutationPromise = traceBatchUpdateMutation.mutateAsync({
+        projectId,
+        traceIds: ids,
+        trace: {
+          workspace_name: workspaceName,
+          project_id: projectId,
+          tags: [newTag],
+        },
+        mergeTags: true,
+      });
+      entityName = "traces";
+    } else if (type === TRACE_DATA_TYPE.spans) {
+      const ids = rows.map((row) => row.id);
+      mutationPromise = spanBatchUpdateMutation.mutateAsync({
+        projectId,
+        spanIds: ids,
+        span: {
+          workspace_name: workspaceName,
+          project_id: projectId,
+          trace_id: (rows[0] as Span).trace_id,
+          tags: [newTag],
+        },
+        mergeTags: true,
+      });
+      entityName = "spans";
+    } else {
+      // threads - use thread_model_id instead of id
+      const threadModelIds = rows.map((row) => (row as Thread).thread_model_id);
+      mutationPromise = threadBatchUpdateMutation.mutateAsync({
+        projectId,
+        threadIds: threadModelIds,
+        thread: {
+          tags: [newTag],
+        },
+        mergeTags: true,
+      });
+      entityName = "threads";
+    }
 
-    Promise.all(promises)
+    mutationPromise
       .then(() => {
         toast({
           title: "Success",
-          description: `Tag "${newTag}" added to ${rows.length} selected ${
-            type === TRACE_DATA_TYPE.traces ? "traces" : "spans"
-          }`,
+          description: `Tag "${newTag}" added to ${rows.length} selected ${entityName}`,
         });
 
         if (onSuccess) {
@@ -104,7 +112,7 @@ const AddTagDialog: React.FunctionComponent<AddTagDialogProps> = ({
         handleClose();
       })
       .catch(() => {
-        // Error handling is already done by the mutation hooks,this just ensures we don't close the dialog on error
+        // Error handling is already done by the mutation hooks
       });
   };
 
@@ -114,34 +122,28 @@ const AddTagDialog: React.FunctionComponent<AddTagDialogProps> = ({
         <DialogHeader>
           <DialogTitle>
             Add tag to {rows.length}{" "}
-            {type === TRACE_DATA_TYPE.traces ? "traces" : "spans"}
+            {type === TRACE_DATA_TYPE.traces
+              ? "traces"
+              : type === TRACE_DATA_TYPE.spans
+                ? "spans"
+                : "threads"}
           </DialogTitle>
         </DialogHeader>
-        {rows.length > MAX_ENTITIES && (
-          <div className="mb-2 text-sm text-destructive">
-            You can only add tags to up to {MAX_ENTITIES} entities at a time.
-            Please select fewer entities.
-          </div>
-        )}
         <div className="grid gap-4 py-4">
           <div className="flex items-center gap-4">
             <Input
               placeholder="New tag"
               value={newTag}
               onChange={(event) => setNewTag(event.target.value)}
               className="col-span-3"
-              disabled={rows.length > MAX_ENTITIES}
             />
           </div>
         </div>
         <DialogFooter>
           <Button variant="outline" onClick={handleClose}>
             Cancel
           </Button>
-          <Button
-            onClick={handleAddTag}
-            disabled={!newTag || rows.length > MAX_ENTITIES}
-          >
+          <Button onClick={handleAddTag} disabled={!newTag}>
             Add tag
           </Button>
         </DialogFooter>
diff --git a/apps/opik-frontend/src/components/pages/TracesPage/ThreadsTab/ThreadsActionsPanel.tsx b/apps/opik-frontend/src/components/pages/TracesPage/ThreadsTab/ThreadsActionsPanel.tsx
diff --git a/apps/opik-frontend/src/components/pages/TracesPage/ThreadsTab/ThreadsTab.tsx b/apps/opik-frontend/src/components/pages/TracesPage/ThreadsTab/ThreadsTab.tsx
