From 0500a87d98c6198807a04268f23a436fb260bec3 Mon Sep 17 00:00:00 2001
From: aaronvg <aaron@boundaryml.com>
Date: Sun, 26 Jan 2025 19:54:51 -0800
Subject: [PATCH] show token usage in the playground card view (#1376)

- Show token usage in card view
- Add contrast in sidebar and in function name and test name

<!-- ELLIPSIS_HIDDEN -->


----

> [!IMPORTANT]
> Enhance playground card view by displaying token usage and improving
UI contrast, with refactoring for better code maintainability.
>
>   - **UI Enhancements**:
> - Display token usage in `ResponseRenderer.tsx` with input/output
token counts.
> - Improve text contrast in `function-test-name.tsx` and
`side-bar/index.tsx`.
>   - **Code Refactoring**:
> - Replace `useAtom` with `useSetAtom` in `preview-toolbar.tsx` for
`setShowEnvDialog`.
> - Add `className` prop to `Loader` in `components.tsx` for flexible
styling.
>   - **Miscellaneous**:
> - Update `root.code-workspace` to use `biomejs.biome` as the default
formatter for `typescriptreact`.
> - Add `ChevronsLeftRight` icon in `ResponseRenderer.tsx` for token
display.
>
> <sup>This description was created by </sup>[<img alt="Ellipsis"
src="https://img.shields.io/badge/Ellipsis-blue?color=175173">](https://www.ellipsis.dev?ref=BoundaryML%2Fbaml&utm_source=github&utm_medium=referral)<sup>
for 68f22c1850dee4e483e191cd1cec665fab8378b3. It will automatically
update as commits are pushed.</sup>


<!-- ELLIPSIS_HIDDEN -->
---
 root.code-workspace                           |  4 +++-
 .../playground-panel/function-test-name.tsx   |  2 +-
 .../playground-panel/preview-toolbar.tsx      |  4 ++--
 .../prompt-preview/components.tsx             |  5 ++--
 .../components/ResponseRenderer.tsx           | 24 ++++++++++++++++++-
 .../test-panel/components/TabularView.tsx     |  5 +++-
 .../playground-panel/side-bar/index.tsx       |  6 ++---
 7 files changed, 39 insertions(+), 11 deletions(-)

diff --git a/root.code-workspace b/root.code-workspace
index 88c503369..ee42923d7 100644
--- a/root.code-workspace
+++ b/root.code-workspace
@@ -44,7 +44,7 @@
       "editor.defaultFormatter": "biomejs.biome"
     },
     "[typescriptreact]": {
-      "editor.defaultFormatter": "esbenp.prettier-vscode"
+      "editor.defaultFormatter": "biomejs.biome"
     },
     "files.associations": {
       "*.baml.j2": "jinja",
@@ -83,6 +83,8 @@
       "titleBar.activeBackground": "#8B0002",
       "titleBar.activeForeground": "#FFFBFB"
     },
+    "biome.lsp.bin": "typescript/node_modules/@biomejs/biome/bin/biome",
+
     "biome.lspBin": "typescript/node_modules/@biomejs/biome/bin/biome",
     "terminal.integrated.scrollback": 1e6
   },
