feat(ui): bulk annotate traces and spans from the selection bar

apps/opik-frontend/src/api/traces/useSpanFeedbackScoreBatchSetMutation.ts

@@ -0,0 +1,63 @@
+import { useMutation, useQueryClient } from "@tanstack/react-query";
+import chunk from "lodash/chunk";
+import { AxiosError } from "axios";
+
+import api, { SPANS_KEY, SPANS_REST_ENDPOINT, TRACE_KEY } from "@/api/api";
+import { FEEDBACK_SCORE_TYPE } from "@/types/traces";
+import { useToast } from "@/ui/use-toast";
+import { extractErrorMessage } from "@/lib/errors";
+import {
+  FeedbackScoreBatchEntry,
+  MAX_FEEDBACK_SCORES_PER_BATCH,
+} from "@/lib/feedback-scores";
+
+type UseSpanFeedbackScoreBatchSetMutationParams = {
+  projectName: string;
+  scores: FeedbackScoreBatchEntry[];
+};
+
+const useSpanFeedbackScoreBatchSetMutation = () => {
+  const queryClient = useQueryClient();
+  const { toast } = useToast();
+
+  return useMutation({
+    mutationFn: async ({
+      projectName,
+      scores,
+    }: UseSpanFeedbackScoreBatchSetMutationParams) => {
+      for (const scoresChunk of chunk(scores, MAX_FEEDBACK_SCORES_PER_BATCH)) {
+        await api.put(`${SPANS_REST_ENDPOINT}feedback-scores`, {
+          scores: scoresChunk.map((score) => ({
+            id: score.id,
+            // The backend groups the batch by project name and derives project_id from it,
+            // falling back to the default project when it is blank, so it has to be sent.
+            project_name: projectName,
+            name: score.name,
+            category_name: score.categoryName,
+            value: score.value,
+            reason: score.reason,
+            source: FEEDBACK_SCORE_TYPE.ui,
+          })),
+        });
+      }
+    },
+    onError: (error: AxiosError) => {
+      toast({
+        title: "Error",
+        description: extractErrorMessage(error),
+        variant: "destructive",
+      });
+    },
+    onSettled: async () => {
+      // Mirror the span branch of useTraceFeedbackScoreSetMutation: refresh the spans list,
+      // its columns/statistics and the trace details panel, which shows the aggregated span
+      // scores. The batch only carries span ids, so the trace cache is invalidated broadly.
+      await queryClient.invalidateQueries({ queryKey: [SPANS_KEY] });
+      await queryClient.invalidateQueries({ queryKey: ["spans-columns"] });
+      await queryClient.invalidateQueries({ queryKey: ["spans-statistic"] });
+      await queryClient.invalidateQueries({ queryKey: [TRACE_KEY] });
+    },
+  });
+};
+
+export default useSpanFeedbackScoreBatchSetMutation;

apps/opik-frontend/src/api/traces/useTraceFeedbackScoreBatchSetMutation.ts

@@ -0,0 +1,70 @@
+import { useMutation, useQueryClient } from "@tanstack/react-query";
+import chunk from "lodash/chunk";
+import { AxiosError } from "axios";
+
+import api, { TRACE_KEY, TRACES_KEY, TRACES_REST_ENDPOINT } from "@/api/api";
+import { FEEDBACK_SCORE_TYPE } from "@/types/traces";
+import { useToast } from "@/ui/use-toast";
+import { extractErrorMessage } from "@/lib/errors";
+import {
+  FeedbackScoreBatchEntry,
+  MAX_FEEDBACK_SCORES_PER_BATCH,
+} from "@/lib/feedback-scores";
+
+type UseTraceFeedbackScoreBatchSetMutationParams = {
+  projectName: string;
+  scores: FeedbackScoreBatchEntry[];
+};
+
+const useTraceFeedbackScoreBatchSetMutation = () => {
+  const queryClient = useQueryClient();
+  const { toast } = useToast();
+
+  return useMutation({
+    mutationFn: async ({
+      projectName,
+      scores,
+    }: UseTraceFeedbackScoreBatchSetMutationParams) => {
+      for (const scoresChunk of chunk(scores, MAX_FEEDBACK_SCORES_PER_BATCH)) {
+        await api.put(`${TRACES_REST_ENDPOINT}feedback-scores`, {
+          scores: scoresChunk.map((score) => ({
+            id: score.id,
+            // The backend groups the batch by project name and derives project_id from it,
+            // falling back to the default project when it is blank, so it has to be sent.
+            project_name: projectName,
+            name: score.name,
+            category_name: score.categoryName,
+            value: score.value,
+            reason: score.reason,
+            source: FEEDBACK_SCORE_TYPE.ui,
+          })),
+        });
+      }
+    },
+    onError: (error: AxiosError) => {
+      toast({
+        title: "Error",
+        description: extractErrorMessage(error),
+        variant: "destructive",
+      });
+    },
+    onSettled: async (data, error, variables) => {
+      // Mirror useTraceFeedbackScoreSetMutation: the scores feed the traces list, its
+      // columns/statistics and the per-trace details panel.
+      await queryClient.invalidateQueries({ queryKey: [TRACES_KEY] });
+      await queryClient.invalidateQueries({ queryKey: ["traces-columns"] });
+      await queryClient.invalidateQueries({ queryKey: ["traces-statistic"] });
+
+      const traceIds = [...new Set(variables.scores.map((score) => score.id))];
+      await Promise.all(
+        traceIds.map((traceId) =>
+          queryClient.invalidateQueries({
+            queryKey: [TRACE_KEY, { traceId }],
+          }),
+        ),
+      );
+    },
+  });
+};
+
+export default useTraceFeedbackScoreBatchSetMutation;

