Merge pull request #1096 from Agenta-AI/Sub-issue-1059/-Add-cost,-lat…

…ency,-tokens-to-playground

Add cost, latency, tokens to playground

agenta-web/src/components/Playground/Views/TestView.tsx

@@ -97,6 +97,11 @@ interface BoxComponentProps {
     inputParams: Parameter[] | null
     testData: GenericObject
     result: string
+    additionalData: {
+        cost: number | null
+        latency: number | null
+        usage: {completion_tokens: number; prompt_tokens: number; total_tokens: number} | null
+    }
     onInputParamChange: (paramName: string, newValue: any) => void
     onRun: () => void
     onAddToTestset: (params: Record<string, string>) => void
@@ -109,6 +114,7 @@ const BoxComponent: React.FC<BoxComponentProps> = ({
     inputParams,
     testData,
     result,
+    additionalData,
     onInputParamChange,
     onRun,
     onAddToTestset,
@@ -160,6 +166,28 @@ const BoxComponent: React.FC<BoxComponentProps> = ({
                     imageSize="large"
                 />
             </Row>
+            {additionalData && (
+                <Space>
+                    <p>
+                        Tokens:{" "}
+                        {additionalData.usage !== null
+                            ? JSON.stringify(additionalData.usage.total_tokens)
+                            : 0}
+                    </p>
+                    <p>
+                        Cost:{" "}
+                        {additionalData.cost !== null
+                            ? `$${additionalData.cost.toFixed(4)}`
+                            : "$0.00"}
+                    </p>
+                    <p>
+                        Latency:{" "}
+                        {additionalData.latency !== null
+                            ? `${Math.round(additionalData.latency * 1000)}ms`
+                            : "0ms"}
+                    </p>
+                </Space>
+            )}
             <Row className={classes.row2} style={{marginBottom: isChatVariant ? 12 : 0}}>
                 <Col span={24} className={classes.row2Col} id={variant.variantId}>
                     <Button
@@ -224,6 +252,13 @@ const App: React.FC<TestViewProps> = ({
     const [params, setParams] = useState<Record<string, string> | null>(null)
     const classes = useStylesApp()
     const rootRef = React.useRef<HTMLDivElement>(null)
+    const [additionalDataList, setAdditionalDataList] = useState<
+        Array<{
+            cost: number | null
+            latency: number | null
+            usage: {completion_tokens: number; prompt_tokens: number; total_tokens: number} | null
+        }>
+    >(testList.map(() => ({cost: null, latency: null, usage: null})))
 
     useEffect(() => {
         setResultsList((prevResultsList) => {
@@ -306,7 +341,12 @@ const App: React.FC<TestViewProps> = ({
                 isChatVariant ? testItem.chat : [],
             )
 
-            setResultForIndex(res, index)
+            setResultForIndex(res.message, index)
+            setAdditionalDataList((prev) => {
+                const newDataList = [...prev]
+                newDataList[index] = {cost: res.cost, latency: res.latency, usage: res.usage}
+                return newDataList
+            })
         } catch (e) {
             setResultForIndex(
                 "The code has resulted in the following error: \n\n --------------------- \n" +
@@ -399,6 +439,7 @@ const App: React.FC<TestViewProps> = ({
                                   ?.content
                             : resultsList[index]
                     }
+                    additionalData={additionalDataList[index]}
                     onInputParamChange={(paramName, value) =>
                         handleInputParamChange(paramName, value, index)
                     }