Merge commit 'a8390532f7' as 'ai-context/workbench-ui'

ai-context/workbench-ui/src/components/token-cost/Controls.tsx

@@ -0,0 +1,39 @@
+import React, { useState } from "react";
+
+import { Plus } from "lucide-react";
+
+import { Button, Box } from "@chakra-ui/react";
+
+import EditDialog from "./EditDialog";
+
+const Controls = ({ onUpdate }) => {
+  const [createOpen, setCreateOpen] = useState(false);
+
+  const onComplete = () => {
+    onUpdate();
+    setCreateOpen(false);
+  };
+
+  return (
+    <Box>
+      <Button
+        mt={5}
+        ml={5}
+        mb={5}
+        variant="solid"
+        colorPalette="primary"
+        onClick={() => setCreateOpen(true)}
+      >
+        <Plus /> Create Model Cost
+      </Button>
+      <EditDialog
+        open={createOpen}
+        onOpenChange={setCreateOpen}
+        create={true}
+        onComplete={() => onComplete()}
+      />
+    </Box>
+  );
+};
+
+export default Controls;

ai-context/workbench-ui/src/components/token-cost/EditDialog.tsx

@@ -0,0 +1,144 @@
+import React, { useEffect, useState } from "react";
+
+import { Trash, SendHorizontal } from "lucide-react";
+
+import { Portal, Button, Dialog, CloseButton } from "@chakra-ui/react";
+
+import { useTokenCosts } from "@trustgraph/react-state";
+import TextField from "../common/TextField";
+
+const EditDialog = ({ open, onOpenChange, model, create }) => {
+  const state = useTokenCosts();
+
+  const [newModel, setNewModel] = useState("");
+  const [input, setInput] = useState(0);
+  const [output, setOutput] = useState(0);
+
+  useEffect(() => {
+    if (!model) return;
+
+    const models = state.tokenCosts.filter((row) => row.model == model);
+
+    if (models.length < 1) return;
+
+    setNewModel(model);
+    setInput(models[0].input_price * 1000000);
+    setOutput(models[0].output_price * 1000000);
+  }, [state.tokenCosts, model, create]);
+
+  const onEdit = () => {
+    // Create is different from edit existing
+    if (create) {
+      // When creating, the order is...
+      // 1) write the prompt template,
+      // 2) get the template index
+      // 3) add this prompt ID to the index if not already there
+
+      state.updateTokenCost({
+        model: newModel,
+        input_price: input,
+        output_price: output,
+        onSuccess: () => onOpenChange(false),
+      });
+    } else {
+      state.updateTokenCost({
+        model: model,
+        input_price: input,
+        output_price: output,
+        onSuccess: () => onOpenChange(false),
+      });
+    }
+  };
+
+  const onDelete = () => {
+    // Shouldn't happen, but can't delete a prompt that hasn't been created
+    // yet
+    if (create) return;
+
+    state.deleteTokenCosts({
+      model: model,
+      onSuccess: () => onOpenChange(false),
+    });
+  };
+
+  return (
+    <Dialog.Root
+      placement="center"
+      size="xl"
+      open={open}
+      onOpenChange={(x) => {
+        onOpenChange(x.open);
+      }}
+    >
+      <Portal>
+        <Dialog.Backdrop />
+        <Dialog.Positioner>
+          <Dialog.Content>
+            <Dialog.Header>
+              {create && <Dialog.Title>Create model cost</Dialog.Title>}
+
+              {!create && (
+                <Dialog.Title>
+                  Edit model cost: <code>{model}</code>
+                </Dialog.Title>
+              )}
+            </Dialog.Header>
+            <Dialog.Body>
+              {create && (
+                <TextField
+                  label="Model ID"
+                  placeholder="Enter a unique model ID"
+                  value={newModel}
+                  onValueChange={(v) => setNewModel(v)}
+                  required={true}
+                />
+              )}
+
+              <TextField
+                label="Input token cost ( $ / 1Mt )"
+                placeholder="Input token cost"
+                value={input}
+                onValueChange={(v) => setInput(v)}
+                required={true}
+              />
+
+              <TextField
+                label="Output token cost ( $ / 1Mt )"
+                placeholder="Output token cost ($/token)"
+                value={output}
+                onValueChange={(v) => setOutput(v)}
+                required={true}
+              />
+            </Dialog.Body>
+            <Dialog.Footer>
+              <Button variant="outline" onClick={() => onOpenChange(false)}>
+                Cancel
+              </Button>
+              {
+                // If a 'create' operation, there's nothing to delete, only
+                // present if an existing prompt exists
+              }
+              {!create && (
+                <Button
+                  variant="solid"
+                  onClick={() => onDelete()}
+                  colorPalette="red"
+                >
+                  <Trash /> Delete
+                </Button>
+              )}
+              <Button onClick={() => onEdit()} colorPalette="primary">
+                <SendHorizontal /> Submit
+              </Button>
+            </Dialog.Footer>
+            <Dialog.CloseTrigger asChild>
+              <CloseButton size="sm" />
+            </Dialog.CloseTrigger>
+          </Dialog.Content>
+        </Dialog.Positioner>
+      </Portal>
+    </Dialog.Root>
+  );
+};
+
+export default EditDialog;

ai-context/workbench-ui/src/components/token-cost/TokenCosts.tsx

@@ -0,0 +1,44 @@
+import { useState } from "react";
+
+import { getCoreRowModel, useReactTable } from "@tanstack/react-table";
+
+import EditDialog from "./EditDialog";
+import Controls from "./Controls";
+import { useTokenCosts } from "@trustgraph/react-state";
+import { columns } from "../../model/token-costs-table";
+import ClickableTable from "../common/ClickableTable";
+
+const TokenCostTable = () => {
+  const state = useTokenCosts();
+
+  const tokenCosts = state.tokenCosts ? state.tokenCosts : [];
+
+  // Initialize React Table with document data and column configuration
+  const table = useReactTable({
+    data: tokenCosts,
+    columns: columns,
+    getCoreRowModel: getCoreRowModel(),
+  });
+
+  const [selected, setSelected] = useState("");
+
+  return (
+    <>
+      <EditDialog
+        open={selected != ""}
+        onOpenChange={() => setSelected("")}
+        create={false}
+        model={selected}
+      />
+
+      <ClickableTable
+        table={table}
+        onClick={(row) => setSelected(row.original.model)}
+      />
+
+      <Controls />
+    </>
+  );
+};
+
+export default TokenCostTable;