explicit mapping in labels (#135)

* explicit mapping in labels

* Update frontend/components/traces/add-label.tsx

Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>

---------

Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>

frontend/app/globals.css

@@ -48,15 +48,15 @@
     --primary-foreground: 16 0% 91%;
 
     --secondary: 240 4% 16%;
-    --secondary-foreground: 0 4% 70%;
+    --secondary-foreground: 215 16% 65%;
 
     --muted: 240 6% 16%;
-    --muted-foreground: 215 20.2% 65.1%;
+    --muted-foreground: 215 16% 65%;
 
     --accent: 232 9% 17%;
     --accent-foreground: 210 100% 100%;
 
-    --destructive: 0 62.8% 30.6%;
+    --destructive: 0 60% 50%;
     --destructive-foreground: 210 40% 98%;
 
     --border: 240 6% 18%;

frontend/components/evaluator/evaluator-editor.tsx

@@ -171,7 +171,7 @@ export function EvaluatorEditor({
                   <div className="flex flex-col space-y-2">
                     <Label className="text-secondary-foreground">Expected output</Label>
                     <div className="flex space-x-1">
-                      {labelClass.valueMap.map((value, index) => (
+                      {Object.keys(labelClass.valueMap).map((value, index) => (
                         <div
                           key={index}
                           className="border rounded-md p-0.5 px-2 text-sm"

frontend/components/traces/add-label-popover.tsx

@@ -1,21 +1,17 @@
 import {
   LabelClass,
   LabelSource,
-  LabelType,
   RegisteredLabelClassForPath,
   Span,
-  SpanLabel
 } from '@/lib/traces/types';
 import { cn, swrFetcher } from '@/lib/utils';
 import { useState } from 'react';
 import useSWR from 'swr';
 import {
-  ArrowDown,
   ChevronDown,
   Loader2,
   MoreVertical,
   Plus,
-  Sparkles,
   Tag,
   X
 } from 'lucide-react';
@@ -31,14 +27,6 @@ import {
   TableRow
 } from '../ui/table';
 import { AddLabel } from './add-label';
-import {
-  Select,
-  SelectContent,
-  SelectItem,
-  SelectTrigger,
-  SelectValue
-} from '../ui/select';
-import { useUserContext } from '@/contexts/user-context';
 import { v4 } from 'uuid';
 import { renderNodeInput } from '@/lib/flow/utils';
 import { PopoverClose } from '@radix-ui/react-popover';
@@ -397,7 +385,7 @@ function AddLabelInstance({
         },
         body: JSON.stringify({
           classId: labelClass.id,
-          value: labelClass.valueMap.findIndex((v) => v === value),
+          value: labelClass.valueMap[value],
           source: source,
           reasoning: reasoning
         })
@@ -491,15 +479,15 @@ function AddLabelInstance({
       <PopoverContent side="bottom" align="end">
         <div className="flex flex-col">
           <div className="flex flex-col space-y-2">
-            {labelClass.valueMap.map((value, index) => (
+            {Object.entries(labelClass.valueMap).map(([key, value], index) => (
               <PopoverClose key={index}>
                 <div
                   onClick={() => {
-                    addLabel(value, labelClass, LabelSource.MANUAL);
+                    addLabel(key, labelClass, LabelSource.MANUAL);
                   }}
                   className="cursor-pointer hover:bg-secondary-foreground/10 p-1 rounded border px-2"
                 >
-                  {value}
+                  {key}
                 </div>
               </PopoverClose>
             ))}

frontend/components/traces/add-label.tsx

@@ -34,16 +34,15 @@ export function AddLabel({ span, onClose }: AddLabelProps) {
     name: '',
     projectId: projectId,
     createdAt: '',
-    labelType: LabelType.BOOLEAN,
-    valueMap: ["false", "true"],
+    labelType: LabelType.CATEGORICAL,
+    valueMap: {
+      "False": 0,
+      "True": 1
+    },
     description: null,
     evaluatorRunnableGraph: null
   });
 
-  const isLabelValueMapValid =
-    labelClass.valueMap.length > 0 &&
-    labelClass.valueMap.every((value) => value.length > 0);
-
   const saveLabel = async () => {
     setIsSaving(true);
 
@@ -76,6 +75,11 @@ export function AddLabel({ span, onClose }: AddLabelProps) {
     onClose();
   };
 
+  const hasDuplicateValues = () => {
+    const values = Object.values(labelClass.valueMap);
+    return values.length !== new Set(values).size;
+  };
+
   return (
     <div className="flex-col items-center space-y-2">
       <div className="flex items-center space-x-2">
@@ -102,77 +106,98 @@ export function AddLabel({ span, onClose }: AddLabelProps) {
           onChange={(e) =>
             setLabelClass({ ...labelClass, description: e.target.value })
           }
-          minRows={3}
+          minRows={1}
         />
       </div>
-      <div className="flex-col space-y-1">
-        <Label>Type</Label>
-        <Select
-          value={selectedType}
-          onValueChange={(labelType) => {
-            setSelectedType(labelType as LabelType);
-
-            if (labelType === LabelType.BOOLEAN) {
-              setLabelClass({ ...labelClass, valueMap: ['false', 'true'] });
-            }
-          }}
-        >
-          <SelectTrigger>
-            <SelectValue placeholder="Select label type" />
-          </SelectTrigger>
-          <SelectContent>
-            <SelectItem value={LabelType.BOOLEAN}>Boolean</SelectItem>
-            <SelectItem value={LabelType.CATEGORICAL}>Categorical</SelectItem>
-          </SelectContent>
-        </Select>
-      </div>
-      {selectedType === LabelType.CATEGORICAL && (
-        <div className="flex flex-col space-y-2">
-          <div className="flex-col space-y-1">
-            <Label>Categorical values</Label>
-          </div>
 
-          {labelClass.valueMap.map((value, index) => (
-            <div key={index} className="flex space-x-2">
-              <Input
-                type="text"
-                placeholder="Categorical value"
-                onChange={(e) =>
-                  setLabelClass({
-                    ...labelClass,
-                    valueMap: labelClass.valueMap.map((value, i) =>
-                      i === index ? e.target.value : value
-                    )
-                  })
-                }
-              />
-              <Button
-                variant="ghost"
-                size="icon"
-                onClick={() =>
-                  setLabelClass({
-                    ...labelClass,
-                    valueMap: labelClass.valueMap.filter((_, i) => i !== index)
-                  })
-                }
-              >
-                <Trash2 size={14} />
-              </Button>
+      <div className="flex flex-col space-y-2">
+        <div className="flex-col space-y-1">
+          <Label>Label values</Label>
+          {hasDuplicateValues() && (
+            <div className="text-sm text-destructive">
+              Duplicate numerical values are not allowed
             </div>
-          ))}
-          <Button
-            variant="secondary"
-            onClick={() =>
-              setLabelClass({
-                ...labelClass,
-                valueMap: [...labelClass.valueMap, '']
-              })
-            }
-          >
-            Add categorical value
-          </Button>
+          )}
         </div>
-      )}
+
+        <table className="w-full">
+          <thead>
+            <tr className="text-sm text-secondary-foreground">
+              <th className="text-left pb-2 w-full">Value</th>
+              <th className="text-left pb-2">Numerical</th>
+              <th className="w-10"></th>
+            </tr>
+          </thead>
+          <tbody>
+            {Object.entries(labelClass.valueMap).map(([key, value], i) => (
+              <tr key={i}>
+                <td className="pr-2 pb-2">
+                  <Input
+                    type="text"
+                    placeholder="Categorical value"
+                    value={key}
+                    onChange={(e) =>
+                      setLabelClass({
+                        ...labelClass,
+                        valueMap: Object.fromEntries(
+                          Object.entries(labelClass.valueMap).map(([k, v], j) =>
+                            j === i ? [e.target.value, v] : [k, v]
+                          )
+                        )
+                      })
+                    }
+                  />
+                </td>
+                <td className="pr-2 pb-2">
+                  <Input
+                    type="number"
+                    className="w-24"
+                    placeholder="#"
+                    value={value}
+                    onChange={(e) => {
+                      setLabelClass({
+                        ...labelClass,
+                        valueMap: Object.fromEntries(
+                          Object.entries(labelClass.valueMap).map(([k, v], j) =>
+                            j === i ? [k, parseInt(e.target.value, 10)] : [k, v]
+                          )
+                        )
+                      });
+                    }}
+                  />
+                </td>
+                <td className="pb-2">
+                  <Button
+                    variant="ghost"
+                    size="icon"
+                    onClick={() =>
+                      setLabelClass({
+                        ...labelClass,
+                        valueMap: Object.fromEntries(
+                          Object.entries(labelClass.valueMap).filter(([_, v], j) => j !== i)
+                        )
+                      })
+                    }
+                  >
+                    <Trash2 size={14} />
+                  </Button>
+                </td>
+              </tr>
+            ))}
+          </tbody>
+        </table>
+        <Button
+          variant="secondary"
+          onClick={() =>
+            setLabelClass({
+              ...labelClass,
+              valueMap: { ...labelClass.valueMap, [""]: 0 }
+            })
+          }
+        >
+          Add label value
+        </Button>
+      </div>
       <div className="flex flex-col space-y-2">
         <div>
           <div className="flex items-center justify-between">
@@ -214,7 +239,7 @@ export function AddLabel({ span, onClose }: AddLabelProps) {
           onClick={async () => {
             await saveLabel();
           }}
-          disabled={!labelClass.name || !isLabelValueMapValid}
+          disabled={!labelClass.name || Object.keys(labelClass.valueMap).length === 0}
         >
           <Loader2
             className={isSaving ? 'animate-spin h-4 w-4 mr-2' : 'hidden'}

frontend/components/traces/span-labels.tsx

@@ -106,7 +106,7 @@ export default function SpanLabels({ spanId }: SpanLabelsProps) {
                     )}
                   </TableCell>
                   <TableCell>{label.labelSource}</TableCell>
-                  <TableCell>{label.valueMap?.[label.value] ?? ''}</TableCell>
+                  <TableCell>{Object.entries(label.valueMap).find(([k, v]) => v === label.value)?.[0] ?? ''}</TableCell>
                   <TableCell>{label.userEmail}</TableCell>
                   <TableCell>
                     <Button
@@ -123,7 +123,7 @@ export default function SpanLabels({ spanId }: SpanLabelsProps) {
             </TableBody>
           </Table>
         ) : (
-          <div className="p-2 text-secondary-foreground bg-secondary text-sm">
+          <div className="p-2 text-secondary-foreground bg-card text-sm">
             No labels
           </div>
         )}

frontend/lib/traces/types.ts

@@ -19,7 +19,7 @@ export type LabelClass = {
   projectId: string;
   createdAt: string;
   labelType: LabelType;
-  valueMap: string[];
+  valueMap: Record<string, number>;
   description: string | null;
   evaluatorRunnableGraph: any | null;
 };
@@ -38,7 +38,7 @@ export type SpanLabel = {
   classId: string;
   createdAt: string;
   value: number;
-  valueMap: string[];
+  valueMap: Record<string, number>;
   status: string;
   userId: string | null;
   jobStatus: string | null;