apps/opik-frontend/src/lib/feedback-scores.tsx

@@ -637,3 +637,18 @@ export const combineExperimentScoresAsMap = (row: {
 
   return result;
 };
+
+// A single entry of the traces/spans batch feedback score payload: the entity id plus the
+// score to set on it.
+export type FeedbackScoreBatchEntry = {
+  id: string;
+  name: string;
+  value: number;
+  categoryName?: string;
+  reason?: string;
+};
+
+// The backend caps a feedback score batch at 1000 items
+// (FeedbackScoreBatchContainer: `@Size(min = 1, max = 1000)`), and annotating a selection
+// produces `rows * scores` entries, so larger payloads have to be split.
+export const MAX_FEEDBACK_SCORES_PER_BATCH = 1000;

apps/opik-frontend/src/v2/pages-shared/traces/AddAnnotationDialog/AddAnnotationDialog.test.tsx

@@ -0,0 +1,184 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { render, screen, fireEvent, waitFor } from "@testing-library/react";
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
+import { ReactNode } from "react";
+import { TooltipProvider } from "@/ui/tooltip";
+import { TRACE_DATA_TYPE } from "@/hooks/useTracesOrSpansList";
+import { Span, Trace } from "@/types/traces";
+import { UpdateFeedbackScoreData } from "@/v2/pages-shared/traces/TraceDetailsPanel/TraceAnnotateViewer/types";
+import AddAnnotationDialog from "./AddAnnotationDialog";
+
+const mockSetTraceFeedbackScores = vi.fn();
+const mockSetSpanFeedbackScores = vi.fn();
+
+vi.mock("@/api/traces/useTraceFeedbackScoreBatchSetMutation", () => ({
+  default: () => ({
+    mutateAsync: mockSetTraceFeedbackScores,
+    isPending: false,
+  }),
+}));
+vi.mock("@/api/traces/useSpanFeedbackScoreBatchSetMutation", () => ({
+  default: () => ({
+    mutateAsync: mockSetSpanFeedbackScores,
+    isPending: false,
+  }),
+}));
+
+vi.mock("@/ui/use-toast", () => ({
+  useToast: () => ({ toast: vi.fn() }),
+}));
+
+type StubFeedbackScoresEditorProps = {
+  onUpdateFeedbackScore: (update: UpdateFeedbackScoreData) => void;
+  onDeleteFeedbackScore: (name: string) => void;
+};
+
+vi.mock(
+  "@/v2/pages-shared/traces/FeedbackScoresEditor/FeedbackScoresEditor",
+  () => {
+    const Stub = ({
+      onUpdateFeedbackScore,
+      onDeleteFeedbackScore,
+    }: StubFeedbackScoresEditorProps) => (
+      <div data-testid="feedback-scores-editor">
+        <button
+          onClick={() =>
+            onUpdateFeedbackScore({
+              name: "Relevance",
+              value: 1,
+              categoryName: "Yes",
+              reason: "looks good",
+            })
+          }
+        >
+          set-score
+        </button>
+        <button onClick={() => onDeleteFeedbackScore("Relevance")}>
+          clear-score
+        </button>
+      </div>
+    );
+    Stub.displayName = "FeedbackScoresEditorStub";
+    const StubHeader = () => null;
+    StubHeader.displayName = "FeedbackScoresEditorHeaderStub";
+    const StubFooter = () => null;
+    StubFooter.displayName = "FeedbackScoresEditorFooterStub";
+    Stub.Header = StubHeader;
+    Stub.Footer = StubFooter;
+    return { default: Stub };
+  },
+);
+
+describe("AddAnnotationDialog", () => {
+  let queryClient: QueryClient;
+
+  beforeEach(() => {
+    queryClient = new QueryClient({
+      defaultOptions: {
+        queries: { retry: false },
+        mutations: { retry: false },
+      },
+    });
+    vi.clearAllMocks();
+    mockSetTraceFeedbackScores.mockResolvedValue(undefined);
+    mockSetSpanFeedbackScores.mockResolvedValue(undefined);
+  });
+
+  const wrapper = ({ children }: { children: ReactNode }) => (
+    <QueryClientProvider client={queryClient}>
+      <TooltipProvider>{children}</TooltipProvider>
+    </QueryClientProvider>
+  );
+
+  const mockTrace: Trace = {
+    id: "row-1",
+    name: "Test Trace",
+    input: { prompt: "test input" },
+    output: { response: "test output" },
+    start_time: "2024-01-01T00:00:00Z",
+    end_time: "2024-01-01T00:00:01Z",
+    duration: 1000,
+    created_at: "2024-01-01T00:00:00Z",
+    last_updated_at: "2024-01-01T00:00:01Z",
+    tags: [],
+    metadata: {},
+    feedback_scores: [],
+    comments: [],
+    project_id: "project-1",
+  };
+
+  const ROWS: Array<Trace | Span> = [mockTrace, { ...mockTrace, id: "row-2" }];
+
+  const renderDialog = (type: TRACE_DATA_TYPE) =>
+    render(
+      <AddAnnotationDialog
+        rows={ROWS}
+        open
+        setOpen={vi.fn()}
+        projectName="project-name"
+        type={type}
+      />,
+      { wrapper },
+    );
+
+  const EXPECTED_SCORES = [
+    {
+      id: "row-1",
+      name: "Relevance",
+      value: 1,
+      categoryName: "Yes",
+      reason: "looks good",
+    },
+    {
+      id: "row-2",
+      name: "Relevance",
+      value: 1,
+      categoryName: "Yes",
+      reason: "looks good",
+    },
+  ];
+
+  it("sends one batch of trace scores, carrying the project name", async () => {
+    renderDialog(TRACE_DATA_TYPE.traces);
+
+    fireEvent.click(screen.getByRole("button", { name: "set-score" }));
+    fireEvent.click(screen.getByTestId("apply-annotation-button"));
+
+    await waitFor(() =>
+      expect(mockSetTraceFeedbackScores).toHaveBeenCalledTimes(1),
+    );
+    expect(mockSetTraceFeedbackScores).toHaveBeenCalledWith({
+      projectName: "project-name",
+      scores: EXPECTED_SCORES,
+    });
+    expect(mockSetSpanFeedbackScores).not.toHaveBeenCalled();
+  });
+
+  it("routes span selections to the spans endpoint, not the traces one", async () => {
+    renderDialog(TRACE_DATA_TYPE.spans);
+
+    fireEvent.click(screen.getByRole("button", { name: "set-score" }));
+    fireEvent.click(screen.getByTestId("apply-annotation-button"));
+
+    await waitFor(() =>
+      expect(mockSetSpanFeedbackScores).toHaveBeenCalledTimes(1),
+    );
+    expect(mockSetSpanFeedbackScores).toHaveBeenCalledWith({
+      projectName: "project-name",
+      scores: EXPECTED_SCORES,
+    });
+    expect(mockSetTraceFeedbackScores).not.toHaveBeenCalled();
+  });
+
+  it("enables apply only while at least one score is set", () => {
+    renderDialog(TRACE_DATA_TYPE.traces);
+
+    expect(screen.getByTestId("apply-annotation-button")).toBeDisabled();
+
+    fireEvent.click(screen.getByRole("button", { name: "set-score" }));
+    expect(screen.getByTestId("apply-annotation-button")).toBeEnabled();
+
+    fireEvent.click(screen.getByRole("button", { name: "clear-score" }));
+    expect(screen.getByTestId("apply-annotation-button")).toBeDisabled();
+  });
+});