[OPIK-619]: playground fe cache (#936)

* [OPIK-619]: cache playground state in localstorage; * [OPIK-619]: add a reset playground button; * [OPIK-619]: lint fixes; * [OPIK-619]: remove console.log; * [OPIK-619]: stop streaming when the location changes; * [OPIK-619]: rename function; --------- Co-authored-by: Sasha <[email protected]>
comet-ml · Dec 19, 2024 · 0500603 · 0500603
1 parent 9ca2b89
commit 0500603
Show file tree

Hide file tree

Showing 10 changed files with 232 additions and 111 deletions.
diff --git a/apps/opik-frontend/src/api/playground/useCompletionProxyStreaming.ts b/apps/opik-frontend/src/api/playground/useCompletionProxyStreaming.ts
@@ -1,4 +1,4 @@
-import { useCallback, useRef } from "react";
+import { useCallback, useEffect, useRef } from "react";
 
 import dayjs from "dayjs";
 import { UsageType } from "@/types/shared";
@@ -7,13 +7,18 @@ import {
   ProviderMessageType,
   ChatCompletionMessageChoiceType,
   ChatCompletionResponse,
-  ChatCompletionErrorMessageType,
+  ChatCompletionProviderErrorMessageType,
   ChatCompletionSuccessMessageType,
-  ChatCompletionProxyErrorMessageType,
+  ChatCompletionOpikErrorMessageType,
 } from "@/types/playground";
-import { safelyParseJSON, snakeCaseObj } from "@/lib/utils";
+import { isValidJsonObject, safelyParseJSON, snakeCaseObj } from "@/lib/utils";
 import { BASE_API_URL } from "@/api/api";
 import { PROVIDER_MODEL_TYPE } from "@/types/providers";
+import { useLocation } from "@tanstack/react-router";
+
+const getNowUtcTimeISOString = (): string => {
+  return dayjs().utc().toISOString();
+};
 
 interface GetCompletionProxyStreamParams {
   model: PROVIDER_MODEL_TYPE | "";
@@ -23,6 +28,21 @@ interface GetCompletionProxyStreamParams {
   workspaceName: string;
 }
 
+const isOpikError = (
+  response: ChatCompletionResponse,
+): response is ChatCompletionOpikErrorMessageType => {
+  return (
+    "errors" in response ||
+    ("code" in response && !isValidJsonObject(response.message))
+  );
+};
+
+const isProviderError = (
+  response: ChatCompletionResponse,
+): response is ChatCompletionProviderErrorMessageType => {
+  return "code" in response && isValidJsonObject(response.message);
+};
+
 const getCompletionProxyStream = async ({
   model,
   messages,
@@ -54,8 +74,8 @@ export interface RunStreamingReturn {
   endTime: string;
   usage: UsageType | null;
   choices: ChatCompletionMessageChoiceType[] | null;
-  platformError: null | string;
-  proxyError: null | string;
+  providerError: null | string;
+  opikError: null | string;
 }
 
 interface UseCompletionProxyStreamingParameters {
@@ -81,15 +101,15 @@ const useCompletionProxyStreaming = ({
   }, []);
 
   const runStreaming = useCallback(async (): Promise<RunStreamingReturn> => {
-    const startTime = dayjs().utc().toISOString();
+    const startTime = getNowUtcTimeISOString();
 
     let accumulatedValue = "";
     let usage = null;
     let choices: ChatCompletionMessageChoiceType[] = [];
 
     // errors
-    let proxyError = null;
-    let platformError = null;
+    let opikError = null;
+    let providerError = null;
 
     try {
       abortControllerRef.current = new AbortController();
@@ -122,25 +142,30 @@ const useCompletionProxyStreaming = ({
       };
 
       const handleAIPlatformErrorMessage = (
-        parsedMessage: ChatCompletionErrorMessageType,
+        parsedMessage: ChatCompletionProviderErrorMessageType,
       ) => {
         const message = safelyParseJSON(parsedMessage?.message);
 
-        platformError = message?.error?.message;
+        providerError = message?.error?.message;
       };
 
-      const handleProxyErrorMessage = (
-        parsedMessage: ChatCompletionProxyErrorMessageType,
+      const handleOpikErrorMessage = (
+        parsedMessage: ChatCompletionOpikErrorMessageType,
       ) => {
-        proxyError = parsedMessage.errors.join(" ");
+        if ("code" in parsedMessage && "message" in parsedMessage) {
+          opikError = parsedMessage.message;
+          return;
+        }
+
+        opikError = parsedMessage.errors.join(" ");
       };
 
       // an analogue of true && reader
       // we need it to wait till the stream is closed
       while (reader) {
         const { done, value } = await reader.read();
 
-        if (done || proxyError || platformError) {
+        if (done || opikError || providerError) {
           break;
         }
 
@@ -151,9 +176,9 @@ const useCompletionProxyStreaming = ({
           const parsed = safelyParseJSON(line) as ChatCompletionResponse;
 
           // handle different message types
-          if ("errors" in parsed) {
-            handleProxyErrorMessage(parsed);
-          } else if ("code" in parsed) {
+          if (isOpikError(parsed)) {
+            handleOpikErrorMessage(parsed);
+          } else if (isProviderError(parsed)) {
             handleAIPlatformErrorMessage(parsed);
           } else {
             handleSuccessMessage(parsed);
@@ -162,10 +187,10 @@ const useCompletionProxyStreaming = ({
       }
       return {
         startTime,
-        endTime: dayjs().utc().toISOString(),
+        endTime: getNowUtcTimeISOString(),
         result: accumulatedValue,
-        platformError,
-        proxyError,
+        providerError,
+        opikError,
         usage,
         choices,
       };
@@ -179,16 +204,23 @@ const useCompletionProxyStreaming = ({
 
       return {
         startTime,
-        endTime: dayjs().utc().toISOString(),
+        endTime: getNowUtcTimeISOString(),
         result: accumulatedValue,
-        platformError,
-        proxyError: proxyError || defaultErrorMessage,
+        providerError,
+        opikError: opikError || defaultErrorMessage,
         usage: null,
         choices,
       };
     }
   }, [messages, model, onAddChunk, configs, workspaceName]);
 
+  const location = useLocation();
+
+  useEffect(() => {
+    // stop streaming whenever the location changes
+    return () => stop();
+  }, [location, stop]);
+
   return { runStreaming, stop };
 };
 

diff --git a/apps/opik-frontend/src/api/playground/useCreateOutputTraceAndSpan.ts b/apps/opik-frontend/src/api/playground/useCreateOutputTraceAndSpan.ts
@@ -42,7 +42,7 @@ const useCreateOutputTraceAndSpan = () => {
       endTime,
       result,
       usage,
-      platformError,
+      providerError,
       choices,
       model,
       providerMessages,
@@ -59,7 +59,7 @@ const useCreateOutputTraceAndSpan = () => {
           startTime,
           endTime,
           input: { messages: providerMessages },
-          output: { output: result || platformError },
+          output: { output: result || providerError },
         });
 
         await createSpanMutateAsync({

diff --git a/apps/opik-frontend/src/components/pages/ConfigurationPage/AIProvidersTab/AIProvidersTab.tsx b/apps/opik-frontend/src/components/pages/ConfigurationPage/AIProvidersTab/AIProvidersTab.tsx
@@ -1,5 +1,4 @@
 import React, { useMemo, useRef, useState } from "react";
-import { keepPreviousData } from "@tanstack/react-query";
 import { ColumnPinningState } from "@tanstack/react-table";
 
 import { convertColumnDataToColumn } from "@/lib/table";
@@ -58,7 +57,6 @@ const AIProvidersTab = () => {
       workspaceName,
     },
     {
-      placeholderData: keepPreviousData,
       refetchInterval: 30000,
     },
   );

diff --git a/.../opik-frontend/src/components/pages/PlaygroundPage/PlaygroundOutputs/PlaygroundOutput.tsx b/.../opik-frontend/src/components/pages/PlaygroundPage/PlaygroundOutputs/PlaygroundOutput.tsx
@@ -6,6 +6,7 @@ import React, {
   useState,
 } from "react";
 import {
+  PlaygroundOutputType,
   PlaygroundMessageType,
   PlaygroundPromptConfigsType,
 } from "@/types/playground";
@@ -14,14 +15,16 @@ import { getAlphabetLetter } from "@/lib/utils";
 import { transformMessageIntoProviderMessage } from "@/lib/playground";
 import PlaygroundOutputLoader from "@/components/pages/PlaygroundPage/PlaygroundOutputs/PlaygroundOutputLoader/PlaygroundOutputLoader";
 import useCreateOutputTraceAndSpan from "@/api/playground/useCreateOutputTraceAndSpan";
-import useAppStore from "@/store/AppStore";
 import { PROVIDER_MODEL_TYPE } from "@/types/providers";
 
 interface PlaygroundOutputProps {
   model: PROVIDER_MODEL_TYPE | "";
+  workspaceName: string;
   messages: PlaygroundMessageType[];
   index: number;
   configs: PlaygroundPromptConfigsType;
+  output: PlaygroundOutputType;
+  onOutputChange: (output: PlaygroundOutputType) => void;
 }
 
 export interface PlaygroundOutputRef {
@@ -30,13 +33,11 @@ export interface PlaygroundOutputRef {
 }
 
 const PlaygroundOutput = forwardRef<PlaygroundOutputRef, PlaygroundOutputProps>(
-  ({ model, messages, index, configs }, ref) => {
-    const workspaceName = useAppStore((state) => state.activeWorkspaceName);
-
+  (
+    { model, messages, index, configs, onOutputChange, output, workspaceName },
+    ref,
+  ) => {
     const [isLoading, setIsLoading] = useState(false);
-    const [error, setError] = useState<string | null>(null);
-
-    const [outputText, setOutputText] = useState<string | null>(null);
 
     const providerMessages = useMemo(() => {
       return messages.map(transformMessageIntoProviderMessage);
@@ -47,24 +48,24 @@ const PlaygroundOutput = forwardRef<PlaygroundOutputRef, PlaygroundOutputProps>(
       model,
       configs,
       messages: providerMessages,
-      onAddChunk: setOutputText,
+      onAddChunk: onOutputChange,
       workspaceName,
     });
 
     const createOutputTraceAndSpan = useCreateOutputTraceAndSpan();
 
     const exposedRun = useCallback(async () => {
-      setError(null);
       setIsLoading(true);
-      setOutputText(null);
+      onOutputChange(null);
 
       const streaming = await runStreaming();
 
       setIsLoading(false);
 
-      const error = streaming.platformError || streaming.proxyError;
+      const error = streaming.providerError || streaming.opikError;
+
       if (error) {
-        setError(error);
+        onOutputChange(error);
       }
 
       createOutputTraceAndSpan({
@@ -79,6 +80,7 @@ const PlaygroundOutput = forwardRef<PlaygroundOutputRef, PlaygroundOutputProps>(
       model,
       configs,
       providerMessages,
+      onOutputChange,
     ]);
 
     useImperativeHandle(
@@ -91,15 +93,11 @@ const PlaygroundOutput = forwardRef<PlaygroundOutputRef, PlaygroundOutputProps>(
     );
 
     const renderContent = () => {
-      if (isLoading && !outputText) {
+      if (isLoading && !output) {
         return <PlaygroundOutputLoader />;
       }
 
-      if (error) {
-        return `Error: ${error}`;
-      }
-
-      return outputText;
+      return output;
     };
 
     return (

diff --git a/...opik-frontend/src/components/pages/PlaygroundPage/PlaygroundOutputs/PlaygroundOutputs.tsx b/...opik-frontend/src/components/pages/PlaygroundPage/PlaygroundOutputs/PlaygroundOutputs.tsx
@@ -1,8 +1,8 @@
-import React, { useRef, useState } from "react";
+import React, { useCallback, useRef, useState } from "react";
 import { Pause, Play } from "lucide-react";
 
 import { Button } from "@/components/ui/button";
-import { PlaygroundPromptType } from "@/types/playground";
+import { PlaygroundOutputType, PlaygroundPromptType } from "@/types/playground";
 
 import PlaygroundOutput, {
   PlaygroundOutputRef,
@@ -11,9 +11,15 @@ import TooltipWrapper from "@/components/shared/TooltipWrapper/TooltipWrapper";
 
 interface PlaygroundOutputsProps {
   prompts: PlaygroundPromptType[];
+  onChange: (id: string, changes: Partial<PlaygroundPromptType>) => void;
+  workspaceName: string;
 }
 
-const PlaygroundOutputs = ({ prompts }: PlaygroundOutputsProps) => {
+const PlaygroundOutputs = ({
+  prompts,
+  onChange,
+  workspaceName,
+}: PlaygroundOutputsProps) => {
   const [isRunning, setIsRunning] = useState(false);
 
   const outputRefs = useRef<Map<number, PlaygroundOutputRef>>(new Map());
@@ -50,6 +56,13 @@ const PlaygroundOutputs = ({ prompts }: PlaygroundOutputsProps) => {
     setIsRunning(false);
   };
 
+  const handleOutputChange = useCallback(
+    (id: string, output: PlaygroundOutputType) => {
+      onChange(id, { output });
+    },
+    [onChange],
+  );
+
   const renderActionButton = () => {
     if (isRunning) {
       return (
@@ -106,10 +119,13 @@ const PlaygroundOutputs = ({ prompts }: PlaygroundOutputsProps) => {
         {prompts?.map((prompt, promptIdx) => (
           <PlaygroundOutput
             key={`output-${prompt.id}`}
+            workspaceName={workspaceName}
             model={prompt.model}
             index={promptIdx}
             messages={prompt.messages}
+            output={prompt.output}
             configs={prompt.configs}
+            onOutputChange={(o) => handleOutputChange(prompt.id, o)}
             ref={(node) => {
               const map = getOutputRefMap();