diff --git a/typescript/playground-common/src/shared/baml-project-panel/playground-panel/function-test-name.tsx b/typescript/playground-common/src/shared/baml-project-panel/playground-panel/function-test-name.tsx
index 7861bb2d3..50c2bf906 100644
--- a/typescript/playground-common/src/shared/baml-project-panel/playground-panel/function-test-name.tsx
+++ b/typescript/playground-common/src/shared/baml-project-panel/playground-panel/function-test-name.tsx
@@ -33,7 +33,7 @@ export const FunctionTestName: React.FC<FunctionTestNameProps> = ({ functionName
   })
 
   return (
-    <div className={`flex w-full items-center space-x-1 text-xs ${selected ? '' : 'text-muted-foreground'}`}>
+    <div className={`flex w-full items-center space-x-1 text-xs ${selected ? '' : 'text-foreground/90'}`}>
       <div
         className='flex items-center cursor-pointer hover:text-primary'
         onClick={() => {
diff --git a/typescript/playground-common/src/shared/baml-project-panel/playground-panel/preview-toolbar.tsx b/typescript/playground-common/src/shared/baml-project-panel/playground-panel/preview-toolbar.tsx
index 365c80f79..24c0c5d25 100644
--- a/typescript/playground-common/src/shared/baml-project-panel/playground-panel/preview-toolbar.tsx
+++ b/typescript/playground-common/src/shared/baml-project-panel/playground-panel/preview-toolbar.tsx
@@ -2,7 +2,7 @@
 
 import { Button } from '@/components/ui/button'
 import { DropdownMenu, DropdownMenuContent, DropdownMenuItem, DropdownMenuTrigger } from '@/components/ui/dropdown-menu'
-import { atom, useAtom, useAtomValue } from 'jotai'
+import { atom, useAtom, useAtomValue, useSetAtom } from 'jotai'
 import { Braces, Bug, BugIcon, ChevronDown, Copy, FileJson, PlayCircle, Settings, Workflow } from 'lucide-react'
 import React from 'react'
 import { ThemeToggle } from '../theme/ThemeToggle'
@@ -43,7 +43,7 @@ export const isClientCallGraphEnabledAtom = atom(false)
 export default function Component() {
   const [renderMode, setRenderMode] = useAtom(renderModeAtom)
   const selections = useAtomValue(selectedItemAtom)
-  const [showEnvDialog, setShowEnvDialog] = useAtom(showEnvDialogAtom)
+  const setShowEnvDialog = useSetAtom(showEnvDialogAtom)
 
   const options: {
     label: string
diff --git a/typescript/playground-common/src/shared/baml-project-panel/playground-panel/prompt-preview/components.tsx b/typescript/playground-common/src/shared/baml-project-panel/playground-panel/prompt-preview/components.tsx
index 131abf1c6..df29777f5 100644
--- a/typescript/playground-common/src/shared/baml-project-panel/playground-panel/prompt-preview/components.tsx
+++ b/typescript/playground-common/src/shared/baml-project-panel/playground-panel/prompt-preview/components.tsx
@@ -1,10 +1,11 @@
 import { Check, Copy, Loader2 } from 'lucide-react'
 import { Button } from '@/components/ui/button'
 import { useState } from 'react'
+import { cn } from '@/lib/utils'
 
-export const Loader: React.FC<{ message?: string }> = ({ message }) => {
+export const Loader: React.FC<{ message?: string; className?: string }> = ({ message, className }) => {
   return (
-    <div className='flex gap-2 justify-center items-center text-gray-500'>
+    <div className={cn('flex gap-2 justify-center items-center text-gray-500', className)}>
       <Loader2 className='animate-spin' />
       {message}
     </div>
diff --git a/typescript/playground-common/src/shared/baml-project-panel/playground-panel/prompt-preview/test-panel/components/ResponseRenderer.tsx b/typescript/playground-common/src/shared/baml-project-panel/playground-panel/prompt-preview/test-panel/components/ResponseRenderer.tsx
index 67d199185..916272a4c 100644
--- a/typescript/playground-common/src/shared/baml-project-panel/playground-panel/prompt-preview/test-panel/components/ResponseRenderer.tsx
+++ b/typescript/playground-common/src/shared/baml-project-panel/playground-panel/prompt-preview/test-panel/components/ResponseRenderer.tsx
@@ -2,7 +2,17 @@
 import { WasmFunctionResponse, WasmTestResponse, WasmLLMFailure, WasmLLMResponse } from '@gloo-ai/baml-schema-wasm-web'
 import { DoneTestStatusType } from '../../../atoms'
 import { useState } from 'react'
-import { AlertCircle, Brain, Check, CheckCircle, ChevronDown, ChevronUp, Clock, Copy } from 'lucide-react'
+import {
+  AlertCircle,
+  Brain,
+  Check,
+  CheckCircle,
+  ChevronDown,
+  ChevronsLeftRight,
+  ChevronUp,
+  Clock,
+  Copy,
+} from 'lucide-react'
 import { Button } from '~/components/ui/button'
 import { Badge } from '~/components/ui/badge'
 import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from '~/components/ui/tooltip'
@@ -132,6 +142,18 @@ const MetadataBadges: React.FC<{ llmResponse: WasmLLMResponse }> = ({ llmRespons
       </TooltipTrigger>
       <TooltipContent>Latency</TooltipContent>
     </Tooltip>
+    <Tooltip delayDuration={0}>
+      <TooltipTrigger asChild>
+        <Badge variant='outline' className='flex items-center space-x-1 font-light text-muted-foreground'>
+          <ChevronsLeftRight className='w-3 h-3' />
+          <span>
+            {llmResponse.input_tokens != null ? Number(llmResponse.input_tokens) : 'unknown'} in |{' '}
+            {llmResponse.output_tokens != null ? Number(llmResponse.output_tokens) : 'unknown'} out
+          </span>
+        </Badge>
+      </TooltipTrigger>
+      <TooltipContent>Tokens (in / out)</TooltipContent>
+    </Tooltip>
   </TooltipProvider>
 )
 
diff --git a/typescript/playground-common/src/shared/baml-project-panel/playground-panel/prompt-preview/test-panel/components/TabularView.tsx b/typescript/playground-common/src/shared/baml-project-panel/playground-panel/prompt-preview/test-panel/components/TabularView.tsx
index 95705e590..db9de28cc 100644
--- a/typescript/playground-common/src/shared/baml-project-panel/playground-panel/prompt-preview/test-panel/components/TabularView.tsx
+++ b/typescript/playground-common/src/shared/baml-project-panel/playground-panel/prompt-preview/test-panel/components/TabularView.tsx
@@ -89,7 +89,10 @@ const ResponseContent = ({
           <ParsedResponseRenderer response={getTestStateResponse(state)} />
 
           {getExplanation(state) && (
-            <div className='mt-2 text-xs text-muted-foreground/80'>{getExplanation(state)}</div>
+            <div className='flex flex-col gap-2 mt-2 text-xs text-muted-foreground/80'>
+              <div>BAML parser fixed the following issues:</div>
+              <pre>{getExplanation(state)}</pre>
+            </div>
           )}
         </>
       )}
diff --git a/typescript/playground-common/src/shared/baml-project-panel/playground-panel/side-bar/index.tsx b/typescript/playground-common/src/shared/baml-project-panel/playground-panel/side-bar/index.tsx
index f7d8a7f98..ab332c407 100644
--- a/typescript/playground-common/src/shared/baml-project-panel/playground-panel/side-bar/index.tsx
+++ b/typescript/playground-common/src/shared/baml-project-panel/playground-panel/side-bar/index.tsx
@@ -11,7 +11,6 @@ import {
   ChevronLeft,
   ChevronRight,
   FlaskConical,
-  Loader,
   Play,
   Search,
   Settings,
@@ -28,6 +27,7 @@ import EnvVars from './env-vars'
 import { ScrollArea } from '@/components/ui/scroll-area'
 import { atomWithStorage } from 'jotai/utils'
 import { vscode } from '../../vscode'
+import { Loader } from '../prompt-preview/components'
 
 interface FunctionData {
   name: string
@@ -117,7 +117,7 @@ export default function CustomSidebar() {
                 </div>
                 <div className='overflow-auto flex-1'>
                   <div className='px-2'>
-                    {searchTerm && filteredFunctions.length > 0 && (
+                    {filteredFunctions.length > 0 && (
                       <Button
                         variant='ghost'
                         size='sm'
@@ -213,7 +213,7 @@ function FunctionItem({ label, tests, isLast = false, isSelected = false, search
       <div
         className={cn(
           'flex relative items-center px-1 py-1 -mx-2 transition-colors cursor-pointer group hover:bg-muted',
-          isSelected ? 'font-bold text-purple-400' : 'text-muted-foreground',
+          isSelected ? 'font-bold text-purple-400' : 'text-muted-primary',
         )}
         onClick={handleClick}
       >