Convert to useSWRMutation

langchain-ai · Nov 27, 2023 · cd1a2d4 · cd1a2d4
1 parent 90632cd
commit cd1a2d4
Showing 1 changed file with 36 additions and 51 deletions.
diff --git a/langserve/playground/src/components/feedback/CorrectnessFeedback.tsx b/langserve/playground/src/components/feedback/CorrectnessFeedback.tsx
@@ -2,52 +2,45 @@ import ThumbsUpIcon from "../../assets/ThumbsUpIcon.svg?react";
 import ThumbsDownIcon from "../../assets/ThumbsDownIcon.svg?react";
 import CircleSpinIcon from "../../assets/CircleSpinIcon.svg?react";
 import { resolveApiUrl } from "../../utils/url";
-import { useEffect, useRef, useState } from "react";
+import { useState } from "react";
 import { cn } from "../../utils/cn";
+import useSWRMutation from "swr/mutation";
 
-export function CorrectnessFeedback(props: { runId?: string }) {
-  const [loading, setLoading] = useState<number | null>(null);
-  const [state, setState] = useState<{ id: string; score: number } | null>(
-    null
-  );
-
-  const isMounted = useRef<boolean>(true);
-  useEffect(() => {
-    isMounted.current = true;
-    return () => {
-      isMounted.current = false;
-    };
-  }, []);
-
-  const sendFeedback = async (payload: {
-    run_id: string;
+const useFeedbackMutation = (runId: string) => {
+  interface FeedbackArguments {
     key: string;
     score: number;
-  }) => {
-    if (isMounted.current) setLoading(payload.score);
+  }
+
+  const [lastArg, setLastArg] = useState<FeedbackArguments | null>(null);
 
-    try {
-      const request = await fetch(
-        resolveApiUrl(state?.id ? `/feedback/${state.id}` : "/feedback"),
-        {
-          method: state?.id ? "PATCH" : "POST",
-          headers: { "Content-Type": "application/json" },
-          body: JSON.stringify(payload),
-        }
-      );
+  const mutation = useSWRMutation(
+    ["feedback", runId],
+    async ([, runId], { arg }: { arg: FeedbackArguments }) => {
+      const payload = { run_id: runId, key: arg.key, score: arg.score };
+      setLastArg(arg);
+
+      const request = await fetch(resolveApiUrl("/feedback"), {
+        method: "POST",
+        headers: { "Content-Type": "application/json" },
+        body: JSON.stringify(payload),
+      });
 
       if (!request.ok) throw new Error(`Failed request ${request.status}`);
       const json: {
         id: string;
         score: number;
       } = await request.json();
 
-      if (isMounted.current)
-        setState(json ?? { id: state?.id, score: payload.score });
-    } finally {
-      if (isMounted.current) setLoading(null);
+      return json;
     }
-  };
+  );
+
+  return { lastArg: mutation.isMutating ? lastArg : null, mutation };
+};
+
+export function CorrectnessFeedback(props: { runId: string }) {
+  const score = useFeedbackMutation(props.runId);
 
   if (props.runId == null) return null;
   return (
@@ -56,20 +49,16 @@ export function CorrectnessFeedback(props: { runId?: string }) {
         type="button"
         className={cn(
           "border focus-within:border-ls-blue focus-within:outline-none bg-background rounded p-1 border-divider-700 hover:bg-divider-500/50 active:bg-divider-500",
-          state?.score === 1 && "text-teal-500"
+          score.mutation.data?.score === 1 && "text-teal-500"
         )}
-        disabled={loading != null}
+        disabled={score.mutation.isMutating}
         onClick={() => {
-          if (props.runId) {
-            sendFeedback({
-              run_id: props.runId,
-              key: "correctness",
-              score: 1,
-            });
+          if (score.mutation.data?.score !== 1) {
+            score.mutation.trigger({ key: "correctness", score: 1 });
           }
         }}
       >
-        {loading === 1 ? (
+        {score.lastArg?.score === 1 ? (
           <CircleSpinIcon className="animate-spin w-4 h-4 text-white/50 fill-white" />
         ) : (
           <ThumbsUpIcon className="w-4 h-4" />
@@ -80,20 +69,16 @@ export function CorrectnessFeedback(props: { runId?: string }) {
         type="button"
         className={cn(
           "border focus-within:border-ls-blue focus-within:outline-none bg-background rounded p-1 border-divider-700 hover:bg-divider-500/50 active:bg-divider-500",
-          state?.score === -1 && "text-red-500"
+          score.mutation.data?.score === -1 && "text-red-500"
         )}
-        disabled={loading != null}
+        disabled={score.mutation.isMutating}
         onClick={() => {
-          if (props.runId) {
-            sendFeedback({
-              run_id: props.runId,
-              key: "correctness",
-              score: -1,
-            });
+          if (score.mutation.data?.score !== -1) {
+            score.mutation.trigger({ key: "correctness", score: -1 });
           }
         }}
       >
-        {loading === -1 ? (
+        {score.lastArg?.score === -1 ? (
           <CircleSpinIcon className="animate-spin w-4 h-4 text-white/50 fill-white" />
         ) : (
           <ThumbsDownIcon className="w-4 h-4" />